🐚 shell of YOLO format processing added

supervision/dataset/core.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Iterator
+from typing import Dict, Iterator, List, Optional, Tuple
 
 import cv2
 import numpy as np
@@ -11,7 +11,11 @@ from supervision.dataset.formats.pascal_voc import (
     detections_to_pascal_voc,
     load_pascal_voc_annotations,
 )
-from supervision.dataset.formats.yolo import load_yolo_annotations, save_yolo_annotations, save_data_yaml
+from supervision.dataset.formats.yolo import (
+    load_yolo_annotations,
+    save_data_yaml,
+    save_yolo_annotations,
+)
 from supervision.dataset.ultils import save_dataset_images
 from supervision.detection.core import Detections
 from supervision.file import list_files_with_extensions
@@ -237,7 +241,9 @@ class DetectionDataset(BaseDataset):
         approximation_percentage: float = 0.75,
     ) -> None:
         if images_directory_path is not None:
-            save_dataset_images(images_directory_path=images_directory_path, images=self.images)
+            save_dataset_images(
+                images_directory_path=images_directory_path, images=self.images
+            )
         if annotations_directory_path is not None:
             save_yolo_annotations(
                 annotations_directory_path=annotations_directory_path,
@@ -245,7 +251,7 @@ class DetectionDataset(BaseDataset):
                 annotations=self.annotations,
                 min_image_area_percentage=min_image_area_percentage,
                 max_image_area_percentage=max_image_area_percentage,
-                approximation_percentage=approximation_percentage
+                approximation_percentage=approximation_percentage,
             )
         if data_yaml_path is not None:
             save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes)

supervision/dataset/formats/yolo.py

@@ -1,14 +1,15 @@
+import os
 from pathlib import Path
-from typing import Dict, List, Tuple, Union, Optional
+from typing import Dict, List, Optional, Tuple, Union
 
 import cv2
-import yaml
 import numpy as np
+import yaml
 
 from supervision.dataset.ultils import approximate_mask_with_polygons
 from supervision.detection.core import Detections
 from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy
-from supervision.file import list_files_with_extensions, read_txt_file
+from supervision.file import list_files_with_extensions, read_txt_file, save_text_file
 
 
 def _parse_box(values: List[str]) -> np.ndarray:
@@ -80,6 +81,11 @@ def _extract_class_names(file_path: str) -> List[str]:
     return names
 
 
+def _image_name_to_annotation_name(image_name: str) -> str:
+    base_name, _ = os.path.splitext(image_name)
+    return base_name + ".txt"
+
+
 def yolo_annotations_to_detections(
     lines: List[str], resolution_wh: Tuple[int, int], with_masks: bool
 ) -> Detections:
@@ -167,7 +173,7 @@ def object_to_pascal_voc(
     xyxy: np.ndarray,
     class_id: int,
     image_shape: Tuple[int, int, int],
-    polygon: Optional[np.ndarray] = None
+    polygon: Optional[np.ndarray] = None,
 ) -> str:
     height, width, _ = image_shape
 
@@ -178,7 +184,7 @@ def detections_to_yolo_annotations(
     min_image_area_percentage: float = 0.0,
     max_image_area_percentage: float = 1.0,
     approximation_percentage: float = 0.75,
-) -> str:
+) -> List[str]:
     annotation = []
     for xyxy, mask, _, class_id, _ in detections:
         if mask is not None:
@@ -191,13 +197,18 @@ def detections_to_yolo_annotations(
             for polygon in polygons:
                 xyxy = polygon_to_xyxy(polygon=polygon)
                 next_object = object_to_pascal_voc(
-                    xyxy=xyxy, class_id=class_id, image_shape=image_shape, polygon=polygon
+                    xyxy=xyxy,
+                    class_id=class_id,
+                    image_shape=image_shape,
+                    polygon=polygon,
                 )
                 annotation.append(next_object)
         else:
-            next_object = object_to_pascal_voc(xyxy=xyxy, class_id=class_id, image_shape=image_shape)
+            next_object = object_to_pascal_voc(
+                xyxy=xyxy, class_id=class_id, image_shape=image_shape
+            )
             annotation.append(next_object)
-    return "\n".join(annotation)
+    return annotation
 
 
 def save_yolo_annotations(
@@ -208,14 +219,23 @@ def save_yolo_annotations(
     max_image_area_percentage: float = 1.0,
     approximation_percentage: float = 0.75,
 ) -> None:
-    pass
+    Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
+    for image_name, image in images:
+        detections = annotations[image_name]
+        yolo_annotations_name = _image_name_to_annotation_name(image_name=image_name)
+        yolo_annotations_path = os.path.join(annotations_directory_path, yolo_annotations_name)
+        lines = detections_to_yolo_annotations(
+            detections=detections,
+            image_shape=image.shape,
+            min_image_area_percentage=min_image_area_percentage,
+            max_image_area_percentage=max_image_area_percentage,
+            approximation_percentage=approximation_percentage,
+        )
+        save_text_file(lines=lines, file_path=yolo_annotations_path)
 
 
 def save_data_yaml(data_yaml_path: str, classes: List[str]) -> None:
-    data = {
-        'nc': len(classes),
-        'names': classes
-    }
+    data = {"nc": len(classes), "names": classes}
     Path(data_yaml_path).parent.mkdir(parents=True, exist_ok=True)
-    with open(data_yaml_path, 'w') as outfile:
+    with open(data_yaml_path, "w") as outfile:
         yaml.dump(data, outfile, default_flow_style=False)

supervision/dataset/ultils.py

@@ -1,6 +1,6 @@
 import os
 from pathlib import Path
-from typing import List, Dict, Optional
+from typing import Dict, List
 
 import cv2
 import numpy as np
@@ -40,9 +40,11 @@ def approximate_mask_with_polygons(
     ]
 
 
-def save_dataset_images(images_directory_path: str, images: Dict[str, np.ndarray]) -> None:
+def save_dataset_images(
+    images_directory_path: str, images: Dict[str, np.ndarray]
+) -> None:
     Path(images_directory_path).mkdir(parents=True, exist_ok=True)
 
     for image_name, image in images.items():
         target_image_path = os.path.join(images_directory_path, image_name)
-        cv2.imwrite(target_image_path, image)
+        cv2.imwrite(target_image_path, image)

supervision/detection/line_counter.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Optional
 
 import cv2
 import numpy as np
@@ -6,7 +6,6 @@ import numpy as np
 from supervision.detection.core import Detections
 from supervision.draw.color import Color
 from supervision.geometry.core import Point, Rect, Vector
-from typing import Optional
 
 
 class LineZone:
@@ -71,6 +70,7 @@ class LineZone:
             else:
                 self.out_count += 1
 
+
 class LineZoneAnnotator:
     def __init__(
         self,
@@ -145,9 +145,16 @@ class LineZoneAnnotator:
             lineType=cv2.LINE_AA,
         )
 
-        in_text = f"{self.custom_in_text}: {line_counter.in_count}" if self.custom_in_text is not None else f"in: {line_counter.in_count}"
-        out_text = f"{self.custom_out_text}: {line_counter.out_count}" if self.custom_out_text is not None else f"out: {line_counter.out_count}"
-
+        in_text = (
+            f"{self.custom_in_text}: {line_counter.in_count}"
+            if self.custom_in_text is not None
+            else f"in: {line_counter.in_count}"
+        )
+        out_text = (
+            f"{self.custom_out_text}: {line_counter.out_count}"
+            if self.custom_out_text is not None
+            else f"out: {line_counter.out_count}"
+        )
 
         (in_text_width, in_text_height), _ = cv2.getTextSize(
             in_text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness

supervision/file.py

@@ -53,3 +53,16 @@ def read_txt_file(file_path: str) -> List[str]:
         lines = [line.rstrip("\n") for line in lines]
 
     return lines
+
+
+def save_text_file(lines: List[str], file_path: str):
+    """
+    Write a list of strings to a text file, each string on a new line.
+
+    Args:
+        lines (List[str]): The list of strings to be written to the file.
+        file_path (str): The path to the text file.
+    """
+    with open(file_path, "w") as file:
+        for line in lines:
+            file.write(line + "\n")

test/dataset/formats/test_yolo.py

@@ -5,7 +5,7 @@ import pytest
 import numpy as np
 
 from supervision.detection.core import Detections
-from supervision.dataset.formats.yolo import yolo_annotations_to_detections, _with_mask
+from supervision.dataset.formats.yolo import yolo_annotations_to_detections, _with_mask, _image_name_to_annotation_name
 
 
 def _mock_simple_mask(resolution_wh: Tuple[int, int], box: List[int]) -> np.array:
@@ -202,3 +202,38 @@ def test_yolo_annotations_to_detections(
         assert np.array_equal(result.xyxy, expected_result.xyxy)
         assert np.array_equal(result.class_id, expected_result.class_id)
         assert (result.mask is None and expected_result.mask is None) or _arrays_almost_equal(result.mask, expected_result.mask)
+
+
+@pytest.mark.parametrize(
+    'image_name, expected_result, exception',
+    [
+        (
+            'image.png',
+            'image.txt',
+            DoesNotRaise()
+        ),  # simple png image
+        (
+            'image.jpeg',
+            'image.txt',
+            DoesNotRaise()
+        ),  # simple jpeg image
+        (
+            'image.jpg',
+            'image.txt',
+            DoesNotRaise()
+        ),  # simple jpg image
+        (
+            'image.000.jpg',
+            'image.000.txt',
+            DoesNotRaise()
+        ),  # jpg image with multiple dots in name
+    ]
+)
+def test_image_name_to_annotation_name(
+    image_name: str,
+    expected_result: Optional[str],
+    exception: Exception
+) -> None:
+    with exception:
+        result = _image_name_to_annotation_name(image_name=image_name)
+        assert result == expected_result