dgenerate-ultralytics-headless 8.3.222__py3-none-any.whl → 8.3.225__py3-none-any.whl

ultralytics/data/base.py

@@ -21,11 +21,10 @@ from ultralytics.utils.patches import imread
 
 
 class BaseDataset(Dataset):
-    """
-    Base dataset class for loading and processing image data.
+    """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.
+    This class provides core functionality for loading images, caching, and preparing data for training and inference in
+    object detection tasks.
 
     Attributes:
         img_path (str): Path to the folder containing images.
@@ -86,8 +85,7 @@ class BaseDataset(Dataset):
         fraction: float = 1.0,
         channels: int = 3,
     ):
-        """
-        Initialize BaseDataset with given configuration and options.
+        """Initialize BaseDataset with given configuration and options.
 
         Args:
             img_path (str | list[str]): Path to the folder containing images or list of image paths.
@@ -148,8 +146,7 @@ class BaseDataset(Dataset):
         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.
+        """Read image files from the specified path.
 
         Args:
             img_path (str | list[str]): Path or list of paths to image directories or files.
@@ -186,8 +183,7 @@ class BaseDataset(Dataset):
         return im_files
 
     def update_labels(self, include_class: list[int] | None) -> None:
-        """
-        Update labels to include only specified classes.
+        """Update labels to include only specified classes.
 
         Args:
             include_class (list[int], optional): List of classes to include. If None, all classes are included.
@@ -210,8 +206,7 @@ class BaseDataset(Dataset):
                 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'.
+        """Load an image from dataset index 'i'.
 
         Args:
             i (int): Index of the image to load.
@@ -286,8 +281,7 @@ class BaseDataset(Dataset):
             np.save(f.as_posix(), imread(self.im_files[i]), allow_pickle=False)
 
     def check_cache_disk(self, safety_margin: float = 0.5) -> bool:
-        """
-        Check if there's enough disk space for caching images.
+        """Check if there's enough disk space for caching images.
 
         Args:
             safety_margin (float): Safety margin factor for disk space calculation.
@@ -322,8 +316,7 @@ class BaseDataset(Dataset):
         return True
 
     def check_cache_ram(self, safety_margin: float = 0.5) -> bool:
-        """
-        Check if there's enough RAM for caching images.
+        """Check if there's enough RAM for caching images.
 
         Args:
             safety_margin (float): Safety margin factor for RAM calculation.
@@ -381,8 +374,7 @@ class BaseDataset(Dataset):
         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.
+        """Get and return label information from the dataset.
 
         Args:
             index (int): Index of the image to retrieve.
@@ -410,8 +402,7 @@ class BaseDataset(Dataset):
         return label
 
     def build_transforms(self, hyp: dict[str, Any] | None = None):
-        """
-        Users can customize augmentations here.
+        """Users can customize augmentations here.
 
         Examples:
             >>> if self.augment:
@@ -424,8 +415,7 @@ class BaseDataset(Dataset):
         raise NotImplementedError
 
     def get_labels(self) -> list[dict[str, Any]]:
-        """
-        Users can customize their own format here.
+        """Users can customize their own format here.
 
         Examples:
             Ensure output is a dictionary with the following keys:

ultralytics/data/build.py

@@ -14,7 +14,7 @@ import numpy as np
 import torch
 import torch.distributed as dist
 from PIL import Image
-from torch.utils.data import dataloader, distributed
+from torch.utils.data import Dataset, dataloader, distributed
 
 from ultralytics.cfg import IterableSimpleNamespace
 from ultralytics.data.dataset import GroundingDataset, YOLODataset, YOLOMultiModalDataset
@@ -35,8 +35,7 @@ from ultralytics.utils.torch_utils import TORCH_2_0
 
 
 class InfiniteDataLoader(dataloader.DataLoader):
-    """
-    Dataloader that reuses workers for infinite iteration.
+    """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.
@@ -94,11 +93,10 @@ class InfiniteDataLoader(dataloader.DataLoader):
 
 
 class _RepeatSampler:
-    """
-    Sampler that repeats forever for infinite iteration.
+    """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.
+    This sampler wraps another sampler and yields its contents indefinitely, allowing for infinite iteration over a
+    dataset without recreating the sampler.
 
     Attributes:
         sampler (Dataset.sampler): The sampler to repeat.
@@ -115,27 +113,26 @@ class _RepeatSampler:
 
 
 class ContiguousDistributedSampler(torch.utils.data.Sampler):
-    """
-    Distributed sampler that assigns contiguous batch-aligned chunks of the dataset to each GPU.
+    """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).
+    [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.
+    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 (torch.utils.data.Dataset): Dataset to sample from. Must implement __len__.
+        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.
         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.
+        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.
 
-    Example:
+    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)
@@ -147,7 +144,14 @@ class ContiguousDistributedSampler(torch.utils.data.Sampler):
         ...         ...
     """
 
-    def __init__(self, dataset, num_replicas=None, batch_size=None, rank=None, shuffle=False):
+    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
@@ -156,7 +160,6 @@ class ContiguousDistributedSampler(torch.utils.data.Sampler):
         if batch_size is None:
             batch_size = getattr(dataset, "batch_size", 1)
 
-        self.dataset = dataset
         self.num_replicas = num_replicas
         self.batch_size = batch_size
         self.rank = rank
@@ -165,7 +168,7 @@ class ContiguousDistributedSampler(torch.utils.data.Sampler):
         self.total_size = len(dataset)
         self.num_batches = math.ceil(self.total_size / self.batch_size)
 
-    def _get_rank_indices(self):
+    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
@@ -184,7 +187,7 @@ class ContiguousDistributedSampler(torch.utils.data.Sampler):
 
         return start_idx, end_idx
 
-    def __iter__(self):
+    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))
@@ -196,14 +199,13 @@ class ContiguousDistributedSampler(torch.utils.data.Sampler):
 
         return iter(indices)
 
-    def __len__(self):
+    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):
-        """
-        Set the epoch for this sampler to ensure different shuffling patterns across epochs.
+    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.
@@ -211,7 +213,7 @@ class ContiguousDistributedSampler(torch.utils.data.Sampler):
         self.epoch = epoch
 
 
-def seed_worker(worker_id: int):
+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)
@@ -227,7 +229,7 @@ def build_yolo_dataset(
     rect: bool = False,
     stride: int = 32,
     multi_modal: bool = False,
-):
+) -> Dataset:
     """Build and return a YOLO dataset based on configuration parameters."""
     dataset = YOLOMultiModalDataset if multi_modal else YOLODataset
     return dataset(
@@ -258,7 +260,7 @@ def build_grounding(
     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,
@@ -288,9 +290,8 @@ def build_dataloader(
     rank: int = -1,
     drop_last: bool = False,
     pin_memory: bool = True,
-):
-    """
-    Create and return an InfiniteDataLoader or DataLoader for training or validation.
+) -> InfiniteDataLoader:
+    """Create and return an InfiniteDataLoader or DataLoader for training or validation.
 
     Args:
         dataset (Dataset): Dataset to load data from.
@@ -336,9 +337,10 @@ def build_dataloader(
     )
 
 
-def check_source(source):
-    """
-    Check the type of input source and return corresponding flag values.
+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.
@@ -385,12 +387,17 @@ def check_source(source):
     return source, webcam, screenshot, from_img, in_memory, tensor
 
 
-def load_inference_source(source=None, batch: int = 1, vid_stride: int = 1, buffer: bool = False, channels: int = 3):
-    """
-    Load an inference source for object detection and apply necessary transformations.
+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 | Path | torch.Tensor | PIL.Image | np.ndarray, optional): The input source for inference.
+        source (str | Path | list | tuple | torch.Tensor | PIL.Image | np.ndarray): 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.

ultralytics/data/converter.py

@@ -21,12 +21,11 @@ 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.
+    """Convert 91-index COCO class IDs to 80-index COCO class IDs.
 
     Returns:
-        (list[int]): A list of 91 class IDs where the index represents the 80-index class ID and the value
-            is the corresponding 91-index class ID.
+        (list[int]): A list of 91 class IDs where the index represents the 80-index class ID and the value is the
+            corresponding 91-index class ID.
     """
     return [
         0,
@@ -124,15 +123,11 @@ def coco91_to_coco80_class() -> list[int]:
 
 
 def coco80_to_coco91_class() -> list[int]:
-    r"""
-    Convert 80-index (val2014) to 91-index (paper).
+    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.
 
-    References:
-        https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
-
     Examples:
         >>> import numpy as np
         >>> a = np.loadtxt("data/coco.names", dtype="str", delimiter="\n")
@@ -143,6 +138,9 @@ def coco80_to_coco91_class() -> list[int]:
 
         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,
@@ -236,8 +234,7 @@ def convert_coco(
     cls91to80: bool = True,
     lvis: bool = False,
 ):
-    """
-    Convert COCO dataset annotations to a YOLO annotation format suitable for training YOLO models.
+    """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.
@@ -348,8 +345,7 @@ def convert_coco(
 
 
 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.
+    """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.
@@ -424,8 +420,7 @@ def convert_segment_masks_to_yolo_seg(masks_dir: str, output_dir: str, classes:
 
 
 def convert_dota_to_yolo_obb(dota_root_path: str):
-    """
-    Convert DOTA dataset annotations to YOLO OBB (Oriented Bounding Box) format.
+    """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.
@@ -517,8 +512,7 @@ def convert_dota_to_yolo_obb(dota_root_path: str):
 
 
 def min_index(arr1: np.ndarray, arr2: np.ndarray):
-    """
-    Find a pair of indexes with the shortest distance between two arrays of 2D points.
+    """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.
@@ -533,14 +527,14 @@ def min_index(arr1: np.ndarray, arr2: np.ndarray):
 
 
 def merge_multi_segment(segments: list[list]):
-    """
-    Merge multiple segments into one list by connecting the coordinates with the minimum distance between each segment.
+    """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,...].
+        segments (list[list]): Original segmentations in COCO's JSON file. Each element is a list of coordinates, like
+            [segmentation1, segmentation2,...].
 
     Returns:
         s (list[np.ndarray]): A list of connected segments represented as NumPy arrays.
@@ -584,14 +578,13 @@ def merge_multi_segment(segments: list[list]):
 
 
 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 or oriented bounding box (OBB) in
-    YOLO format. Generate segmentation data using SAM auto-annotator as needed.
+    """Convert existing object detection dataset (bounding boxes) to segmentation dataset or oriented bounding box (OBB)
+    in YOLO format. Generate 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.
+        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.
 
@@ -648,12 +641,11 @@ def yolo_bbox2segment(im_dir: str | Path, save_dir: str | Path | None = None, sa
 
 
 def create_synthetic_coco_dataset():
-    """
-    Create a synthetic COCO dataset with random images based on filenames from label lists.
+    """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.
+    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
@@ -704,11 +696,10 @@ def create_synthetic_coco_dataset():
 
 
 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.
+    """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.
+    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.
@@ -756,12 +747,11 @@ def convert_to_multispectral(path: str | Path, n_channels: int = 10, replace: bo
 
 
 async def convert_ndjson_to_yolo(ndjson_path: str | Path, output_path: str | Path | None = None) -> Path:
-    """
-    Convert NDJSON dataset format to Ultralytics YOLO11 dataset structure.
+    """Convert NDJSON dataset format to Ultralytics YOLO11 dataset structure.
 
-    This function converts datasets stored in NDJSON (Newline Delimited JSON) format to the standard YOLO
-    format with separate directories for images and labels. It supports parallel processing for efficient
-    conversion of large datasets and can download images from URLs if they don't exist locally.
+    This function converts datasets stored in NDJSON (Newline Delimited JSON) format to the standard YOLO format with
+    separate directories for images and labels. It supports parallel processing for efficient conversion of large
+    datasets and can download images from URLs if they don't exist locally.
 
     The NDJSON format consists of:
     - First line: Dataset metadata with class names and configuration
@@ -769,8 +759,8 @@ async def convert_ndjson_to_yolo(ndjson_path: str | Path, output_path: str | Pat
 
     Args:
         ndjson_path (Union[str, Path]): Path to the input NDJSON file containing dataset information.
-        output_path (Optional[Union[str, Path]], optional): Directory where the converted YOLO dataset
-            will be saved. If None, uses the parent directory of the NDJSON file. Defaults to None.
+        output_path (Optional[Union[str, Path]], optional): Directory where the converted YOLO dataset will be saved. If
+            None, uses the parent directory of the NDJSON file. Defaults to None.
 
     Returns:
         (Path): Path to the generated data.yaml file that can be used for YOLO training.