zhao.zhu/yolov26_3d
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

单目3D初始代码

Browse Source
This commit is contained in:
zhao.zhu
2026-06-24 09:35:46 +08:00
commit 04a5895b6b
1153 changed files with 340700 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
ultralytics/data/__init__.py Executable file
View File

@@ -0,0 +1,28 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from .base import BaseDataset
from .build import build_dataloader, build_grounding, build_yolo_dataset, load_inference_source
from .dataset import (
ClassificationDataset,
GroundingDataset,
SemanticDataset,
YOLOConcatDataset,
YOLODataset,
YOLOGroundDataset,
YOLOMultiModalDataset,
)
__all__ = (
"BaseDataset",
"ClassificationDataset",
"GroundingDataset",
"SemanticDataset",
"YOLOConcatDataset",
"YOLODataset",
"YOLOGroundDataset",
"YOLOMultiModalDataset",
"build_dataloader",
"build_grounding",
"build_yolo_dataset",
"load_inference_source",
)

66
ultralytics/data/annotator.py Executable file
View File

@@ -0,0 +1,66 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
from pathlib import Path
from ultralytics import SAM, YOLO
def auto_annotate(
data: str | Path,
det_model: str = "yolo26x.pt",
sam_model: str = "sam_b.pt",
device: str = "",
conf: float = 0.25,
iou: float = 0.45,
imgsz: int = 640,
max_det: int = 300,
classes: list[int] | None = None,
output_dir: str | Path | None = None,
) -> None:
"""Automatically annotate images using a YOLO object detection model and a SAM segmentation model.
This function processes images in a specified directory, detects objects using a YOLO model, and then generates
segmentation masks using a SAM model. The resulting annotations are saved as text files in YOLO format.
Args:
data (str | Path): Path to a folder containing images to be annotated.
det_model (str): Path or name of the pre-trained YOLO detection model.
sam_model (str): Path or name of the pre-trained SAM segmentation model.
device (str): Device to run the models on (e.g., 'cpu', 'cuda', '0'). Empty string for auto-selection.
conf (float): Confidence threshold for detection model.
iou (float): IoU threshold for filtering overlapping boxes in detection results.
imgsz (int): Input image resize dimension.
max_det (int): Maximum number of detections per image.
classes (list[int], optional): Filter predictions to specified class IDs, returning only relevant detections.
output_dir (str | Path, optional): Directory to save the annotated results. If None, creates a default directory
based on the input data path.
Examples:
>>> from ultralytics.data.annotator import auto_annotate
>>> auto_annotate(data="ultralytics/assets", det_model="yolo26n.pt", sam_model="mobile_sam.pt")
"""
det_model = YOLO(det_model)
sam_model = SAM(sam_model)
data = Path(data)
if not output_dir:
output_dir = data.parent / f"{data.stem}_auto_annotate_labels"
Path(output_dir).mkdir(exist_ok=True, parents=True)
det_results = det_model(
data, stream=True, device=device, conf=conf, iou=iou, imgsz=imgsz, max_det=max_det, classes=classes
)
for result in det_results:
if class_ids := result.boxes.cls.int().tolist(): # Extract class IDs from detection results
boxes = result.boxes.xyxy # Boxes object for bbox outputs
sam_results = sam_model(result.orig_img, bboxes=boxes, verbose=False, save=False, device=device)
segments = sam_results[0].masks.xyn
with open(f"{Path(output_dir) / Path(result.path).stem}.txt", "w", encoding="utf-8") as f:
for i, s in enumerate(segments):
if s.any():
segment = map(str, s.reshape(-1).tolist())
f.write(f"{class_ids[i]} " + " ".join(segment) + "\n")

2817
ultralytics/data/augment.py Executable file
View File

File diff suppressed because it is too large Load Diff

454
ultralytics/data/base.py Executable file
View File

@@ -0,0 +1,454 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import glob
import math
import os
import random
from copy import deepcopy
from multiprocessing.pool import ThreadPool
from pathlib import Path
from typing import Any
import cv2
import numpy as np
from torch.utils.data import Dataset
from ultralytics.data.utils import FORMATS_HELP_MSG, HELP_URL, IMG_FORMATS, check_file_speeds
from ultralytics.utils import DEFAULT_CFG, LOCAL_RANK, LOGGER, NUM_THREADS, TQDM
from ultralytics.utils.patches import imread
class BaseDataset(Dataset):
"""Base dataset class for loading and processing image data.
This class provides core functionality for loading images, caching, and preparing data for training and inference in
object detection tasks.
Attributes:
img_path (str | list[str]): Path to the folder containing images.
imgsz (int): Target image size for resizing.
augment (bool): Whether to apply data augmentation.
single_cls (bool): Whether to treat all objects as a single class.
prefix (str): Prefix to print in log messages.
fraction (float): Fraction of dataset to utilize.
channels (int): Number of channels in the images (1 for grayscale, 3 for color). Color images loaded with OpenCV
are in BGR channel order.
cv2_flag (int): OpenCV flag for reading images.
im_files (list[str]): List of image file paths.
labels (list[dict]): List of label data dictionaries.
ni (int): Number of images in the dataset.
rect (bool): Whether to use rectangular training.
batch_size (int): Size of batches.
stride (int): Stride used in the model.
pad (float): Padding value.
buffer (list): Buffer for mosaic images.
max_buffer_length (int): Maximum buffer size.
ims (list): List of loaded images.
im_hw0 (list): List of original image dimensions (h, w).
im_hw (list): List of resized image dimensions (h, w).
npy_files (list[Path]): List of numpy file paths.
cache (str | None): Cache setting ('ram', 'disk', or None for no caching).
transforms (callable): Image transformation function.
batch_shapes (np.ndarray): Batch shapes for rectangular training.
batch (np.ndarray): Batch index of each image.
Methods:
get_img_files: Read image files from the specified path.
update_labels: Update labels to include only specified classes.
load_image: Load an image from the dataset.
cache_images: Cache images to memory or disk.
cache_images_to_disk: Save an image as an *.npy file for faster loading.
check_cache_disk: Check image caching requirements vs available disk space.
check_cache_ram: Check image caching requirements vs available memory.
set_rectangle: Sort images by aspect ratio and set batch shapes for rectangular training.
get_image_and_label: Get and return label information from the dataset.
update_labels_info: Custom label format method to be implemented by subclasses.
build_transforms: Build transformation pipeline to be implemented by subclasses.
get_labels: Get labels method to be implemented by subclasses.
"""
def __init__(
self,
img_path: str | list[str],
imgsz: int = 640,
cache: bool | str = False,
augment: bool = True,
hyp: dict[str, Any] = DEFAULT_CFG,
prefix: str = "",
rect: bool = False,
batch_size: int = 16,
stride: int = 32,
pad: float = 0.5,
single_cls: bool = False,
classes: list[int] | None = None,
fraction: float = 1.0,
channels: int = 3,
use_yuv444: bool = False,
):
"""Initialize BaseDataset with given configuration and options.
Args:
img_path (str | list[str]): Path to the folder containing images or list of image paths.
imgsz (int): Image size for resizing.
cache (bool | str): Cache images to RAM or disk during training.
augment (bool): If True, data augmentation is applied.
hyp (dict[str, Any]): Hyperparameters to apply data augmentation.
prefix (str): Prefix to print in log messages.
rect (bool): If True, rectangular training is used.
batch_size (int): Size of batches.
stride (int): Stride used in the model.
pad (float): Padding value.
single_cls (bool): If True, single class training is used.
classes (list[int], optional): List of included classes.
fraction (float): Fraction of dataset to utilize.
channels (int): Number of channels in the images (1 for grayscale, 3 for color). Color images loaded with
OpenCV are in BGR channel order.
use_yuv444 (bool): If True, convert YUV444 images to BGR after loading.
"""
super().__init__()
self.img_path = img_path
self.imgsz = imgsz
self.augment = augment
self.single_cls = single_cls
self.prefix = prefix
self.fraction = fraction
self.channels = channels
self.use_yuv444 = use_yuv444
self.cv2_flag = cv2.IMREAD_GRAYSCALE if channels == 1 else cv2.IMREAD_COLOR
self.im_files = self.get_img_files(self.img_path)
self.labels = self.get_labels()
self.update_labels(include_class=classes) # single_cls and include_class
self.ni = len(self.labels) # number of images
self.rect = rect
self.batch_size = batch_size
self.stride = stride
self.pad = pad
if self.rect:
assert self.batch_size is not None
self.set_rectangle()
# Buffer thread for mosaic images
self.buffer = [] # buffer size = batch size
self.max_buffer_length = min((self.ni, self.batch_size * 8, 1000)) if self.augment else 0
# Cache images (options are cache = True, False, None, "ram", "disk")
self.ims, self.im_hw0, self.im_hw = [None] * self.ni, [None] * self.ni, [None] * self.ni
self.npy_files = [Path(f).with_suffix(".npy") for f in self.im_files]
self.cache = cache.lower() if isinstance(cache, str) else "ram" if cache is True else None
if self.cache == "ram" and self.check_cache_ram():
if hyp.deterministic:
LOGGER.warning(
"cache='ram' may produce non-deterministic training results. "
"Consider cache='disk' as a deterministic alternative if your disk space allows."
)
self.cache_images()
elif self.cache == "disk" and self.check_cache_disk():
self.cache_images()
# Transforms
self.transforms = self.build_transforms(hyp=hyp)
def get_img_files(self, img_path: str | list[str]) -> list[str]:
"""Read image files from the specified path.
Args:
img_path (str | list[str]): Path or list of paths to image directories or files.
Returns:
(list[str]): List of image file paths.
Raises:
FileNotFoundError: If no images are found or the path doesn't exist.
"""
try:
f = [] # image files
for p in img_path if isinstance(img_path, list) else [img_path]:
p = Path(p) # os-agnostic
if p.is_dir(): # dir
f += glob.glob(str(p / "**" / "*.*"), recursive=True)
# F = list(p.rglob('*.*')) # pathlib
elif p.is_file(): # file
with open(p, encoding="utf-8") as t:
t = t.read().strip().splitlines()
parent = str(p.parent) + os.sep
f += [x.replace("./", parent) if x.startswith("./") else x for x in t] # local to global path
# F += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
else:
raise FileNotFoundError(f"{self.prefix}{p} does not exist")
im_files = sorted(x.replace("/", os.sep) for x in f if x.rpartition(".")[-1].lower() in IMG_FORMATS)
# self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS]) # pathlib
assert im_files, f"{self.prefix}No images found in {img_path}. {FORMATS_HELP_MSG}"
except Exception as e:
raise FileNotFoundError(f"{self.prefix}Error loading data from {img_path}\n{HELP_URL}") from e
if self.fraction < 1:
im_files = im_files[: round(len(im_files) * self.fraction)] # retain a fraction of the dataset
check_file_speeds(im_files, prefix=self.prefix) # check image read speeds
return im_files
def update_labels(self, include_class: list[int] | None) -> None:
"""Update labels to include only specified classes.
Args:
include_class (list[int], optional): List of classes to include. If None, all classes are included.
"""
include_class_array = np.array(include_class).reshape(1, -1)
for i in range(len(self.labels)):
if include_class is not None:
cls = self.labels[i]["cls"]
bboxes = self.labels[i]["bboxes"]
segments = self.labels[i]["segments"]
keypoints = self.labels[i]["keypoints"]
j = (cls == include_class_array).any(1)
self.labels[i]["cls"] = cls[j]
self.labels[i]["bboxes"] = bboxes[j]
if segments:
self.labels[i]["segments"] = [segments[si] for si, idx in enumerate(j) if idx]
if keypoints is not None:
self.labels[i]["keypoints"] = keypoints[j]
if self.single_cls:
self.labels[i]["cls"][:, 0] = 0
def load_image(self, i: int, rect_mode: bool = True) -> tuple[np.ndarray, tuple[int, int], tuple[int, int]]:
"""Load an image from dataset index 'i'.
Args:
i (int): Index of the image to load.
rect_mode (bool): Whether to use rectangular resizing.
Returns:
im (np.ndarray): Loaded image as a NumPy array.
hw_original (tuple[int, int]): Original image dimensions in (height, width) format.
hw_resized (tuple[int, int]): Resized image dimensions in (height, width) format.
Raises:
FileNotFoundError: If the image file is not found.
"""
im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i]
if im is None: # not cached in RAM
if fn.exists(): # load npy
try:
im = np.load(fn)
except Exception as e:
LOGGER.warning(f"{self.prefix}Removing corrupt *.npy image file {fn} due to: {e}")
Path(fn).unlink(missing_ok=True)
im = imread(f, flags=self.cv2_flag) # BGR
else: # read image
im = imread(f, flags=self.cv2_flag) # BGR (or YUV444 if use_yuv444=True)
if im is None:
raise FileNotFoundError(f"Image Not Found {f}")
# Convert YUV444 to BGR if needed
if self.use_yuv444 and im.ndim == 3:
# YUV444 to BGR conversion using BT.601 full range
y, u, v = im[:, :, 0].astype(np.float32), im[:, :, 1].astype(np.float32), im[:, :, 2].astype(np.float32)
u -= 128.0
v -= 128.0
r = y + 1.402 * v
g = y - 0.344136 * u - 0.714136 * v
b = y + 1.772 * u
im = np.stack([b, g, r], axis=2).clip(0, 255).astype(np.uint8)
h0, w0 = im.shape[:2] # orig hw
if rect_mode: # resize long side to imgsz while maintaining aspect ratio
max_dim = max(self.imgsz) if isinstance(self.imgsz, (list, tuple)) else self.imgsz
r = max_dim / max(h0, w0) # ratio
if r != 1: # if sizes are not equal
w, h = (min(math.ceil(w0 * r), max_dim), min(math.ceil(h0 * r), max_dim))
im = cv2.resize(im, (w, h), interpolation=cv2.INTER_LINEAR)
elif isinstance(self.imgsz, (list, tuple)): # resize to rectangular target
target_w, target_h = (self.imgsz[1], self.imgsz[0]) if len(self.imgsz) == 2 else (self.imgsz[0], self.imgsz[0])
if not (h0 == target_h and w0 == target_w):
im = cv2.resize(im, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
elif not (h0 == w0 == self.imgsz): # resize by stretching image to square imgsz
im = cv2.resize(im, (self.imgsz, self.imgsz), interpolation=cv2.INTER_LINEAR)
if im.ndim == 2:
im = im[..., None]
# Add to buffer if training with augmentations
if self.augment:
self.ims[i], self.im_hw0[i], self.im_hw[i] = im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized
self.buffer.append(i)
if 1 < len(self.buffer) >= self.max_buffer_length: # prevent empty buffer
j = self.buffer.pop(0)
if self.cache != "ram":
self.ims[j], self.im_hw0[j], self.im_hw[j] = None, None, None
return im, (h0, w0), im.shape[:2]
return self.ims[i], self.im_hw0[i], self.im_hw[i]
def cache_images(self) -> None:
"""Cache images to memory or disk for faster training."""
b, gb = 0, 1 << 30 # bytes of cached images, bytes per gigabytes
fcn, storage = (self.cache_images_to_disk, "Disk") if self.cache == "disk" else (self.load_image, "RAM")
with ThreadPool(NUM_THREADS) as pool:
results = pool.imap(fcn, range(self.ni))
pbar = TQDM(enumerate(results), total=self.ni, disable=LOCAL_RANK > 0)
for i, x in pbar:
if self.cache == "disk":
b += self.npy_files[i].stat().st_size
else: # 'ram'
self.ims[i], self.im_hw0[i], self.im_hw[i] = x # im, hw_orig, hw_resized = load_image(self, i)
b += self.ims[i].nbytes
pbar.desc = f"{self.prefix}Caching images ({b / gb:.1f}GB {storage})"
pbar.close()
def cache_images_to_disk(self, i: int) -> None:
"""Save an image as an *.npy file for faster loading."""
f = self.npy_files[i]
if not f.exists():
np.save(f.as_posix(), imread(self.im_files[i], flags=self.cv2_flag), allow_pickle=False)
def check_cache_disk(self, safety_margin: float = 0.5) -> bool:
"""Check if there's enough disk space for caching images.
Args:
safety_margin (float): Safety margin factor for disk space calculation.
Returns:
(bool): True if there's enough disk space, False otherwise.
"""
import shutil
b, gb = 0, 1 << 30 # bytes of cached images, bytes per gigabytes
n = min(self.ni, 30) # extrapolate from 30 random images
for _ in range(n):
im_file = random.choice(self.im_files)
im = imread(im_file)
if im is None:
continue
b += im.nbytes
if not os.access(Path(im_file).parent, os.W_OK):
self.cache = None
LOGGER.warning(f"{self.prefix}Skipping caching images to disk, directory not writable")
return False
disk_required = b * self.ni / n * (1 + safety_margin) # bytes required to cache dataset to disk
total, _used, free = shutil.disk_usage(Path(self.im_files[0]).parent)
if disk_required > free:
self.cache = None
LOGGER.warning(
f"{self.prefix}{disk_required / gb:.1f}GB disk space required, "
f"with {int(safety_margin * 100)}% safety margin but only "
f"{free / gb:.1f}/{total / gb:.1f}GB free, not caching images to disk"
)
return False
return True
def check_cache_ram(self, safety_margin: float = 0.5) -> bool:
"""Check if there's enough RAM for caching images.
Args:
safety_margin (float): Safety margin factor for RAM calculation.
Returns:
(bool): True if there's enough RAM, False otherwise.
"""
b, gb = 0, 1 << 30 # bytes of cached images, bytes per gigabytes
n = min(self.ni, 30) # extrapolate from 30 random images
for _ in range(n):
im = imread(random.choice(self.im_files)) # sample image
if im is None:
continue
ratio = self.imgsz / max(im.shape[0], im.shape[1]) # max(h, w) # ratio
b += im.nbytes * ratio**2
mem_required = b * self.ni / n * (1 + safety_margin) # GB required to cache dataset into RAM
mem = __import__("psutil").virtual_memory()
if mem_required > mem.available:
self.cache = None
LOGGER.warning(
f"{self.prefix}{mem_required / gb:.1f}GB RAM required to cache images "
f"with {int(safety_margin * 100)}% safety margin but only "
f"{mem.available / gb:.1f}/{mem.total / gb:.1f}GB available, not caching images"
)
return False
return True
def set_rectangle(self) -> None:
"""Sort images by aspect ratio and set batch shapes for rectangular training."""
bi = np.floor(np.arange(self.ni) / self.batch_size).astype(int) # batch index
nb = bi[-1] + 1 # number of batches
s = np.array([x.pop("shape") for x in self.labels]) # hw
ar = s[:, 0] / s[:, 1] # aspect ratio
irect = ar.argsort()
self.im_files = [self.im_files[i] for i in irect]
self.labels = [self.labels[i] for i in irect]
ar = ar[irect]
# Set training image shapes
shapes = [[1, 1]] * nb
for i in range(nb):
ari = ar[bi == i]
mini, maxi = ari.min(), ari.max()
if maxi < 1:
shapes[i] = [maxi, 1]
elif mini > 1:
shapes[i] = [1, 1 / mini]
self.batch_shapes = np.ceil(np.array(shapes) * self.imgsz / self.stride + self.pad).astype(int) * self.stride
self.batch = bi # batch index of image
def __getitem__(self, index: int) -> dict[str, Any]:
"""Return transformed label information for given index."""
return self.transforms(self.get_image_and_label(index))
def get_image_and_label(self, index: int) -> dict[str, Any]:
"""Get and return label information from the dataset.
Args:
index (int): Index of the image to retrieve.
Returns:
(dict[str, Any]): Label dictionary with image and metadata.
"""
label = deepcopy(self.labels[index]) # requires deepcopy() https://github.com/ultralytics/ultralytics/pull/1948
label.pop("shape", None) # shape is for rect, remove it
label["img"], label["ori_shape"], label["resized_shape"] = self.load_image(index)
label["ratio_pad"] = (
label["resized_shape"][0] / label["ori_shape"][0],
label["resized_shape"][1] / label["ori_shape"][1],
) # for evaluation
if self.rect:
label["rect_shape"] = self.batch_shapes[self.batch[index]]
return self.update_labels_info(label)
def __len__(self) -> int:
"""Return the length of the labels list for the dataset."""
return len(self.labels)
def update_labels_info(self, label: dict[str, Any]) -> dict[str, Any]:
"""Customize your label format here."""
return label
def build_transforms(self, hyp: dict[str, Any] | None = None):
"""Users can customize augmentations here.
Examples:
>>> if self.augment:
... # Training transforms
... return Compose([])
>>> else:
... # Val transforms
... return Compose([])
"""
raise NotImplementedError
def get_labels(self) -> list[dict[str, Any]]:
"""Users can customize their own format here.
Examples:
Ensure output is a dictionary with the following keys:
>>> dict(
... im_file=im_file,
... shape=shape, # format: (height, width)
... cls=cls,
... bboxes=bboxes, # xywh
... segments=segments, # xy
... keypoints=keypoints, # xy
... normalized=True, # or False
... bbox_format="xyxy", # or xywh, ltwh
... )
"""
raise NotImplementedError

444
ultralytics/data/build.py Executable file
View File

@@ -0,0 +1,444 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import math
import os
import random
from collections.abc import Iterator
from pathlib import Path
from typing import Any
from urllib.parse import urlsplit
import numpy as np
import torch
import torch.distributed as dist
from PIL import Image
from torch.utils.data import Dataset, dataloader, distributed
from ultralytics.cfg import IterableSimpleNamespace
from ultralytics.data.dataset import GroundingDataset, YOLODataset, YOLOGroundDataset, YOLOMultiModalDataset
from ultralytics.data.loaders import (
LOADERS,
LoadImagesAndVideos,
LoadPilAndNumpy,
LoadScreenshots,
LoadStreams,
LoadTensor,
SourceTypes,
autocast_list,
)
from ultralytics.data.utils import IMG_FORMATS, VID_FORMATS
from ultralytics.utils import RANK, colorstr
from ultralytics.utils.checks import check_file
from ultralytics.utils.torch_utils import TORCH_2_0
class InfiniteDataLoader(dataloader.DataLoader):
"""DataLoader that reuses workers for infinite iteration.
This dataloader extends the PyTorch DataLoader to provide infinite recycling of workers, which improves efficiency
for training loops that need to iterate through the dataset multiple times without recreating workers.
Attributes:
batch_sampler (_RepeatSampler): A sampler that repeats indefinitely.
iterator (Iterator): The iterator from the parent DataLoader.
Methods:
__len__: Return the length of the batch sampler's sampler.
__iter__: Yield batches from the underlying iterator.
__del__: Ensure workers are properly terminated.
reset: Reset the iterator, useful when modifying dataset settings during training.
Examples:
Create an infinite DataLoader for training
>>> dataset = YOLODataset(...)
>>> dataloader = InfiniteDataLoader(dataset, batch_size=16, shuffle=True)
>>> for batch in dataloader: # Infinite iteration
>>> train_step(batch)
"""
def __init__(self, *args: Any, **kwargs: Any):
"""Initialize the InfiniteDataLoader with the same arguments as DataLoader."""
if not TORCH_2_0:
kwargs.pop("prefetch_factor", None) # not supported by earlier versions
super().__init__(*args, **kwargs)
object.__setattr__(self, "batch_sampler", _RepeatSampler(self.batch_sampler))
self.iterator = super().__iter__()
def __len__(self) -> int:
"""Return the length of the batch sampler's sampler."""
return len(self.batch_sampler.sampler)
def __iter__(self) -> Iterator:
"""Create an iterator that yields indefinitely from the underlying iterator."""
for _ in range(len(self)):
yield next(self.iterator)
def __del__(self):
"""Ensure that workers are properly terminated when the DataLoader is deleted."""
try:
if not hasattr(self.iterator, "_workers"):
return
for w in self.iterator._workers: # force terminate
if w.is_alive():
w.terminate()
self.iterator._shutdown_workers() # cleanup
except Exception:
pass
def reset(self):
"""Reset the iterator to allow modifications to the dataset during training."""
self.iterator = self._get_iterator()
class _RepeatSampler:
"""Sampler that repeats forever for infinite iteration.
This sampler wraps another sampler and yields its contents indefinitely, allowing for infinite iteration over a
dataset without recreating the sampler.
Attributes:
sampler (torch.utils.data.Sampler): The sampler to repeat.
"""
def __init__(self, sampler: Any):
"""Initialize the _RepeatSampler with a sampler to repeat indefinitely."""
self.sampler = sampler
def __iter__(self) -> Iterator:
"""Iterate over the sampler indefinitely, yielding its contents."""
while True:
yield from iter(self.sampler)
class ContiguousDistributedSampler(torch.utils.data.Sampler):
"""Distributed sampler that assigns contiguous batch-aligned chunks of the dataset to each GPU.
Unlike PyTorch's DistributedSampler which distributes samples in a round-robin fashion (GPU 0 gets indices
[0,2,4,...], GPU 1 gets [1,3,5,...]), this sampler gives each GPU contiguous batches of the dataset (GPU 0 gets
batches [0,1,2,...], GPU 1 gets batches [k,k+1,...], etc.). This preserves any ordering or grouping in the original
dataset, which is critical when samples are organized by similarity (e.g., images sorted by size to enable efficient
batching without padding when using rect=True).
The sampler handles uneven batch counts by distributing remainder batches to the first few ranks, ensuring all
samples are covered exactly once across all GPUs.
Args:
dataset (Dataset): Dataset to sample from. Must implement __len__.
num_replicas (int, optional): Number of distributed processes. Defaults to world size.
batch_size (int, optional): Batch size used by dataloader. Defaults to dataset.batch_size or 1.
rank (int, optional): Rank of current process. Defaults to current rank.
shuffle (bool, optional): Whether to shuffle indices within each rank's chunk. Defaults to False. When True,
shuffling is deterministic and controlled by set_epoch() for reproducibility.
Examples:
>>> # For validation with size-grouped images
>>> sampler = ContiguousDistributedSampler(val_dataset, batch_size=32, shuffle=False)
>>> loader = DataLoader(val_dataset, batch_size=32, sampler=sampler)
>>> # For training with shuffling
>>> sampler = ContiguousDistributedSampler(train_dataset, batch_size=32, shuffle=True)
>>> for epoch in range(num_epochs):
... sampler.set_epoch(epoch)
... for batch in loader:
... ...
"""
def __init__(
self,
dataset: Dataset,
num_replicas: int | None = None,
batch_size: int | None = None,
rank: int | None = None,
shuffle: bool = False,
) -> None:
"""Initialize the sampler with dataset and distributed training parameters."""
if num_replicas is None:
num_replicas = dist.get_world_size() if dist.is_initialized() else 1
if rank is None:
rank = dist.get_rank() if dist.is_initialized() else 0
if batch_size is None:
batch_size = getattr(dataset, "batch_size", 1)
self.num_replicas = num_replicas
self.rank = rank
self.epoch = 0
self.shuffle = shuffle
self.total_size = len(dataset)
# ensure all ranks have a sample if batch size >= total size; degenerates to round-robin sampler
self.batch_size = 1 if batch_size >= self.total_size else batch_size
self.num_batches = math.ceil(self.total_size / self.batch_size)
def _get_rank_indices(self) -> tuple[int, int]:
"""Calculate the start and end sample indices for this rank."""
# Calculate which batches this rank handles
batches_per_rank_base = self.num_batches // self.num_replicas
remainder = self.num_batches % self.num_replicas
# This rank gets an extra batch if rank < remainder
batches_for_this_rank = batches_per_rank_base + (1 if self.rank < remainder else 0)
# Calculate starting batch: base position + number of extra batches given to earlier ranks
start_batch = self.rank * batches_per_rank_base + min(self.rank, remainder)
end_batch = start_batch + batches_for_this_rank
# Convert batch indices to sample indices
start_idx = start_batch * self.batch_size
end_idx = min(end_batch * self.batch_size, self.total_size)
return start_idx, end_idx
def __iter__(self) -> Iterator:
"""Generate indices for this rank's contiguous chunk of the dataset."""
start_idx, end_idx = self._get_rank_indices()
indices = list(range(start_idx, end_idx))
if self.shuffle:
g = torch.Generator()
g.manual_seed(self.epoch)
indices = [indices[i] for i in torch.randperm(len(indices), generator=g).tolist()]
return iter(indices)
def __len__(self) -> int:
"""Return the number of samples in this rank's chunk."""
start_idx, end_idx = self._get_rank_indices()
return end_idx - start_idx
def set_epoch(self, epoch: int) -> None:
"""Set the epoch for this sampler to ensure different shuffling patterns across epochs.
Args:
epoch (int): Epoch number to use as the random seed for shuffling.
"""
self.epoch = epoch
def seed_worker(worker_id: int) -> None:
"""Set dataloader worker seed for reproducibility across worker processes."""
worker_seed = torch.initial_seed() % 2**32
np.random.seed(worker_seed)
random.seed(worker_seed)
def build_yolo_dataset(
cfg: IterableSimpleNamespace,
img_path: str,
batch: int,
data: dict[str, Any],
mode: str = "train",
rect: bool = False,
stride: int = 32,
multi_modal: bool = False,
) -> Dataset:
"""Build and return a YOLO dataset based on configuration parameters."""
# Detect ground dataset by presence of class_map
if "class_map" in data:
dataset = YOLOGroundDataset
elif multi_modal:
dataset = YOLOMultiModalDataset
else:
dataset = YOLODataset
return dataset(
img_path=img_path,
imgsz=cfg.imgsz,
batch_size=batch,
augment=mode == "train", # augmentation
hyp=cfg, # TODO: probably add a get_hyps_from_cfg function
rect=cfg.rect or rect, # rectangular batches
cache=cfg.cache or None,
single_cls=cfg.single_cls or False,
stride=stride,
pad=0.0 if mode == "train" else 0.5,
prefix=colorstr(f"{mode}: "),
task=cfg.task,
classes=cfg.classes,
data=data,
fraction=cfg.fraction if mode == "train" else 1.0,
)
def build_grounding(
cfg: IterableSimpleNamespace,
img_path: str,
json_file: str,
batch: int,
mode: str = "train",
rect: bool = False,
stride: int = 32,
max_samples: int = 80,
) -> Dataset:
"""Build and return a GroundingDataset based on configuration parameters."""
return GroundingDataset(
img_path=img_path,
json_file=json_file,
max_samples=max_samples,
imgsz=cfg.imgsz,
batch_size=batch,
augment=mode == "train", # augmentation
hyp=cfg, # TODO: probably add a get_hyps_from_cfg function
rect=cfg.rect or rect, # rectangular batches
cache=cfg.cache or None,
single_cls=cfg.single_cls or False,
stride=stride,
pad=0.0 if mode == "train" else 0.5,
prefix=colorstr(f"{mode}: "),
task=cfg.task,
classes=cfg.classes,
fraction=cfg.fraction if mode == "train" else 1.0,
)
def build_dataloader(
dataset,
batch: int,
workers: int,
shuffle: bool = True,
rank: int = -1,
drop_last: bool = False,
pin_memory: bool = True,
) -> InfiniteDataLoader:
"""Create and return an InfiniteDataLoader for training or validation.
Args:
dataset (Dataset): Dataset to load data from.
batch (int): Batch size for the dataloader.
workers (int): Number of worker processes for data loading.
shuffle (bool, optional): Whether to shuffle the dataset.
rank (int, optional): Process rank in distributed training. -1 for single-GPU training.
drop_last (bool, optional): Whether to drop the last incomplete batch.
pin_memory (bool, optional): Whether to use pinned memory for dataloader.
Returns:
(InfiniteDataLoader): A dataloader that can be used for training or validation.
Examples:
Create a dataloader for training
>>> dataset = YOLODataset(...)
>>> dataloader = build_dataloader(dataset, batch=16, workers=4, shuffle=True)
"""
batch = min(batch, len(dataset))
nd = torch.cuda.device_count() # number of CUDA devices
nw = min(os.cpu_count() // max(nd, 1), workers) # number of workers
sampler = (
None
if rank == -1
else distributed.DistributedSampler(dataset, shuffle=shuffle)
if shuffle
else ContiguousDistributedSampler(dataset)
)
generator = torch.Generator()
generator.manual_seed(6148914691236517205 + RANK)
return InfiniteDataLoader(
dataset=dataset,
batch_size=batch,
shuffle=shuffle and sampler is None,
num_workers=nw,
sampler=sampler,
prefetch_factor=4 if nw > 0 else None, # increase over default 2
pin_memory=nd > 0 and pin_memory,
collate_fn=getattr(dataset, "collate_fn", None),
worker_init_fn=seed_worker,
generator=generator,
drop_last=drop_last and len(dataset) % batch != 0,
)
def check_source(
source: str | int | Path | list | tuple | np.ndarray | Image.Image | torch.Tensor,
) -> tuple[Any, bool, bool, bool, bool, bool]:
"""Check the type of input source and return corresponding flag values.
Args:
source (str | int | Path | list | tuple | np.ndarray | PIL.Image | torch.Tensor): The input source to check.
Returns:
source (str | int | Path | list | tuple | np.ndarray | PIL.Image | torch.Tensor): The processed source.
webcam (bool): Whether the source is a webcam.
screenshot (bool): Whether the source is a screenshot.
from_img (bool): Whether the source is an image or list of images.
in_memory (bool): Whether the source is an in-memory object.
tensor (bool): Whether the source is a torch.Tensor.
Examples:
Check a file path source
>>> source, webcam, screenshot, from_img, in_memory, tensor = check_source("image.jpg")
Check a webcam source
>>> source, webcam, screenshot, from_img, in_memory, tensor = check_source(0)
"""
webcam, screenshot, from_img, in_memory, tensor = False, False, False, False, False
if isinstance(source, (str, int, Path)): # int for local usb camera
source = str(source)
source_lower = source.lower()
is_url = source_lower.startswith(("https://", "http://", "rtsp://", "rtmp://", "tcp://"))
is_file = (urlsplit(source_lower).path if is_url else source_lower).rpartition(".")[-1] in (
IMG_FORMATS | VID_FORMATS
)
webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
screenshot = source_lower == "screen"
if is_url and is_file:
source = check_file(source) # download
elif isinstance(source, LOADERS):
in_memory = True
elif isinstance(source, (list, tuple)):
source = autocast_list(source) # convert all list elements to PIL or np arrays
from_img = True
elif isinstance(source, (Image.Image, np.ndarray)):
from_img = True
elif isinstance(source, torch.Tensor):
tensor = True
else:
raise TypeError("Unsupported image type. For supported types see https://docs.ultralytics.com/modes/predict")
return source, webcam, screenshot, from_img, in_memory, tensor
def load_inference_source(
source: str | int | Path | list | tuple | np.ndarray | Image.Image | torch.Tensor,
batch: int = 1,
vid_stride: int = 1,
buffer: bool = False,
channels: int = 3,
):
"""Load an inference source for object detection and apply necessary transformations.
Args:
source (str | int | Path | list | tuple | np.ndarray | PIL.Image | torch.Tensor): The input source for
inference.
batch (int, optional): Batch size for dataloaders.
vid_stride (int, optional): The frame interval for video sources.
buffer (bool, optional): Whether stream frames will be buffered.
channels (int, optional): The number of input channels for the model.
Returns:
(Dataset): A dataset object for the specified input source with attached source_type attribute.
Examples:
Load an image source for inference
>>> dataset = load_inference_source("image.jpg", batch=1)
Load a video stream source
>>> dataset = load_inference_source("rtsp://example.com/stream", vid_stride=2)
"""
source, stream, screenshot, from_img, in_memory, tensor = check_source(source)
source_type = source.source_type if in_memory else SourceTypes(stream, screenshot, from_img, tensor)
# DataLoader
if tensor:
dataset = LoadTensor(source)
elif in_memory:
dataset = source
elif stream:
dataset = LoadStreams(source, vid_stride=vid_stride, buffer=buffer, channels=channels)
elif screenshot:
dataset = LoadScreenshots(source, channels=channels)
elif from_img:
dataset = LoadPilAndNumpy(source, channels=channels)
else:
dataset = LoadImagesAndVideos(source, batch=batch, vid_stride=vid_stride, channels=channels)
# Attach source types to the dataset
setattr(dataset, "source_type", source_type)
return dataset

948
ultralytics/data/converter.py Executable file
View File

@@ -0,0 +1,948 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import asyncio
import hashlib
import json
import random
import shutil
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
import cv2
import numpy as np
from PIL import Image
from ultralytics.utils import ASSETS_URL, DATASETS_DIR, LOGGER, NUM_THREADS, TQDM, YAML
from ultralytics.utils.checks import check_file
from ultralytics.utils.downloads import download, zip_directory
from ultralytics.utils.files import increment_path
def coco91_to_coco80_class() -> list[int]:
"""Convert 91-index COCO class IDs to 80-index COCO class IDs.
Returns:
(list[int | None]): A list of 91 elements where the index represents the 91-index class ID and the value is the
corresponding 80-index class ID, or None if there is no mapping.
"""
return [
0,
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
None,
11,
12,
13,
14,
15,
16,
17,
18,
19,
20,
21,
22,
23,
None,
24,
25,
None,
None,
26,
27,
28,
29,
30,
31,
32,
33,
34,
35,
36,
37,
38,
39,
None,
40,
41,
42,
43,
44,
45,
46,
47,
48,
49,
50,
51,
52,
53,
54,
55,
56,
57,
58,
59,
None,
60,
None,
None,
61,
None,
62,
63,
64,
65,
66,
67,
68,
69,
70,
71,
72,
None,
73,
74,
75,
76,
77,
78,
79,
None,
]
def coco80_to_coco91_class() -> list[int]:
r"""Convert 80-index (val2014) to 91-index (paper).
Returns:
(list[int]): A list of 80 class IDs where each value is the corresponding 91-index class ID.
Examples:
>>> import numpy as np
>>> a = np.loadtxt("data/coco.names", dtype="str", delimiter="\n")
>>> b = np.loadtxt("data/coco_paper.names", dtype="str", delimiter="\n")
Convert the darknet to COCO format
>>> x1 = [list(a[i] == b).index(True) + 1 for i in range(80)]
Convert the COCO to darknet format
>>> x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)]
References:
https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
"""
return [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
11,
13,
14,
15,
16,
17,
18,
19,
20,
21,
22,
23,
24,
25,
27,
28,
31,
32,
33,
34,
35,
36,
37,
38,
39,
40,
41,
42,
43,
44,
46,
47,
48,
49,
50,
51,
52,
53,
54,
55,
56,
57,
58,
59,
60,
61,
62,
63,
64,
65,
67,
70,
72,
73,
74,
75,
76,
77,
78,
79,
80,
81,
82,
84,
85,
86,
87,
88,
89,
90,
]
def convert_coco(
labels_dir: str = "../coco/annotations/",
save_dir: str = "coco_converted/",
use_segments: bool = False,
use_keypoints: bool = False,
cls91to80: bool = True,
lvis: bool = False,
):
"""Convert COCO dataset annotations to a YOLO annotation format suitable for training YOLO models.
Args:
labels_dir (str, optional): Path to directory containing COCO dataset annotation files.
save_dir (str, optional): Path to directory to save results to.
use_segments (bool, optional): Whether to include segmentation masks in the output.
use_keypoints (bool, optional): Whether to include keypoint annotations in the output.
cls91to80 (bool, optional): Whether to map 91 COCO class IDs to the corresponding 80 COCO class IDs.
lvis (bool, optional): Whether to convert data in lvis dataset way.
Examples:
>>> from ultralytics.data.converter import convert_coco
Convert COCO annotations to YOLO format
>>> convert_coco("coco/annotations/", use_segments=True, use_keypoints=False, cls91to80=False)
Convert LVIS annotations to YOLO format
>>> convert_coco("lvis/annotations/", use_segments=True, use_keypoints=False, cls91to80=False, lvis=True)
"""
# Create dataset directory
save_dir = increment_path(save_dir) # increment if save directory already exists
for p in save_dir / "labels", save_dir / "images":
p.mkdir(parents=True, exist_ok=True) # make dir
# Convert classes
coco80 = coco91_to_coco80_class()
# Import json
for json_file in sorted(Path(labels_dir).resolve().glob("*.json")):
lname = "" if lvis else json_file.stem.replace("instances_", "")
fn = Path(save_dir) / "labels" / lname # folder name
fn.mkdir(parents=True, exist_ok=True)
if lvis:
# NOTE: create folders for both train and val in advance,
# since LVIS val set contains images from COCO 2017 train in addition to the COCO 2017 val split.
(fn / "train2017").mkdir(parents=True, exist_ok=True)
(fn / "val2017").mkdir(parents=True, exist_ok=True)
with open(json_file, encoding="utf-8") as f:
data = json.load(f)
# Create image dict
images = {f"{x['id']:d}": x for x in data["images"]}
# Create image-annotations dict
annotations = defaultdict(list)
for ann in data["annotations"]:
annotations[ann["image_id"]].append(ann)
image_txt = []
# Write labels file
for img_id, anns in TQDM(annotations.items(), desc=f"Annotations {json_file}"):
img = images[f"{img_id:d}"]
h, w = img["height"], img["width"]
f = str(Path(img["coco_url"]).relative_to("http://images.cocodataset.org")) if lvis else img["file_name"]
if lvis:
image_txt.append(str(Path("./images") / f))
bboxes = []
segments = []
keypoints = []
for ann in anns:
if ann.get("iscrowd", False):
continue
# The COCO box format is [top left x, top left y, width, height]
box = np.array(ann["bbox"], dtype=np.float64)
box[:2] += box[2:] / 2 # xy top-left corner to center
box[[0, 2]] /= w # normalize x
box[[1, 3]] /= h # normalize y
if box[2] <= 0 or box[3] <= 0: # if w <= 0 and h <= 0
continue
cls = coco80[ann["category_id"] - 1] if cls91to80 else ann["category_id"] - 1 # class
box = [cls, *box.tolist()]
if box not in bboxes:
bboxes.append(box)
if use_segments and ann.get("segmentation") is not None:
if len(ann["segmentation"]) == 0:
segments.append([])
continue
elif len(ann["segmentation"]) > 1:
s = merge_multi_segment(ann["segmentation"])
s = (np.concatenate(s, axis=0) / np.array([w, h])).reshape(-1).tolist()
else:
s = [j for i in ann["segmentation"] for j in i] # all segments concatenated
s = (np.array(s).reshape(-1, 2) / np.array([w, h])).reshape(-1).tolist()
s = [cls, *s]
segments.append(s)
if use_keypoints and ann.get("keypoints") is not None:
keypoints.append(
box + (np.array(ann["keypoints"]).reshape(-1, 3) / np.array([w, h, 1])).reshape(-1).tolist()
)
# Write
with open((fn / f).with_suffix(".txt"), "a", encoding="utf-8") as file:
for i in range(len(bboxes)):
if use_keypoints:
line = (*(keypoints[i]),) # cls, box, keypoints
else:
line = (
*(segments[i] if use_segments and len(segments[i]) > 0 else bboxes[i]),
) # cls, box or segments
file.write(("%g " * len(line)).rstrip() % line + "\n")
if lvis:
filename = Path(save_dir) / json_file.name.replace("lvis_v1_", "").replace(".json", ".txt")
with open(filename, "a", encoding="utf-8") as f:
f.writelines(f"{line}\n" for line in image_txt)
LOGGER.info(f"{'LVIS' if lvis else 'COCO'} data converted successfully.\nResults saved to {save_dir.resolve()}")
def convert_segment_masks_to_yolo_seg(masks_dir: str, output_dir: str, classes: int):
"""Convert a dataset of segmentation mask images to the YOLO segmentation format.
This function takes the directory containing the binary format mask images and converts them into YOLO segmentation
format. The converted masks are saved in the specified output directory.
Args:
masks_dir (str): The path to the directory where all mask images (png, jpg) are stored.
output_dir (str): The path to the directory where the converted YOLO segmentation masks will be stored.
classes (int): Total number of classes in the dataset, e.g., 80 for COCO.
Examples:
>>> from ultralytics.data.converter import convert_segment_masks_to_yolo_seg
The classes here is the total classes in the dataset, for COCO dataset we have 80 classes
>>> convert_segment_masks_to_yolo_seg("path/to/masks_directory", "path/to/output/directory", classes=80)
Notes:
The expected directory structure for the masks is:
- masks
├─ mask_image_01.png or mask_image_01.jpg
├─ mask_image_02.png or mask_image_02.jpg
├─ mask_image_03.png or mask_image_03.jpg
└─ mask_image_04.png or mask_image_04.jpg
After execution, the labels will be organized in the following structure:
- output_dir
├─ mask_yolo_01.txt
├─ mask_yolo_02.txt
├─ mask_yolo_03.txt
└─ mask_yolo_04.txt
"""
pixel_to_class_mapping = {i + 1: i for i in range(classes)}
for mask_path in Path(masks_dir).iterdir():
if mask_path.suffix in {".png", ".jpg"}:
mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE) # Read the mask image in grayscale
img_height, img_width = mask.shape # Get image dimensions
LOGGER.info(f"Processing {mask_path} imgsz = {img_height} x {img_width}")
unique_values = np.unique(mask) # Get unique pixel values representing different classes
yolo_format_data = []
for value in unique_values:
if value == 0:
continue # Skip background
class_index = pixel_to_class_mapping.get(value, -1)
if class_index == -1:
LOGGER.warning(f"Unknown class for pixel value {value} in file {mask_path}, skipping.")
continue
# Create a binary mask for the current class and find contours
contours, _ = cv2.findContours(
(mask == value).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
) # Find contours
for contour in contours:
if len(contour) >= 3: # YOLO requires at least 3 points for a valid segmentation
contour = contour.squeeze() # Remove single-dimensional entries
yolo_format = [class_index]
for point in contour:
# Normalize the coordinates
yolo_format.append(round(point[0] / img_width, 6)) # Rounding to 6 decimal places
yolo_format.append(round(point[1] / img_height, 6))
yolo_format_data.append(yolo_format)
# Save Ultralytics YOLO format data to file
output_path = Path(output_dir) / f"{mask_path.stem}.txt"
with open(output_path, "w", encoding="utf-8") as file:
for item in yolo_format_data:
line = " ".join(map(str, item))
file.write(line + "\n")
LOGGER.info(f"Processed and stored at {output_path} imgsz = {img_height} x {img_width}")
def convert_dota_to_yolo_obb(dota_root_path: str):
"""Convert DOTA dataset annotations to YOLO OBB (Oriented Bounding Box) format.
The function processes images in the 'train' and 'val' folders of the DOTA dataset. For each image, it reads the
associated label from the original labels directory and writes new labels in YOLO OBB format to a new directory.
Args:
dota_root_path (str): The root directory path of the DOTA dataset.
Examples:
>>> from ultralytics.data.converter import convert_dota_to_yolo_obb
>>> convert_dota_to_yolo_obb("path/to/DOTA")
Notes:
The directory structure assumed for the DOTA dataset:
- DOTA
├─ images
│ ├─ train
│ └─ val
└─ labels
├─ train_original
└─ val_original
After execution, the function will organize the labels into:
- DOTA
└─ labels
├─ train
└─ val
"""
dota_root_path = Path(dota_root_path)
# Class names to indices mapping
class_mapping = {
"plane": 0,
"ship": 1,
"storage-tank": 2,
"baseball-diamond": 3,
"tennis-court": 4,
"basketball-court": 5,
"ground-track-field": 6,
"harbor": 7,
"bridge": 8,
"large-vehicle": 9,
"small-vehicle": 10,
"helicopter": 11,
"roundabout": 12,
"soccer-ball-field": 13,
"swimming-pool": 14,
"container-crane": 15,
"airport": 16,
"helipad": 17,
}
def convert_label(image_name: str, image_width: int, image_height: int, orig_label_dir: Path, save_dir: Path):
"""Convert a single image's DOTA annotation to YOLO OBB format and save it to a specified directory."""
orig_label_path = orig_label_dir / f"{image_name}.txt"
save_path = save_dir / f"{image_name}.txt"
with orig_label_path.open("r") as f, save_path.open("w") as g:
lines = f.readlines()
for line in lines:
parts = line.strip().split()
if len(parts) < 9:
continue
class_name = parts[8]
class_idx = class_mapping[class_name]
coords = [float(p) for p in parts[:8]]
normalized_coords = [
coords[i] / image_width if i % 2 == 0 else coords[i] / image_height for i in range(8)
]
formatted_coords = [f"{coord:.6g}" for coord in normalized_coords]
g.write(f"{class_idx} {' '.join(formatted_coords)}\n")
for phase in {"train", "val"}:
image_dir = dota_root_path / "images" / phase
orig_label_dir = dota_root_path / "labels" / f"{phase}_original"
save_dir = dota_root_path / "labels" / phase
save_dir.mkdir(parents=True, exist_ok=True)
image_paths = list(image_dir.iterdir())
for image_path in TQDM(image_paths, desc=f"Processing {phase} images"):
if image_path.suffix != ".png":
continue
image_name_without_ext = image_path.stem
img = cv2.imread(str(image_path))
h, w = img.shape[:2]
convert_label(image_name_without_ext, w, h, orig_label_dir, save_dir)
def min_index(arr1: np.ndarray, arr2: np.ndarray):
"""Find a pair of indexes with the shortest distance between two arrays of 2D points.
Args:
arr1 (np.ndarray): A NumPy array of shape (N, 2) representing N 2D points.
arr2 (np.ndarray): A NumPy array of shape (M, 2) representing M 2D points.
Returns:
(tuple[int, int]): A tuple (idx1, idx2) where idx1 is the index in arr1 and idx2 is the index in arr2 of the
pair with the shortest distance.
"""
dis = ((arr1[:, None, :] - arr2[None, :, :]) ** 2).sum(-1)
return np.unravel_index(np.argmin(dis, axis=None), dis.shape)
def merge_multi_segment(segments: list[list]):
"""Merge multiple segments into one list by connecting the coordinates with the minimum distance between each
segment.
This function connects these coordinates with a thin line to merge all segments into one.
Args:
segments (list[list]): Original segmentations in COCO's JSON file. Each element is a list of coordinates, like
[segmentation1, segmentation2,...].
Returns:
(list[np.ndarray]): A list of connected segments represented as NumPy arrays.
"""
s = []
segments = [np.array(i).reshape(-1, 2) for i in segments]
idx_list = [[] for _ in range(len(segments))]
# Record the indexes with min distance between each segment
for i in range(1, len(segments)):
idx1, idx2 = min_index(segments[i - 1], segments[i])
idx_list[i - 1].append(idx1)
idx_list[i].append(idx2)
# Use two round to connect all the segments
for k in range(2):
# Forward connection
if k == 0:
for i, idx in enumerate(idx_list):
# Middle segments have two indexes, reverse the index of middle segments
if len(idx) == 2 and idx[0] > idx[1]:
idx = idx[::-1]
segments[i] = segments[i][::-1, :]
segments[i] = np.roll(segments[i], -idx[0], axis=0)
segments[i] = np.concatenate([segments[i], segments[i][:1]])
# Deal with the first segment and the last one
if i in {0, len(idx_list) - 1}:
s.append(segments[i])
else:
idx = [0, idx[1] - idx[0]]
s.append(segments[i][idx[0] : idx[1] + 1])
else:
for i in range(len(idx_list) - 1, -1, -1):
if i not in {0, len(idx_list) - 1}:
idx = idx_list[i]
nidx = abs(idx[1] - idx[0])
s.append(segments[i][nidx:])
return s
def yolo_bbox2segment(im_dir: str | Path, save_dir: str | Path | None = None, sam_model: str = "sam_b.pt", device=None):
"""Convert existing object detection dataset (bounding boxes) to segmentation dataset in YOLO format.
Generates segmentation data using SAM auto-annotator as needed.
Args:
im_dir (str | Path): Path to image directory to convert.
save_dir (str | Path, optional): Path to save the generated labels, labels will be saved into `labels-segment`
in the same directory level of `im_dir` if save_dir is None.
sam_model (str): Segmentation model to use for intermediate segmentation data.
device (int | str, optional): The specific device to run SAM models.
Notes:
The input directory structure assumed for dataset:
- im_dir
├─ 001.jpg
├─ ...
└─ NNN.jpg
- labels
├─ 001.txt
├─ ...
└─ NNN.txt
"""
from ultralytics import SAM
from ultralytics.data import YOLODataset
from ultralytics.utils.ops import xywh2xyxy
# NOTE: add placeholder to pass class index check
dataset = YOLODataset(im_dir, data=dict(names=list(range(1000)), channels=3))
if len(dataset.labels[0]["segments"]) > 0: # if it's segment data
LOGGER.info("Segmentation labels detected, no need to generate new ones!")
return
LOGGER.info("Detection labels detected, generating segment labels by SAM model!")
sam_model = SAM(sam_model)
for label in TQDM(dataset.labels, total=len(dataset.labels), desc="Generating segment labels"):
h, w = label["shape"]
boxes = label["bboxes"]
if len(boxes) == 0: # skip empty labels
continue
boxes[:, [0, 2]] *= w
boxes[:, [1, 3]] *= h
im = cv2.imread(label["im_file"])
sam_results = sam_model(im, bboxes=xywh2xyxy(boxes), verbose=False, save=False, device=device)
label["segments"] = sam_results[0].masks.xyn
save_dir = Path(save_dir) if save_dir else Path(im_dir).parent / "labels-segment"
save_dir.mkdir(parents=True, exist_ok=True)
for label in dataset.labels:
texts = []
lb_name = Path(label["im_file"]).with_suffix(".txt").name
txt_file = save_dir / lb_name
cls = label["cls"]
for i, s in enumerate(label["segments"]):
if len(s) == 0:
continue
line = (int(cls[i]), *s.reshape(-1))
texts.append(("%g " * len(line)).rstrip() % line)
with open(txt_file, "a", encoding="utf-8") as f:
f.writelines(text + "\n" for text in texts)
LOGGER.info(f"Generated segment labels saved in {save_dir}")
def create_synthetic_coco_dataset():
"""Create a synthetic COCO dataset with random images based on filenames from label lists.
This function downloads COCO labels, reads image filenames from label list files, creates synthetic images for
train2017 and val2017 subsets, and organizes them in the COCO dataset structure. It uses multithreading to generate
images efficiently.
Examples:
>>> from ultralytics.data.converter import create_synthetic_coco_dataset
>>> create_synthetic_coco_dataset()
Notes:
- Requires internet connection to download label files.
- Generates random RGB images of varying sizes (480x480 to 640x640 pixels).
- Existing test2017 directory is removed as it's not needed.
- Reads image filenames from train2017.txt and val2017.txt files.
"""
def create_synthetic_image(image_file: Path):
"""Generate a synthetic image with random size and color for dataset augmentation or testing purposes."""
if not image_file.exists():
size = (random.randint(480, 640), random.randint(480, 640))
Image.new(
"RGB",
size=size,
color=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)),
).save(image_file)
# Download labels
dir = DATASETS_DIR / "coco"
download([f"{ASSETS_URL}/coco2017labels-segments.zip"], dir=dir.parent)
# Create synthetic images
shutil.rmtree(dir / "labels" / "test2017", ignore_errors=True) # Remove test2017 directory as not needed
with ThreadPoolExecutor(max_workers=NUM_THREADS) as executor:
for subset in {"train2017", "val2017"}:
subset_dir = dir / "images" / subset
subset_dir.mkdir(parents=True, exist_ok=True)
# Read image filenames from label list file
label_list_file = dir / f"{subset}.txt"
if label_list_file.exists():
with open(label_list_file, encoding="utf-8") as f:
image_files = [dir / line.strip() for line in f]
# Submit all tasks
futures = [executor.submit(create_synthetic_image, image_file) for image_file in image_files]
for _ in TQDM(as_completed(futures), total=len(futures), desc=f"Generating images for {subset}"):
pass # The actual work is done in the background
else:
LOGGER.warning(f"Labels file {label_list_file} does not exist. Skipping image creation for {subset}.")
LOGGER.info("Synthetic COCO dataset created successfully.")
def convert_to_multispectral(path: str | Path, n_channels: int = 10, replace: bool = False, zip: bool = False):
"""Convert RGB images to multispectral images by interpolating across wavelength bands.
This function takes RGB images and interpolates them to create multispectral images with a specified number of
channels. It can process either a single image or a directory of images.
Args:
path (str | Path): Path to an image file or directory containing images to convert.
n_channels (int): Number of spectral channels to generate in the output image.
replace (bool): Whether to replace the original image file with the converted one.
zip (bool): Whether to zip the converted images into a zip file.
Examples:
Convert a single image
>>> convert_to_multispectral("path/to/image.jpg", n_channels=10)
Convert a dataset
>>> convert_to_multispectral("coco8", n_channels=10)
"""
from scipy.interpolate import interp1d
from ultralytics.data.utils import IMG_FORMATS
path = Path(path)
if path.is_dir():
# Process directory
im_files = [f for ext in (IMG_FORMATS - {"tif", "tiff"}) for f in path.rglob(f"*.{ext}")]
for im_path in im_files:
try:
convert_to_multispectral(im_path, n_channels)
if replace:
im_path.unlink()
except Exception as e:
LOGGER.info(f"Error converting {im_path}: {e}")
if zip:
zip_directory(path)
else:
# Process a single image
output_path = path.with_suffix(".tiff")
img = cv2.cvtColor(cv2.imread(str(path)), cv2.COLOR_BGR2RGB)
# Interpolate all pixels at once
rgb_wavelengths = np.array([650, 510, 475]) # R, G, B wavelengths (nm)
target_wavelengths = np.linspace(450, 700, n_channels)
f = interp1d(rgb_wavelengths.T, img, kind="linear", bounds_error=False, fill_value="extrapolate")
multispectral = f(target_wavelengths)
cv2.imwritemulti(str(output_path), np.clip(multispectral, 0, 255).astype(np.uint8).transpose(2, 0, 1))
LOGGER.info(f"Converted {output_path}")
async def convert_ndjson_to_yolo(ndjson_path: str | Path, output_path: str | Path | None = None) -> Path:
"""Convert NDJSON dataset format to Ultralytics YOLO dataset structure.
This function converts datasets stored in NDJSON (Newline Delimited JSON) format to the standard YOLO format. For
detection/segmentation/pose/obb tasks, it creates separate directories for images and labels. For classification
tasks, it creates the ImageNet-style {split}/{class_name}/ folder structure. It supports parallel processing for
efficient conversion of large datasets and can download images from URLs.
The NDJSON format consists of:
- First line: Dataset metadata with class names, task type, and configuration
- Subsequent lines: Individual image records with annotations and optional URLs
Args:
ndjson_path (str | Path): Path to the input NDJSON file containing dataset information.
output_path (str | Path | None, optional): Directory where the converted YOLO dataset will be saved. If None,
uses the DATASETS_DIR directory. Defaults to None.
Returns:
(Path): Path to the generated data.yaml file (detection) or dataset directory (classification).
Examples:
Convert a local NDJSON file:
>>> yaml_path = await convert_ndjson_to_yolo("dataset.ndjson")
>>> print(f"Dataset converted to: {yaml_path}")
Convert with custom output directory:
>>> yaml_path = await convert_ndjson_to_yolo("dataset.ndjson", output_path="./converted_datasets")
Use with YOLO training
>>> from ultralytics import YOLO
>>> model = YOLO("yolo26n.pt")
>>> model.train(data="https://github.com/ultralytics/assets/releases/download/v0.0.0/coco8-ndjson.ndjson")
"""
from ultralytics.utils.checks import check_requirements
check_requirements("aiohttp")
import aiohttp
ndjson_path = Path(check_file(ndjson_path))
output_path = Path(output_path or DATASETS_DIR)
with open(ndjson_path) as f:
lines = [json.loads(line.strip()) for line in f if line.strip()]
dataset_record, image_records = lines[0], lines[1:]
# Hash semantic content only (excludes signed URLs which change on every export)
_h = hashlib.sha256()
for r in lines:
_h.update(json.dumps({k: v for k, v in r.items() if k != "url"}, sort_keys=True).encode())
_hash = _h.hexdigest()[:16]
# Early exit if dataset content unchanged (hash stored in data.yaml)
dataset_dir = output_path / ndjson_path.stem
yaml_path = dataset_dir / "data.yaml"
if yaml_path.is_file():
try:
if YAML.load(yaml_path).get("hash") == _hash:
return yaml_path
except Exception:
pass
splits = {record["split"] for record in image_records}
# Check if this is a classification dataset
is_classification = dataset_record.get("task") == "classify"
class_names = {int(k): v for k, v in dataset_record.get("class_names", {}).items()}
len(class_names)
# Validate required fields before downloading images
task = dataset_record.get("task", "detect")
if not is_classification:
if "train" not in splits:
raise ValueError(f"Dataset missing required 'train' split. Found splits: {sorted(splits)}")
if "val" not in splits and "test" not in splits:
raise ValueError(f"Dataset missing required 'val' split. Found splits: {sorted(splits)}")
if task == "pose" and "kpt_shape" not in dataset_record:
raise ValueError("Pose dataset missing required 'kpt_shape'. See https://docs.ultralytics.com/datasets/pose/")
# Check if dataset already exists (enables image reuse across split changes)
_reuse = dataset_dir.exists()
if _reuse:
yaml_path.unlink(missing_ok=True) # Invalidate hash before destructive ops (crash safety)
if not is_classification:
shutil.rmtree(dataset_dir / "labels", ignore_errors=True)
dataset_dir.mkdir(parents=True, exist_ok=True)
data_yaml = None
if not is_classification:
# Detection/segmentation/pose/obb: prepare YAML and create base structure
data_yaml = dict(dataset_record)
data_yaml["names"] = class_names
data_yaml.pop("class_names", None)
data_yaml.pop("type", None) # Remove NDJSON-specific fields
for split in sorted(splits):
(dataset_dir / "images" / split).mkdir(parents=True, exist_ok=True)
(dataset_dir / "labels" / split).mkdir(parents=True, exist_ok=True)
data_yaml[split] = f"images/{split}"
async def process_record(session, semaphore, record):
"""Process single image record with async session."""
async with semaphore:
split, original_name = record["split"], record["file"]
annotations = record.get("annotations", {})
if is_classification:
# Classification: place image in {split}/{class_name}/ folder
class_ids = annotations.get("classification", [])
class_id = class_ids[0] if class_ids else 0
class_name = class_names.get(class_id, str(class_id))
image_path = dataset_dir / split / class_name / original_name
else:
# Detection: write label file and place image in images/{split}/
image_path = dataset_dir / "images" / split / original_name
label_path = dataset_dir / "labels" / split / f"{Path(original_name).stem}.txt"
lines_to_write = []
for key in annotations.keys():
lines_to_write = [" ".join(map(str, item)) for item in annotations[key]]
break
label_path.write_text("\n".join(lines_to_write) + "\n" if lines_to_write else "")
# Reuse existing image from another split dir (avoids redownload on resplit) or download
if not image_path.exists():
if _reuse:
for s in ("train", "val", "test"):
if s == split:
continue
candidate = (
(dataset_dir / s / class_name / original_name)
if is_classification
else (dataset_dir / "images" / s / original_name)
)
if candidate.exists():
image_path.parent.mkdir(parents=True, exist_ok=True)
candidate.rename(image_path)
break
if not image_path.exists() and (http_url := record.get("url")):
image_path.parent.mkdir(parents=True, exist_ok=True)
# Retry with exponential backoff (3 attempts: 0s, 2s, 4s delays)
for attempt in range(3):
try:
async with session.get(http_url, timeout=aiohttp.ClientTimeout(total=30)) as response:
response.raise_for_status()
image_path.write_bytes(await response.read())
return True
except Exception as e:
if attempt < 2: # Don't sleep after last attempt
await asyncio.sleep(2**attempt) # 1s, 2s backoff
else:
LOGGER.warning(f"Failed to download {http_url} after 3 attempts: {e}")
return False
return True
# Process all images with async downloads (limit connections for small datasets)
semaphore = asyncio.Semaphore(min(128, len(image_records)))
async with aiohttp.ClientSession() as session:
pbar = TQDM(
total=len(image_records),
desc=f"Converting {ndjson_path.name} → {dataset_dir} ({len(image_records)} images)",
)
async def tracked_process(record):
result = await process_record(session, semaphore, record)
pbar.update(1)
return result
results = await asyncio.gather(*[tracked_process(record) for record in image_records])
pbar.close()
# Validate images were downloaded successfully
success_count = sum(1 for r in results if r)
if success_count == 0:
raise RuntimeError(f"Failed to download any images from {ndjson_path}. Check network connection and URLs.")
if success_count < len(image_records):
LOGGER.warning(f"Downloaded {success_count}/{len(image_records)} images from {ndjson_path}")
# Remove orphaned images no longer in the dataset (prevents stale background images in training)
if _reuse:
expected_paths = set()
for r in image_records:
s, name = r["split"], r["file"]
if is_classification:
ann = r.get("annotations", {})
cids = ann.get("classification", [])
cid = cids[0] if cids else 0
expected_paths.add(dataset_dir / s / class_names.get(cid, str(cid)) / name)
else:
expected_paths.add(dataset_dir / "images" / s / name)
img_root = dataset_dir if is_classification else (dataset_dir / "images")
for p in img_root.rglob("*"):
if p.is_file() and p not in expected_paths:
p.unlink()
if is_classification:
# Classification: return dataset directory (check_cls_dataset expects a directory path)
return dataset_dir
else:
# Detection: write data.yaml with hash for future change detection
data_yaml["hash"] = _hash
YAML.save(yaml_path, data_yaml)
return yaml_path

1865
ultralytics/data/dataset.py Executable file
View File

File diff suppressed because it is too large Load Diff

731
ultralytics/data/ground3d_augment.py Executable file
View File

@@ -0,0 +1,731 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
"""Ground 3D detection data utilities.
Functions for on-the-fly label parsing, calibration reading, and virtual camera augmentation
for joint 2D+3D ground detection training.
"""
import json
import os
from functools import lru_cache
from pathlib import Path
import cv2
import numpy as np
def parse_ground_3d_label_file(lb_file, class_map, difficulty_weights, face_3d_classes, complete_3d_classes, min_wh=2.0):
"""Parse a ground 3D label file on-the-fly. Returns 2D labels dict + 3D array.
Reads label files with variable column counts:
- 6-col: [class_name, x, y, w, h, difficulty] — 2D only
- 19-col: [class_name, x, y, w, h, ...3D(13cols)..., difficulty] — complete 3D (no face)
- 51-col: [class_name, x, y, w, h, ...3D+faces(45cols)..., difficulty] — face 3D
The 48-dim internal format per object:
[0]: class_id, [1-4]: 2D bbox (xywh norm)
[5-7]: 3D center (x3d, y3d, z3d), [8-10]: dims (l, h, w)
[11]: rot_y, [12-13]: 3D box center projection (xc, yc), [14]: alpha
[15-22]: front face, [23-30]: rear face, [31-38]: left face, [39-46]: right face
[47]: difficulty level
Args:
lb_file (str): Path to label file.
class_map (dict): Mapping from class names to class IDs.
difficulty_weights (list): Loss weights for difficulty levels [easy, normal, medium, hard].
face_3d_classes (set): Class IDs with 4-face annotations (51-col).
complete_3d_classes (set): Class IDs with whole-box 3D only (19-col).
min_wh (float): Minimum box width/height in normalized coords for filtering.
Returns:
lb_2d (dict): Dict with cls (n,1), bboxes (n,4), difficulties (n,1),
segments, keypoints, normalized, bbox_format.
lb_3d (np.ndarray): 3D portion shape (n, 42). Objects without 3D GT keep NaN values;
images without labels return an empty `(0, 42)` array.
"""
empty_2d = {
"cls": np.zeros((0, 1), dtype=np.float32),
"bboxes": np.zeros((0, 4), dtype=np.float32),
"difficulties": np.zeros((0, 1), dtype=np.float32),
"difficulty_levels": np.zeros((0, 1), dtype=np.int64),
"segments": [],
"keypoints": None,
"normalized": True,
"bbox_format": "xywh",
}
empty_3d = np.full((0, 42), np.nan, dtype=np.float32)
if not os.path.isfile(lb_file):
return empty_2d, empty_3d
# Read file once
with open(lb_file, encoding="utf-8") as f:
lines = f.read().strip().splitlines()
if not lines:
return empty_2d, empty_3d
labels = []
# Cache method lookups to avoid repeated attribute access
class_map_get = class_map.get
face_3d_classes_contains = face_3d_classes.__contains__
complete_3d_classes_contains = complete_3d_classes.__contains__
for line in lines:
parts = line.split() # split() without strip() is faster
if not parts:
continue
cls_name = parts[0]
cls_id = class_map_get(cls_name)
if cls_id is None:
continue # skip unknown classes
ncols = len(parts)
if ncols == 6:
# 2D only: [class, x, y, w, h, difficulty]
temp = np.full(48, np.nan, dtype=np.float32)
temp[0] = cls_id
temp[1:5] = [float(parts[i]) for i in range(1, 5)]
temp[47] = float(parts[5])
labels.append(temp)
elif ncols == 19 and complete_3d_classes_contains(cls_id):
# Complete 3D without face: 19 cols, difficulty at index 18
# Use indices 1-13 + 16 (skip 14,15 which are unused), map to temp[1:15]
temp = np.full(48, np.nan, dtype=np.float32)
temp[0] = cls_id
useful_indices = list(range(1, 14)) + [16]
temp[1:15] = [float(parts[i]) for i in useful_indices]
temp[47] = float(parts[18])
labels.append(temp)
elif ncols == 51 and face_3d_classes_contains(cls_id):
# Face 3D: 51 cols, difficulty at index 50
# [cls_id] + indices 1-13 + 16 → temp[0:15], then indices 18-49 → temp[15:47], difficulty → temp[47]
temp_list = [cls_id] + [float(parts[i]) for i in (list(range(1, 14)) + [16])]
temp_list.extend(float(parts[i]) for i in range(18, 50))
temp_list.append(float(parts[50])) # difficulty
labels.append(np.array(temp_list, dtype=np.float32))
elif ncols == 7:
# 2D with two difficulty columns: [class, x, y, w, h, diff1, diff2]
# In joint 2D&3D training, diff2 is truncation, so use only diff1 for difficulty supervision.
temp = np.full(48, np.nan, dtype=np.float32)
temp[0] = cls_id
temp[1:5] = [float(parts[i]) for i in range(1, 5)]
temp[47] = float(parts[5])
labels.append(temp)
else:
raise ValueError(f"Unexpected number of columns ({ncols}) in label file {lb_file}: '{line}'")
if not labels:
return empty_2d, empty_3d
lb = np.stack(labels, axis=0) # (n, 48)
nl = len(lb)
# Validate
if nl > 0:
assert lb.shape[1] == 48, f"labels require 48 columns, got {lb.shape[1]}"
# Remove duplicates
_, idx = np.unique(lb, axis=0, return_index=True)
if len(idx) < nl:
lb = lb[idx]
if len(lb) == 0:
return empty_2d, empty_3d
# Split into 2D and 3D portions
cls = lb[:, 0:1] # (n, 1)
bboxes = lb[:, 1:5] # (n, 4) xywh normalized
# Difficulty → loss weights
dw = difficulty_weights
raw_diff = lb[:, 47].astype(int).clip(0, len(dw) - 1)
difficulties = np.array([dw[d] for d in raw_diff], dtype=np.float32).reshape(-1, 1)
difficulty_levels = raw_diff.astype(np.int64).reshape(-1, 1)
# 3D portion: columns 5-46 (42 dims)
# [x3d, y3d, z3d, l, h, w, rot_y, xc, yc, alpha, front(8), rear(8), left(8), right(8)]
labels_3d = lb[:, 5:47] # (n, 42)
lb_2d = {
"cls": cls,
"bboxes": bboxes,
"difficulties": difficulties,
"difficulty_levels": difficulty_levels,
"segments": [],
"keypoints": None,
"normalized": True,
"bbox_format": "xywh",
}
return lb_2d, labels_3d
@lru_cache(maxsize=256)
def _read_label_root_camera4_cached(calib_path_str):
"""Cached helper to read a label-root clip-level camera4.json file."""
import math
with open(calib_path_str, encoding="utf-8") as f:
payload = json.load(f)
required = ("focal_u", "focal_v", "cu", "cv")
if any(key not in payload for key in required):
return None
calib = dict(payload)
calib["focal_u"] = float(payload["focal_u"])
calib["focal_v"] = float(payload["focal_v"])
calib["cu"] = float(payload["cu"])
calib["cv"] = float(payload["cv"])
calib["distort_coeffs"] = list(payload.get("distort_coeffs", []))
if "pitch" in payload:
calib["pitch"] = math.radians(float(payload["pitch"]))
for angle_key in ("roll", "yaw"):
if angle_key in payload:
calib[angle_key] = math.radians(float(payload[angle_key]))
return calib
def read_calib_from_path(img_path, image_root=None, extra_calib_candidates=None):
"""Read clip-level camera4.json from the label-root calibration folder.
Args:
img_path (str): Path to the image file.
image_root (str | Path | None): Unused compatibility arg kept for existing call sites.
extra_calib_candidates (Iterable[str | Path] | None): Label-root per-frame calibration candidates. The loader
resolves each candidate's sibling `L2_calib/camera4.json` and ignores all other layouts.
Returns:
dict | None: Calibration dict with keys: focal_u, focal_v, cu, cv, pitch (radians), distort_coeffs.
Returns None if calibration file not found.
"""
_ = img_path, image_root
for candidate in extra_calib_candidates or ():
candidate = Path(candidate).resolve()
camera4_path = candidate if candidate.name == "camera4.json" else candidate.parent / "L2_calib" / "camera4.json"
if camera4_path.exists():
return _read_label_root_camera4_cached(str(camera4_path))
return None
def compute_vanishing_point_x(raw_calib, ori_w):
"""Compute vanishing point X from calibration."""
if raw_calib is None:
return ori_w / 2
return raw_calib.get("cu", ori_w / 2)
def compute_vanishing_point_y(raw_calib, ori_h):
"""Compute vanishing point Y from calibration."""
if raw_calib is None:
return ori_h / 2
cv_orig = raw_calib.get("cv", ori_h / 2)
pitch = raw_calib.get("pitch", 0.0)
focal_v = raw_calib.get("focal_v", ori_h)
return cv_orig - focal_v * np.tan(pitch) if pitch != 0 else cv_orig
def compute_centered_roi_bounds(ori_w, ori_h, roi_w, roi_h, center_x, center_y):
"""Compute ROI bounds centered on the requested crop center."""
crop_x1 = int(max(0, min(center_x - roi_w / 2, ori_w - roi_w)))
crop_y1 = int(max(0, min(center_y - roi_h / 2, ori_h - roi_h)))
return crop_x1, crop_y1, crop_x1 + roi_w, crop_y1 + roi_h
def adjust_calib_for_roi_crop(raw_calib, ori_w, ori_h, crop_bounds=None):
"""Shift intrinsics into ROI crop coordinates before resize."""
crop_x1, crop_y1, crop_x2, crop_y2 = crop_bounds or (0, 0, ori_w, ori_h)
cu = raw_calib.get("cu", ori_w / 2) if raw_calib else ori_w / 2
cv = raw_calib.get("cv", ori_h / 2) if raw_calib else ori_h / 2
focal_u = raw_calib.get("focal_u", ori_w) if raw_calib else ori_w
focal_v = raw_calib.get("focal_v", ori_h) if raw_calib else ori_h
distort_coeffs = raw_calib.get("distort_coeffs", []) if raw_calib else []
return {
"focal_u": focal_u,
"focal_v": focal_v,
"cu": cu - crop_x1,
"cv": cv - crop_y1,
"src_w": crop_x2 - crop_x1,
"src_h": crop_y2 - crop_y1,
"distort_coeffs": distort_coeffs,
}
def build_final_resized_calib(focal_u, focal_v, cu, cv, src_w, src_h, target_w, target_h, virtual_fx, distort_coeffs=None):
"""Build final calibration after ROI crop and direct resize."""
scale_x = target_w / src_w
scale_y = target_h / src_h
fx_final = focal_u * scale_x
return {
"fx": fx_final,
"fy": focal_v * scale_y,
"cx": cu * scale_x,
"cy": cv * scale_y,
"distort_coeffs": distort_coeffs if distort_coeffs is not None else [],
"depth_scale": fx_final / virtual_fx,
}
def pack_labels_to_48(lb_2d, lb_3d):
"""Pack 2D and 3D labels into the internal augmentation representation."""
bboxes = lb_2d["bboxes"]
n = len(bboxes)
if n == 0:
return np.zeros((0, 49), dtype=np.float32)
labels_48 = np.full((n, 49), np.nan, dtype=np.float32)
labels_48[:, 0] = lb_2d["cls"].reshape(-1)
labels_48[:, 1:5] = bboxes
labels_48[:, 47] = lb_2d["difficulties"].reshape(-1)
labels_48[:, 48] = lb_2d.get("difficulty_levels", np.zeros((n, 1), dtype=np.int64)).reshape(-1)
if lb_3d is not None and len(lb_3d):
labels_48[:, 5:47] = lb_3d
return labels_48
def unpack_labels_from_48(labels_48):
"""Unpack the internal 48-dim representation into 2D and 3D labels."""
lb_2d = {
"cls": np.zeros((0, 1), dtype=np.float32),
"bboxes": np.zeros((0, 4), dtype=np.float32),
"difficulties": np.zeros((0, 1), dtype=np.float32),
"difficulty_levels": np.zeros((0, 1), dtype=np.int64),
"segments": [],
"keypoints": None,
"normalized": True,
"bbox_format": "xywh",
}
if len(labels_48) == 0:
return lb_2d, None
lb_2d["cls"] = labels_48[:, 0:1]
lb_2d["bboxes"] = labels_48[:, 1:5]
lb_2d["difficulties"] = labels_48[:, 47:48]
lb_2d["difficulty_levels"] = labels_48[:, 48:49].astype(np.int64) if labels_48.shape[1] > 48 else np.zeros((len(labels_48), 1), dtype=np.int64)
return lb_2d, labels_48[:, 5:47]
def _handle_cut_labels_42(labels, outside_mask, still_inside_mask):
"""Handle cut-in/cut-out updates for ROI-remapped 42-dim labels."""
if len(labels) == 0:
return
partial_mask = ~(still_inside_mask | outside_mask)
if not np.any(partial_mask):
return
rot_y = labels[partial_mask, 6]
is_cut_in = (rot_y >= -np.pi) & (rot_y <= 0)
partial_indices = np.where(partial_mask)[0]
def _invalidate_face(face_indices, face_offset):
labels[np.ix_(face_indices, np.arange(face_offset, face_offset + 6))] = -1
labels[face_indices, face_offset + 6] = 0
labels[face_indices, face_offset + 7] = 0
cut_in_idx = partial_indices[is_cut_in]
if len(cut_in_idx):
for face_offset in (18, 26, 34):
_invalidate_face(cut_in_idx, face_offset)
labels[cut_in_idx, 16] = 1
labels[cut_in_idx, 17] = 1
cut_out_idx = partial_indices[~is_cut_in]
if len(cut_out_idx):
for face_offset in (10, 26, 34):
_invalidate_face(cut_out_idx, face_offset)
labels[cut_out_idx, 24] = 1
labels[cut_out_idx, 25] = 1
def remap_labels_to_roi(lb_2d, lb_3d, ori_w, ori_h, crop_bounds):
"""Shift boxes and UV coordinates from original image space into ROI-normalized space."""
bboxes = lb_2d["bboxes"]
if len(bboxes) == 0:
return lb_2d, lb_3d
crop_x1, crop_y1, crop_x2, crop_y2 = crop_bounds
roi_width = crop_x2 - crop_x1
roi_height = crop_y2 - crop_y1
bboxes = bboxes.copy()
x1 = (bboxes[:, 0] - bboxes[:, 2] / 2) * ori_w
y1 = (bboxes[:, 1] - bboxes[:, 3] / 2) * ori_h
x2 = (bboxes[:, 0] + bboxes[:, 2] / 2) * ori_w
y2 = (bboxes[:, 1] + bboxes[:, 3] / 2) * ori_h
x1_roi = x1 - crop_x1
y1_roi = y1 - crop_y1
x2_roi = x2 - crop_x1
y2_roi = y2 - crop_y1
still_inside = (x1_roi >= 0) & (y1_roi >= 0) & (x2_roi < roi_width) & (y2_roi < roi_height)
outside = (
((x1_roi < 0) & (x2_roi < 0))
| ((x1_roi >= roi_width) & (x2_roi >= roi_width))
| ((y1_roi < 0) & (y2_roi < 0))
| ((y1_roi >= roi_height) & (y2_roi >= roi_height))
)
if lb_3d is not None and len(lb_3d) > 0:
lb_3d = lb_3d.copy()
_handle_cut_labels_42(lb_3d, outside, still_inside)
x1_roi = np.clip(x1_roi, 0, roi_width - 1)
y1_roi = np.clip(y1_roi, 0, roi_height - 1)
x2_roi = np.clip(x2_roi, 0, roi_width - 1)
y2_roi = np.clip(y2_roi, 0, roi_height - 1)
bboxes[:, 0] = (x1_roi + x2_roi) * 0.5 / roi_width
bboxes[:, 1] = (y1_roi + y2_roi) * 0.5 / roi_height
bboxes[:, 2] = (x2_roi - x1_roi) / roi_width
bboxes[:, 3] = (y2_roi - y1_roi) / roi_height
keep = ~outside
lb_2d = {
**lb_2d,
"bboxes": bboxes[keep],
"cls": lb_2d["cls"][keep],
"difficulties": lb_2d["difficulties"][keep],
"difficulty_levels": lb_2d["difficulty_levels"][keep],
}
if lb_3d is not None and len(lb_3d) > 0:
lb_3d = lb_3d[keep]
for xi, yi in [(7, 8), (14, 15), (22, 23), (30, 31), (38, 39)]:
valid = ~np.isnan(lb_3d[:, xi]) & (lb_3d[:, xi] != -1)
if np.any(valid):
lb_3d[valid, xi] = (lb_3d[valid, xi] * ori_w - crop_x1) / roi_width
lb_3d[valid, yi] = (lb_3d[valid, yi] * ori_h - crop_y1) / roi_height
return lb_2d, lb_3d
def normalize_roi_depth(lb_3d, fx_final, virtual_fx):
"""Normalize ROI z3d targets to the canonical virtual focal length."""
if lb_3d is None or len(lb_3d) == 0:
return lb_3d
lb_3d = lb_3d.copy()
z3d_scale = virtual_fx / fx_final
mask = ~np.isnan(lb_3d[:, 2]) & (lb_3d[:, 2] > 0)
lb_3d[mask, 2] *= z3d_scale
for col in [12, 20, 28, 36]:
mask = ~np.isnan(lb_3d[:, col]) & (lb_3d[:, col] != -1.0) & (lb_3d[:, col] > 0)
lb_3d[mask, col] *= z3d_scale
return lb_3d
def compute_simul_calib(calib_params, ori_img_size, target_size, crop_center_x, crop_center_y, target_fx, augment=False):
"""Compute virtual camera calibration parameters from original fisheye calibration.
Uses OpenCV to compute optimal new camera matrix after undistortion (no black margins),
then crops a region while maintaining target aspect ratio.
Port from yolov5-3d/utils/dataloaders3d_ground.py:698-824.
Args:
calib_params (dict): Original calibration with focal_u, focal_v, cu, cv, distort_coeffs.
ori_img_size (tuple): Original image size (width, height).
target_size (tuple): Target size (width, height) e.g., (960, 480).
crop_center_x (float): Crop center X in distorted image (typically image center).
crop_center_y (float): Crop center Y in distorted image (typically vanishing point Y).
target_fx (float): Target focal length x for virtual camera.
augment (bool): If True, randomly select crop size between min and max.
Returns:
dict: Virtual camera calibration with fx, fy, cx, cy, crop_bounds, scale, K_undistorted, fx_to_target_scale.
"""
import math
import random
fx_orig = calib_params["focal_u"]
fy_orig = calib_params["focal_v"]
cx_orig = calib_params["cu"]
cy_orig = calib_params["cv"]
distort_coeffs = calib_params.get("distort_coeffs", [])
ori_w, ori_h = ori_img_size
target_w, target_h = target_size
K_orig = np.array([[fx_orig, 0, cx_orig], [0, fy_orig, cy_orig], [0, 0, 1]], dtype=np.float64)
D = np.array(distort_coeffs[:4], dtype=np.float64) if len(distort_coeffs) >= 4 else np.zeros(4, dtype=np.float64)
# Optimal new camera matrix (no black margins)
K_new = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(K_orig, D, (ori_w, ori_h), np.eye(3), balance=0.0)
fx_undist = K_new[0, 0]
fy_undist = K_new[1, 1]
cx_undist = K_new[0, 2]
cy_undist = K_new[1, 2]
# Undistort crop center point
dist_point = np.array([[[crop_center_x, crop_center_y]]], dtype=np.float32)
undist_point = cv2.fisheye.undistortPoints(dist_point, K_orig, D, P=K_new)
cx_undist_crop = undist_point[0, 0, 0]
cy_undist_crop = undist_point[0, 0, 1]
# Max crop dimensions centered on undistorted crop center
max_w = min(cx_undist_crop * 2, (ori_w - cx_undist_crop) * 2)
max_h = min(cy_undist_crop * 2, (ori_h - cy_undist_crop) * 2)
# GCD approach for exact aspect ratio with integer coordinates
gcd = math.gcd(target_w, target_h)
ratio_w = target_w // gcd
ratio_h = target_h // gcd
k_from_w = int(max_w / ratio_w)
k_from_h = int(max_h / ratio_h)
k_max = min(k_from_w, k_from_h)
k_min = max(target_w // ratio_w, target_h // ratio_h)
if augment and k_max > k_min:
k = random.randint(k_min, k_max)
else:
k = k_max
crop_w = k * ratio_w
crop_h = k * ratio_h
crop_x1 = int(cx_undist_crop - crop_w / 2)
crop_y1 = int(cy_undist_crop - crop_h / 2)
crop_x2 = crop_x1 + crop_w
crop_y2 = crop_y1 + crop_h
scale_x = target_w / crop_w
scale_y = target_h / crop_h
scaled_fx = scale_x * fx_undist
fx_to_target_scale = target_fx / scaled_fx
return {
"fx": scaled_fx,
"fy": scale_y * fy_undist,
"cx": (cx_undist - crop_x1) * scale_x,
"cy": (cy_undist - crop_y1) * scale_y,
"distort_coeffs": [],
"depth_scale": scaled_fx / target_fx,
"crop_bounds": (crop_x1, crop_y1, crop_x2, crop_y2),
"scale": (scale_x, scale_y),
"K_undistorted": K_new,
"K_orig": K_orig,
"D": D,
"fx_to_target_scale": fx_to_target_scale,
}
def apply_simul_transform(img, labels_48, simul_calib, calib_params, target_size, augment=False):
"""Apply fisheye undistortion + crop + resize to image and 48-dim labels.
Port from yolov5-3d/utils/dataloaders3d_ground.py:826-1039.
Args:
img (np.ndarray): Input image (H, W, 3) BGR — distorted fisheye image.
labels_48 (np.ndarray): Label array (N, 48) in 48-dim format.
simul_calib (dict): Pre-computed virtual camera calibration from compute_simul_calib().
calib_params (dict): Original calibration dict.
target_size (tuple): Target size (width, height).
augment (bool): If True, use random interpolation for resize.
Returns:
img_transformed (np.ndarray): Transformed image.
labels_transformed (np.ndarray): Transformed labels (M, 48), M <= N.
"""
import random
h_orig, w_orig = img.shape[:2]
target_w, target_h = target_size
K_orig = simul_calib["K_orig"]
D = simul_calib["D"]
K_new = simul_calib["K_undistorted"]
# Step 1: Undistort full image
img_undistorted = cv2.fisheye.undistortImage(img, K_orig, D, Knew=K_new)
# Step 2: Crop
crop_x1, crop_y1, crop_x2, crop_y2 = simul_calib["crop_bounds"]
img_cropped = img_undistorted[crop_y1:crop_y2, crop_x1:crop_x2]
# Step 3: Resize
if augment:
interp = random.choice([cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_AREA, cv2.INTER_CUBIC, cv2.INTER_LANCZOS4])
else:
interp = cv2.INTER_LINEAR
img_transformed = cv2.resize(img_cropped, target_size, interpolation=interp)
# Step 4: Transform labels
labels_transformed = labels_48.copy()
if len(labels_transformed) == 0:
return img_transformed, labels_transformed
scale_x, scale_y = simul_calib["scale"]
fx_to_target_scale = simul_calib["fx_to_target_scale"]
# Collect all 2D points to undistort in batch
all_points = []
point_map = [] # (label_idx, field_type, col_indices)
for i in range(len(labels_transformed)):
# Bbox corners from xywhn
xc = labels_transformed[i, 1] * w_orig
yc = labels_transformed[i, 2] * h_orig
bw = labels_transformed[i, 3] * w_orig
bh = labels_transformed[i, 4] * h_orig
x1, y1 = xc - bw / 2, yc - bh / 2
x2, y2 = xc + bw / 2, yc + bh / 2
all_points.extend([[x1, y1], [x2, y2]])
point_map.extend([(i, "bbox_tl"), (i, "bbox_br")])
# Whole box UV (dims 12-13 in 48-dim = xc, yc normalized)
if not np.isnan(labels_transformed[i, 12]):
ux = labels_transformed[i, 12] * w_orig
uy = labels_transformed[i, 13] * h_orig
all_points.append([ux, uy])
point_map.append((i, "whole_uv"))
# Face UVs: front(19,20), rear(27,28), left(35,36), right(43,44) in 48-dim
for face_name, uv_cols in [("front", (19, 20)), ("rear", (27, 28)), ("left", (35, 36)), ("right", (43, 44))]:
if not np.isnan(labels_transformed[i, uv_cols[0]]) and labels_transformed[i, uv_cols[0]] != -1:
fu = labels_transformed[i, uv_cols[0]] * w_orig
fv = labels_transformed[i, uv_cols[1]] * h_orig
all_points.append([fu, fv])
point_map.append((i, f"{face_name}_uv"))
if not all_points:
return img_transformed, labels_transformed
# Batch undistort all points
pts_dist = np.array(all_points, dtype=np.float32).reshape(-1, 1, 2)
pts_undist = cv2.fisheye.undistortPoints(pts_dist, K_orig, D, P=K_new).reshape(-1, 2)
# Apply crop + resize to undistorted points
pts_transformed = np.zeros_like(pts_undist)
pts_transformed[:, 0] = (pts_undist[:, 0] - crop_x1) * scale_x
pts_transformed[:, 1] = (pts_undist[:, 1] - crop_y1) * scale_y
# Write back transformed coordinates
outside = np.zeros(len(labels_transformed), dtype=bool)
still_inside = np.ones(len(labels_transformed), dtype=bool)
for idx, (label_i, field) in enumerate(point_map):
px, py = pts_transformed[idx]
if field == "bbox_tl":
labels_transformed[label_i, 1] = px # temp store x1
elif field == "bbox_br":
x1_t = labels_transformed[label_i, 1]
# Get the tl point from previous entry
tl_idx = idx - 1
y1_t = pts_transformed[tl_idx, 1]
# Check if fully inside before clipping
if x1_t < 0 or y1_t < 0 or px > target_w or py > target_h:
still_inside[label_i] = False
# Clip to image bounds
x1_c = np.clip(x1_t, 0, target_w)
x2_c = np.clip(px, 0, target_w)
y1_c = np.clip(y1_t, 0, target_h)
y2_c = np.clip(py, 0, target_h)
bw_new = x2_c - x1_c
bh_new = y2_c - y1_c
if bw_new <= 0 or bh_new <= 0:
outside[label_i] = True
continue
# Convert back to xywhn
labels_transformed[label_i, 1] = (x1_c + x2_c) / 2 / target_w
labels_transformed[label_i, 2] = (y1_c + y2_c) / 2 / target_h
labels_transformed[label_i, 3] = bw_new / target_w
labels_transformed[label_i, 4] = bh_new / target_h
elif field == "whole_uv":
labels_transformed[label_i, 12] = px / target_w
labels_transformed[label_i, 13] = py / target_h
elif field.endswith("_uv"):
face = field.split("_")[0]
uv_map = {"front": (19, 20), "rear": (27, 28), "left": (35, 36), "right": (43, 44)}
cols = uv_map[face]
labels_transformed[label_i, cols[0]] = px / target_w
labels_transformed[label_i, cols[1]] = py / target_h
# Scale z3d by fx_to_target_scale
labels_transformed = _scale_z3d(labels_transformed, fx_to_target_scale)
# Handle partial visibility (cut-in/cut-out)
_handle_cut_labels(labels_transformed, outside, still_inside)
# Remove outside boxes
labels_transformed = labels_transformed[~outside]
return img_transformed, labels_transformed
def _scale_z3d(labels, scale):
"""Scale z3d coordinates for depth normalization.
Port from yolov5-3d/utils/dataloaders3d_ground.py:1041-1080.
"""
if len(labels) == 0 or scale == 1.0:
return labels
labels_scaled = labels.copy()
# Whole z3d (dim 7 in 48-dim)
labels_scaled[:, 7] *= scale
# Face z3d: front(17), rear(25), left(33), right(41)
# Note: only scale if not NaN and not -1 (invalid), to preserve missing/invalid indicators
for col in [17, 25, 33, 41]:
mask = ~np.isnan(labels_scaled[:, col]) & (labels_scaled[:, col] != -1.0)
labels_scaled[mask, col] *= scale
return labels_scaled
def _handle_cut_labels(labels, outside_mask, still_inside_mask):
"""Handle partial visibility for objects partially outside the image.
For objects with bbox partially outside, mark as cut-in or cut-out based on rotation angle.
Cut-in (approaching, rot_y in [-pi, 0]): keep front face only.
Cut-out (leaving, rot_y > 0): keep rear face only.
Port from yolov5-3d/utils/dataloaders3d_ground.py:1000-1037.
Args:
labels (np.ndarray): Label array (N, 48) in 48-dim format.
outside_mask (np.ndarray): Boolean mask (N,) — True for fully outside boxes.
still_inside_mask (np.ndarray): Boolean mask (N,) — True for fully inside boxes.
"""
if len(labels) == 0:
return
partial_mask = ~(still_inside_mask | outside_mask)
if not np.any(partial_mask):
return
rot_y = labels[partial_mask, 11] # dim11: rot_y in 48-dim
is_cut_in = (rot_y >= -np.pi) & (rot_y <= 0)
partial_indices = np.where(partial_mask)[0]
def _invalidate_face(face_indices, face_offset):
labels[np.ix_(face_indices, np.arange(face_offset, face_offset + 6))] = -1
labels[face_indices, face_offset + 6] = 0
labels[face_indices, face_offset + 7] = 0
# Cut-in: keep front face, invalidate others
cut_in_idx = partial_indices[is_cut_in]
if len(cut_in_idx):
for face_offset in (23, 31, 39):
_invalidate_face(cut_in_idx, face_offset)
labels[cut_in_idx, 21] = 1 # front face score = 1
labels[cut_in_idx, 22] = 1
# Cut-out: keep rear face, invalidate others
cut_out_idx = partial_indices[~is_cut_in]
if len(cut_out_idx):
for face_offset in (15, 31, 39):
_invalidate_face(cut_out_idx, face_offset)
labels[cut_out_idx, 29] = 1 # rear face score = 1
labels[cut_out_idx, 30] = 1

695
ultralytics/data/loaders.py Executable file
View File

@@ -0,0 +1,695 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import glob
import math
import os
import time
import urllib
from dataclasses import dataclass
from pathlib import Path
from threading import Thread
from typing import Any
import cv2
import numpy as np
import torch
from PIL import Image, ImageOps
from ultralytics.data.utils import FORMATS_HELP_MSG, IMG_FORMATS, VID_FORMATS
from ultralytics.utils import IS_COLAB, IS_KAGGLE, LOGGER, ops
from ultralytics.utils.checks import check_requirements
from ultralytics.utils.patches import imread
@dataclass
class SourceTypes:
"""Class to represent various types of input sources for predictions.
This class uses dataclass to define boolean flags for different types of input sources that can be used for making
predictions with YOLO models.
Attributes:
stream (bool): Flag indicating if the input source is a video stream.
screenshot (bool): Flag indicating if the input source is a screenshot.
from_img (bool): Flag indicating if the input source is an image file.
tensor (bool): Flag indicating if the input source is a tensor.
Examples:
>>> source_types = SourceTypes(stream=True, screenshot=False, from_img=False)
>>> print(source_types.stream)
True
>>> print(source_types.from_img)
False
"""
stream: bool = False
screenshot: bool = False
from_img: bool = False
tensor: bool = False
class LoadStreams:
"""Stream Loader for various types of video streams.
Supports RTSP, RTMP, HTTP, and TCP streams. This class handles the loading and processing of multiple video streams
simultaneously, making it suitable for real-time video analysis tasks.
Attributes:
sources (list[str]): The source input paths or URLs for the video streams.
vid_stride (int): Video frame-rate stride.
buffer (bool): Whether to buffer input streams.
running (bool): Flag to indicate if the streaming thread is running.
mode (str): Set to 'stream' indicating real-time capture.
imgs (list[list[np.ndarray]]): List of image frames for each stream.
fps (list[float]): List of FPS for each stream.
frames (list[int]): List of total frames for each stream.
threads (list[Thread]): List of threads for each stream.
shape (list[tuple[int, int, int]]): List of shapes for each stream.
caps (list[cv2.VideoCapture]): List of cv2.VideoCapture objects for each stream.
bs (int): Batch size for processing.
cv2_flag (int): OpenCV flag for image reading (grayscale or color/BGR).
Methods:
update: Read stream frames in daemon thread.
close: Close stream loader and release resources.
__iter__: Returns an iterator object for the class.
__next__: Returns source paths, transformed, and original images for processing.
__len__: Return the length of the sources object.
Examples:
>>> stream_loader = LoadStreams("rtsp://example.com/stream1.mp4")
>>> for sources, imgs, _ in stream_loader:
... # Process the images
... pass
>>> stream_loader.close()
Notes:
- The class uses threading to efficiently load frames from multiple streams simultaneously.
- It automatically handles YouTube links, converting them to the best available stream URL.
- The class implements a buffer system to manage frame storage and retrieval.
"""
def __init__(self, sources: str = "file.streams", vid_stride: int = 1, buffer: bool = False, channels: int = 3):
"""Initialize stream loader for multiple video sources, supporting various stream types.
Args:
sources (str): Path to streams file or single stream URL.
vid_stride (int): Video frame-rate stride.
buffer (bool): Whether to buffer input streams.
channels (int): Number of image channels (1 for grayscale, 3 for color).
"""
torch.backends.cudnn.benchmark = True # faster for fixed-size inference
self.buffer = buffer # buffer input streams
self.running = True # running flag for Thread
self.mode = "stream"
self.vid_stride = vid_stride # video frame-rate stride
self.cv2_flag = cv2.IMREAD_GRAYSCALE if channels == 1 else cv2.IMREAD_COLOR # grayscale or color (BGR)
sources = Path(sources).read_text().rsplit() if os.path.isfile(sources) else [sources]
n = len(sources)
self.bs = n
self.fps = [0] * n # frames per second
self.frames = [0] * n
self.threads = [None] * n
self.caps = [None] * n # video capture objects
self.imgs = [[] for _ in range(n)] # images
self.shape = [[] for _ in range(n)] # image shapes
self.sources = [ops.clean_str(x).replace(os.sep, "_") for x in sources] # clean source names for later
for i, s in enumerate(sources): # index, source
# Start thread to read frames from video stream
st = f"{i + 1}/{n}: {s}... "
if urllib.parse.urlparse(s).hostname in {"www.youtube.com", "youtube.com", "youtu.be"}: # YouTube video
# YouTube format i.e. 'https://www.youtube.com/watch?v=Jsn8D3aC840' or 'https://youtu.be/Jsn8D3aC840'
s = get_best_youtube_url(s)
s = int(s) if s.isnumeric() else s # i.e. s = '0' local webcam
if s == 0 and (IS_COLAB or IS_KAGGLE):
raise NotImplementedError(
"'source=0' webcam not supported in Colab and Kaggle notebooks. "
"Try running 'source=0' in a local environment."
)
self.caps[i] = cv2.VideoCapture(s) # store video capture object
if not self.caps[i].isOpened():
raise ConnectionError(f"{st}Failed to open {s}")
w = int(self.caps[i].get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(self.caps[i].get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = self.caps[i].get(cv2.CAP_PROP_FPS) # warning: may return 0 or nan
self.frames[i] = max(int(self.caps[i].get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float(
"inf"
) # infinite stream fallback
self.fps[i] = max((fps if math.isfinite(fps) else 0) % 100, 0) or 30 # 30 FPS fallback
success, im = self.caps[i].read() # guarantee first frame
im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)[..., None] if self.cv2_flag == cv2.IMREAD_GRAYSCALE else im
if not success or im is None:
raise ConnectionError(f"{st}Failed to read images from {s}")
self.imgs[i].append(im)
self.shape[i] = im.shape
self.threads[i] = Thread(target=self.update, args=([i, self.caps[i], s]), daemon=True)
LOGGER.info(f"{st}Success ✅ ({self.frames[i]} frames of shape {w}x{h} at {self.fps[i]:.2f} FPS)")
self.threads[i].start()
LOGGER.info("") # newline
def update(self, i: int, cap: cv2.VideoCapture, stream: str):
"""Read stream frames in daemon thread and update image buffer."""
n, f = 0, self.frames[i] # frame number, frame array
while self.running and cap.isOpened() and n < (f - 1):
if len(self.imgs[i]) < 30: # keep a <=30-image buffer
n += 1
cap.grab() # .read() = .grab() followed by .retrieve()
if n % self.vid_stride == 0:
success, im = cap.retrieve()
im = (
cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)[..., None] if self.cv2_flag == cv2.IMREAD_GRAYSCALE else im
)
if not success:
im = np.zeros(self.shape[i], dtype=np.uint8)
LOGGER.warning("Video stream unresponsive, please check your IP camera connection.")
cap.open(stream) # re-open stream if signal was lost
if self.buffer:
self.imgs[i].append(im)
else:
self.imgs[i] = [im]
else:
time.sleep(0.01) # wait until the buffer is empty
def close(self):
"""Terminate stream loader, stop threads, and release video capture resources."""
self.running = False # stop flag for Thread
for thread in self.threads:
if thread.is_alive():
thread.join(timeout=5) # Add timeout
for cap in self.caps: # Iterate through the stored VideoCapture objects
try:
cap.release() # release video capture
except Exception as e:
LOGGER.warning(f"Could not release VideoCapture object: {e}")
def __iter__(self):
"""Return an iterator object and reset the frame counter."""
self.count = -1
return self
def __next__(self) -> tuple[list[str], list[np.ndarray], list[str]]:
"""Return the next batch of frames from multiple video streams for processing."""
self.count += 1
images = []
for i, x in enumerate(self.imgs):
# Wait until a frame is available in each buffer
while not x:
if not self.threads[i].is_alive():
self.close()
raise StopIteration
time.sleep(1 / min(self.fps))
x = self.imgs[i]
if not x:
LOGGER.warning(f"Waiting for stream {i}")
# Get and remove the first frame from imgs buffer
if self.buffer:
images.append(x.pop(0))
# Get the last frame, and clear the rest from the imgs buffer
else:
images.append(x.pop(-1) if x else np.zeros(self.shape[i], dtype=np.uint8))
x.clear()
return self.sources, images, [""] * self.bs
def __len__(self) -> int:
"""Return the number of video streams in the LoadStreams object."""
return self.bs # 1E12 frames = 32 streams at 30 FPS for 30 years
class LoadScreenshots:
"""Ultralytics screenshot dataloader for capturing and processing screen images.
This class manages the loading of screenshot images for processing with YOLO. It is suitable for use with `yolo
predict source=screen`.
Attributes:
screen (int): The screen number to capture.
left (int): The left coordinate for screen capture area.
top (int): The top coordinate for screen capture area.
width (int): The width of the screen capture area.
height (int): The height of the screen capture area.
mode (str): Set to 'stream' indicating real-time capture.
frame (int): Counter for captured frames.
sct (mss.mss): Screen capture object from `mss` library.
bs (int): Batch size, set to 1.
fps (int): Frames per second, set to 30.
monitor (dict[str, int]): Monitor configuration details.
cv2_flag (int): OpenCV flag for image reading (grayscale or color/BGR).
Methods:
__iter__: Returns an iterator object.
__next__: Captures the next screenshot and returns it.
Examples:
>>> loader = LoadScreenshots("0 100 100 640 480") # screen 0, top-left (100,100), 640x480
>>> for sources, imgs, info in loader:
... print(f"Captured frame: {imgs[0].shape}")
"""
def __init__(self, source: str, channels: int = 3):
"""Initialize screenshot capture with specified screen and region parameters.
Args:
source (str): Screen capture source string in format "screen_num left top width height".
channels (int): Number of image channels (1 for grayscale, 3 for color).
"""
check_requirements("mss")
import mss
source, *params = source.split()
self.screen, left, top, width, height = 0, None, None, None, None # default to full screen 0
if len(params) == 1:
self.screen = int(params[0])
elif len(params) == 4:
left, top, width, height = (int(x) for x in params)
elif len(params) == 5:
self.screen, left, top, width, height = (int(x) for x in params)
self.mode = "stream"
self.frame = 0
self.sct = mss.mss()
self.bs = 1
self.fps = 30
self.cv2_flag = cv2.IMREAD_GRAYSCALE if channels == 1 else cv2.IMREAD_COLOR # grayscale or color (BGR)
# Parse monitor shape
monitor = self.sct.monitors[self.screen]
self.top = monitor["top"] if top is None else (monitor["top"] + top)
self.left = monitor["left"] if left is None else (monitor["left"] + left)
self.width = width or monitor["width"]
self.height = height or monitor["height"]
self.monitor = {"left": self.left, "top": self.top, "width": self.width, "height": self.height}
def __iter__(self):
"""Return an iterator object for the screenshot capture."""
return self
def __next__(self) -> tuple[list[str], list[np.ndarray], list[str]]:
"""Capture and return the next screenshot as a numpy array using the mss library."""
im0 = np.asarray(self.sct.grab(self.monitor))[:, :, :3] # BGRA to BGR
im0 = cv2.cvtColor(im0, cv2.COLOR_BGR2GRAY)[..., None] if self.cv2_flag == cv2.IMREAD_GRAYSCALE else im0
s = f"screen {self.screen} (LTWH): {self.left},{self.top},{self.width},{self.height}: "
self.frame += 1
return [str(self.screen)], [im0], [s] # screen, img, string
class LoadImagesAndVideos:
"""A class for loading and processing images and videos for YOLO object detection.
This class manages the loading and pre-processing of image and video data from various sources, including single
image files, video files, and lists of image and video paths.
Attributes:
files (list[str]): List of image and video file paths.
nf (int): Total number of files (images and videos).
video_flag (list[bool]): Flags indicating whether a file is a video (True) or an image (False).
mode (str): Current mode, 'image' or 'video'.
vid_stride (int): Stride for video frame-rate.
bs (int): Batch size.
cap (cv2.VideoCapture): Video capture object for OpenCV.
frame (int): Frame counter for video.
frames (int): Total number of frames in the video.
count (int): Counter for iteration, initialized at 0 during __iter__().
ni (int): Number of images.
cv2_flag (int): OpenCV flag for image reading (grayscale or color/BGR).
Methods:
__init__: Initialize the LoadImagesAndVideos object.
__iter__: Returns an iterator object for VideoStream or ImageFolder.
__next__: Returns the next batch of images or video frames along with their paths and metadata.
_new_video: Creates a new video capture object for the given path.
__len__: Returns the number of batches in the object.
Examples:
>>> loader = LoadImagesAndVideos("path/to/data", batch=32, vid_stride=1)
>>> for paths, imgs, info in loader:
... # Process batch of images or video frames
... pass
Notes:
- Supports various image formats including HEIC.
- Handles both local files and directories.
- Can read from a text file containing paths to images and videos.
"""
def __init__(self, path: str | Path | list, batch: int = 1, vid_stride: int = 1, channels: int = 3):
"""Initialize dataloader for images and videos, supporting various input formats.
Args:
path (str | Path | list): Path to images/videos, directory, or list of paths.
batch (int): Batch size for processing.
vid_stride (int): Video frame-rate stride.
channels (int): Number of image channels (1 for grayscale, 3 for color).
"""
parent = None
if isinstance(path, str) and Path(path).suffix in {".txt", ".csv"}: # txt/csv file with source paths
parent, content = Path(path).parent, Path(path).read_text()
path = content.splitlines() if Path(path).suffix == ".txt" else content.split(",") # list of sources
path = [p.strip() for p in path]
files = []
for p in sorted(path) if isinstance(path, (list, tuple)) else [path]:
a = str(Path(p).absolute()) # do not use .resolve() https://github.com/ultralytics/ultralytics/issues/2912
if "*" in a:
files.extend(sorted(glob.glob(a, recursive=True))) # glob
elif os.path.isdir(a):
files.extend(sorted(glob.glob(os.path.join(a, "*.*")))) # dir
elif os.path.isfile(a):
files.append(a) # files (absolute or relative to CWD)
elif parent and (parent / p).is_file():
files.append(str((parent / p).absolute())) # files (relative to *.txt file parent)
else:
raise FileNotFoundError(f"{p} does not exist")
# Define files as images or videos
images, videos = [], []
for f in files:
suffix = f.rpartition(".")[-1].lower() # Get file extension without the dot and lowercase
if suffix in IMG_FORMATS:
images.append(f)
elif suffix in VID_FORMATS:
videos.append(f)
ni, nv = len(images), len(videos)
self.files = images + videos
self.nf = ni + nv # number of files
self.ni = ni # number of images
self.video_flag = [False] * ni + [True] * nv
self.mode = "video" if ni == 0 else "image" # default to video if no images
self.vid_stride = vid_stride # video frame-rate stride
self.bs = batch
self.cv2_flag = cv2.IMREAD_GRAYSCALE if channels == 1 else cv2.IMREAD_COLOR # grayscale or color (BGR)
if any(videos):
self._new_video(videos[0]) # new video
else:
self.cap = None
if self.nf == 0:
raise FileNotFoundError(f"No images or videos found in {p}. {FORMATS_HELP_MSG}")
def __iter__(self):
"""Iterate through image/video files, yielding source paths, images, and metadata."""
self.count = 0
return self
def __next__(self) -> tuple[list[str], list[np.ndarray], list[str]]:
"""Return the next batch of images or video frames with their paths and metadata."""
paths, imgs, info = [], [], []
while len(imgs) < self.bs:
if self.count >= self.nf: # end of file list
if imgs:
return paths, imgs, info # return last partial batch
else:
raise StopIteration
path = self.files[self.count]
if self.video_flag[self.count]:
self.mode = "video"
if not self.cap or not self.cap.isOpened():
self._new_video(path)
success = False
for _ in range(self.vid_stride):
success = self.cap.grab()
if not success:
break # end of video or failure
if success:
success, im0 = self.cap.retrieve()
im0 = (
cv2.cvtColor(im0, cv2.COLOR_BGR2GRAY)[..., None]
if self.cv2_flag == cv2.IMREAD_GRAYSCALE
else im0
)
if success:
self.frame += 1
paths.append(path)
imgs.append(im0)
info.append(f"video {self.count + 1}/{self.nf} (frame {self.frame}/{self.frames}) {path}: ")
if self.frame == self.frames: # end of video
self.count += 1
self.cap.release()
else:
# Move to the next file if the current video ended or failed to open
self.count += 1
if self.cap:
self.cap.release()
if self.count < self.nf:
self._new_video(self.files[self.count])
else:
# Handle image files
self.mode = "image"
im0 = imread(path, flags=self.cv2_flag) # BGR
if im0 is None:
LOGGER.warning(f"Image Read Error {path}")
else:
paths.append(path)
imgs.append(im0)
info.append(f"image {self.count + 1}/{self.nf} {path}: ")
self.count += 1 # move to the next file
if self.count >= self.ni: # end of image list
break
return paths, imgs, info
def _new_video(self, path: str):
"""Create a new video capture object for the given path and initialize video-related attributes."""
self.frame = 0
self.cap = cv2.VideoCapture(path)
self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))
if not self.cap.isOpened():
raise FileNotFoundError(f"Failed to open video {path}")
self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT) / self.vid_stride)
def __len__(self) -> int:
"""Return the number of batches in the dataset."""
return math.ceil(self.nf / self.bs) # number of batches
class LoadPilAndNumpy:
"""Load images from PIL and Numpy arrays for batch processing.
This class manages loading and pre-processing of image data from both PIL and Numpy formats. It performs basic
validation and format conversion to ensure that the images are in the required format for downstream processing.
Attributes:
paths (list[str]): List of image paths or autogenerated filenames.
im0 (list[np.ndarray]): List of images stored as Numpy arrays.
mode (str): Type of data being processed, set to 'image'.
bs (int): Batch size, equivalent to the length of `im0`.
Methods:
_single_check: Validate and format a single image to a Numpy array.
Examples:
>>> from PIL import Image
>>> import numpy as np
>>> pil_img = Image.new("RGB", (100, 100))
>>> np_img = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
>>> loader = LoadPilAndNumpy([pil_img, np_img])
>>> paths, images, _ = next(iter(loader))
>>> print(f"Loaded {len(images)} images")
Loaded 2 images
"""
def __init__(self, im0: Image.Image | np.ndarray | list, channels: int = 3):
"""Initialize a loader for PIL and Numpy images, converting inputs to a standardized format.
Args:
im0 (PIL.Image.Image | np.ndarray | list): Single image or list of images in PIL or numpy format.
channels (int): Number of image channels (1 for grayscale, 3 for color).
"""
if not isinstance(im0, list):
im0 = [im0]
# use `image{i}.jpg` when Image.filename returns an empty path.
self.paths = [getattr(im, "filename", "") or f"image{i}.jpg" for i, im in enumerate(im0)]
pil_flag = "L" if channels == 1 else "RGB" # grayscale or RGB
self.im0 = [self._single_check(im, pil_flag) for im in im0]
self.mode = "image"
self.bs = len(self.im0)
@staticmethod
def _single_check(im: Image.Image | np.ndarray, flag: str = "RGB") -> np.ndarray:
"""Validate and format an image to a NumPy array.
Notes:
- PIL inputs are converted to NumPy and returned in OpenCV-compatible BGR order for color images.
- NumPy inputs are returned as-is (no channel-order conversion is applied).
"""
assert isinstance(im, (Image.Image, np.ndarray)), f"Expected PIL/np.ndarray image type, but got {type(im)}"
if isinstance(im, Image.Image):
im = np.asarray(im.convert(flag))
# Add a new axis if grayscale; convert RGB -> BGR for OpenCV compatibility.
im = im[..., None] if flag == "L" else im[..., ::-1]
im = np.ascontiguousarray(im) # contiguous
elif im.ndim == 2: # grayscale in numpy form
im = im[..., None]
return im
def __len__(self) -> int:
"""Return the length of the 'im0' attribute, representing the number of loaded images."""
return len(self.im0)
def __next__(self) -> tuple[list[str], list[np.ndarray], list[str]]:
"""Return the next batch of images, paths, and metadata for processing."""
if self.count == 1: # loop only once as it's batch inference
raise StopIteration
self.count += 1
return self.paths, self.im0, [""] * self.bs
def __iter__(self):
"""Iterate through PIL/numpy images, yielding paths, raw images, and metadata for processing."""
self.count = 0
return self
class LoadTensor:
"""A class for loading and processing tensor data for object detection tasks.
This class handles the loading and pre-processing of image data from PyTorch tensors, preparing them for further
processing in object detection pipelines.
Attributes:
im0 (torch.Tensor): The input tensor containing the image(s) with shape (B, C, H, W).
bs (int): Batch size, inferred from the shape of `im0`.
mode (str): Current processing mode, set to 'image'.
paths (list[str]): List of image paths or auto-generated filenames.
Methods:
_single_check: Validates and formats an input tensor.
Examples:
>>> import torch
>>> tensor = torch.rand(1, 3, 640, 640)
>>> loader = LoadTensor(tensor)
>>> paths, images, info = next(iter(loader))
>>> print(f"Processed {len(images)} images")
"""
def __init__(self, im0: torch.Tensor) -> None:
"""Initialize LoadTensor object for processing torch.Tensor image data.
Args:
im0 (torch.Tensor): Input tensor with shape (B, C, H, W).
"""
self.im0 = self._single_check(im0)
self.bs = self.im0.shape[0]
self.mode = "image"
self.paths = [getattr(im, "filename", f"image{i}.jpg") for i, im in enumerate(im0)]
@staticmethod
def _single_check(im: torch.Tensor, stride: int = 32) -> torch.Tensor:
"""Validate and format a single image tensor, ensuring correct shape and normalization."""
s = (
f"torch.Tensor inputs should be BCHW i.e. shape(1, 3, 640, 640) "
f"divisible by stride {stride}. Input shape{tuple(im.shape)} is incompatible."
)
if len(im.shape) != 4:
if len(im.shape) != 3:
raise ValueError(s)
LOGGER.warning(s)
im = im.unsqueeze(0)
if im.shape[2] % stride or im.shape[3] % stride:
raise ValueError(s)
if im.max() > 1.0 + torch.finfo(im.dtype).eps: # torch.float32 eps is 1.2e-07
LOGGER.warning(
f"torch.Tensor inputs should be normalized 0.0-1.0 but max value is {im.max()}. Dividing input by 255."
)
im = im.float() / 255.0
return im
def __iter__(self):
"""Yield an iterator object for iterating through tensor image data."""
self.count = 0
return self
def __next__(self) -> tuple[list[str], torch.Tensor, list[str]]:
"""Yield the next batch of tensor images and metadata for processing."""
if self.count == 1:
raise StopIteration
self.count += 1
return self.paths, self.im0, [""] * self.bs
def __len__(self) -> int:
"""Return the batch size of the tensor input."""
return self.bs
def autocast_list(source: list[Any]) -> list[Image.Image | np.ndarray]:
"""Convert a list of sources into a list of numpy arrays or PIL images for Ultralytics prediction."""
files = []
for im in source:
if isinstance(im, (str, Path)): # filename or uri
files.append(
ImageOps.exif_transpose(Image.open(urllib.request.urlopen(im) if str(im).startswith("http") else im))
)
elif isinstance(im, (Image.Image, np.ndarray)): # PIL or np Image
files.append(im)
else:
raise TypeError(
f"type {type(im).__name__} is not a supported Ultralytics prediction source type. \n"
f"See https://docs.ultralytics.com/modes/predict for supported source types."
)
return files
def get_best_youtube_url(url: str, method: str = "pytube") -> str | None:
"""Retrieve the URL of the best quality MP4 video stream from a given YouTube video.
Args:
url (str): The URL of the YouTube video.
method (str): The method to use for extracting video info. Options are "pytube", "pafy", and "yt-dlp".
Returns:
(str | None): The URL of the best quality MP4 video stream, or None if no suitable stream is found.
Examples:
>>> url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ"
>>> best_url = get_best_youtube_url(url)
>>> print(best_url)
https://rr4---sn-q4flrnek.googlevideo.com/videoplayback?expire=...
Notes:
- Requires additional libraries based on the chosen method: pytubefix, pafy, or yt-dlp.
- The function prioritizes streams with at least 1080p resolution when available.
- For the "yt-dlp" method, it looks for formats with video codec, no audio, and *.mp4 extension.
"""
if method == "pytube":
# Switched from pytube to pytubefix to resolve https://github.com/pytube/pytube/issues/1954
check_requirements("pytubefix>=6.5.2")
from pytubefix import YouTube
streams = YouTube(url).streams.filter(file_extension="mp4", only_video=True)
streams = sorted(streams, key=lambda s: s.resolution, reverse=True) # sort streams by resolution
for stream in streams:
if stream.resolution and int(stream.resolution[:-1]) >= 1080: # check if resolution is at least 1080p
return stream.url
elif method == "pafy":
check_requirements(("pafy", "youtube_dl==2020.12.2"))
import pafy
return pafy.new(url).getbestvideo(preftype="mp4").url
elif method == "yt-dlp":
check_requirements("yt-dlp")
import yt_dlp
with yt_dlp.YoutubeDL({"quiet": True}) as ydl:
info_dict = ydl.extract_info(url, download=False) # extract info
for f in reversed(info_dict.get("formats", [])): # reversed because best is usually last
# Find a format with video codec, no audio, *.mp4 extension at least 1920x1080 size
good_size = (f.get("width") or 0) >= 1920 or (f.get("height") or 0) >= 1080
if good_size and f["vcodec"] != "none" and f["acodec"] == "none" and f["ext"] == "mp4":
return f.get("url")
# Define constants
LOADERS = (LoadStreams, LoadPilAndNumpy, LoadImagesAndVideos, LoadScreenshots)

18
ultralytics/data/scripts/download_weights.sh Executable file
View File

@@ -0,0 +1,18 @@
#!/bin/bash
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
# Download latest models from https://github.com/ultralytics/assets/releases
# Example usage: bash ultralytics/data/scripts/download_weights.sh
# parent
# └── weights
# ├── yolov8n.pt ← downloads here
# ├── yolov8s.pt
# └── ...
python << EOF
from ultralytics.utils.downloads import attempt_download_asset
assets = [f"yolov8{size}{suffix}.pt" for size in "nsmlx" for suffix in ("", "-cls", "-seg", "-pose")]
for x in assets:
attempt_download_asset(f"weights/{x}")
EOF

61
ultralytics/data/scripts/get_coco.sh Executable file
View File

@@ -0,0 +1,61 @@
#!/bin/bash
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
# Download COCO 2017 dataset https://cocodataset.org
# Example usage: bash data/scripts/get_coco.sh
# parent
# ├── ultralytics
# └── datasets
# └── coco ← downloads here
# Arguments (optional) Usage: bash data/scripts/get_coco.sh --train --val --test --segments
if [ "$#" -gt 0 ]; then
for opt in "$@"; do
case "${opt}" in
--train) train=true ;;
--val) val=true ;;
--test) test=true ;;
--segments) segments=true ;;
--sama) sama=true ;;
esac
done
else
train=true
val=true
test=false
segments=false
sama=false
fi
# Download/unzip labels
d='../datasets' # unzip directory
url=https://github.com/ultralytics/assets/releases/download/v0.0.0/
if [ "$segments" == "true" ]; then
f='coco2017labels-segments.zip' # 169 MB
elif [ "$sama" == "true" ]; then
f='coco2017labels-segments-sama.zip' # 199 MB https://www.sama.com/sama-coco-dataset/
else
f='coco2017labels.zip' # 46 MB
fi
echo 'Downloading' $url$f ' ...'
curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
# Download/unzip images
d='../datasets/coco/images' # unzip directory
url=http://images.cocodataset.org/zips/
if [ "$train" == "true" ]; then
f='train2017.zip' # 19G, 118k images
echo 'Downloading' $url$f '...'
curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
fi
if [ "$val" == "true" ]; then
f='val2017.zip' # 1G, 5k images
echo 'Downloading' $url$f '...'
curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
fi
if [ "$test" == "true" ]; then
f='test2017.zip' # 7G, 41k images (optional)
echo 'Downloading' $url$f '...'
curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
fi
wait # finish background tasks

18
ultralytics/data/scripts/get_coco128.sh Executable file
View File

@@ -0,0 +1,18 @@
#!/bin/bash
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
# Download COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017)
# Example usage: bash data/scripts/get_coco128.sh
# parent
# ├── ultralytics
# └── datasets
# └── coco128 ← downloads here
# Download/unzip images and labels
d='../datasets' # unzip directory
url=https://github.com/ultralytics/assets/releases/download/v0.0.0/
f='coco128.zip' # or 'coco128-segments.zip', 68 MB
echo 'Downloading' $url$f ' ...'
curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
wait # finish background tasks

52
ultralytics/data/scripts/get_imagenet.sh Executable file
View File

@@ -0,0 +1,52 @@
#!/bin/bash
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
# Download ILSVRC2012 ImageNet dataset https://image-net.org
# Example usage: bash data/scripts/get_imagenet.sh
# parent
# ├── ultralytics
# └── datasets
# └── imagenet ← downloads here
# Arguments (optional) Usage: bash data/scripts/get_imagenet.sh --train --val
if [ "$#" -gt 0 ]; then
for opt in "$@"; do
case "${opt}" in
--train) train=true ;;
--val) val=true ;;
esac
done
else
train=true
val=true
fi
# Make dir
d='../datasets/imagenet' # unzip directory
mkdir -p $d && cd $d
# Download/unzip train
if [ "$train" == "true" ]; then
wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_train.tar # download 138G, 1281167 images
mkdir train && mv ILSVRC2012_img_train.tar train/ && cd train
tar -xf ILSVRC2012_img_train.tar && rm -f ILSVRC2012_img_train.tar
find . -name "*.tar" | while read NAME; do
mkdir -p "${NAME%.tar}"
tar -xf "${NAME}" -C "${NAME%.tar}"
rm -f "${NAME}"
done
cd ..
fi
# Download/unzip val
if [ "$val" == "true" ]; then
wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar # download 6.3G, 50000 images
mkdir val && mv ILSVRC2012_img_val.tar val/ && cd val && tar -xf ILSVRC2012_img_val.tar
wget -qO- https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh | bash # move into subdirs
fi
# Delete corrupted image (optional: PNG under JPEG name that may cause dataloaders to fail)
# rm train/n04266014/n04266014_10835.JPEG
# TFRecords (optional)
# wget https://raw.githubusercontent.com/tensorflow/models/master/research/slim/datasets/imagenet_lsvrc_2015_synsets.txt

138
ultralytics/data/split.py Executable file
View File

@@ -0,0 +1,138 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import random
import shutil
from pathlib import Path
from ultralytics.data.utils import IMG_FORMATS, img2label_paths
from ultralytics.utils import DATASETS_DIR, LOGGER, TQDM
def split_classify_dataset(source_dir: str | Path, train_ratio: float = 0.8) -> Path:
"""Split classification dataset into train and val directories in a new directory.
Creates a new directory '{source_dir}_split' with train/val subdirectories, preserving the original class structure
with an 80/20 split by default.
Directory structure:
Before:
caltech/
├── class1/
│ ├── img1.jpg
│ ├── img2.jpg
│ └── ...
├── class2/
│ ├── img1.jpg
│ └── ...
└── ...
After:
caltech_split/
├── train/
│ ├── class1/
│ │ ├── img1.jpg
│ │ └── ...
│ ├── class2/
│ │ ├── img1.jpg
│ │ └── ...
│ └── ...
└── val/
├── class1/
│ ├── img2.jpg
│ └── ...
├── class2/
│ └── ...
└── ...
Args:
source_dir (str | Path): Path to classification dataset root directory.
train_ratio (float): Ratio for train split, between 0 and 1.
Returns:
(Path): Path to the created split directory.
Examples:
Split dataset with default 80/20 ratio
>>> split_classify_dataset("path/to/caltech")
Split with custom ratio
>>> split_classify_dataset("path/to/caltech", 0.75)
"""
source_path = Path(source_dir)
split_path = Path(f"{source_path}_split")
train_path, val_path = split_path / "train", split_path / "val"
# Create directory structure
split_path.mkdir(exist_ok=True)
train_path.mkdir(exist_ok=True)
val_path.mkdir(exist_ok=True)
# Process class directories
class_dirs = [d for d in source_path.iterdir() if d.is_dir()]
total_images = sum(len(list(d.glob("*.*"))) for d in class_dirs)
stats = f"{len(class_dirs)} classes, {total_images} images"
LOGGER.info(f"Splitting {source_path} ({stats}) into {train_ratio:.0%} train, {1 - train_ratio:.0%} val...")
for class_dir in class_dirs:
# Create class directories
(train_path / class_dir.name).mkdir(exist_ok=True)
(val_path / class_dir.name).mkdir(exist_ok=True)
# Split and copy files
image_files = list(class_dir.glob("*.*"))
random.shuffle(image_files)
split_idx = int(len(image_files) * train_ratio)
for img in image_files[:split_idx]:
shutil.copy2(img, train_path / class_dir.name / img.name)
for img in image_files[split_idx:]:
shutil.copy2(img, val_path / class_dir.name / img.name)
LOGGER.info(f"Split complete in {split_path} ✅")
return split_path
def autosplit(
path: Path = DATASETS_DIR / "coco8/images",
weights: tuple[float, float, float] = (0.9, 0.1, 0.0),
annotated_only: bool = False,
) -> None:
"""Automatically split a dataset into train/val/test splits and save the resulting splits into autosplit_*.txt
files.
Args:
path (Path): Path to images directory.
weights (tuple[float, float, float]): Train, validation, and test split fractions.
annotated_only (bool): If True, only images with an associated txt file are used.
Examples:
Split images with default weights
>>> from ultralytics.data.split import autosplit
>>> autosplit()
Split with custom weights and annotated images only
>>> autosplit(path="path/to/images", weights=(0.8, 0.15, 0.05), annotated_only=True)
"""
path = Path(path) # images dir
files = sorted(x for x in path.rglob("*.*") if x.suffix[1:].lower() in IMG_FORMATS) # image files only
n = len(files) # number of files
random.seed(0) # for reproducibility
indices = random.choices([0, 1, 2], weights=weights, k=n) # assign each image to a split
txt = ["autosplit_train.txt", "autosplit_val.txt", "autosplit_test.txt"] # 3 txt files
for x in txt:
if (path.parent / x).exists():
(path.parent / x).unlink() # remove existing
LOGGER.info(f"Autosplitting images from {path}" + ", using *.txt labeled images only" * annotated_only)
for i, img in TQDM(zip(indices, files), total=n):
if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check label
with open(path.parent / txt[i], "a", encoding="utf-8") as f:
f.write(f"./{img.relative_to(path.parent).as_posix()}" + "\n") # add image to txt file
if __name__ == "__main__":
split_classify_dataset("caltech101")

344
ultralytics/data/split_dota.py Executable file
View File

@@ -0,0 +1,344 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import itertools
from glob import glob
from math import ceil
from pathlib import Path
from typing import Any
import cv2
import numpy as np
from PIL import Image
from ultralytics.data.utils import exif_size, img2label_paths
from ultralytics.utils import TQDM
from ultralytics.utils.checks import check_requirements
def bbox_iof(polygon1: np.ndarray, bbox2: np.ndarray, eps: float = 1e-6) -> np.ndarray:
"""Calculate Intersection over Foreground (IoF) between polygons and bounding boxes.
Args:
polygon1 (np.ndarray): Polygon coordinates with shape (N, 8).
bbox2 (np.ndarray): Bounding boxes with shape (M, 4).
eps (float, optional): Small value to prevent division by zero.
Returns:
(np.ndarray): IoF scores with shape (N, M).
Notes:
Polygon format: [x1, y1, x2, y2, x3, y3, x4, y4].
Bounding box format: [x_min, y_min, x_max, y_max].
"""
check_requirements("shapely>=2.0.0")
from shapely.geometry import Polygon
polygon1 = polygon1.reshape(-1, 4, 2)
lt_point = np.min(polygon1, axis=-2) # left-top
rb_point = np.max(polygon1, axis=-2) # right-bottom
bbox1 = np.concatenate([lt_point, rb_point], axis=-1)
lt = np.maximum(bbox1[:, None, :2], bbox2[..., :2])
rb = np.minimum(bbox1[:, None, 2:], bbox2[..., 2:])
wh = np.clip(rb - lt, 0, np.inf)
h_overlaps = wh[..., 0] * wh[..., 1]
left, top, right, bottom = (bbox2[..., i] for i in range(4))
polygon2 = np.stack([left, top, right, top, right, bottom, left, bottom], axis=-1).reshape(-1, 4, 2)
sg_polys1 = [Polygon(p) for p in polygon1]
sg_polys2 = [Polygon(p) for p in polygon2]
overlaps = np.zeros(h_overlaps.shape)
for p in zip(*np.nonzero(h_overlaps)):
overlaps[p] = sg_polys1[p[0]].intersection(sg_polys2[p[-1]]).area
unions = np.array([p.area for p in sg_polys1], dtype=np.float32)
unions = unions[..., None]
unions = np.clip(unions, eps, np.inf)
outputs = overlaps / unions
if outputs.ndim == 1:
outputs = outputs[..., None]
return outputs
def load_yolo_dota(data_root: str, split: str = "train") -> list[dict[str, Any]]:
"""Load DOTA dataset annotations and image information.
Args:
data_root (str): Data root directory.
split (str, optional): The split data set, could be 'train' or 'val'.
Returns:
(list[dict[str, Any]]): List of annotation dictionaries containing image information.
Notes:
The directory structure assumed for the DOTA dataset:
- data_root
- images
- train
- val
- labels
- train
- val
"""
assert split in {"train", "val"}, f"Split must be 'train' or 'val', not {split}."
im_dir = Path(data_root) / "images" / split
assert im_dir.exists(), f"Can't find {im_dir}, please check your data root."
im_files = glob(str(Path(data_root) / "images" / split / "*"))
lb_files = img2label_paths(im_files)
annos = []
for im_file, lb_file in zip(im_files, lb_files):
w, h = exif_size(Image.open(im_file))
with open(lb_file, encoding="utf-8") as f:
lb = [x.split() for x in f.read().strip().splitlines() if len(x)]
lb = np.array(lb, dtype=np.float32)
annos.append(dict(ori_size=(h, w), label=lb, filepath=im_file))
return annos
def get_windows(
im_size: tuple[int, int],
crop_sizes: tuple[int, ...] = (1024,),
gaps: tuple[int, ...] = (200,),
im_rate_thr: float = 0.6,
eps: float = 0.01,
) -> np.ndarray:
"""Get the coordinates of sliding windows for image cropping.
Args:
im_size (tuple[int, int]): Original image size, (H, W).
crop_sizes (tuple[int, ...], optional): Crop size of windows.
gaps (tuple[int, ...], optional): Gap between crops.
im_rate_thr (float, optional): Threshold for the ratio of image area within a window to the total window area.
eps (float, optional): Epsilon value for math operations.
Returns:
(np.ndarray): Array of window coordinates of shape (N, 4) where each row is [x_start, y_start, x_stop, y_stop].
"""
h, w = im_size
windows = []
for crop_size, gap in zip(crop_sizes, gaps):
assert crop_size > gap, f"invalid crop_size gap pair [{crop_size} {gap}]"
step = crop_size - gap
xn = 1 if w <= crop_size else ceil((w - crop_size) / step + 1)
xs = [step * i for i in range(xn)]
if len(xs) > 1 and xs[-1] + crop_size > w:
xs[-1] = w - crop_size
yn = 1 if h <= crop_size else ceil((h - crop_size) / step + 1)
ys = [step * i for i in range(yn)]
if len(ys) > 1 and ys[-1] + crop_size > h:
ys[-1] = h - crop_size
start = np.array(list(itertools.product(xs, ys)), dtype=np.int64)
stop = start + crop_size
windows.append(np.concatenate([start, stop], axis=1))
windows = np.concatenate(windows, axis=0)
im_in_wins = windows.copy()
im_in_wins[:, 0::2] = np.clip(im_in_wins[:, 0::2], 0, w)
im_in_wins[:, 1::2] = np.clip(im_in_wins[:, 1::2], 0, h)
im_areas = (im_in_wins[:, 2] - im_in_wins[:, 0]) * (im_in_wins[:, 3] - im_in_wins[:, 1])
win_areas = (windows[:, 2] - windows[:, 0]) * (windows[:, 3] - windows[:, 1])
im_rates = im_areas / win_areas
if not (im_rates > im_rate_thr).any():
max_rate = im_rates.max()
im_rates[abs(im_rates - max_rate) < eps] = 1
return windows[im_rates > im_rate_thr]
def get_window_obj(anno: dict[str, Any], windows: np.ndarray, iof_thr: float = 0.7) -> list[np.ndarray]:
"""Get objects for each window based on IoF threshold."""
h, w = anno["ori_size"]
label = anno["label"]
if len(label):
label[:, 1::2] *= w
label[:, 2::2] *= h
iofs = bbox_iof(label[:, 1:], windows)
# Unnormalized and misaligned coordinates
return [(label[iofs[:, i] >= iof_thr]) for i in range(len(windows))] # window_anns
else:
return [np.zeros((0, 9), dtype=np.float32) for _ in range(len(windows))] # window_anns
def crop_and_save(
anno: dict[str, Any],
windows: np.ndarray,
window_objs: list[np.ndarray],
im_dir: str,
lb_dir: str,
allow_background_images: bool = True,
) -> None:
"""Crop images and save new labels for each window.
Args:
anno (dict[str, Any]): Annotation dict, including 'filepath', 'label', 'ori_size' as its keys.
windows (np.ndarray): Array of windows coordinates with shape (N, 4).
window_objs (list[np.ndarray]): A list of labels inside each window.
im_dir (str): The output directory path of images.
lb_dir (str): The output directory path of labels.
allow_background_images (bool, optional): Whether to include background images without labels.
Notes:
The directory structure assumed for the DOTA dataset:
- data_root
- images
- train
- val
- labels
- train
- val
"""
im = cv2.imread(anno["filepath"])
name = Path(anno["filepath"]).stem
for i, window in enumerate(windows):
x_start, y_start, x_stop, y_stop = window.tolist()
new_name = f"{name}__{x_stop - x_start}__{x_start}___{y_start}"
patch_im = im[y_start:y_stop, x_start:x_stop]
ph, pw = patch_im.shape[:2]
label = window_objs[i]
if len(label) or allow_background_images:
cv2.imwrite(str(Path(im_dir) / f"{new_name}.jpg"), patch_im)
if len(label):
label[:, 1::2] -= x_start
label[:, 2::2] -= y_start
label[:, 1::2] /= pw
label[:, 2::2] /= ph
with open(Path(lb_dir) / f"{new_name}.txt", "w", encoding="utf-8") as f:
for lb in label:
formatted_coords = [f"{coord:.6g}" for coord in lb[1:]]
f.write(f"{int(lb[0])} {' '.join(formatted_coords)}\n")
def split_images_and_labels(
data_root: str,
save_dir: str,
split: str = "train",
crop_sizes: tuple[int, ...] = (1024,),
gaps: tuple[int, ...] = (200,),
) -> None:
"""Split both images and labels for a given dataset split.
Args:
data_root (str): Root directory of the dataset.
save_dir (str): Directory to save the split dataset.
split (str, optional): The split data set, could be 'train' or 'val'.
crop_sizes (tuple[int, ...], optional): Tuple of crop sizes.
gaps (tuple[int, ...], optional): Tuple of gaps between crops.
Notes:
The directory structure assumed for the DOTA dataset:
- data_root
- images
- split
- labels
- split
and the output directory structure is:
- save_dir
- images
- split
- labels
- split
"""
im_dir = Path(save_dir) / "images" / split
im_dir.mkdir(parents=True, exist_ok=True)
lb_dir = Path(save_dir) / "labels" / split
lb_dir.mkdir(parents=True, exist_ok=True)
annos = load_yolo_dota(data_root, split=split)
for anno in TQDM(annos, total=len(annos), desc=split):
windows = get_windows(anno["ori_size"], crop_sizes, gaps)
window_objs = get_window_obj(anno, windows)
crop_and_save(anno, windows, window_objs, str(im_dir), str(lb_dir))
def split_trainval(
data_root: str, save_dir: str, crop_size: int = 1024, gap: int = 200, rates: tuple[float, ...] = (1.0,)
) -> None:
"""Split train and val sets of DOTA dataset with multiple scaling rates.
Args:
data_root (str): Root directory of the dataset.
save_dir (str): Directory to save the split dataset.
crop_size (int, optional): Base crop size.
gap (int, optional): Base gap between crops.
rates (tuple[float, ...], optional): Scaling rates for crop_size and gap.
Notes:
The directory structure assumed for the DOTA dataset:
- data_root
- images
- train
- val
- labels
- train
- val
and the output directory structure is:
- save_dir
- images
- train
- val
- labels
- train
- val
"""
crop_sizes, gaps = [], []
for r in rates:
crop_sizes.append(int(crop_size / r))
gaps.append(int(gap / r))
for split in {"train", "val"}:
split_images_and_labels(data_root, save_dir, split, crop_sizes, gaps)
def split_test(
data_root: str, save_dir: str, crop_size: int = 1024, gap: int = 200, rates: tuple[float, ...] = (1.0,)
) -> None:
"""Split test set of DOTA dataset, labels are not included within this set.
Args:
data_root (str): Root directory of the dataset.
save_dir (str): Directory to save the split dataset.
crop_size (int, optional): Base crop size.
gap (int, optional): Base gap between crops.
rates (tuple[float, ...], optional): Scaling rates for crop_size and gap.
Notes:
The directory structure assumed for the DOTA dataset:
- data_root
- images
- test
and the output directory structure is:
- save_dir
- images
- test
"""
crop_sizes, gaps = [], []
for r in rates:
crop_sizes.append(int(crop_size / r))
gaps.append(int(gap / r))
save_dir = Path(save_dir) / "images" / "test"
save_dir.mkdir(parents=True, exist_ok=True)
im_dir = Path(data_root) / "images" / "test"
assert im_dir.exists(), f"Can't find {im_dir}, please check your data root."
im_files = glob(str(im_dir / "*"))
for im_file in TQDM(im_files, total=len(im_files), desc="test"):
w, h = exif_size(Image.open(im_file))
windows = get_windows((h, w), crop_sizes=crop_sizes, gaps=gaps)
im = cv2.imread(im_file)
name = Path(im_file).stem
for window in windows:
x_start, y_start, x_stop, y_stop = window.tolist()
new_name = f"{name}__{x_stop - x_start}__{x_start}___{y_start}"
patch_im = im[y_start:y_stop, x_start:x_stop]
cv2.imwrite(str(save_dir / f"{new_name}.jpg"), patch_im)
if __name__ == "__main__":
split_trainval(data_root="DOTAv2", save_dir="DOTAv2-split")
split_test(data_root="DOTAv2", save_dir="DOTAv2-split")

883
ultralytics/data/utils.py Executable file
View File

@@ -0,0 +1,883 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import json
import os
import random
import subprocess
import time
import zipfile
from multiprocessing.pool import ThreadPool
from pathlib import Path
from tarfile import is_tarfile
from typing import Any
import cv2
import numpy as np
from PIL import Image, ImageOps
from ultralytics.nn.autobackend import check_class_names
from ultralytics.utils import (
ASSETS_URL,
DATASETS_DIR,
LOGGER,
NUM_THREADS,
ROOT,
SETTINGS_FILE,
TQDM,
YAML,
clean_url,
colorstr,
emojis,
is_dir_writeable,
)
from ultralytics.utils.checks import check_file, check_font, is_ascii
from ultralytics.utils.downloads import download, safe_download, unzip_file
from ultralytics.utils.ops import segments2boxes
HELP_URL = "See https://docs.ultralytics.com/datasets for dataset formatting guidance."
IMG_FORMATS = {
"avif",
"bmp",
"dng",
"heic",
"heif",
"jp2",
"jpeg",
"jpeg2000",
"jpg",
"mpo",
"png",
"tif",
"tiff",
"webp",
}
VID_FORMATS = {"asf", "avi", "gif", "m4v", "mkv", "mov", "mp4", "mpeg", "mpg", "ts", "wmv", "webm"} # videos
FORMATS_HELP_MSG = f"Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}"
def img2label_paths(img_paths: list[str]) -> list[str]:
"""Convert image paths to label paths by replacing 'images' with 'labels' and extension with '.txt'."""
sa, sb = f"{os.sep}images{os.sep}", f"{os.sep}labels{os.sep}" # /images/, /labels/ substrings
return [sb.join(x.rsplit(sa, 1)).rsplit(".", 1)[0] + ".txt" for x in img_paths]
def check_file_speeds(
files: list[str], threshold_ms: float = 10, threshold_mb: float = 50, max_files: int = 5, prefix: str = ""
):
"""Check dataset file access speed and provide performance feedback.
This function tests the access speed of dataset files by measuring ping (stat call) time and read speed. It samples
up to `max_files` files from the provided list and warns if access times exceed the threshold.
Args:
files (list[str]): List of file paths to check for access speed.
threshold_ms (float, optional): Threshold in milliseconds for ping time warnings.
threshold_mb (float, optional): Threshold in megabytes per second for read speed warnings.
max_files (int, optional): The maximum number of files to check.
prefix (str, optional): Prefix string to add to log messages.
Examples:
>>> from pathlib import Path
>>> image_files = list(Path("dataset/images").glob("*.jpg"))
>>> check_file_speeds(image_files, threshold_ms=15)
"""
if not files:
LOGGER.warning(f"{prefix}Image speed checks: No files to check")
return
# Sample files (max 5)
files = random.sample(files, min(max_files, len(files)))
# Test ping (stat time)
ping_times = []
file_sizes = []
read_speeds = []
for f in files:
try:
# Measure ping (stat call)
start = time.perf_counter()
file_size = os.stat(f).st_size
ping_times.append((time.perf_counter() - start) * 1000) # ms
file_sizes.append(file_size)
# Measure read speed
start = time.perf_counter()
with open(f, "rb") as file_obj:
_ = file_obj.read()
read_time = time.perf_counter() - start
if read_time > 0: # Avoid division by zero
read_speeds.append(file_size / (1 << 20) / read_time) # MB/s
except Exception:
pass
if not ping_times:
LOGGER.warning(f"{prefix}Image speed checks: failed to access files")
return
# Calculate stats with uncertainties
avg_ping = np.mean(ping_times)
std_ping = np.std(ping_times, ddof=1) if len(ping_times) > 1 else 0
size_msg = f", size: {np.mean(file_sizes) / (1 << 10):.1f} KB"
ping_msg = f"ping: {avg_ping:.1f}±{std_ping:.1f} ms"
if read_speeds:
avg_speed = np.mean(read_speeds)
std_speed = np.std(read_speeds, ddof=1) if len(read_speeds) > 1 else 0
speed_msg = f", read: {avg_speed:.1f}±{std_speed:.1f} MB/s"
else:
speed_msg = ""
if avg_ping < threshold_ms or avg_speed < threshold_mb:
LOGGER.info(f"{prefix}Fast image access ✅ ({ping_msg}{speed_msg}{size_msg})")
else:
LOGGER.warning(
f"{prefix}Slow image access detected ({ping_msg}{speed_msg}{size_msg}). "
f"Use local storage instead of remote/mounted storage for better performance. "
f"See https://docs.ultralytics.com/guides/model-training-tips/"
)
def get_hash(paths: list[str]) -> str:
"""Return a single hash value of a list of paths (files or dirs)."""
size = 0
for p in paths:
try:
size += os.stat(p).st_size
except OSError:
continue
h = __import__("hashlib").sha256(str(size).encode()) # hash sizes
h.update("".join(paths).encode()) # hash paths
return h.hexdigest() # return hash
def exif_size(img: Image.Image) -> tuple[int, int]:
"""Return exif-corrected PIL size."""
s = img.size # (width, height)
if img.format == "JPEG": # only support JPEG images
try:
if exif := img.getexif():
rotation = exif.get(274, None) # the EXIF key for the orientation tag is 274
if rotation in {6, 8}: # rotation 270 or 90
s = s[1], s[0]
except Exception:
pass
return s
def verify_image(args: tuple) -> tuple:
"""Verify one image."""
(im_file, cls), prefix = args
# Number (found, corrupt), message
nf, nc, msg = 0, 0, ""
try:
im = Image.open(im_file)
im.verify() # PIL verify
shape = exif_size(im) # image size
shape = (shape[1], shape[0]) # hw
assert (shape[0] > 9) & (shape[1] > 9), f"image size {shape} <10 pixels"
assert im.format.lower() in IMG_FORMATS, f"Invalid image format {im.format}. {FORMATS_HELP_MSG}"
if im.format.lower() in {"jpg", "jpeg"}:
with open(im_file, "rb") as f:
f.seek(-2, 2)
if f.read() != b"\xff\xd9": # corrupt JPEG
ImageOps.exif_transpose(Image.open(im_file)).save(im_file, "JPEG", subsampling=0, quality=100)
msg = f"{prefix}{im_file}: corrupt JPEG restored and saved"
nf = 1
except Exception as e:
nc = 1
msg = f"{prefix}{im_file}: ignoring corrupt image/label: {e}"
return (im_file, cls), nf, nc, msg
def verify_image_label(args: tuple) -> list:
"""Verify one image-label pair."""
im_file, lb_file, prefix, keypoint, num_cls, nkpt, ndim, single_cls = args
# Number (missing, found, empty, corrupt), message, segments, keypoints
nm, nf, ne, nc, msg, segments, keypoints = 0, 0, 0, 0, "", [], None
try:
# Verify images
im = Image.open(im_file)
im.verify() # PIL verify
shape = exif_size(im) # image size
shape = (shape[1], shape[0]) # hw
assert (shape[0] > 9) & (shape[1] > 9), f"image size {shape} <10 pixels"
assert im.format.lower() in IMG_FORMATS, f"invalid image format {im.format}. {FORMATS_HELP_MSG}"
if im.format.lower() in {"jpg", "jpeg"}:
with open(im_file, "rb") as f:
f.seek(-2, 2)
if f.read() != b"\xff\xd9": # corrupt JPEG
ImageOps.exif_transpose(Image.open(im_file)).save(im_file, "JPEG", subsampling=0, quality=100)
msg = f"{prefix}{im_file}: corrupt JPEG restored and saved"
# Verify labels
if os.path.isfile(lb_file):
nf = 1 # label found
with open(lb_file, encoding="utf-8") as f:
lb = [x.split() for x in f.read().strip().splitlines() if len(x)]
if any(len(x) > 6 for x in lb) and (not keypoint): # is segment
classes = np.array([x[0] for x in lb], dtype=np.float32)
segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in lb] # (cls, xy1...)
lb = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1) # (cls, xywh)
lb = np.array(lb, dtype=np.float32)
if nl := len(lb):
if keypoint:
assert lb.shape[1] == (5 + nkpt * ndim), f"labels require {(5 + nkpt * ndim)} columns each"
points = lb[:, 5:].reshape(-1, ndim)[:, :2]
else:
assert lb.shape[1] == 5, f"labels require 5 columns, {lb.shape[1]} columns detected"
points = lb[:, 1:]
# Coordinate points check with 1% tolerance
assert points.max() <= 1.01, f"non-normalized or out of bounds coordinates {points[points > 1.01]}"
assert lb.min() >= -0.01, f"negative class labels or coordinate {lb[lb < -0.01]}"
# All labels
max_cls = 0 if single_cls else lb[:, 0].max() # max label count
assert max_cls < num_cls, (
f"Label class {int(max_cls)} exceeds dataset class count {num_cls}. "
f"Possible class labels are 0-{num_cls - 1}"
)
_, i = np.unique(lb, axis=0, return_index=True)
if len(i) < nl: # duplicate row check
lb = lb[i] # remove duplicates
if segments:
segments = [segments[x] for x in i]
msg = f"{prefix}{im_file}: {nl - len(i)} duplicate labels removed"
else:
ne = 1 # label empty
lb = np.zeros((0, (5 + nkpt * ndim) if keypoint else 5), dtype=np.float32)
else:
nm = 1 # label missing
lb = np.zeros((0, (5 + nkpt * ndim) if keypoint else 5), dtype=np.float32)
if keypoint:
keypoints = lb[:, 5:].reshape(-1, nkpt, ndim)
if ndim == 2:
kpt_mask = np.where((keypoints[..., 0] < 0) | (keypoints[..., 1] < 0), 0.0, 1.0).astype(np.float32)
keypoints = np.concatenate([keypoints, kpt_mask[..., None]], axis=-1) # (nl, nkpt, 3)
lb = lb[:, :5]
return im_file, lb, shape, segments, keypoints, nm, nf, ne, nc, msg
except Exception as e:
nc = 1
msg = f"{prefix}{im_file}: ignoring corrupt image/label: {e}"
return [None, None, None, None, None, nm, nf, ne, nc, msg]
def verify_image_label_ground(args: tuple) -> list:
"""Verify one image-label pair for ground dataset with difficulty scores."""
im_file, lb_file, prefix, class_map = args
# Number (missing, found, empty, corrupt), message, segments
nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, "", np.zeros((0, 1000, 2), dtype=np.float32)
try:
# Verify images
im = Image.open(im_file)
im.verify() # PIL verify
shape = exif_size(im) # image size
shape = (shape[1], shape[0]) # hw
assert (shape[0] > 9) & (shape[1] > 9), f"image size {shape} <10 pixels"
assert im.format.lower() in IMG_FORMATS, f"invalid image format {im.format}. {FORMATS_HELP_MSG}"
if im.format.lower() in {"jpg", "jpeg"}:
with open(im_file, "rb") as f:
f.seek(-2, 2)
if f.read() != b"\xff\xd9": # corrupt JPEG
ImageOps.exif_transpose(Image.open(im_file)).save(im_file, "JPEG", subsampling=0, quality=100)
msg = f"{prefix}{im_file}: corrupt JPEG restored and saved"
# Verify labels
if os.path.isfile(lb_file):
nf = 1 # label found
with open(lb_file, encoding="utf-8") as f:
ori_lb = [x.split() for x in f.read().strip().splitlines() if len(x)]
lb = []
for x in ori_lb:
# Map class name to class ID
if x[0] in class_map:
class_id = class_map[x[0]]
else:
continue # Skip unknown classes
# Parse coordinates and difficulties
coords = [float(y) for y in x[1:5]] # x, y, w, h
difficulty = float(x[5]) # diff2 is truncation; use only diff1 for difficulty supervision
# Append [class_id, x, y, w, h, difficulty]
lb.append([class_id] + coords + [difficulty])
lb = np.array(lb, dtype=np.float32)
if nl := len(lb):
assert lb.shape[1] == 6, f"labels require 6 columns (class, x, y, w, h, difficulty), {lb.shape[1]} columns detected"
assert (lb[:, :5] >= 0).all(), f"negative label values {lb[:, :5][lb[:, :5] < 0]}"
assert (lb[:, 1:5] <= 1).all(), f"non-normalized or out of bounds coordinates {lb[:, 1:5][lb[:, 1:5] > 1]}"
_, i = np.unique(lb, axis=0, return_index=True)
if len(i) < nl: # duplicate row check
lb = lb[i] # remove duplicates
msg = f"{prefix}{im_file}: {nl - len(i)} duplicate labels removed"
else:
ne = 1 # label empty
lb = np.zeros((0, 6), dtype=np.float32)
else:
nm = 1 # label missing
lb = np.zeros((0, 6), dtype=np.float32)
return im_file, lb, shape, segments, nm, nf, ne, nc, msg
except Exception as e:
nc = 1
msg = f"{prefix}{im_file}: ignoring corrupt image/label: {e}"
return [None, None, None, None, nm, nf, ne, nc, msg]
def visualize_image_annotations(image_path: str, txt_path: str, label_map: dict[int, str]):
"""Visualize YOLO annotations (bounding boxes and class labels) on an image.
This function reads an image and its corresponding annotation file in YOLO format, then draws bounding boxes around
detected objects and labels them with their respective class names. The bounding box colors are assigned based on
the class ID, and the text color is dynamically adjusted for readability, depending on the background color's
luminance.
Args:
image_path (str): Path to the image file to annotate. The file must be readable by PIL.
txt_path (str): Path to the annotation file in YOLO format, which should contain one line per object.
label_map (dict[int, str]): A dictionary that maps class IDs (integers) to class labels (strings).
Examples:
>>> label_map = {0: "cat", 1: "dog", 2: "bird"} # Should include all annotated classes
>>> visualize_image_annotations("path/to/image.jpg", "path/to/annotations.txt", label_map)
"""
import matplotlib.pyplot as plt
from ultralytics.utils.plotting import colors
img = np.array(Image.open(image_path))
img_height, img_width = img.shape[:2]
annotations = []
with open(txt_path, encoding="utf-8") as file:
for line in file:
class_id, x_center, y_center, width, height = map(float, line.split())
x = (x_center - width / 2) * img_width
y = (y_center - height / 2) * img_height
w = width * img_width
h = height * img_height
annotations.append((x, y, w, h, int(class_id)))
_, ax = plt.subplots(1) # Plot the image and annotations
for x, y, w, h, label in annotations:
color = tuple(c / 255 for c in colors(label, False)) # Get and normalize an RGB color for Matplotlib
rect = plt.Rectangle((x, y), w, h, linewidth=2, edgecolor=color, facecolor="none") # Create a rectangle
ax.add_patch(rect)
luminance = 0.2126 * color[0] + 0.7152 * color[1] + 0.0722 * color[2] # Formula for luminance
ax.text(x, y - 5, label_map[label], color="white" if luminance < 0.5 else "black", backgroundcolor=color)
ax.imshow(img)
plt.show()
def polygon2mask(
imgsz: tuple[int, int], polygons: list[np.ndarray], color: int = 1, downsample_ratio: int = 1
) -> np.ndarray:
"""Convert a list of polygons to a binary mask of the specified image size.
Args:
imgsz (tuple[int, int]): The size of the image as (height, width).
polygons (list[np.ndarray]): A list of polygons. Each polygon is a 1D array of coordinates with length M, where
M % 2 = 0 (alternating x, y values).
color (int, optional): The color value to fill in the polygons on the mask.
downsample_ratio (int, optional): Factor by which to downsample the mask.
Returns:
(np.ndarray): A binary mask of the specified image size with the polygons filled in.
"""
mask = np.zeros(imgsz, dtype=np.uint8)
polygons = np.asarray(polygons, dtype=np.int32)
polygons = polygons.reshape((polygons.shape[0], -1, 2))
cv2.fillPoly(mask, polygons, color=color)
nh, nw = (imgsz[0] // downsample_ratio, imgsz[1] // downsample_ratio)
# Note: fillPoly first then resize is trying to keep the same loss calculation method when mask-ratio=1
return cv2.resize(mask, (nw, nh))
def polygons2masks(
imgsz: tuple[int, int], polygons: list[np.ndarray], color: int, downsample_ratio: int = 1
) -> np.ndarray:
"""Convert a list of polygons to a set of binary masks of the specified image size.
Args:
imgsz (tuple[int, int]): The size of the image as (height, width).
polygons (list[np.ndarray]): A list of polygons. Each polygon is an array of coordinates that can be reshaped to
(-1, 2) as (x, y) point pairs.
color (int): The color value to fill in the polygons on the masks.
downsample_ratio (int, optional): Factor by which to downsample each mask.
Returns:
(np.ndarray): A set of binary masks of the specified image size with the polygons filled in.
"""
return np.array([polygon2mask(imgsz, [x.reshape(-1)], color, downsample_ratio) for x in polygons])
def polygons2masks_overlap(
imgsz: tuple[int, int], segments: list[np.ndarray], downsample_ratio: int = 1
) -> tuple[np.ndarray, np.ndarray]:
"""Return a downsampled overlap mask and sorted area indices."""
masks = np.zeros(
(imgsz[0] // downsample_ratio, imgsz[1] // downsample_ratio),
dtype=np.int32 if len(segments) > 255 else np.uint8,
)
areas = []
ms = []
for segment in segments:
mask = polygon2mask(
imgsz,
[segment.reshape(-1)],
downsample_ratio=downsample_ratio,
color=1,
)
ms.append(mask.astype(masks.dtype))
areas.append(mask.sum())
areas = np.asarray(areas)
index = np.argsort(-areas)
ms = np.array(ms)[index]
for i in range(len(segments)):
mask = ms[i] * (i + 1)
masks = masks + mask
masks = np.clip(masks, a_min=0, a_max=i + 1)
return masks, index
def find_dataset_yaml(path: Path) -> Path:
"""Find and return the YAML file associated with a Detect, Segment or Pose dataset.
This function searches for a YAML file at the root level of the provided directory first, and if not found, it
performs a recursive search. It prefers YAML files that have the same stem as the provided path.
Args:
path (Path): The directory path to search for the YAML file.
Returns:
(Path): The path of the found YAML file.
"""
files = list(path.glob("*.yaml")) or list(path.rglob("*.yaml")) # try root level first and then recursive
assert files, f"No YAML file found in '{path.resolve()}'"
if len(files) > 1:
files = [f for f in files if f.stem == path.stem] # prefer *.yaml files that match
assert len(files) == 1, f"Expected 1 YAML file in '{path.resolve()}', but found {len(files)}.\n{files}"
return files[0]
def check_det_dataset(dataset: str, autodownload: bool = True) -> dict[str, Any]:
"""Download, verify, and/or unzip a dataset if not found locally.
This function checks the availability of a specified dataset, and if not found, it has the option to download and
unzip the dataset. It then reads and parses the accompanying YAML data, ensuring key requirements are met and also
resolves paths related to the dataset.
Args:
dataset (str): Path to the dataset or dataset descriptor (like a YAML file).
autodownload (bool, optional): Whether to automatically download the dataset if not found.
Returns:
(dict[str, Any]): Parsed dataset information and paths.
"""
file = check_file(dataset)
# Download (optional)
extract_dir = ""
if zipfile.is_zipfile(file) or is_tarfile(file):
new_dir = safe_download(file, dir=DATASETS_DIR, unzip=True, delete=False)
file = find_dataset_yaml(DATASETS_DIR / new_dir)
extract_dir, autodownload = file.parent, False
# Read YAML
data = YAML.load(file, append_filename=True) # dictionary
# Checks
for k in "train", "val":
if k not in data:
if k != "val" or "validation" not in data:
raise SyntaxError(
emojis(f"{dataset} '{k}:' key missing ❌.\n'train' and 'val' are required in all data YAMLs.")
)
LOGGER.warning("renaming data YAML 'validation' key to 'val' to match YOLO format.")
data["val"] = data.pop("validation") # replace 'validation' key with 'val' key
# Support class_map for flexible class name to ID mapping (e.g., ground 2D detection)
if "class_map" in data:
class_map = data["class_map"]
max_id = max(class_map.values())
names = [""] * (max_id + 1)
for name, idx in class_map.items():
if not names[idx]: # Use first name for each ID
names[idx] = name
data["names"] = names
data["nc"] = len(names)
elif "names" not in data and "nc" not in data:
raise SyntaxError(emojis(f"{dataset} key missing ❌.\n either 'names' or 'nc' are required in all data YAMLs."))
if "names" in data and "nc" in data and len(data["names"]) != data["nc"]:
raise SyntaxError(emojis(f"{dataset} 'names' length {len(data['names'])} and 'nc: {data['nc']}' must match."))
if "names" not in data:
data["names"] = [f"class_{i}" for i in range(data["nc"])]
else:
data["nc"] = len(data["names"])
data["names"] = check_class_names(data["names"])
data["channels"] = data.get("channels", 3) # get image channels, default to 3
# Resolve paths
path = Path(extract_dir or data.get("path") or Path(data.get("yaml_file", "")).parent) # dataset root
if not path.exists() and not path.is_absolute():
path = (DATASETS_DIR / path).resolve() # path relative to DATASETS_DIR
# Set paths
data["path"] = path # download scripts
for k in "train", "val", "test", "minival":
if data.get(k): # prepend path
if isinstance(data[k], str):
x = (path / data[k]).resolve()
if not x.exists() and data[k].startswith("../"):
x = (path / data[k][3:]).resolve()
data[k] = str(x)
else:
data[k] = [str((path / x).resolve()) for x in data[k]]
# Parse YAML
val, s = (data.get(x) for x in ("val", "download"))
if val:
val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path
if not all(x.exists() for x in val):
name = clean_url(dataset) # dataset name with URL auth stripped
LOGGER.info("")
m = f"Dataset '{name}' images not found, missing path '{next(x for x in val if not x.exists())}'"
if s and autodownload:
LOGGER.warning(m)
else:
m += f"\nNote dataset download directory is '{DATASETS_DIR}'. You can update this in '{SETTINGS_FILE}'"
raise FileNotFoundError(m)
t = time.time()
r = None # success
if s.startswith("http") and s.endswith(".zip"): # URL
safe_download(url=s, dir=DATASETS_DIR, delete=True)
elif s.startswith("bash "): # bash script
LOGGER.info(f"Running {s} ...")
subprocess.run(s.split(), check=True)
else: # python script
exec(s, {"yaml": data})
dt = f"({round(time.time() - t, 1)}s)"
s = f"success ✅ {dt}, saved to {colorstr('bold', DATASETS_DIR)}" if r in {0, None} else f"failure {dt} ❌"
LOGGER.info(f"Dataset download {s}\n")
check_font("Arial.ttf" if is_ascii(data["names"]) else "Arial.Unicode.ttf") # download fonts
return data # dictionary
def check_cls_dataset(dataset: str | Path, split: str = "") -> dict[str, Any]:
"""Check a classification dataset such as Imagenet.
This function accepts a `dataset` name and attempts to retrieve the corresponding dataset information. If the
dataset is not found locally, it attempts to download the dataset from the internet and save it locally.
Args:
dataset (str | Path): The name of the dataset.
split (str, optional): The split of the dataset. Either 'val', 'test', or ''.
Returns:
(dict[str, Any]): A dictionary containing the following keys:
- 'train' (Path): The directory path containing the training set of the dataset.
- 'val' (Path): The directory path containing the validation set of the dataset.
- 'test' (Path): The directory path containing the test set of the dataset.
- 'nc' (int): The number of classes in the dataset.
- 'names' (dict[int, str]): A dictionary of class names in the dataset.
"""
# Download (optional if dataset=https://file.zip is passed directly)
if str(dataset).startswith(("http:/", "https:/")):
dataset = safe_download(dataset, dir=DATASETS_DIR, unzip=True, delete=False)
elif str(dataset).endswith((".zip", ".tar", ".gz")):
file = check_file(dataset)
dataset = safe_download(file, dir=DATASETS_DIR, unzip=True, delete=False)
dataset = Path(dataset)
data_dir = (dataset if dataset.is_dir() else (DATASETS_DIR / dataset)).resolve()
if not data_dir.is_dir():
if data_dir.suffix != "":
raise ValueError(
f'Classification datasets must be a directory (data="path/to/dir") not a file (data="{dataset}"), '
"See https://docs.ultralytics.com/datasets/classify/"
)
LOGGER.info("")
LOGGER.warning(f"Dataset not found, missing path {data_dir}, attempting download...")
t = time.time()
if str(dataset) == "imagenet":
subprocess.run(["bash", str(ROOT / "data/scripts/get_imagenet.sh")], check=True)
else:
download(f"{ASSETS_URL}/{dataset}.zip", dir=data_dir.parent)
LOGGER.info(f"Dataset download success ✅ ({time.time() - t:.1f}s), saved to {colorstr('bold', data_dir)}\n")
train_set = data_dir / "train"
if not train_set.is_dir():
LOGGER.warning(f"Dataset 'split=train' not found at {train_set}")
if image_files := list(data_dir.rglob("*.jpg")) + list(data_dir.rglob("*.png")):
from ultralytics.data.split import split_classify_dataset
LOGGER.info(f"Found {len(image_files)} images in subdirectories. Attempting to split...")
data_dir = split_classify_dataset(data_dir, train_ratio=0.8)
train_set = data_dir / "train"
else:
LOGGER.error(f"No images found in {data_dir} or its subdirectories.")
val_set = (
data_dir / "val"
if (data_dir / "val").exists()
else data_dir / "validation"
if (data_dir / "validation").exists()
else data_dir / "valid"
if (data_dir / "valid").exists()
else None
) # data/test or data/val
test_set = data_dir / "test" if (data_dir / "test").exists() else None # data/val or data/test
if split == "val" and not val_set:
LOGGER.warning("Dataset 'split=val' not found, using 'split=test' instead.")
val_set = test_set
elif split == "test" and not test_set:
LOGGER.warning("Dataset 'split=test' not found, using 'split=val' instead.")
test_set = val_set
nc = len([x for x in (data_dir / "train").glob("*") if x.is_dir()]) # number of classes
names = [x.name for x in (data_dir / "train").iterdir() if x.is_dir()] # class names list
names = dict(enumerate(sorted(names)))
# Print to console
for k, v in {"train": train_set, "val": val_set, "test": test_set}.items():
prefix = f"{colorstr(f'{k}:')} {v}..."
if v is None:
LOGGER.info(prefix)
else:
files = [path for path in v.rglob("*.*") if path.suffix[1:].lower() in IMG_FORMATS]
nf = len(files) # number of files
nd = len({file.parent for file in files}) # number of directories
if nf == 0:
if k == "train":
raise FileNotFoundError(f"{dataset} '{k}:' no training images found")
else:
LOGGER.warning(f"{prefix} found {nf} images in {nd} classes (no images found)")
elif nd != nc:
LOGGER.error(f"{prefix} found {nf} images in {nd} classes (requires {nc} classes, not {nd})")
else:
LOGGER.info(f"{prefix} found {nf} images in {nd} classes ✅ ")
return {"train": train_set, "val": val_set, "test": test_set, "nc": nc, "names": names, "channels": 3}
class HUBDatasetStats:
"""A class for generating HUB dataset JSON and `-hub` dataset directory.
Args:
path (str): Path to data.yaml or data.zip (with data.yaml inside data.zip).
task (str): Dataset task. Options are 'detect', 'segment', 'pose', 'classify', 'obb'.
autodownload (bool): Attempt to download dataset if not found locally.
Attributes:
task (str): Dataset task type.
hub_dir (Path): Directory path for HUB dataset files.
im_dir (Path): Directory path for compressed images.
stats (dict): Statistics dictionary containing dataset information.
data (dict): Dataset configuration data.
Methods:
get_json: Return dataset JSON for Ultralytics HUB.
process_images: Compress images for Ultralytics HUB.
Examples:
>>> from ultralytics.data.utils import HUBDatasetStats
>>> stats = HUBDatasetStats("path/to/coco8.zip", task="detect") # detect dataset
>>> stats = HUBDatasetStats("path/to/coco8-seg.zip", task="segment") # segment dataset
>>> stats = HUBDatasetStats("path/to/coco8-pose.zip", task="pose") # pose dataset
>>> stats = HUBDatasetStats("path/to/dota8.zip", task="obb") # OBB dataset
>>> stats = HUBDatasetStats("path/to/imagenet10.zip", task="classify") # classification dataset
>>> stats.get_json(save=True)
>>> stats.process_images()
Notes:
Download *.zip files from https://github.com/ultralytics/hub/tree/main/example_datasets
i.e. https://github.com/ultralytics/hub/raw/main/example_datasets/coco8.zip for coco8.zip.
"""
def __init__(self, path: str = "coco8.yaml", task: str = "detect", autodownload: bool = False):
"""Initialize class."""
path = Path(path).resolve()
LOGGER.info(f"Starting HUB dataset checks for {path}....")
self.task = task # detect, segment, pose, classify, obb
if self.task == "classify":
unzip_dir = unzip_file(path)
data = check_cls_dataset(unzip_dir)
data["path"] = unzip_dir
else: # detect, segment, pose, obb
_, data_dir, yaml_path = self._unzip(Path(path))
try:
# Load YAML with checks
data = YAML.load(yaml_path)
data["path"] = "" # strip path since YAML should be in dataset root for all HUB datasets
YAML.save(yaml_path, data)
data = check_det_dataset(yaml_path, autodownload) # dict
data["path"] = data_dir # YAML path should be set to '' (relative) or parent (absolute)
except Exception as e:
raise Exception("error/HUB/dataset_stats/init") from e
self.hub_dir = Path(f"{data['path']}-hub")
self.im_dir = self.hub_dir / "images"
self.stats = {"nc": len(data["names"]), "names": list(data["names"].values())} # statistics dictionary
self.data = data
@staticmethod
def _unzip(path: Path) -> tuple[bool, str, Path]:
"""Unzip data.zip."""
if not str(path).endswith(".zip"): # path is data.yaml
return False, None, path
unzip_dir = unzip_file(path, path=path.parent)
assert unzip_dir.is_dir(), (
f"Error unzipping {path}, {unzip_dir} not found. path/to/abc.zip MUST unzip to path/to/abc/"
)
return True, str(unzip_dir), find_dataset_yaml(unzip_dir) # zipped, data_dir, yaml_path
def _hub_ops(self, f: str):
"""Save a compressed image for HUB previews."""
compress_one_image(f, self.im_dir / Path(f).name) # save to dataset-hub
def get_json(self, save: bool = False, verbose: bool = False) -> dict:
"""Return dataset JSON for Ultralytics HUB."""
def _round(labels):
"""Update labels to integer class and 4 decimal place floats."""
if self.task == "detect":
coordinates = labels["bboxes"]
elif self.task in {"segment", "obb"}: # Segment and OBB use segments. OBB segments are normalized xyxyxyxy
coordinates = [x.flatten() for x in labels["segments"]]
elif self.task == "pose":
n, nk, nd = labels["keypoints"].shape
coordinates = np.concatenate((labels["bboxes"], labels["keypoints"].reshape(n, nk * nd)), 1)
else:
raise ValueError(f"Undefined dataset task={self.task}.")
zipped = zip(labels["cls"], coordinates)
return [[int(c[0]), *(round(float(x), 4) for x in points)] for c, points in zipped]
for split in "train", "val", "test":
self.stats[split] = None # predefine
path = self.data.get(split)
# Check split
if path is None: # no split
continue
files = [f for f in Path(path).rglob("*.*") if f.suffix[1:].lower() in IMG_FORMATS] # image files in split
if not files: # no images
continue
# Get dataset statistics
if self.task == "classify":
from torchvision.datasets import ImageFolder # scope for faster 'import ultralytics'
dataset = ImageFolder(self.data[split])
x = np.zeros(len(dataset.classes)).astype(int)
for im in dataset.imgs:
x[im[1]] += 1
self.stats[split] = {
"instance_stats": {"total": len(dataset), "per_class": x.tolist()},
"image_stats": {"total": len(dataset), "unlabelled": 0, "per_class": x.tolist()},
"labels": [{Path(k).name: v} for k, v in dataset.imgs],
}
else:
from ultralytics.data import YOLODataset
dataset = YOLODataset(img_path=self.data[split], data=self.data, task=self.task)
x = np.array(
[
np.bincount(label["cls"].astype(int).flatten(), minlength=self.data["nc"])
for label in TQDM(dataset.labels, total=len(dataset), desc="Statistics")
]
) # shape(128x80)
self.stats[split] = {
"instance_stats": {"total": int(x.sum()), "per_class": x.sum(0).tolist()},
"image_stats": {
"total": len(dataset),
"unlabelled": int(np.all(x == 0, 1).sum()),
"per_class": (x > 0).sum(0).tolist(),
},
"labels": [{Path(k).name: _round(v)} for k, v in zip(dataset.im_files, dataset.labels)],
}
# Save, print and return
if save:
self.hub_dir.mkdir(parents=True, exist_ok=True) # makes dataset-hub/
stats_path = self.hub_dir / "stats.json"
LOGGER.info(f"Saving {stats_path.resolve()}...")
with open(stats_path, "w", encoding="utf-8") as f:
json.dump(self.stats, f) # save stats.json
if verbose:
LOGGER.info(json.dumps(self.stats, indent=2, sort_keys=False))
return self.stats
def process_images(self) -> Path:
"""Compress images for Ultralytics HUB."""
from ultralytics.data import YOLODataset # ClassificationDataset
self.im_dir.mkdir(parents=True, exist_ok=True) # makes dataset-hub/images/
for split in "train", "val", "test":
if self.data.get(split) is None:
continue
dataset = YOLODataset(img_path=self.data[split], data=self.data)
with ThreadPool(NUM_THREADS) as pool:
for _ in TQDM(pool.imap(self._hub_ops, dataset.im_files), total=len(dataset), desc=f"{split} images"):
pass
LOGGER.info(f"Done. All images saved to {self.im_dir}")
return self.im_dir
def compress_one_image(f: str, f_new: str | None = None, max_dim: int = 1920, quality: int = 50):
"""Compress a single image file to reduced size while preserving its aspect ratio and quality using either the
Python Imaging Library (PIL) or OpenCV library. If the input image is smaller than the maximum dimension, it
will not be resized.
Args:
f (str): The path to the input image file.
f_new (str, optional): The path to the output image file. If not specified, the input file will be overwritten.
max_dim (int, optional): The maximum dimension (width or height) of the output image.
quality (int, optional): The image compression quality as a percentage.
Examples:
>>> from pathlib import Path
>>> from ultralytics.data.utils import compress_one_image
>>> for f in Path("path/to/dataset").rglob("*.jpg"):
>>> compress_one_image(f)
"""
try: # use PIL
Image.MAX_IMAGE_PIXELS = None # Fix DecompressionBombError, allow optimization of image > ~178.9 million pixels
im = Image.open(f)
if im.mode in {"RGBA", "LA"}: # Convert to RGB if needed (for JPEG)
im = im.convert("RGB")
r = max_dim / max(im.height, im.width) # ratio
if r < 1.0: # image too large
im = im.resize((int(im.width * r), int(im.height * r)))
im.save(f_new or f, "JPEG", quality=quality, optimize=True) # save
except Exception as e: # use OpenCV
LOGGER.warning(f"HUB ops PIL failure {f}: {e}")
im = cv2.imread(f)
im_height, im_width = im.shape[:2]
r = max_dim / max(im_height, im_width) # ratio
if r < 1.0: # image too large
im = cv2.resize(im, (int(im_width * r), int(im_height * r)), interpolation=cv2.INTER_AREA)
cv2.imwrite(str(f_new or f), im)
def load_dataset_cache_file(path: Path) -> dict:
"""Load an Ultralytics *.cache dictionary from path."""
import gc
gc.disable() # reduce pickle load time https://github.com/ultralytics/ultralytics/pull/1585
cache = np.load(str(path), allow_pickle=True).item() # load dict
gc.enable()
return cache
def save_dataset_cache_file(prefix: str, path: Path, x: dict, version: str):
"""Save an Ultralytics dataset *.cache dictionary x to path."""
x["version"] = version # add cache version
if is_dir_writeable(path.parent):
if path.exists():
path.unlink() # remove *.cache file if exists
with open(str(path), "wb") as file: # context manager here fixes windows async np.save bug
np.save(file, x)
LOGGER.info(f"{prefix}New cache created: {path}")
else:
LOGGER.warning(f"{prefix}Cache directory {path.parent} is not writable, cache not saved.")