diff --git a/tests/test_cvat2slowfast.py b/tests/test_cvat2slowfast.py index 813abeb..8b9841d 100644 --- a/tests/test_cvat2slowfast.py +++ b/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) diff --git a/tests/test_cvat2ultralytics.py b/tests/test_cvat2ultralytics.py index 202e95e..0ea3d03 100644 --- a/tests/test_cvat2ultralytics.py +++ b/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): diff --git a/tests/test_tracks_extractor.py b/tests/test_tracks_extractor.py index 87dd994..564c1f2 100644 --- a/tests/test_tracks_extractor.py +++ b/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()) diff --git a/tests/utils.py b/tests/utils.py index 24d2414..ed0aff2 100644 --- a/tests/utils.py +++ b/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):