Skip to content

Commit

Permalink
Merge pull request #1636 from patel-zeel/feat/expand-image-formats
Browse files Browse the repository at this point in the history
Allow "tiff" and more extensions in `DetectionDataset.from_yolo` function
  • Loading branch information
SkalskiP authored Jan 8, 2025
2 parents d4c6fd7 + 6aedf48 commit 651ed9d
Showing 1 changed file with 21 additions and 4 deletions.
25 changes: 21 additions & 4 deletions supervision/dataset/formats/yolo.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -153,7 +153,18 @@ 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=[
"bmp",
"dng",
"jpg",
"jpeg",
"mpo",
"png",
"tif",
"tiff",
"webp",
],
)
]

Expand All @@ -167,10 +178,16 @@ 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 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
Expand Down

0 comments on commit 651ed9d

Please sign in to comment.