Merge pull request #9 from DagsHub/bug/context-path-image

Bug: Handle YOLO annotation setups where "image" is part of the main path and not the annotations

Author: kbolashev

dagshub_annotation_converter/converters/yolo.py

@@ -180,17 +180,17 @@ def export_to_fs(
     export_path = Path(export_dir)
 
     for filename, anns in grouped_annotations.items():
-        annotation_filepath = replace_folder(
-            Path(filename), context.image_dir_name, context.label_dir_name, context.label_extension
+        annotation_filepath = export_path / context.path / filename
+        out_path = replace_folder(
+            annotation_filepath, context.image_dir_name, context.label_dir_name, context.label_extension
         )
-        if annotation_filepath is None:
+        if out_path is None:
             logger.warning(f"Couldn't generate annotation file path for image file [{filename}]")
             continue
-        annotation_filename = export_path / context.path / annotation_filepath
-        annotation_filename.parent.mkdir(parents=True, exist_ok=True)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
         annotation_content = annotations_to_string(anns, context)
         if annotation_content is not None:
-            with open(annotation_filename, "w") as f:
+            with open(out_path, "w") as f:
                 f.write(annotation_content)
 
     guessed_train_path, guessed_val_path, guessed_test_path = _guess_train_val_test_split(

tests/fs_export/yolo/test_fs_export.py

@@ -15,7 +15,7 @@
 
 
 def test_bbox_export(tmp_path):
-    ctx = YoloContext(annotation_type="bbox", path="data")
+    ctx = YoloContext(annotation_type="bbox", path=Path("data"))
     ctx.categories.add(name="cat")
     ctx.categories.add(name="dog")
     annotations = [
@@ -53,7 +53,7 @@ def test_bbox_export(tmp_path):
 
 
 def test_segmentation_export(tmp_path):
-    ctx = YoloContext(annotation_type="segmentation", path="data")
+    ctx = YoloContext(annotation_type="segmentation", path=Path("data"))
     ctx.categories.add(name="cat")
     ctx.categories.add(name="dog")
     annotations = [
@@ -85,7 +85,7 @@ def test_segmentation_export(tmp_path):
 
 
 def test_pose_export(tmp_path):
-    ctx = YoloContext(annotation_type="pose", path="data")
+    ctx = YoloContext(annotation_type="pose", path=Path("data"))
     ctx.categories.add(name="cat")
     ctx.categories.add(name="dog")
     ctx.keypoints_in_annotation = 2
@@ -118,7 +118,7 @@ def test_pose_export(tmp_path):
 
 
 def test_not_exporting_wrong_annotations(tmp_path):
-    ctx = YoloContext(annotation_type="bbox", path="data")
+    ctx = YoloContext(annotation_type="bbox", path=Path("data"))
     ctx.categories.add(name="cat")
     ctx.categories.add(name="dog")
     annotations = [
@@ -173,3 +173,41 @@ def test__get_common_folder_with_part(paths, prefix, expected):
         expected = Path(expected)
 
     assert actual == expected
+
+
+def test_export_with_image_in_path(tmp_path):
+    ctx = YoloContext(annotation_type="bbox", path=Path("data/images"))
+    ctx.categories.add(name="cat")
+    ctx.categories.add(name="dog")
+    annotations = [
+        IRBBoxImageAnnotation(
+            filename="cats/1.jpg",
+            categories={"cat": 1.0},
+            top=0.0,
+            left=0.0,
+            width=0.5,
+            height=0.5,
+            image_width=100,
+            image_height=200,
+            coordinate_style=CoordinateStyle.NORMALIZED,
+        ),
+        IRBBoxImageAnnotation(
+            filename="dogs/2.jpg",
+            categories={"dog": 1.0},
+            top=0.5,
+            left=0.5,
+            width=0.5,
+            height=0.5,
+            image_width=100,
+            image_height=200,
+            coordinate_style=CoordinateStyle.NORMALIZED,
+        ),
+    ]
+
+    p = export_to_fs(ctx, annotations, export_dir=tmp_path)
+
+    assert p == tmp_path / "yolo_dagshub.yaml"
+
+    assert (tmp_path / "yolo_dagshub.yaml").exists()
+    assert (tmp_path / "data" / "labels" / "cats" / "1.txt").exists()
+    assert (tmp_path / "data" / "labels" / "dogs" / "2.txt").exists()