AnyImage/tests/test_worker.py

import importlib.util
import json
import os
import sys
import tempfile
import unittest
from contextlib import nullcontext
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import patch

import numpy as np


PROJECT_ROOT = Path(__file__).resolve().parents[1]


def load_worker():
    specification = importlib.util.spec_from_file_location(
        "snapform_worker",
        PROJECT_ROOT / "src" / "snapform" / "worker.py",
    )
    module = importlib.util.module_from_spec(specification)
    specification.loader.exec_module(module)
    return module


class WorkerTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.worker = load_worker()

    def test_foreground_preserves_source_rgb_and_ben2_alpha(self):
        alpha = SimpleNamespace(size=(2, 1))
        source_rgba = SimpleNamespace(size=(2, 1), applied_alpha=None)
        source_rgba.putalpha = lambda value: setattr(
            source_rgba,
            "applied_alpha",
            value,
        )
        source = SimpleNamespace(convert=lambda _mode: source_rgba)
        foreground_rgba = SimpleNamespace(
            getchannel=lambda _name: alpha,
        )
        foreground = SimpleNamespace(convert=lambda _mode: foreground_rgba)

        result = self.worker.preserve_source_rgb(source, foreground)

        self.assertIs(result, source_rgba)
        self.assertIs(result.applied_alpha, alpha)

    def test_collect_images_is_sorted_and_filters_suffixes(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / "b.PNG").touch()
            (root / "a.jpg").touch()
            (root / "notes.txt").touch()

            images = self.worker.collect_images(root)

        self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"])

    def test_ai_tasks_support_cooperative_cancellation(self):
        with self.assertRaises(self.worker.TaskCancelled):
            self.worker.check_cancelled(lambda: True)

        self.worker.check_cancelled(lambda: False)

    def test_reconstruct_parser_accepts_background_model_directory(self):
        args = self.worker.parser().parse_args(
            [
                "reconstruct",
                "--input",
                "input.png",
                "--model-dir",
                "da3",
                "--output-dir",
                "result",
                "--export-format",
                "mini_npz",
                "--status",
                "status.json",
                "--remove-background-model-dir",
                "ben2",
                "--process-res",
                "1024",
                "--process-res-method",
                "upper_bound_resize",
                "--conf-thresh-percentile",
                "40",
                "--num-max-points",
                "1000",
                "--show-cameras",
                "0",
                "--gs-views-interval",
                "1",
                "--video-frames",
                "16",
                "--device",
                "auto",
                "--use-ray-pose",
                "1",
                "--ref-view-strategy",
                "first",
            ]
        )
        self.assertEqual(args.remove_background_model_dir, "ben2")

    def test_prepares_a_single_image_file(self):
        with tempfile.TemporaryDirectory() as directory:
            image_path = Path(directory) / "input.png"
            image_path.touch()

            images, temporary_directory = self.worker.prepare_input(
                image_path,
                16,
                directory,
            )

        self.assertEqual(images, [image_path])
        self.assertIsNone(temporary_directory)

    def test_prepares_a_frame_directory(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / "frame-0001.jpg").touch()
            (root / "frame-0002.png").touch()
            (root / "notes.txt").touch()

            images, temporary_directory = self.worker.prepare_input(
                root,
                16,
                directory,
            )

        self.assertEqual(
            [path.name for path in images],
            ["frame-0001.jpg", "frame-0002.png"],
        )
        self.assertIsNone(temporary_directory)

    def test_prepares_first_frames_of_a_frame_directory(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            for index in range(1, 6):
                (root / f"frame-{index:04d}.jpg").touch()

            images, temporary_directory = self.worker.prepare_input(
                root,
                3,
                directory,
            )

        self.assertEqual(
            [path.name for path in images],
            ["frame-0001.jpg", "frame-0002.jpg", "frame-0003.jpg"],
        )
        self.assertIsNone(temporary_directory)

    def test_prepares_first_frames_after_offset_in_a_frame_directory(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            for index in range(1, 6):
                (root / f"frame-{index:04d}.jpg").touch()

            images, temporary_directory = self.worker.prepare_input(
                root,
                3,
                directory,
                start_frame=2,
            )

        self.assertEqual(
            [path.name for path in images],
            ["frame-0002.jpg", "frame-0003.jpg", "frame-0004.jpg"],
        )
        self.assertIsNone(temporary_directory)

    def test_extract_video_frames_reads_frames_sequentially(self):
        class FakeCapture:
            def __init__(self):
                self.position = 0

            def isOpened(self):
                return True

            def read(self):
                if self.position >= 100:
                    return False, None
                frame = self.position
                self.position += 1
                return True, frame

            def release(self):
                return None

        class FakeCv2:
            def __init__(self):
                self.captures = []

            def VideoCapture(self, _path):
                capture = FakeCapture()
                self.captures.append(capture)
                return capture

            def imwrite(self, path, frame):
                Path(path).write_text(str(frame), encoding="utf-8")
                return True

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            fake_cv2 = FakeCv2()
            image_paths = self.worker.extract_video_frames(
                root / "clip.mkv",
                root / "out",
                10,
                fake_cv2,
            )

        self.assertEqual(
            [path.name for path in image_paths],
            [f"frame.{index:04d}.jpg" for index in range(1, 11)],
        )
        self.assertEqual(len(fake_cv2.captures), 1)
        self.assertEqual(fake_cv2.captures[0].position, 10)

    def test_extract_video_frames_starts_after_offset(self):
        class FakeCapture:
            def __init__(self):
                self.position = 0

            def isOpened(self):
                return True

            def read(self):
                if self.position >= 100:
                    return False, None
                frame = self.position
                self.position += 1
                return True, frame

            def release(self):
                return None

        class FakeCv2:
            def __init__(self):
                self.captures = []

            def VideoCapture(self, _path):
                capture = FakeCapture()
                self.captures.append(capture)
                return capture

            def imwrite(self, path, frame):
                Path(path).write_text(str(frame), encoding="utf-8")
                return True

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            fake_cv2 = FakeCv2()
            image_paths = self.worker.extract_video_frames(
                root / "clip.mkv",
                root / "out",
                10,
                fake_cv2,
                start_frame=71,
            )
            extracted = [
                Path(path).read_text(encoding="utf-8")
                for path in image_paths
            ]

        self.assertEqual(extracted, [str(frame) for frame in range(70, 80)])
        self.assertEqual(fake_cv2.captures[0].position, 80)

    def test_extract_video_frames_tolerates_short_video_tail(self):
        class ShortCapture:
            def __init__(self):
                self.position = 0

            def isOpened(self):
                return True

            def read(self):
                if self.position >= 99:
                    return False, None
                frame = self.position
                self.position += 1
                return True, frame

            def release(self):
                return None

        class FakeCv2:
            def VideoCapture(self, _path):
                return ShortCapture()

            def imwrite(self, path, frame):
                Path(path).write_text(str(frame), encoding="utf-8")
                return True

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            image_paths = self.worker.extract_video_frames(
                root / "clip.mkv",
                root / "out",
                10,
                FakeCv2(),
            )

        self.assertEqual(
            [path.name for path in image_paths],
            [f"frame.{index:04d}.jpg" for index in range(1, 11)],
        )

    def test_npz_is_complete_requires_readable_archive(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            valid_path = root / "valid.npz"
            np.savez_compressed(
                valid_path,
                depth=np.zeros((4, 4)),
            )
            self.assertTrue(self.worker.npz_is_complete(valid_path))

            empty_path = root / "empty.npz"
            empty_path.write_bytes(b"PK\x05\x06" + b"\x00" * 18)
            self.assertFalse(self.worker.npz_is_complete(empty_path))

    def test_wait_for_exported_paths_waits_for_complete_npz(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            npz_path = root / "exports" / "mini_npz" / "results.npz"
            npz_path.parent.mkdir(parents=True)
            npz_path.touch()
            with patch.object(
                self.worker,
                "npz_is_complete",
                side_effect=[False, True],
            ):
                result = self.worker.wait_for_exported_paths(
                    root,
                    "mini_npz",
                    timeout=2.0,
                )
            self.assertEqual(result, [npz_path])

    def test_move_exported_file_retries_windows_file_lock(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source = root / "source.npz"
            destination = root / "destination.npz"
            source.touch()
            path_type = type(source)
            original_replace = path_type.replace
            calls = 0

            def locked_once(path, target):
                nonlocal calls
                calls += 1
                if calls == 1:
                    error = PermissionError("file is in use")
                    error.winerror = 32
                    raise error
                return original_replace(path, target)

            with patch.object(path_type, "replace", new=locked_once), patch.object(
                self.worker.time,
                "sleep",
            ) as sleep:
                result = self.worker.move_exported_file(
                    source,
                    destination,
                )

            self.assertEqual(result, destination)
            self.assertTrue(destination.is_file())
            sleep.assert_called_once_with(0.1)

    def test_rejects_unsupported_input_file(self):
        with tempfile.TemporaryDirectory() as directory:
            input_path = Path(directory) / "input.txt"
            input_path.touch()

            with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
                self.worker.prepare_input(input_path, 16, directory)

    def test_download_model_uses_actual_request_without_preflight(self):
        calls = []
        fake_hugging_face = ModuleType("huggingface_hub")
        fake_hugging_face.snapshot_download = lambda **options: calls.append(
            options
        )
        args = SimpleNamespace(
            repository="depth-anything/DA3-SMALL",
            model_dir="model",
            status="status.json",
            allow_pattern=[],
        )

        with tempfile.TemporaryDirectory() as directory, patch.dict(
            sys.modules,
            {"huggingface_hub": fake_hugging_face},
        ), patch.object(
            self.worker,
            "create_download_progress_class",
            return_value=object,
        ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
            self.worker,
            "_verify_model_complete",
            return_value=None,
        ):
            args.model_dir = str(Path(directory) / "model")
            args.status = str(Path(directory) / "status.json")
            self.worker.download_model(args)

        self.assertEqual(len(calls), 1)

    def test_download_model_refuses_success_when_files_missing(self):
        fake_hugging_face = ModuleType("huggingface_hub")
        fake_hugging_face.snapshot_download = lambda **options: None
        args = SimpleNamespace(
            repository="PramaLLC/BEN2",
            model_dir="model",
            status="status.json",
            allow_pattern=[],
        )

        with tempfile.TemporaryDirectory() as directory, patch.dict(
            sys.modules,
            {"huggingface_hub": fake_hugging_face},
        ), patch.object(
            self.worker,
            "create_download_progress_class",
            return_value=object,
        ), patch.dict(os.environ, {"HF_ENDPOINT": ""}):
            args.model_dir = str(Path(directory) / "model")
            args.status = str(Path(directory) / "status.json")
            with self.assertRaisesRegex(
                RuntimeError,
                "All download sources failed",
            ):
                self.worker.download_model(args)

    def test_download_model_reports_all_failed_sources(self):
        fake_hugging_face = ModuleType("huggingface_hub")

        def fake_snapshot_download(**options):
            raise OSError("direct down")

        fake_hugging_face.snapshot_download = fake_snapshot_download
        args = SimpleNamespace(
            repository="depth-anything/DA3-SMALL",
            model_dir="model",
            status="status.json",
            allow_pattern=[],
        )

        with tempfile.TemporaryDirectory() as directory, patch.dict(
            sys.modules,
            {"huggingface_hub": fake_hugging_face},
        ), patch.object(
            self.worker,
            "create_download_progress_class",
            return_value=object,
        ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
            self.worker,
            "download_from_r2",
            side_effect=OSError("mirror down"),
        ):
            args.r2_directory = "DA3Small"
            args.r2_file = ["config.json|100|abc"]
            args.model_dir = str(Path(directory) / "model")
            args.status = str(Path(directory) / "status.json")
            with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"):
                self.worker.download_model(args)

    def test_download_from_r2_downloads_declared_files(self):
        args = SimpleNamespace(
            repository="depth-anything/DA3-SMALL",
            model_dir="model",
            status="status.json",
            allow_pattern=["config.json"],
            r2_directory="DA3Small",
            r2_file=["config.json|100|abc"],
        )
        downloaded = []

        with tempfile.TemporaryDirectory() as directory, patch.object(
            self.worker,
            "_download_r2_file",
            side_effect=lambda *args, **kwargs: downloaded.append(args[1]),
        ), patch.object(self.worker, "write_status"):
            args.model_dir = str(Path(directory) / "model")
            args.status = str(Path(directory) / "status.json")
            self.worker.download_from_r2(args, args.model_dir)

        self.assertEqual(downloaded, ["config.json"])

    def test_ben2_background_removal_writes_fast_transparent_png(self):
        inference_options = []
        save_options = []

        class FakeImage:
            size = (1, 1)

            def __enter__(self):
                return self

            def __exit__(self, *_args):
                return None

            def convert(self, mode):
                self.mode = mode
                return self

            def getchannel(self, _name):
                return self

            def putalpha(self, _alpha):
                return None

            def save(self, path, **options):
                save_options.append(options)
                Path(path).touch()

        class FakeModel:
            @classmethod
            def from_pretrained(cls, model_dir):
                self.assertEqual(model_dir, "model")
                return cls()

            def to(model_self, device):
                self.assertEqual(device.type, "cpu")
                return model_self

            def eval(self):
                return None

            def inference(self, image, **options):
                inference_options.append((image, options))
                return FakeImage()

        fake_torch = ModuleType("torch")
        fake_torch.cuda = SimpleNamespace(
            is_available=lambda: False,
            synchronize=lambda: None,
        )
        fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
        fake_torch.inference_mode = nullcontext
        fake_ben2 = ModuleType("ben2_onnx")
        fake_ben2.create_session = lambda _model_dir, _device: object()
        fake_ben2.infer = lambda _model, image: (
            inference_options.append((image, {})) or FakeImage()
        )
        fake_pil = ModuleType("PIL")
        fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())

        with tempfile.TemporaryDirectory() as directory, patch.dict(
            sys.modules,
            {
                "torch": fake_torch,
                "ben2_onnx": fake_ben2,
                "PIL": fake_pil,
            },
        ):
            root = Path(directory)
            input_path = root / "input.png"
            input_path.touch()
            args = SimpleNamespace(
                input=str(input_path),
                model_dir="model",
                output_dir=str(root / "result"),
                status=str(root / "status.json"),
                device="cpu",
            )
            self.worker.remove_background(args)
            status = json.loads(Path(args.status).read_text(encoding="utf-8"))

            self.assertTrue((root / "result" / "foreground.png").is_file())

        self.assertEqual(
            inference_options[0][1],
            {},
        )
        self.assertEqual(
            save_options[0],
            {"format": "PNG", "compress_level": 1},
        )
        self.assertIn("Background removed in", status["message"])

    def test_ben2_background_removal_processes_each_frame(self):
        inference_options = []

        class FakeImage:
            size = (1, 1)

            def __enter__(self):
                return self

            def __exit__(self, *_args):
                return None

            def convert(self, mode):
                self.mode = mode
                return self

            def getchannel(self, _name):
                return self

            def putalpha(self, _alpha):
                return None

            def save(self, path, **options):
                Path(path).touch()

        class FakeModel:
            @classmethod
            def from_pretrained(cls, model_dir):
                return cls()

            def to(model_self, device):
                return model_self

            def eval(self):
                return None

            def inference(self, image, **options):
                inference_options.append((image, options))
                return FakeImage()

        fake_torch = ModuleType("torch")
        fake_torch.cuda = SimpleNamespace(
            is_available=lambda: False,
            synchronize=lambda: None,
        )
        fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
        fake_torch.inference_mode = nullcontext
        fake_ben2 = ModuleType("ben2_onnx")
        fake_ben2.create_session = lambda _model_dir, _device: object()
        fake_ben2.infer = lambda _model, image: (
            inference_options.append((image, {})) or FakeImage()
        )
        fake_pil = ModuleType("PIL")
        fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())

        with tempfile.TemporaryDirectory() as directory, patch.dict(
            sys.modules,
            {
                "torch": fake_torch,
                "ben2_onnx": fake_ben2,
                "PIL": fake_pil,
            },
        ):
            root = Path(directory)
            input_directory = root / "frames"
            input_directory.mkdir()
            (input_directory / "frame-0001.jpg").touch()
            (input_directory / "frame-0002.jpg").touch()
            args = SimpleNamespace(
                input=str(input_directory),
                model_dir="model",
                output_dir=str(root / "result"),
                status=str(root / "status.json"),
                device="cpu",
                video_frames=8,
            )
            self.worker.remove_background(args)

            result_directory = Path(args.output_dir)
            self.assertTrue(
                (result_directory / "foreground.0001.png").is_file()
            )
            self.assertTrue(
                (result_directory / "foreground.0002.png").is_file()
            )
            self.assertTrue(
                (result_directory / "frame.0001.jpg").is_file()
            )
            self.assertTrue(
                (result_directory / "frame.0002.jpg").is_file()
            )

        self.assertEqual(len(inference_options), 2)
        for _image, options in inference_options:
            self.assertEqual(options, {})

    def test_remove_background_parser_accepts_video_frames(self):
        args = self.worker.parser().parse_args(
            [
                "remove-background",
                "--input",
                "input.mp4",
                "--model-dir",
                "ben2",
                "--output-dir",
                "result",
                "--status",
                "status.json",
                "--device",
                "auto",
                "--video-frames",
                "24",
            ]
        )

        self.assertEqual(args.command, "remove-background")
        self.assertEqual(args.video_frames, 24)

    def test_extract_frames_parser_accepts_frame_range(self):
        args = self.worker.parser().parse_args(
            [
                "extract-frames",
                "--input",
                "frames",
                "--output-dir",
                "result",
                "--status",
                "status.json",
                "--video-frames",
                "24",
                "--start-frame",
                "70",
            ]
        )

        self.assertEqual(args.command, "extract-frames")
        self.assertEqual(args.video_frames, 24)
        self.assertEqual(args.start_frame, 70)

    def test_extract_frames_copies_selected_frames(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            input_directory = root / "frames"
            input_directory.mkdir()
            for index in range(1, 6):
                (input_directory / f"frame-{index:04d}.jpg").touch()
            args = SimpleNamespace(
                input=str(input_directory),
                output_dir=str(root / "result"),
                status=str(root / "status.json"),
                video_frames=3,
                start_frame=2,
            )

            self.worker.extract_frames(args)

            result_directory = Path(args.output_dir)
            self.assertEqual(
                sorted(
                    path.name
                    for path in result_directory.iterdir()
                    if path.is_file()
                ),
                ["frame.0001.jpg", "frame.0002.jpg", "frame.0003.jpg"],
            )
            status = json.loads(
                Path(args.status).read_text(encoding="utf-8")
            )
            self.assertEqual(status["progress"], 1.0)
            self.assertIn("3 frame", status["message"])

    def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self):
        inference_options = []
        mesh_calls = []

        prediction = SimpleNamespace(
            depth=np.full((1, 2, 2), 2.0, dtype=np.float32),
            conf=np.ones((1, 2, 2), dtype=np.float32),
            extrinsics=np.eye(4, dtype=np.float32)[None, :3],
            intrinsics=np.eye(3, dtype=np.float32)[None],
            processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
        )
        fake_da3 = ModuleType("da3_onnx")
        fake_da3.create_session = lambda _model_dir, _device: object()

        def fake_infer(_session, images, **options):
            inference_options.append((images, options))
            return prediction

        fake_da3.infer = fake_infer
        fake_runtime = ModuleType("onnx_runtime")
        fake_runtime.provider_label = lambda _session: "CPU"

        def fake_mesh_export(
            prediction,
            images,
            output_directory,
            confidence_percentile,
            filter_edges,
            max_edge_length,
        ):
            mesh_calls.append(
                (
                    prediction,
                    images,
                    confidence_percentile,
                    filter_edges,
                    max_edge_length,
                )
            )
            output_path = Path(output_directory) / "mesh.glb"
            output_path.touch()
            return output_path

        def fake_point_cloud(_prediction, output_directory, *_args):
            (Path(output_directory) / "scene.glb").touch()

        def fake_depth_preview(_prediction, output_directory):
            target = Path(output_directory) / "depth_vis"
            target.mkdir()
            (target / "0000.jpg").touch()

        def fake_mini_npz(_prediction, output_directory):
            np.savez_compressed(
                Path(output_directory) / "results.npz",
                depth=np.zeros((1, 2, 2)),
            )

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            input_path = root / "input.png"
            input_path.touch()
            with patch.dict(
                sys.modules,
                {
                    "da3_onnx": fake_da3,
                    "onnx_runtime": fake_runtime,
                },
            ), patch.object(
                self.worker,
                "export_textured_mesh_glb",
                side_effect=fake_mesh_export,
            ), patch.object(
                self.worker,
                "export_da3_point_cloud",
                side_effect=fake_point_cloud,
            ), patch.object(
                self.worker,
                "export_da3_depth_preview",
                side_effect=fake_depth_preview,
            ), patch.object(
                self.worker,
                "export_da3_mini_npz",
                side_effect=fake_mini_npz,
            ):
                debug_formats = (
                    "glb",
                    "depth_vis",
                    "mini_npz",
                    "mesh",
                )
                for export_format in debug_formats:
                    args = SimpleNamespace(
                        input=str(input_path),
                        video_frames=16,
                        output_dir=str(root / export_format),
                        export_format=export_format,
                        status=str(root / f"{export_format}.json"),
                        device="cpu",
                        model_dir="model",
                        process_res=518,
                        process_res_method="lower_bound_resize",
                        use_ray_pose=1,
                        ref_view_strategy="middle",
                        conf_thresh_percentile=35.0,
                        num_max_points=12345,
                        show_cameras=0,
                        gs_views_interval=2,
                        mesh_filter_edges=True,
                        mesh_max_edge_length=0.2,
                    )
                    self.worker.reconstruct(args)
                    status = json.loads(
                        Path(args.status).read_text(encoding="utf-8")
                    )
                    self.assertEqual(status["progress"], 1.0)
                    self.assertIn(export_format, status["message"])
                    result_directory = Path(args.output_dir)
                    if export_format == "depth_vis":
                        self.assertTrue((result_directory / "0000.jpg").is_file())
                        self.assertFalse((result_directory / "depth_vis").exists())

        self.assertEqual(len(inference_options), 4)
        self.assertEqual(len(mesh_calls), 1)
        for images, options in inference_options:
            self.assertEqual(images, [str(input_path)])
            self.assertEqual(options["process_res"], 518)
            self.assertEqual(
                options["process_res_method"],
                "lower_bound_resize",
            )
            self.assertTrue(options["use_ray_pose"])
            self.assertEqual(options["ref_view_strategy"], "middle")
        self.assertEqual(mesh_calls[0][1], [input_path])
        self.assertEqual(mesh_calls[0][2:], (35.0, True, 0.2))

    def test_reconstruct_copies_frames_for_multi_frame_input(self):
        fake_da3 = ModuleType("da3_onnx")
        fake_da3.create_session = lambda _model_dir, _device: object()
        fake_da3.infer = lambda *_args, **_kwargs: SimpleNamespace(
            depth=np.full((2, 2, 2), 2.0, dtype=np.float32),
            conf=np.ones((2, 2, 2), dtype=np.float32),
            extrinsics=np.repeat(np.eye(4, dtype=np.float32)[None, :3], 2, axis=0),
            intrinsics=np.repeat(np.eye(3, dtype=np.float32)[None], 2, axis=0),
            processed_images=np.zeros((2, 2, 2, 3), dtype=np.uint8),
        )
        fake_runtime = ModuleType("onnx_runtime")
        fake_runtime.provider_label = lambda _session: "CPU"

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            input_directory = root / "frames"
            input_directory.mkdir()
            (input_directory / "frame-0001.jpg").touch()
            (input_directory / "frame-0002.jpg").touch()
            with patch.dict(
                sys.modules,
                {
                    "da3_onnx": fake_da3,
                    "onnx_runtime": fake_runtime,
                },
            ):
                args = SimpleNamespace(
                    input=str(input_directory),
                    video_frames=16,
                    output_dir=str(root / "result"),
                    export_format="mini_npz",
                    status=str(root / "status.json"),
                    device="cpu",
                    model_dir="model",
                    process_res=518,
                    process_res_method="upper_bound_resize",
                    use_ray_pose=1,
                    ref_view_strategy="first",
                    conf_thresh_percentile=40.0,
                    num_max_points=1000,
                    show_cameras=0,
                    gs_views_interval=1,
                    remove_background_model_dir="",
                )
                self.worker.reconstruct(args)

            result_directory = Path(args.output_dir)
            self.assertTrue(
                (result_directory / "frame.0001.jpg").is_file()
            )
            self.assertTrue(
                (result_directory / "frame.0002.jpg").is_file()
            )
            self.assertTrue(
                (result_directory / "results.npz").is_file()
            )
            self.assertFalse(
                (result_directory / "exports").exists()
            )
            status = json.loads(Path(args.status).read_text(encoding="utf-8"))
            self.assertEqual(status["progress"], 1.0)

    def test_mesh_confidence_uses_percentile(self):
        depth = np.ones((1, 2, 2), dtype=np.float32)
        confidence = np.array([[[1.0, 2.0], [3.0, 4.0]]])

        threshold = self.worker.confidence_percentile_threshold(
            depth,
            confidence,
            50.0,
        )

        self.assertEqual(threshold, 2.5)

    def test_zero_mesh_confidence_disables_confidence_filter(self):
        depth = np.ones((2, 2), dtype=np.float32)
        confidence = np.full((2, 2), np.nan, dtype=np.float32)
        x, y = np.meshgrid([0.0, 1.0], [0.0, 1.0])
        points = np.stack((x, y, depth), axis=-1)

        threshold = self.worker.confidence_percentile_threshold(
            depth,
            confidence,
            0.0,
        )
        faces = self.worker.mesh_triangle_faces(
            points,
            depth,
            confidence,
            threshold,
            filter_edges=False,
            max_edge_length=0.0,
        )

        self.assertIsNone(threshold)
        self.assertEqual(len(faces), 2)

    def test_mesh_filters_triangles_independently(self):
        depth = np.ones((2, 2), dtype=np.float32)
        confidence = np.ones((2, 2), dtype=np.float32)
        confidence[1, 0] = 0.0
        x, y = np.meshgrid(
            np.arange(2, dtype=np.float32),
            np.arange(2, dtype=np.float32),
        )
        points = np.stack((x, y, depth), axis=-1)

        faces = self.worker.mesh_triangle_faces(
            points,
            depth,
            confidence,
            confidence_threshold=0.5,
            filter_edges=False,
            max_edge_length=0.0,
        )

        np.testing.assert_array_equal(faces, [[0, 3, 1]])

    def test_mesh_background_mask_combines_with_confidence(self):
        depth = np.ones((2, 3), dtype=np.float32)
        confidence = np.ones((2, 3), dtype=np.float32)
        x, y = np.meshgrid(
            np.arange(3, dtype=np.float32),
            np.arange(2, dtype=np.float32),
        )
        points = np.stack((x, y, depth), axis=-1)
        foreground = np.ones((2, 3), dtype=bool)
        foreground[:, 0] = False

        faces = self.worker.mesh_triangle_faces(
            points,
            depth,
            confidence,
            confidence_threshold=None,
            filter_edges=False,
            max_edge_length=0.0,
            vertex_mask=foreground,
        )

        np.testing.assert_array_equal(faces, [[1, 4, 5], [1, 5, 2]])

    def test_depth_edge_filter_ignores_smooth_slope(self):
        smooth = np.tile(
            np.linspace(2.0, 2.04, 5, dtype=np.float32),
            (3, 1),
        )
        discontinuous = smooth.copy()
        discontinuous[:, 3:] += 1.0

        self.assertFalse(self.worker.depth_edge_mask(smooth).any())
        self.assertTrue(self.worker.depth_edge_mask(discontinuous)[:, 2:4].all())

    def test_higher_edge_culling_allows_less_stretch(self):
        limits = [
            self.worker.edge_culling_stretch_limit(value)
            for value in (0.0, 50.0, 100.0)
        ]

        self.assertEqual(limits, [10000.0, 2.0, 1.0])

    def test_stretched_edge_culling_uses_camera_space_edge_ratio(self):
        flat_points = np.asarray(
            [
                [[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]],
                [[0.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
            ],
            dtype=np.float32,
        )
        stretched_points = flat_points.copy()
        stretched_points[0, 1] = (2.0, 0.0, 2.0)
        triangles = np.asarray([[0, 2, 3], [0, 3, 1]], dtype=np.int32)

        flat_mask = self.worker.triangle_edge_stretch_mask(
            flat_points,
            triangles,
            1.0,
        )
        loose_mask = self.worker.triangle_edge_stretch_mask(
            stretched_points,
            triangles,
            2.0,
        )
        strict_mask = self.worker.triangle_edge_stretch_mask(
            stretched_points,
            triangles,
            1.0,
        )

        self.assertTrue(flat_mask.all())
        self.assertTrue(loose_mask.all())
        self.assertFalse(strict_mask.all())

    def test_mesh_faces_apply_confidence_threshold_and_face_direction(self):
        points = np.array(
            [
                [[0.0, 0.0, 1.0], [1.0, 0.0, 1.0], [2.0, 0.0, 1.0]],
                [[0.0, 1.0, 1.0], [1.0, 1.0, 1.0], [2.0, 1.0, 1.0]],
            ],
            dtype=np.float32,
        )
        depth = np.ones((2, 3), dtype=np.float32)
        confidence = np.ones((2, 3), dtype=np.float32)
        confidence[:, 0] = 0.0

        faces = self.worker.mesh_triangle_faces(
            points,
            depth,
            confidence,
            confidence_threshold=0.5,
            filter_edges=False,
            max_edge_length=0.0,
        )

        np.testing.assert_array_equal(faces, [[1, 4, 5], [1, 5, 2]])
        vertices = points.reshape(-1, 3)[faces[0]]
        normal = np.cross(
            vertices[1] - vertices[0],
            vertices[2] - vertices[0],
        )
        self.assertLess(normal[2], 0.0)

    def test_mesh_percentile_reduces_face_count(self):
        depth = np.ones((3, 3), dtype=np.float32)
        confidence = np.arange(1, 10, dtype=np.float32).reshape(3, 3)
        x, y = np.meshgrid(
            np.arange(3, dtype=np.float32),
            np.arange(3, dtype=np.float32),
        )
        points = np.stack((x, y, depth), axis=-1)
        face_counts = []
        for percentile in (0.0, 50.0, 100.0):
            threshold = self.worker.confidence_percentile_threshold(
                depth,
                confidence,
                percentile,
            )
            faces = self.worker.mesh_triangle_faces(
                points,
                depth,
                confidence,
                confidence_threshold=threshold,
                filter_edges=False,
                max_edge_length=0.0,
            )
            face_counts.append(len(faces))

        self.assertEqual(face_counts, [8, 2, 0])

    def test_mesh_compaction_removes_unreferenced_vertices(self):
        vertices = np.arange(18, dtype=np.float32).reshape(6, 3)
        faces = np.asarray([[1, 3, 5]], dtype=np.int32)

        compact_vertices, compact_faces = self.worker.compact_mesh(
            vertices,
            faces,
        )

        np.testing.assert_array_equal(
            compact_vertices,
            vertices[[1, 3, 5]],
        )
        np.testing.assert_array_equal(compact_faces, [[0, 1, 2]])

    def test_surface_background_faces_are_removed_before_export(self):
        faces = np.asarray(
            ((0, 1, 4), (0, 4, 3), (1, 2, 5), (1, 5, 4)),
            dtype=np.int32,
        )
        foreground = np.asarray(
            ((False, True, True), (False, True, True)),
            dtype=bool,
        )

        filtered = self.worker.remove_background_faces(faces, foreground)

        np.testing.assert_array_equal(filtered, faces[2:])

    def test_image_mesh_data_keeps_camera_geometry_and_uv(self):
        class FakeImage:
            def __enter__(self):
                return self

            def __exit__(self, *_args):
                return None

            def convert(self, mode):
                self.mode = mode
                return self

            def save(self, path, **_options):
                Path(path).write_bytes(b"texture")

        fake_pil = ModuleType("PIL")
        fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
        depth = np.full((1, 3, 3), 2.0, dtype=np.float32)
        confidence = np.ones((1, 3, 3), dtype=np.float32)
        intrinsics = np.asarray(
            [[[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]]],
            dtype=np.float32,
        )
        prediction = SimpleNamespace(
            depth=depth,
            conf=confidence,
            intrinsics=intrinsics,
        )
        progress = []

        with tempfile.TemporaryDirectory() as directory, patch.dict(
            sys.modules,
            {"PIL": fake_pil},
        ), patch.object(
            self.worker,
            "compact_mesh",
            side_effect=lambda vertices, faces: (
                vertices,
                faces,
            ),
        ):
            result_path, texture_path = self.worker.export_image_mesh_data(
                prediction,
                Path(directory) / "input.png",
                Path(directory) / "result",
                progress_callback=lambda step, total, label: progress.append(
                    (step, total, label)
                ),
            )
            with np.load(result_path) as result:
                self.assertEqual(result["vertices"].shape, (9, 3))
                self.assertEqual(result["faces"].shape, (8, 3))
                self.assertEqual(result["uv"].shape, (9, 2))
                self.assertEqual(float(result["reference_depth"]), 2.0)
                self.assertNotIn("postprocess_method", result.files)
                self.assertEqual(result["point_depth"].shape, (9,))
                self.assertEqual(
                    result["face_confidence_percentile"].shape,
                    (8,),
                )
                self.assertEqual(
                    set(result.files),
                    {
                        "vertices",
                        "faces",
                        "uv",
                        "intrinsics",
                        "image_size",
                        "reference_depth",
                        "point_depth",
                        "face_confidence_percentile",
                        "face_depth_discontinuity",
                    },
                )
                np.testing.assert_allclose(result["uv"][0], [0.0, 1.0])
                np.testing.assert_allclose(result["uv"][-1], [1.0, 0.0])
            self.assertTrue(texture_path.is_file())
            self.assertEqual([item[0] for item in progress], list(range(1, 7)))
            self.assertTrue(all(item[1] == 6 for item in progress))

    def test_prediction_attributes_keep_debug_structure(self):
        prediction = SimpleNamespace(
            depth=np.ones((1, 2, 3), dtype=np.float32),
            conf=np.arange(6, dtype=np.float32).reshape(1, 2, 3),
            intrinsics=np.eye(3, dtype=np.float32)[None],
            extrinsics=np.eye(4, dtype=np.float32)[None],
        )

        with tempfile.TemporaryDirectory() as directory:
            output_path = self.worker.export_prediction_attributes(
                prediction,
                directory,
            )
            with np.load(output_path) as attributes:
                self.assertEqual(attributes["depth"].shape, (1, 2, 3))
                self.assertEqual(attributes["conf"].shape, (1, 2, 3))
                self.assertEqual(
                    attributes["confidence_percentile"].shape,
                    (1, 2, 3),
                )


if __name__ == "__main__":
    unittest.main()