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

ultralytics/data/loaders.py

@@ -1,5 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 import glob
 import math
 import os
@@ -8,6 +10,7 @@ import urllib
 from dataclasses import dataclass
 from pathlib import Path
 from threading import Thread
+from typing import Any
 
 import cv2
 import numpy as np
@@ -22,11 +25,10 @@ from ultralytics.utils.patches import imread
 
 @dataclass
 class SourceTypes:
-    """
-    Class to represent various types of input sources for predictions.
+    """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.
+    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.
@@ -49,24 +51,23 @@ class SourceTypes:
 
 
 class LoadStreams:
-    """
-    Stream Loader for various types of video streams.
+    """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.
+    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.
+        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.
+        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 RGB).
 
@@ -90,8 +91,15 @@ class LoadStreams:
         - The class implements a buffer system to manage frame storage and retrieval.
     """
 
-    def __init__(self, sources="file.streams", vid_stride=1, buffer=False, channels=3):
-        """Initialize stream loader for multiple video sources, supporting various stream types."""
+    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 RGB).
+        """
         torch.backends.cudnn.benchmark = True  # faster for fixed-size inference
         self.buffer = buffer  # buffer input streams
         self.running = True  # running flag for Thread
@@ -143,7 +151,7 @@ class LoadStreams:
             self.threads[i].start()
         LOGGER.info("")  # newline
 
-    def update(self, i, cap, stream):
+    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):
@@ -167,7 +175,7 @@ class LoadStreams:
                 time.sleep(0.01)  # wait until the buffer is empty
 
     def close(self):
-        """Terminates stream loader, stops threads, and releases video capture resources."""
+        """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():
@@ -177,22 +185,21 @@ class LoadStreams:
                 cap.release()  # release video capture
             except Exception as e:
                 LOGGER.warning(f"Could not release VideoCapture object: {e}")
-        cv2.destroyAllWindows()
 
     def __iter__(self):
-        """Iterates through YOLO image feed and re-opens unresponsive streams."""
+        """Iterate through YOLO image feed and re-open unresponsive streams."""
         self.count = -1
         return self
 
-    def __next__(self):
-        """Returns the next batch of frames from multiple video streams for processing."""
+    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() or cv2.waitKey(1) == ord("q"):  # q to quit
+                if not self.threads[i].is_alive():
                     self.close()
                     raise StopIteration
                 time.sleep(1 / min(self.fps))
@@ -211,17 +218,16 @@ class LoadStreams:
 
         return self.sources, images, [""] * self.bs
 
-    def __len__(self):
+    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.
+    """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`.
+    This class manages the loading of screenshot images for processing with YOLO. It is suitable for use with `yolo
+    predict source=screen`.
 
     Attributes:
         source (str): The source input indicating which screen to capture.
@@ -235,7 +241,7 @@ class LoadScreenshots:
         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.
+        monitor (dict[str, int]): Monitor configuration details.
         cv2_flag (int): OpenCV flag for image reading (grayscale or RGB).
 
     Methods:
@@ -248,10 +254,15 @@ class LoadScreenshots:
         ...     print(f"Captured frame: {im.shape}")
     """
 
-    def __init__(self, source, channels=3):
-        """Initialize screenshot capture with specified screen and region parameters."""
+    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 RGB).
+        """
         check_requirements("mss")
-        import mss  # noqa
+        import mss
 
         source, *params = source.split()
         self.screen, left, top, width, height = 0, None, None, None, None  # default to full screen 0
@@ -277,11 +288,11 @@ class LoadScreenshots:
         self.monitor = {"left": self.left, "top": self.top, "width": self.width, "height": self.height}
 
     def __iter__(self):
-        """Yields the next screenshot image from the specified screen or region for processing."""
+        """Yield the next screenshot image from the specified screen or region for processing."""
         return self
 
-    def __next__(self):
-        """Captures and returns the next screenshot as a numpy array using the mss library."""
+    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}: "
@@ -291,16 +302,15 @@ class LoadScreenshots:
 
 
 class LoadImagesAndVideos:
-    """
-    A class for loading and processing images and videos for YOLO object detection.
+    """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.
+    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.
+        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).
+        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.
@@ -330,12 +340,20 @@ class LoadImagesAndVideos:
         - Can read from a text file containing paths to images and videos.
     """
 
-    def __init__(self, path, batch=1, vid_stride=1, channels=3):
-        """Initialize dataloader for images and videos, supporting various input formats."""
+    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 RGB).
+        """
         parent = None
-        if isinstance(path, str) and Path(path).suffix == ".txt":  # *.txt file with img/vid/dir on each line
-            parent = Path(path).parent
-            path = Path(path).read_text().splitlines()  # list of sources
+        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
@@ -353,7 +371,7 @@ class LoadImagesAndVideos:
         # Define files as images or videos
         images, videos = [], []
         for f in files:
-            suffix = f.split(".")[-1].lower()  # Get file extension without the dot and lowercase
+            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:
@@ -376,12 +394,12 @@ class LoadImagesAndVideos:
             raise FileNotFoundError(f"No images or videos found in {p}. {FORMATS_HELP_MSG}")
 
     def __iter__(self):
-        """Iterates through image/video files, yielding source paths, images, and metadata."""
+        """Iterate through image/video files, yielding source paths, images, and metadata."""
         self.count = 0
         return self
 
-    def __next__(self):
-        """Returns the next batch of images or video frames with their paths and metadata."""
+    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
@@ -427,11 +445,11 @@ class LoadImagesAndVideos:
             else:
                 # Handle image files (including HEIC)
                 self.mode = "image"
-                if path.split(".")[-1].lower() == "heic":
+                if path.rpartition(".")[-1].lower() == "heic":
                     # Load HEIC image using Pillow with pillow-heif
-                    check_requirements("pillow-heif")
+                    check_requirements("pi-heif")
 
-                    from pillow_heif import register_heif_opener
+                    from pi_heif import register_heif_opener
 
                     register_heif_opener()  # Register HEIF opener with Pillow
                     with Image.open(path) as img:
@@ -450,8 +468,8 @@ class LoadImagesAndVideos:
 
         return paths, imgs, info
 
-    def _new_video(self, path):
-        """Creates a new video capture object for the given path and initializes video-related attributes."""
+    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))
@@ -459,21 +477,20 @@ class LoadImagesAndVideos:
             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):
-        """Returns the number of files (images and videos) in the dataset."""
+    def __len__(self) -> int:
+        """Return the number of files (images and videos) 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.
+    """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.
+        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`.
 
@@ -491,8 +508,13 @@ class LoadPilAndNumpy:
         Loaded 2 images
     """
 
-    def __init__(self, im0, channels=3):
-        """Initializes a loader for PIL and Numpy images, converting inputs to a standardized format."""
+    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 RGB).
+        """
         if not isinstance(im0, list):
             im0 = [im0]
         # use `image{i}.jpg` when Image.filename returns an empty path.
@@ -503,7 +525,7 @@ class LoadPilAndNumpy:
         self.bs = len(self.im0)
 
     @staticmethod
-    def _single_check(im, flag="RGB"):
+    def _single_check(im: Image.Image | np.ndarray, flag: str = "RGB") -> np.ndarray:
         """Validate and format an image to numpy array, ensuring RGB order and contiguous memory."""
         assert isinstance(im, (Image.Image, np.ndarray)), f"Expected PIL/np.ndarray image type, but got {type(im)}"
         if isinstance(im, Image.Image):
@@ -515,35 +537,34 @@ class LoadPilAndNumpy:
             im = im[..., None]
         return im
 
-    def __len__(self):
-        """Returns the length of the 'im0' attribute, representing the number of loaded images."""
+    def __len__(self) -> int:
+        """Return the length of the 'im0' attribute, representing the number of loaded images."""
         return len(self.im0)
 
-    def __next__(self):
-        """Returns the next batch of images, paths, and metadata for processing."""
+    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):
-        """Iterates through PIL/numpy images, yielding paths, raw images, and metadata for processing."""
+        """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.
+    """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.
+    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.
+        paths (list[str]): List of image paths or auto-generated filenames.
 
     Methods:
         _single_check: Validates and formats an input tensor.
@@ -556,16 +577,20 @@ class LoadTensor:
         >>> print(f"Processed {len(images)} images")
     """
 
-    def __init__(self, im0) -> None:
-        """Initialize LoadTensor object for processing torch.Tensor image data."""
+    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, stride=32):
-        """Validates and formats a single image tensor, ensuring correct shape and normalization."""
+    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."
@@ -586,24 +611,24 @@ class LoadTensor:
         return im
 
     def __iter__(self):
-        """Yields an iterator object for iterating through tensor image data."""
+        """Yield an iterator object for iterating through tensor image data."""
         self.count = 0
         return self
 
-    def __next__(self):
-        """Yields the next batch of tensor images and metadata for processing."""
+    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):
-        """Returns the batch size of the tensor input."""
+    def __len__(self) -> int:
+        """Return the batch size of the tensor input."""
         return self.bs
 
 
-def autocast_list(source):
-    """Merges a list of sources into a list of numpy arrays or PIL images for Ultralytics prediction."""
+def autocast_list(source: list[Any]) -> list[Image.Image | np.ndarray]:
+    """Merge 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
@@ -619,14 +644,12 @@ def autocast_list(source):
     return files
 
 
-def get_best_youtube_url(url, method="pytube"):
-    """
-    Retrieves the URL of the best quality MP4 video stream from a given YouTube video.
+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".
-            Defaults to "pytube".
 
     Returns:
         (str | None): The URL of the best quality MP4 video stream, or None if no suitable stream is found.
@@ -655,7 +678,7 @@ def get_best_youtube_url(url, method="pytube"):
 
     elif method == "pafy":
         check_requirements(("pafy", "youtube_dl==2020.12.2"))
-        import pafy  # noqa
+        import pafy
 
         return pafy.new(url).getbestvideo(preftype="mp4").url
 

ultralytics/data/split.py

@@ -1,5 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 import random
 import shutil
 from pathlib import Path
@@ -8,12 +10,11 @@ from ultralytics.data.utils import IMG_FORMATS, img2label_paths
 from ultralytics.utils import DATASETS_DIR, LOGGER, TQDM
 
 
-def split_classify_dataset(source_dir, train_ratio=0.8):
-    """
-    Split dataset into train and val directories in a new directory.
+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.
+    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:
@@ -46,13 +47,17 @@ def split_classify_dataset(source_dir, train_ratio=0.8):
                 └── ...
 
     Args:
-        source_dir (str | Path): Path to Caltech dataset root directory.
+        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 dataset with default 80/20 ratio
         >>> split_classify_dataset("path/to/caltech")
-        >>> # Split with custom ratio
+
+        Split with custom ratio
         >>> split_classify_dataset("path/to/caltech", 0.75)
     """
     source_path = Path(source_dir)
@@ -90,18 +95,26 @@ def split_classify_dataset(source_dir, train_ratio=0.8):
     return split_path
 
 
-def autosplit(path=DATASETS_DIR / "coco8/images", weights=(0.9, 0.1, 0.0), annotated_only=False):
-    """
-    Automatically split a dataset into train/val/test splits and save the resulting splits into autosplit_*.txt files.
+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, optional): Path to images directory.
-        weights (list | tuple, optional): Train, validation, and test split fractions.
-        annotated_only (bool, optional): If True, only images with an associated txt file are used.
+        path (Path): Path to images directory.
+        weights (tuple): 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
@@ -122,4 +135,4 @@ def autosplit(path=DATASETS_DIR / "coco8/images", weights=(0.9, 0.1, 0.0), annot
 
 
 if __name__ == "__main__":
-    split_classify_dataset("../datasets/caltech101")
+    split_classify_dataset("caltech101")