Fix bbox offset

test_local_turn_features.py

@@ -15,13 +15,14 @@ from types import SimpleNamespace
 
 import numpy as np
 import torch
-from PIL import Image
+from PIL import Image, ImageDraw
 
 from easycom_head import TurnPredictionHead, checkpoint_head_config
 from extract_turn_local_features import (
     FrozenQwenLocalEncoder, ego4d_clip_boxes, extract_local,
     validate_source_timeline,
     model_config_digest,
+    box_coordinate_offset, frame_crops, parse_args, preview_boxes,
 )
 from local_turn_features import (
     LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes, load_local_features,
@@ -82,6 +83,73 @@ class LocalTests(unittest.TestCase):
         actual = ego4d_clip_boxes({"persons": [wearer, person]})
         self.assertEqual(actual, {3: [{"id": "x", "xyxy": [1., 2., 4., 6.]}]})
 
+    def test_dataset_specific_coordinate_defaults_and_explicit_overrides(self):
+        required = ["--data-root", "/raw", "--feature-root", "/global",
+                    "--target-root", "/targets", "--output-root", "/local", "--samples", "4"]
+        self.assertEqual(parse_args(required).box_coordinate_space, "encoded")
+        self.assertEqual(parse_args(required + ["--dataset", "ego4d"]).box_coordinate_space, "camera")
+        self.assertEqual(parse_args(required + ["--box-coordinate-space", "camera"]).box_coordinate_space, "camera")
+        # An appended correction overrides an old LOCAL_ARGS array in the shell.
+        self.assertEqual(parse_args(required + ["--box-coordinate-space", "camera",
+                                                "--box-coordinate-space", "encoded"]).box_coordinate_space, "encoded")
+
+    def test_frame_crops_subtract_each_clips_banner_before_encoding(self):
+        expected = (100, 80, 160, 140)
+        image = Image.new("RGB", (1920, 1080))
+        ImageDraw.Draw(image).rectangle(expected, fill="red")
+        clips, annotations = [], []
+        for width in (2123, 2124, 2203, 1920):
+            offset = width - 1920
+            clips.append(SimpleNamespace(crop=(offset, 0, 1920, 1080), dense_start=float(len(clips))))
+            annotations.append({0: [{"id": "3", "xyxy": [100 + offset, 80, 160 + offset, 140]}]})
+        media = SimpleNamespace(clips=clips,
+                                _source_frame_for_time=lambda t: (int(t), 0, 0.),
+                                frame_at_or_before=lambda t: image.copy())
+        args = SimpleNamespace(box_coordinate_space="encoded", crop_expansion=1.)
+        for index, clip in enumerate(clips):
+            with self.subTest(crop=clip.crop):
+                _, crops, records = frame_crops(media, annotations, float(index), args)
+                self.assertEqual(records[0]["box"], expected)
+                self.assertEqual(records[0]["coordinate_offset"], [clip.crop[0], 0])
+                self.assertEqual(records[0]["annotation_box"], annotations[index][0][0]["xyxy"])
+                self.assertEqual(crops[0].size, (60, 60))
+                self.assertEqual(crops[0].getpixel((30, 30)), (255, 0, 0))
+
+    def test_cropped_camera_annotations_and_uncropped_ego4d_are_not_shifted_twice(self):
+        clip = SimpleNamespace(crop=(204, 12, 1920, 1080), dense_start=0.)
+        self.assertEqual(box_coordinate_offset(clip, "encoded"), (204, 12))
+        self.assertEqual(box_coordinate_offset(clip, "camera"), (0, 0))
+        self.assertEqual(box_coordinate_offset(SimpleNamespace(crop=None), "encoded"), (0, 0))
+        with self.assertRaisesRegex(ValueError, "coordinate space"):
+            box_coordinate_offset(clip, "typo")
+        media = SimpleNamespace(clips=[clip], _source_frame_for_time=lambda t: (0, 0, 0.),
+                                frame_at_or_before=lambda t: Image.new("RGB", (1920, 1080)))
+        annotations = [{0: [{"id": "3", "xyxy": [100, 80, 160, 140]}]}]
+        args = SimpleNamespace(box_coordinate_space="camera", crop_expansion=1.)
+        _, _, records = frame_crops(media, annotations, .1, args)
+        self.assertEqual(records[0]["box"], (100, 80, 160, 140))
+        # Encoded boxes translate on BOTH axes, then clip at the visible edge.
+        annotations = [{0: [{"id": "3", "xyxy": [190, 4, 230, 50]}]}]
+        args.box_coordinate_space = "encoded"
+        _, _, records = frame_crops(media, annotations, .1, args)
+        self.assertEqual(records[0]["box"], (0, 0, 26, 38))
+
+    def test_preview_records_the_same_transformation_used_for_extraction(self):
+        clip = SimpleNamespace(path=Path("source.mp4"), width=2124, height=1080,
+                               crop=(204, 0, 1920, 1080), dense_start=0.)
+        media = SimpleNamespace(clips=[clip], _source_frame_for_time=lambda t: (0, 0, 0.),
+                                frame_at_or_before=lambda t: Image.new("RGB", (1920, 1080)))
+        annotations = [{0: [{"id": "3", "xyxy": [304, 80, 364, 140]}]}]
+        with tempfile.TemporaryDirectory() as directory:
+            args = SimpleNamespace(box_coordinate_space="encoded", crop_expansion=1.,
+                                   video_fps=2., output_root=Path(directory))
+            preview_boxes(media, annotations, args, "Session_4", .2)
+            folder = Path(directory) / "Session_4_preview"
+            audit = json.loads((folder / "boxes.json").read_text())
+            self.assertEqual(audit[0]["coordinate_offset"], [204, 0])
+            self.assertEqual(audit[0]["boxes"][0]["box"], [100, 80, 160, 140])
+            self.assertTrue((folder / audit[0]["image"]).is_file())
+
     def test_sidecar_alignment_contract_and_modalities(self):
         reference, side, _ = fixture()
         with tempfile.TemporaryDirectory() as directory: