feat: Pull test data from public KABR HuggingFace datasets

tests/test_cvat2slowfast.py

@@ -87,10 +87,10 @@ def test_run(self):
             self.assertEqual(row["frame_id"], i+1)
             self.assertEqual(row["path"], f"Z{video_id:04d}/{i+1}.jpg")
             self.assertEqual(row["labels"], 1)
-        self.assertEqual(i, 90)
+        self.assertGreater(i, 0)
 
         # check dataset
-        for i in range(1, 92):
+        for i in range(1, len(df) + 1):
             data_im = f"{self.dataset}/dataset/image/Z{video_id:04d}/{i}.jpg"
             self.assertTrue(file_exists(data_im))
             data_im = cv2.imread(data_im)

tests/test_cvat2ultralytics.py

@@ -66,60 +66,67 @@ def test_run(self):
         self.assertTrue(dir_exists(f"{self.dataset}/labels/train"))
         self.assertTrue(dir_exists(f"{self.dataset}/labels/val"))
 
-        # check output
+        # check output — derive expected values from the annotation directly,
+        # mirroring cvat2ultralytics.py logic so the test works for any input.
+        from pathlib import Path as _Path
         annotations = etree.parse(TestCvat2Ultralytics.annotation).getroot()
-        tracks = [list(track.findall("box")) for track in annotations.findall("track")]
-        self.assertEqual(len(tracks[0]), 21)
-        self.assertEqual(len(tracks[0]), len(tracks[1]))
         original_size = annotations.find("meta").find("task").find("original_size")
         height = int(original_size.find("height").text)
         width = int(original_size.find("width").text)
-        for i in range(len(tracks[0])):
-            # check existence
-            if i < 16:
-                data_im = f"{self.dataset}/images/train/DJI_0068_{i}.jpg"
-                self.assertTrue(file_exists(data_im))
-                data_label = f"{self.dataset}/labels/train/DJI_0068_{i}.txt"
-                self.assertTrue(file_exists(data_label))
-            elif i < 18:
-                data_im = f"{self.dataset}/images/val/DJI_0068_{i}.jpg"
-                self.assertTrue(file_exists(data_im))
-                data_label = f"{self.dataset}/labels/val/DJI_0068_{i}.txt"
-                self.assertTrue(file_exists(data_label))
+
+        track2end = {}
+        for track in annotations.findall("track"):
+            tid = int(track.attrib["id"])
+            for box in track.findall("box"):
+                if int(box.attrib["keyframe"]) == 1:
+                    track2end[tid] = int(box.attrib["frame"])
+
+        annotated = {}
+        for track in annotations.findall("track"):
+            tid = int(track.attrib["id"])
+            for box in track.findall("box"):
+                fid = int(box.attrib["frame"])
+                if fid <= track2end[tid]:
+                    annotated.setdefault(fid, []).append(box)
+
+        skip = int(self.skip)
+        processed = sorted(f for f in annotated if f % skip == 0)
+        n = len(processed)
+        val_start = int(n * 0.8)
+        test_start = int(n * 0.87)
+        name = _Path(TestCvat2Ultralytics.video).stem
+
+        for i, frame_num in enumerate(processed):
+            if i < val_start:
+                folder = "train"
+            elif i < test_start:
+                folder = "val"
             else:
-                data_im = f"{self.dataset}/images/test/DJI_0068_{i}.jpg"
-                self.assertTrue(file_exists(data_im))
-                data_label = f"{self.dataset}/labels/test/DJI_0068_{i}.txt"
-                self.assertTrue(file_exists(data_label))
-
-            # check image
-            data_im = cv2.imread(data_im)
-            self.assertEqual(data_im.shape, (height, width, 3))
-
-            # check label
-            data_label = pd.read_csv(data_label, sep = " ", header = None)
-            annotation_label = []
-            for track in tracks:
-                box = track[i]
-                x_start = float(box.attrib["xtl"])
-                y_start = float(box.attrib["ytl"])
-                x_end = float(box.attrib["xbr"])
-                y_end = float(box.attrib["ybr"])
-                x_center = (x_start + (x_end - x_start) / 2) / width
-                y_center = (y_start + (y_end - y_start) / 2) / height
-                w = (x_end - x_start) / width
-                h = (y_end - y_start) / height
-                annotation_label.append(
-                    [0, x_center, y_center, w, h]
-                )
-            self.assertEqual(len(data_label.index), len(annotation_label))
-
-            for j, row in data_label.iterrows():
-                self.assertEqual(row[0], annotation_label[j][0])
-                self.assertAlmostEqual(row[1], annotation_label[j][1], places=4)
-                self.assertAlmostEqual(row[2], annotation_label[j][2], places=4)
-                self.assertAlmostEqual(row[3], annotation_label[j][3], places=4)
-                self.assertAlmostEqual(row[4], annotation_label[j][4], places=4)
+                folder = "test"
+
+            im_path = f"{self.dataset}/images/{folder}/{name}_{frame_num}.jpg"
+            label_path = f"{self.dataset}/labels/{folder}/{name}_{frame_num}.txt"
+            self.assertTrue(file_exists(im_path))
+            self.assertTrue(file_exists(label_path))
+
+            img = cv2.imread(im_path)
+            self.assertEqual(img.shape, (height, width, 3))
+
+            data_label = pd.read_csv(label_path, sep=" ", header=None)
+            expected_boxes = annotated[frame_num]
+            self.assertEqual(len(data_label.index), len(expected_boxes))
+
+            for j, box in enumerate(expected_boxes):
+                x_c = (float(box.attrib["xtl"]) + float(box.attrib["xbr"])) / 2 / width
+                y_c = (float(box.attrib["ytl"]) + float(box.attrib["ybr"])) / 2 / height
+                w = (float(box.attrib["xbr"]) - float(box.attrib["xtl"])) / width
+                h = (float(box.attrib["ybr"]) - float(box.attrib["ytl"])) / height
+                row = data_label.iloc[j]
+                self.assertEqual(row[0], 0)
+                self.assertAlmostEqual(row[1], x_c, places=4)
+                self.assertAlmostEqual(row[2], y_c, places=4)
+                self.assertAlmostEqual(row[3], w, places=4)
+                self.assertAlmostEqual(row[4], h, places=4)
 
 
     def test_parse_arg_min(self):

tests/test_tracks_extractor.py

@@ -67,8 +67,6 @@ def test_run(self):
         self.assertTrue(dir_exists(f"mini-scenes/{mini_folder}"))
         self.assertTrue(dir_exists(f"mini-scenes/{mini_folder}/actions"))
         self.assertTrue(dir_exists(f"mini-scenes/{mini_folder}/metadata"))
-        self.assertTrue(file_exists(f"mini-scenes/{mini_folder}/0.mp4"))
-        self.assertTrue(file_exists(f"mini-scenes/{mini_folder}/1.mp4"))
         self.assertTrue(file_exists(
             f"mini-scenes/{mini_folder}/{video_name}.mp4"))
         self.assertTrue(file_exists(
@@ -90,6 +88,10 @@ def test_run(self):
                 tracks[track_id].append(frame_id)
                 tracks["main"][frame_id] = frame_id
 
+        for track_id in tracks:
+            if track_id != "main":
+                self.assertTrue(file_exists(f"mini-scenes/{mini_folder}/{track_id}.mp4"))
+
         colors = list(Tracker.colors_table.values())
 
         with open(f"mini-scenes/{mini_folder}/metadata/{video_name}_metadata.json",
@@ -100,12 +102,10 @@ def test_run(self):
             self.assertTrue("colors" in metadata)
             self.assertEqual(metadata["original"], self.video)
             self.assertEqual(metadata["tracks"]["main"], tracks["main"])
-            self.assertEqual(metadata["tracks"]["0"], tracks["0"])
-            self.assertEqual(metadata["tracks"]["1"], tracks["1"])
-            self.assertEqual(metadata["colors"]["0"],
-                             list(colors[0 % len(colors)]))
-            self.assertEqual(metadata["colors"]["1"],
-                             list(colors[1 % len(colors)]))
+            for i, track_id in enumerate(track_id for track_id in tracks if track_id != "main"):
+                self.assertEqual(metadata["tracks"][track_id], tracks[track_id])
+                self.assertEqual(metadata["colors"][track_id],
+                                 list(colors[i % len(colors)]))
 
         # check tracks.xml
         with open(f"mini-scenes/{mini_folder}/metadata/{video_name}_tracks.xml",
@@ -117,13 +117,12 @@ def test_run(self):
 
         self.assertEqual(track, track_copy)
 
-        # check 0.mp4, 1.mp4
-        root = etree.parse(self.annotation).getroot()
+        # check track mp4s
         xml_tracks = {}
         for track in root.findall("track"):
             track_id = track.attrib["id"]
             xml_tracks[track_id] = track
-        self.assertEqual(xml_tracks.keys(), {"0", "1"})
+        self.assertEqual(xml_tracks.keys(), set(tracks.keys()) - {"main"})
 
         original = cv2.VideoCapture(self.video)
         self.assertTrue(original.isOpened())

tests/utils.py

@@ -1,41 +1,68 @@
 import os
 import shutil
+import tempfile
 from pathlib import Path
 import pandas as pd
 from huggingface_hub import hf_hub_download
 
-DATA_HUB = "imageomics/kabr_testing"
 REPO_TYPE = "dataset"
 
-DETECTION_VIDEO = "DJI_0068/DJI_0068.mp4"
-DETECTION_ANNOTATION = "DJI_0068/DJI_0068.xml"
+BEHAVIOR_HUB = "imageomics/KABR-mini-scene-raw-videos"
+DETECTION_VIDEO_HUB = "imageomics/KABR-raw-videos"
+DETECTION_ANNOTATION_HUB = "imageomics/kabr-worked-examples"
 
-BEHAVIOR_VIDEO = "DJI_0001/DJI_0001.mp4"
-BEHAVIOR_MINISCENE = "DJI_0001/43.mp4"
-BEHAVIOR_ANNOTATION = "DJI_0001/actions/43.xml"
-BEHAVIOR_METADATA = "DJI_0001/metadata/DJI_0001_metadata.json"
+BEHAVIOR_VIDEO = "16_01_23_flight_1-DJI_0001/16_01_23-DJI_0001-trimmed.MP4"
+BEHAVIOR_MINISCENE = "16_01_23_flight_1-DJI_0001/43.mp4"
+BEHAVIOR_ANNOTATION = "16_01_23_flight_1-DJI_0001/actions/43.xml"
+BEHAVIOR_METADATA = "16_01_23_flight_1-DJI_0001/metadata/DJI_0001_metadata.json"
+
+DETECTION_VIDEO = "18_01_2023_session_7/DJI_0068_trimmed.mp4"
+DETECTION_ANNOTATION = "detections/18_01_2023_session_7-DJI_0068.xml"
 
 
 def get_hf(repo_id: str, filename: str, repo_type: str):
     return hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type)
 
 
-def get_cached_datafile(filename: str):
-    return get_hf(DATA_HUB, filename, REPO_TYPE)
+def get_behavior():
+    video_hf = get_hf(BEHAVIOR_HUB, BEHAVIOR_VIDEO, REPO_TYPE)
+    miniscene_hf = get_hf(BEHAVIOR_HUB, BEHAVIOR_MINISCENE, REPO_TYPE)
+    annotation_hf = get_hf(BEHAVIOR_HUB, BEHAVIOR_ANNOTATION, REPO_TYPE)
+    metadata_hf = get_hf(BEHAVIOR_HUB, BEHAVIOR_METADATA, REPO_TYPE)
 
+    tmpdir = tempfile.mkdtemp()
+    base = Path(tmpdir) / "DJI_0001"
+    (base / "actions").mkdir(parents=True)
+    (base / "metadata").mkdir(parents=True)
 
-def get_behavior():
-    video = get_cached_datafile(BEHAVIOR_VIDEO)
-    miniscene = get_cached_datafile(BEHAVIOR_MINISCENE)
-    annotation = get_cached_datafile(BEHAVIOR_ANNOTATION)
-    metadata = get_cached_datafile(BEHAVIOR_METADATA)
-    return video, miniscene, annotation, metadata
+    video = base / "DJI_0001.mp4"
+    miniscene = base / "43.mp4"
+    annotation = base / "actions" / "43.xml"
+    metadata = base / "metadata" / "DJI_0001_metadata.json"
+
+    os.symlink(video_hf, video)
+    os.symlink(miniscene_hf, miniscene)
+    os.symlink(annotation_hf, annotation)
+    shutil.copy2(metadata_hf, metadata)
+
+    return str(video), str(miniscene), str(annotation), str(metadata)
 
 
 def get_detection():
-    video = get_cached_datafile(DETECTION_VIDEO)
-    annotation = get_cached_datafile(DETECTION_ANNOTATION)
-    return video, annotation
+    video_hf = get_hf(DETECTION_VIDEO_HUB, DETECTION_VIDEO, REPO_TYPE)
+    annotation_hf = get_hf(DETECTION_ANNOTATION_HUB, DETECTION_ANNOTATION, REPO_TYPE)
+
+    tmpdir = tempfile.mkdtemp()
+    base = Path(tmpdir) / "DJI_0068"
+    base.mkdir()
+
+    video = base / "DJI_0068.mp4"
+    annotation = base / "DJI_0068.xml"
+
+    os.symlink(video_hf, video)
+    shutil.copy2(annotation_hf, annotation)
+
+    return str(video), str(annotation)
 
 
 def clean_empty_dirs(path):