From f616bda40f5e5fabff8ad7bfad172c3f02edfd06 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 17 May 2023 00:58:17 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9A=20shell=20of=20YOLO=20format=20pro?= =?UTF-8?q?cessing=20added?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/dataset/core.py | 14 +++++--- supervision/dataset/formats/yolo.py | 48 +++++++++++++++++++-------- supervision/dataset/ultils.py | 8 +++-- supervision/detection/line_counter.py | 17 +++++++--- supervision/file.py | 13 ++++++++ test/dataset/formats/test_yolo.py | 37 ++++++++++++++++++++- 6 files changed, 110 insertions(+), 27 deletions(-) diff --git a/supervision/dataset/core.py b/supervision/dataset/core.py index 460fd2be..66e2da1e 100644 --- a/supervision/dataset/core.py +++ b/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) diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index 5e6eeda1..eced89fb 100644 --- a/supervision/dataset/formats/yolo.py +++ b/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) diff --git a/supervision/dataset/ultils.py b/supervision/dataset/ultils.py index f201f274..44a09407 100644 --- a/supervision/dataset/ultils.py +++ b/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) \ No newline at end of file + cv2.imwrite(target_image_path, image) diff --git a/supervision/detection/line_counter.py b/supervision/detection/line_counter.py index 23142967..45a9f7ee 100644 --- a/supervision/detection/line_counter.py +++ b/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 diff --git a/supervision/file.py b/supervision/file.py index cc3012d4..af2ffc79 100644 --- a/supervision/file.py +++ b/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") diff --git a/test/dataset/formats/test_yolo.py b/test/dataset/formats/test_yolo.py index 5be8c879..00240e34 100644 --- a/test/dataset/formats/test_yolo.py +++ b/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