first commit
This commit is contained in:
@@ -0,0 +1,339 @@
|
||||
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
|
||||
Reference in New Issue
Block a user