Jerry
340 lines
11 KiB
Python
340 lines
11 KiB
Python
import os
|
|
import sys
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
|
|
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
if PROJECT_ROOT not in sys.path:
|
|
sys.path.insert(0, PROJECT_ROOT)
|
|
|
|
|
|
import tools.model_view_tools as model_view_tools_module
|
|
import tools.reference_image as reference_image_module
|
|
from tools.model_view_tools import get_visual_feedback
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_visual_feedback_uses_latest_uploaded_image_for_both_llm_steps(
|
|
monkeypatch, tmp_path
|
|
):
|
|
image_path = tmp_path / "query_image_001.png"
|
|
image_path.write_bytes(b"fake-png-bytes")
|
|
render_path = tmp_path / "render.png"
|
|
render_path.write_bytes(b"fake-render")
|
|
|
|
class FakeContext:
|
|
def retrieve_full_messages(self):
|
|
return [
|
|
SimpleNamespace(
|
|
role="user",
|
|
content=[
|
|
{"type": "text", "text": "Inspect this part"},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": "data:image/png;base64,abcd",
|
|
"local_path": str(image_path),
|
|
},
|
|
},
|
|
],
|
|
)
|
|
]
|
|
|
|
monkeypatch.setattr(
|
|
reference_image_module,
|
|
"get_current_context",
|
|
lambda: FakeContext(),
|
|
)
|
|
monkeypatch.setattr(
|
|
model_view_tools_module, "_require_simplecad_renderer", lambda: None
|
|
)
|
|
monkeypatch.setattr(
|
|
model_view_tools_module, "print_tool_output", lambda *args, **kwargs: None
|
|
)
|
|
|
|
class FakeSketchPad:
|
|
async def set_item(self, key, value, ttl=None, summary=None, tags=None):
|
|
return key
|
|
|
|
monkeypatch.setattr(
|
|
model_view_tools_module,
|
|
"get_current_sketch_pad",
|
|
lambda: FakeSketchPad(),
|
|
)
|
|
|
|
captured = {}
|
|
|
|
async def fake_question_generator(user_query, code, query_image_path):
|
|
captured["question_image_path"] = (
|
|
str(query_image_path.path) if query_image_path else None
|
|
)
|
|
return "Checklist"
|
|
|
|
async def fake_visual_feedback_generator(
|
|
questions, multi_view_results, query_image_path
|
|
):
|
|
captured["visual_image_path"] = (
|
|
str(query_image_path.path) if query_image_path else None
|
|
)
|
|
captured["multi_view_results"] = str(multi_view_results.path)
|
|
return "Looks correct\nPASS"
|
|
|
|
monkeypatch.setattr(
|
|
model_view_tools_module,
|
|
"question_generator",
|
|
fake_question_generator,
|
|
)
|
|
monkeypatch.setattr(
|
|
model_view_tools_module,
|
|
"visual_feedback_generator",
|
|
fake_visual_feedback_generator,
|
|
)
|
|
|
|
from SimpleLLMFunc.type import ImgPath
|
|
|
|
monkeypatch.setattr(
|
|
model_view_tools_module,
|
|
"render_multi_view_model",
|
|
lambda model_path, output_path: ImgPath(render_path, detail="high"),
|
|
)
|
|
|
|
result = await get_visual_feedback(
|
|
user_query="Inspect this part",
|
|
code="result = None",
|
|
model_path="./part/model.stl",
|
|
)
|
|
|
|
expected_image_path = str(image_path.resolve())
|
|
assert captured["question_image_path"] == expected_image_path
|
|
assert captured["visual_image_path"] == expected_image_path
|
|
assert captured["multi_view_results"] == str(render_path.resolve())
|
|
assert "Model path: ./part/model.stl" in result
|
|
|
|
|
|
def test_camera_relative_light_rig_tracks_camera_direction() -> None:
|
|
for view_dir in (
|
|
np.array([1.0, 0.0, 0.0]),
|
|
np.array([-1.0, 0.0, 0.0]),
|
|
np.array([0.0, 0.0, 1.0]),
|
|
np.array([1.0, 1.0, -1.0]),
|
|
):
|
|
normalized_view = view_dir / np.linalg.norm(view_dir)
|
|
light_dirs, light_weights, ambient = (
|
|
model_view_tools_module._camera_relative_light_rig(normalized_view)
|
|
)
|
|
|
|
assert ambient > 0.0
|
|
assert len(light_dirs) == len(light_weights) >= 3
|
|
assert np.dot(light_dirs[0], normalized_view) > 0.85
|
|
assert all(
|
|
abs(np.linalg.norm(light_dir) - 1.0) < 1e-6 for light_dir in light_dirs
|
|
)
|
|
|
|
|
|
def test_camera_relative_shading_keeps_front_faces_bright_across_views() -> None:
|
|
base_color = np.array([0.72, 0.76, 0.81], dtype=float)
|
|
front_face_brightness: list[float] = []
|
|
|
|
for view_dir in (
|
|
np.array([1.0, 0.0, 0.0]),
|
|
np.array([-1.0, 0.0, 0.0]),
|
|
np.array([0.0, 1.0, 0.0]),
|
|
np.array([1.0, 1.0, -1.0]),
|
|
):
|
|
normalized_view = view_dir / np.linalg.norm(view_dir)
|
|
normals = np.stack([normalized_view, -normalized_view])
|
|
shaded = model_view_tools_module._shade_normals_camera_relative(
|
|
normals,
|
|
base_color,
|
|
normalized_view,
|
|
)
|
|
front_face_brightness.append(float(np.mean(shaded[0, :3])))
|
|
assert float(np.mean(shaded[0, :3])) > float(np.mean(shaded[1, :3]))
|
|
|
|
assert min(front_face_brightness) > 0.55
|
|
assert max(front_face_brightness) - min(front_face_brightness) < 0.12
|
|
|
|
|
|
def test_surface_shading_normals_keep_planar_triangles_consistent() -> None:
|
|
class FakeVector:
|
|
def __init__(self, x: float, y: float, z: float) -> None:
|
|
self.x = x
|
|
self.y = y
|
|
self.z = z
|
|
|
|
class FakeCadFace:
|
|
def geomType(self) -> str:
|
|
return "PLANE"
|
|
|
|
def normalAt(self, location=None):
|
|
return FakeVector(0.0, 0.0, 1.0)
|
|
|
|
class FakeFace:
|
|
cq_face = FakeCadFace()
|
|
|
|
tri_pts = np.array(
|
|
[
|
|
[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0]],
|
|
[[0.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0]],
|
|
],
|
|
dtype=float,
|
|
)
|
|
|
|
normals = model_view_tools_module._compute_surface_shading_normals(
|
|
FakeFace(), tri_pts
|
|
)
|
|
|
|
assert normals.shape == (2, 3)
|
|
assert np.allclose(normals[0], [0.0, 0.0, 1.0])
|
|
assert np.allclose(normals[1], [0.0, 0.0, 1.0])
|
|
|
|
|
|
def test_direct_cad_renderer_is_used_for_brep_like_formats() -> None:
|
|
assert model_view_tools_module._should_use_direct_cad_renderer("part.step") is True
|
|
assert model_view_tools_module._should_use_direct_cad_renderer("part.stp") is True
|
|
assert model_view_tools_module._should_use_direct_cad_renderer("part.brep") is True
|
|
assert model_view_tools_module._should_use_direct_cad_renderer("part.bin") is True
|
|
assert model_view_tools_module._should_use_direct_cad_renderer("part.stl") is False
|
|
|
|
|
|
def test_prefer_cad_native_model_path_uses_step_over_stl(tmp_path) -> None:
|
|
stl_path = tmp_path / "model.stl"
|
|
step_path = tmp_path / "model.step"
|
|
stl_path.write_text("solid", encoding="utf-8")
|
|
step_path.write_text("step", encoding="utf-8")
|
|
|
|
selected = model_view_tools_module._prefer_cad_native_model_path(str(stl_path))
|
|
|
|
assert selected == str(step_path.resolve())
|
|
|
|
|
|
def test_prefer_cad_native_model_path_keeps_stl_without_step(tmp_path) -> None:
|
|
stl_path = tmp_path / "model.stl"
|
|
stl_path.write_text("solid", encoding="utf-8")
|
|
|
|
selected = model_view_tools_module._prefer_cad_native_model_path(str(stl_path))
|
|
|
|
assert selected == str(stl_path.resolve())
|
|
|
|
|
|
def test_load_renderable_shapes_flattens_step_compounds(monkeypatch) -> None:
|
|
class FakeCadShape:
|
|
def __init__(self, shape_type: str, solids=None) -> None:
|
|
self._shape_type = shape_type
|
|
self._solids = list(solids or [])
|
|
|
|
def ShapeType(self) -> str:
|
|
return self._shape_type
|
|
|
|
def Solids(self):
|
|
return list(self._solids)
|
|
|
|
class FakeWrappedSolid:
|
|
def __init__(self, obj) -> None:
|
|
self.obj = obj
|
|
|
|
class FakeWorkplane:
|
|
def __init__(self, values) -> None:
|
|
self._values = values
|
|
|
|
def vals(self):
|
|
return list(self._values)
|
|
|
|
solid_a = FakeCadShape("Solid")
|
|
solid_b = FakeCadShape("Solid")
|
|
compound = FakeCadShape("Compound", solids=[solid_a, solid_b])
|
|
|
|
fake_cq = SimpleNamespace(
|
|
importers=SimpleNamespace(
|
|
importShape=lambda import_type, model_path: FakeWorkplane([compound])
|
|
)
|
|
)
|
|
|
|
monkeypatch.setattr(model_view_tools_module, "cq", fake_cq)
|
|
monkeypatch.setattr(model_view_tools_module, "ScadSolid", FakeWrappedSolid)
|
|
monkeypatch.setattr(
|
|
model_view_tools_module, "_require_simplecad_renderer", lambda: None
|
|
)
|
|
|
|
result = model_view_tools_module._load_renderable_shapes("part.step")
|
|
|
|
assert [wrapped.obj for wrapped in result] == [solid_a, solid_b]
|
|
|
|
|
|
def test_feature_edge_mask_detects_normal_and_depth_discontinuities() -> None:
|
|
mask = np.ones((8, 8), dtype=bool)
|
|
depth = np.zeros((8, 8), dtype=float)
|
|
normals = np.zeros((8, 8, 3), dtype=float)
|
|
normals[:, :4] = np.array([0.0, 0.0, 1.0])
|
|
normals[:, 4:] = np.array([1.0, 0.0, 0.0])
|
|
|
|
edge_mask = model_view_tools_module._compute_feature_edge_mask(
|
|
mask,
|
|
depth,
|
|
normals,
|
|
depth_jump_threshold=10.0,
|
|
normal_cos_threshold=0.95,
|
|
)
|
|
|
|
assert edge_mask[:, 3:5].any()
|
|
|
|
depth[:, 4:] = 4.0
|
|
normals[:, :] = np.array([0.0, 0.0, 1.0])
|
|
edge_mask = model_view_tools_module._compute_feature_edge_mask(
|
|
mask,
|
|
depth,
|
|
normals,
|
|
depth_jump_threshold=1.0,
|
|
normal_cos_threshold=0.95,
|
|
)
|
|
|
|
assert edge_mask[:, 3:5].any()
|
|
|
|
|
|
def test_direct_rasterizer_supersamples_and_preserves_target_size() -> None:
|
|
triangles = [
|
|
np.array([[0.0, 0.0, 1.0], [1.0, 0.0, 1.0], [1.0, 1.0, 1.0]], dtype=float),
|
|
np.array([[0.0, 0.0, 1.0], [1.0, 1.0, 1.0], [0.0, 1.0, 1.0]], dtype=float),
|
|
]
|
|
normals = [
|
|
np.array([0.0, 0.0, 1.0], dtype=float),
|
|
np.array([0.0, 0.0, 1.0], dtype=float),
|
|
]
|
|
|
|
image = model_view_tools_module._rasterize_projected_triangles(
|
|
triangles,
|
|
normals,
|
|
image_size=(48, 48),
|
|
zoom=4.0,
|
|
background_rgb=np.array([255, 255, 255], dtype=np.uint8),
|
|
fill_rgb=np.array([180, 190, 200], dtype=np.uint8),
|
|
outline_rgb=np.array([0, 0, 0], dtype=np.uint8),
|
|
)
|
|
|
|
pixels = np.asarray(image)
|
|
unique_colors = np.unique(pixels.reshape(-1, 3), axis=0)
|
|
|
|
assert image.size == (48, 48)
|
|
assert len(unique_colors) > 3
|
|
|
|
|
|
def test_axis_triad_overlay_draws_small_corner_marker() -> None:
|
|
image = Image.new("RGB", (320, 320), "white")
|
|
annotated = model_view_tools_module._add_axis_triad_overlay(
|
|
image,
|
|
np.array([1.0, 1.0, 1.0]) / np.sqrt(3.0),
|
|
)
|
|
|
|
original = np.asarray(image)
|
|
updated = np.asarray(annotated)
|
|
diff = np.abs(updated.astype(int) - original.astype(int)).sum(axis=2)
|
|
changed_pixels = np.argwhere(diff > 0)
|
|
|
|
assert changed_pixels.size > 0
|
|
assert int(changed_pixels[:, 0].max()) > 220
|
|
assert int(changed_pixels[:, 1].max()) < 120
|