dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__py3-none-any.whl

ultralytics/data/base.py

@@ -1,5 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 import glob
 import math
 import os
@@ -7,7 +9,7 @@ import random
 from copy import deepcopy
 from multiprocessing.pool import ThreadPool
 from pathlib import Path
-from typing import Optional
+from typing import Any
 
 import cv2
 import numpy as np
@@ -19,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.
@@ -32,9 +33,10 @@ class BaseDataset(Dataset):
         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 RGB).
         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.
+        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.
@@ -45,9 +47,11 @@ class BaseDataset(Dataset):
         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.
+        npy_files (list[Path]): List of numpy file paths.
         cache (str): Cache images to RAM or disk during training.
         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.
@@ -66,39 +70,38 @@ class BaseDataset(Dataset):
 
     def __init__(
         self,
-        img_path,
-        imgsz=640,
-        cache=False,
-        augment=True,
-        hyp=DEFAULT_CFG,
-        prefix="",
-        rect=False,
-        batch_size=16,
-        stride=32,
-        pad=0.5,
-        single_cls=False,
-        classes=None,
-        fraction=1.0,
-        channels=3,
+        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,
     ):
-        """
-        Initialize BaseDataset with given configuration and options.
+        """Initialize BaseDataset with given configuration and options.
 
         Args:
-            img_path (str): Path to the folder containing images.
-            imgsz (int, optional): Image size for resizing.
-            cache (bool | str, optional): Cache images to RAM or disk during training.
-            augment (bool, optional): If True, data augmentation is applied.
-            hyp (dict, optional): Hyperparameters to apply data augmentation.
-            prefix (str, optional): Prefix to print in log messages.
-            rect (bool, optional): If True, rectangular training is used.
-            batch_size (int, optional): Size of batches.
-            stride (int, optional): Stride used in the model.
-            pad (float, optional): Padding value.
-            single_cls (bool, optional): If True, single class training is used.
-            classes (list, optional): List of included classes.
-            fraction (float, optional): Fraction of dataset to utilize.
-            channels (int, optional): Number of channels in the images (1 for grayscale, 3 for RGB).
+            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 RGB).
         """
         super().__init__()
         self.img_path = img_path
@@ -142,15 +145,14 @@ class BaseDataset(Dataset):
         # Transforms
         self.transforms = self.build_transforms(hyp=hyp)
 
-    def get_img_files(self, img_path):
-        """
-        Read image files from the specified path.
+    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.
+            img_path (str | list[str]): Path or list of paths to image directories or files.
 
         Returns:
-            (List[str]): List of image file paths.
+            (list[str]): List of image file paths.
 
         Raises:
             FileNotFoundError: If no images are found or the path doesn't exist.
@@ -170,7 +172,7 @@ class BaseDataset(Dataset):
                         # 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.split(".")[-1].lower() in IMG_FORMATS)
+            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:
@@ -180,12 +182,11 @@ class BaseDataset(Dataset):
         check_file_speeds(im_files, prefix=self.prefix)  # check image read speeds
         return im_files
 
-    def update_labels(self, include_class: Optional[list]):
-        """
-        Update labels to include only specified classes.
+    def update_labels(self, include_class: list[int] | None) -> None:
+        """Update labels to include only specified classes.
 
         Args:
-            include_class (list, optional): List of classes to include. If None, all classes are included.
+            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)):
@@ -204,18 +205,17 @@ class BaseDataset(Dataset):
             if self.single_cls:
                 self.labels[i]["cls"][:, 0] = 0
 
-    def load_image(self, i, rect_mode=True):
-        """
-        Load an image from dataset index 'i'.
+    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, optional): Whether to use rectangular resizing.
+            rect_mode (bool): Whether to use rectangular resizing.
 
         Returns:
-            (np.ndarray): Loaded image as a NumPy array.
-            (Tuple[int, int]): Original image dimensions in (height, width) format.
-            (Tuple[int, int]): Resized image dimensions in (height, width) format.
+            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.
@@ -258,7 +258,7 @@ class BaseDataset(Dataset):
 
         return self.ims[i], self.im_hw0[i], self.im_hw[i]
 
-    def cache_images(self):
+    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")
@@ -274,18 +274,17 @@ class BaseDataset(Dataset):
                 pbar.desc = f"{self.prefix}Caching images ({b / gb:.1f}GB {storage})"
             pbar.close()
 
-    def cache_images_to_disk(self, i):
+    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]), allow_pickle=False)
 
-    def check_cache_disk(self, safety_margin=0.5):
-        """
-        Check if there's enough disk space for caching images.
+    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, optional): Safety margin factor for disk space calculation.
+            safety_margin (float): Safety margin factor for disk space calculation.
 
         Returns:
             (bool): True if there's enough disk space, False otherwise.
@@ -302,10 +301,10 @@ class BaseDataset(Dataset):
             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 writeable")
+                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)
+        total, _used, free = shutil.disk_usage(Path(self.im_files[0]).parent)
         if disk_required > free:
             self.cache = None
             LOGGER.warning(
@@ -316,12 +315,11 @@ class BaseDataset(Dataset):
             return False
         return True
 
-    def check_cache_ram(self, safety_margin=0.5):
-        """
-        Check if there's enough RAM for caching images.
+    def check_cache_ram(self, safety_margin: float = 0.5) -> bool:
+        """Check if there's enough RAM for caching images.
 
         Args:
-            safety_margin (float, optional): Safety margin factor for RAM calculation.
+            safety_margin (float): Safety margin factor for RAM calculation.
 
         Returns:
             (bool): True if there's enough RAM, False otherwise.
@@ -346,7 +344,7 @@ class BaseDataset(Dataset):
             return False
         return True
 
-    def set_rectangle(self):
+    def set_rectangle(self) -> None:
         """Set the shape of bounding boxes for YOLO detections as rectangles."""
         bi = np.floor(np.arange(self.ni) / self.batch_size).astype(int)  # batch index
         nb = bi[-1] + 1  # number of batches
@@ -371,19 +369,18 @@ class BaseDataset(Dataset):
         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):
+    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):
-        """
-        Get and return label information from the dataset.
+    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): Label dictionary with image and metadata.
+            (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
@@ -396,17 +393,16 @@ class BaseDataset(Dataset):
             label["rect_shape"] = self.batch_shapes[self.batch[index]]
         return self.update_labels_info(label)
 
-    def __len__(self):
+    def __len__(self) -> int:
         """Return the length of the labels list for the dataset."""
         return len(self.labels)
 
-    def update_labels_info(self, label):
+    def update_labels_info(self, label: dict[str, Any]) -> dict[str, Any]:
         """Custom your label format here."""
         return label
 
-    def build_transforms(self, hyp=None):
-        """
-        Users can customize augmentations here.
+    def build_transforms(self, hyp: dict[str, Any] | None = None):
+        """Users can customize augmentations here.
 
         Examples:
             >>> if self.augment:
@@ -418,9 +414,8 @@ class BaseDataset(Dataset):
         """
         raise NotImplementedError
 
-    def get_labels(self):
-        """
-        Users can customize their own format here.
+    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: