From f4850b870f28f4ccbcc74a76c1eee8c7955c4b94 Mon Sep 17 00:00:00 2001 From: patel-zeel Date: Thu, 31 Oct 2024 16:45:19 +0530 Subject: [PATCH 1/3] replace cv2 with PIL for speed --- supervision/dataset/formats/yolo.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index 1f9033fc5..bae938c23 100644 --- a/supervision/dataset/formats/yolo.py +++ b/supervision/dataset/formats/yolo.py @@ -2,8 +2,8 @@ from pathlib import Path from typing import TYPE_CHECKING, Dict, List, Optional, Tuple -import cv2 import numpy as np +from PIL import Image from supervision.config import ORIENTED_BOX_COORDINATES from supervision.dataset.utils import approximate_mask_with_polygons @@ -153,7 +153,7 @@ def load_yolo_annotations( image_paths = [ str(path) for path in list_files_with_extensions( - directory=images_directory_path, extensions=["jpg", "jpeg", "png"] + directory=images_directory_path, extensions=["*"] ) ] @@ -167,10 +167,15 @@ def load_yolo_annotations( annotations[image_path] = Detections.empty() continue - image = cv2.imread(image_path) + # PIL is much faster than cv2 for checking image shape and mode: https://github.com/roboflow/supervision/issues/1554 + image = Image.open(image_path) lines = read_txt_file(file_path=annotation_path, skip_empty=True) - h, w, _ = image.shape + w, h = image.size resolution_wh = (w, h) + if image.mode != "RGB": + raise ValueError( + f"Images must be 'RGB', but {image_path} is '{image.mode}'." + ) with_masks = _with_mask(lines=lines) with_masks = force_masks if force_masks else with_masks From 2ed66bfffc20889a8c60dcb52f26dac1bd1ca0a0 Mon Sep 17 00:00:00 2001 From: patel-zeel Date: Tue, 7 Jan 2025 22:16:05 +0530 Subject: [PATCH 2/3] expand support to grayscale similar to cv2 --- supervision/dataset/formats/yolo.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index bae938c23..602715945 100644 --- a/supervision/dataset/formats/yolo.py +++ b/supervision/dataset/formats/yolo.py @@ -173,9 +173,13 @@ def load_yolo_annotations( w, h = image.size resolution_wh = (w, h) if image.mode != "RGB": - raise ValueError( - f"Images must be 'RGB', but {image_path} is '{image.mode}'." - ) + if image.mode == "L": + image = image.convert("RGB") + else: + raise ValueError( + f"Images must be 'RGB' or 'grayscale', \ + but {image_path} mode is '{image.mode}'." + ) with_masks = _with_mask(lines=lines) with_masks = force_masks if force_masks else with_masks From d33c159a267b6544e4a507b4b1a1f479aaf60197 Mon Sep 17 00:00:00 2001 From: patel-zeel Date: Wed, 8 Jan 2025 15:02:05 +0530 Subject: [PATCH 3/3] explicit extensions, simplified code, speed improvement --- supervision/dataset/formats/yolo.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index 602715945..7e5a9d454 100644 --- a/supervision/dataset/formats/yolo.py +++ b/supervision/dataset/formats/yolo.py @@ -153,7 +153,18 @@ def load_yolo_annotations( image_paths = [ str(path) for path in list_files_with_extensions( - directory=images_directory_path, extensions=["*"] + directory=images_directory_path, + extensions=[ + "bmp", + "dng", + "jpg", + "jpeg", + "mpo", + "png", + "tif", + "tiff", + "webp", + ], ) ] @@ -172,14 +183,11 @@ def load_yolo_annotations( lines = read_txt_file(file_path=annotation_path, skip_empty=True) w, h = image.size resolution_wh = (w, h) - if image.mode != "RGB": - if image.mode == "L": - image = image.convert("RGB") - else: - raise ValueError( - f"Images must be 'RGB' or 'grayscale', \ - but {image_path} mode is '{image.mode}'." - ) + if image.mode not in ("RGB", "L"): + raise ValueError( + f"Images must be 'RGB' or 'grayscale', \ + but {image_path} mode is '{image.mode}'." + ) with_masks = _with_mask(lines=lines) with_masks = force_masks if force_masks else with_masks