validate extracted start-label first frame

Author: Copilot

poseinterface/clips.py

@@ -21,6 +21,24 @@ def _cliplabels_suffix_error(name: str) -> str:
     return f"Input file must end with '_cliplabels.json', got {name}"
 
 
+def _extract_start_labels(clip_labels: dict) -> dict:
+    start_images = [img for img in clip_labels["images"] if img["id"] == 0]
+    if len(start_images) != 1:
+        raise ValueError(
+            "Clip labels must contain exactly one image with id 0"
+        )
+
+    return {
+        "images": start_images,
+        "annotations": [
+            annot
+            for annot in clip_labels["annotations"]
+            if annot["image_id"] == 0
+        ],
+        "categories": clip_labels["categories"],
+    }
+
+
 def extract_clip(
     video_path: str | Path,
     start_frame: int,
@@ -228,19 +246,7 @@ def extract_start_frames_label(cliplabels_path: str | Path) -> Path:
     with open(cliplabels_path) as f:
         clip_labels = json.load(f)
 
-    # Extract only the first frame's data (frame with id=0)
-    start_labels = {}
-    start_labels["images"] = [
-        img for img in clip_labels["images"] if img["id"] == 0
-    ]
-
-    # Extract only annotations for the first frame (image_id=0)
-    start_labels["annotations"] = [
-        annot for annot in clip_labels["annotations"] if annot["image_id"] == 0
-    ]
-
-    # Keep categories unchanged
-    start_labels["categories"] = clip_labels["categories"]
+    start_labels = _extract_start_labels(clip_labels)
 
     # Generate output path by replacing _cliplabels.json with _startlabels.json
     output_path = cliplabels_path.parent / cliplabels_path.name.replace(
@@ -329,21 +335,7 @@ def extract_start_frames_label_s3(
         response = s3_client.get_object(Bucket=bucket_name, Key=key)
         clip_labels = json.loads(response["Body"].read().decode("utf-8"))
 
-        # Extract only the first frame's data (frame with id=0)
-        start_labels = {}
-        start_labels["images"] = [
-            img for img in clip_labels["images"] if img["id"] == 0
-        ]
-
-        # Extract only annotations for the first frame (image_id=0)
-        start_labels["annotations"] = [
-            annot
-            for annot in clip_labels["annotations"]
-            if annot["image_id"] == 0
-        ]
-
-        # Keep categories unchanged
-        start_labels["categories"] = clip_labels["categories"]
+        start_labels = _extract_start_labels(clip_labels)
 
         # Generate output key by replacing _cliplabels.json with
         # _startlabels.json

tests/test_unit/test_clips.py

@@ -7,6 +7,7 @@
 
 from poseinterface.clips import (
     _extract_cliplabels,
+    _extract_start_labels,
     extract_clip,
     main,
     parse_args,
@@ -157,6 +158,29 @@ def test_extract_clip_invalid(
         extract_clip(video_path, start_frame, duration)
 
 
+def test_extract_start_labels(video_labels):
+    """Test start labels contain only the first frame and its annotations."""
+    start_labels = _extract_start_labels(video_labels)
+
+    assert start_labels["images"] == [{"id": 0}]
+    assert start_labels["annotations"] == [
+        {"image_id": 0, "id": j} for j in range(5)
+    ]
+    assert start_labels["categories"] == video_labels["categories"]
+
+
+def test_extract_start_labels_requires_first_frame(video_labels):
+    """Test start-label extraction validates the presence of frame 0."""
+    clip_labels = {
+        **video_labels,
+        "images": [{"id": 1}],
+        "annotations": [{"image_id": 1, "id": 1}],
+    }
+
+    with pytest.raises(ValueError, match="exactly one image with id 0"):
+        _extract_start_labels(clip_labels)
+
+
 def test_parse_args():
     """Check arguments are parsed with correct types"""
     video_path = "foo.mp4"