ultralytics 8.1.29__py3-none-any.whl → 8.3.63__py3-none-any.whl

ultralytics/data/loaders.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import glob
 import math
@@ -15,14 +15,32 @@ import requests
 import torch
 from PIL import Image
 
-from ultralytics.data.utils import IMG_FORMATS, VID_FORMATS
-from ultralytics.utils import LOGGER, is_colab, is_kaggle, ops
+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."""
+    """
+    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.
+
+    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
@@ -32,38 +50,47 @@ class SourceTypes:
 
 class LoadStreams:
     """
-    Stream Loader for various types of video streams, Supports RTSP, RTMP, HTTP, and TCP 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.
 
     Attributes:
-        sources (str): The source input paths or URLs for the video streams.
-        vid_stride (int): Video frame-rate stride, defaults to 1.
-        buffer (bool): Whether to buffer input streams, defaults to False.
+        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 of image frames for each stream.
-        fps (list): List of FPS for each stream.
-        frames (list): List of total frames for each stream.
-        threads (list): List of threads for each stream.
-        shape (list): List of shapes for each stream.
-        caps (list): 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.
 
     Methods:
-        __init__: Initialize the stream loader.
         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.
 
-    Example:
-         ```bash
-         yolo predict source='rtsp://example.com/media.mp4'
-         ```
+    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="file.streams", vid_stride=1, buffer=False):
-        """Initialize instance variables and check for consistent input stream shapes."""
+        """Initialize stream loader for multiple video sources, supporting various stream types."""
         torch.backends.cudnn.benchmark = True  # faster for fixed-size inference
         self.buffer = buffer  # buffer input streams
         self.running = True  # running flag for Thread
@@ -83,11 +110,11 @@ class LoadStreams:
         for i, s in enumerate(sources):  # index, source
             # Start thread to read frames from video stream
             st = f"{i + 1}/{n}: {s}... "
-            if urlparse(s).hostname in ("www.youtube.com", "youtube.com", "youtu.be"):  # if source is YouTube video
-                # YouTube format i.e. 'https://www.youtube.com/watch?v=Zgi9g1ksQHc' or 'https://youtu.be/LNwODJXcvt4'
+            if urlparse(s).hostname in {"www.youtube.com", "youtube.com", "youtu.be"}:  # if source is 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 = eval(s) if s.isnumeric() else s  # i.e. s = '0' local webcam
-            if s == 0 and (is_colab() or is_kaggle()):
+            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."
@@ -114,7 +141,7 @@ class LoadStreams:
         LOGGER.info("")  # newline
 
     def update(self, i, cap, stream):
-        """Read stream `i` frames in daemon thread."""
+        """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
@@ -134,7 +161,7 @@ class LoadStreams:
                 time.sleep(0.01)  # wait until the buffer is empty
 
     def close(self):
-        """Close stream loader and release resources."""
+        """Terminates stream loader, stops threads, and releases video capture resources."""
         self.running = False  # stop flag for Thread
         for thread in self.threads:
             if thread.is_alive():
@@ -152,7 +179,7 @@ class LoadStreams:
         return self
 
     def __next__(self):
-        """Returns source paths, transformed and original images for processing."""
+        """Returns the next batch of frames from multiple video streams for processing."""
         self.count += 1
 
         images = []
@@ -179,16 +206,16 @@ class LoadStreams:
         return self.sources, images, [""] * self.bs
 
     def __len__(self):
-        """Return the length of the sources object."""
+        """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:
     """
-    YOLOv8 screenshot dataloader.
+    Ultralytics screenshot dataloader for capturing and processing screen images.
 
-    This class manages the loading of screenshot images for processing with YOLOv8.
-    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.
@@ -201,15 +228,21 @@ class LoadScreenshots:
         frame (int): Counter for captured frames.
         sct (mss.mss): Screen capture object from `mss` library.
         bs (int): Batch size, set to 1.
-        monitor (dict): Monitor configuration details.
+        fps (int): Frames per second, set to 30.
+        monitor (Dict[str, int]): Monitor configuration details.
 
     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 source, im, im0s, vid_cap, s in loader:
+        ...     print(f"Captured frame: {im.shape}")
     """
 
     def __init__(self, source):
-        """Source = [screen_number left top width height] (pixels)."""
+        """Initialize screenshot capture with specified screen and region parameters."""
         check_requirements("mss")
         import mss  # noqa
 
@@ -236,11 +269,11 @@ class LoadScreenshots:
         self.monitor = {"left": self.left, "top": self.top, "width": self.width, "height": self.height}
 
     def __iter__(self):
-        """Returns an iterator of the object."""
+        """Yields the next screenshot image from the specified screen or region for processing."""
         return self
 
     def __next__(self):
-        """mss screen capture: get raw pixels from the screen as np array."""
+        """Captures and returns the next screenshot as a numpy array using the mss library."""
         im0 = np.asarray(self.sct.grab(self.monitor))[:, :, :3]  # BGRA to BGR
         s = f"screen {self.screen} (LTWH): {self.left},{self.top},{self.width},{self.height}: "
 
@@ -250,29 +283,45 @@ class LoadScreenshots:
 
 class LoadImagesAndVideos:
     """
-    YOLOv8 image/video dataloader.
+    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 for YOLOv8. It supports loading from
-    various formats, 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): 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): 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, defaults to 1.
-        bs (int): Batch size, set to 1 for this class.
+        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__()`.
+        count (int): Counter for iteration, initialized at 0 during __iter__().
+        ni (int): Number of images.
 
     Methods:
-        _new_video(path): Create a new cv2.VideoCapture object for a given video path.
+        __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, batch=1, vid_stride=1):
-        """Initialize the Dataloader and raise FileNotFoundError if file not found."""
+        """Initialize dataloader for images and videos, supporting various input formats."""
         parent = None
         if isinstance(path, str) and Path(path).suffix == ".txt":  # *.txt file with img/vid/dir on each line
             parent = Path(path).parent
@@ -291,15 +340,21 @@ class LoadImagesAndVideos:
             else:
                 raise FileNotFoundError(f"{p} does not exist")
 
-        images = [x for x in files if x.split(".")[-1].lower() in IMG_FORMATS]
-        videos = [x for x in files if x.split(".")[-1].lower() in VID_FORMATS]
+        # 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
+            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 = "image"
+        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
         if any(videos):
@@ -307,22 +362,19 @@ class LoadImagesAndVideos:
         else:
             self.cap = None
         if self.nf == 0:
-            raise FileNotFoundError(
-                f"No images or videos found in {p}. "
-                f"Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}"
-            )
+            raise FileNotFoundError(f"No images or videos found in {p}. {FORMATS_HELP_MSG}")
 
     def __iter__(self):
-        """Returns an iterator object for VideoStream or ImageFolder."""
+        """Iterates 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 along with their paths and metadata."""
+        """Returns 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 len(imgs) > 0:
+                if imgs:
                     return paths, imgs, info  # return last partial batch
                 else:
                     raise StopIteration
@@ -333,6 +385,7 @@ class LoadImagesAndVideos:
                 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:
@@ -356,13 +409,25 @@ class LoadImagesAndVideos:
                     if self.count < self.nf:
                         self._new_video(self.files[self.count])
             else:
+                # Handle image files (including HEIC)
                 self.mode = "image"
-                im0 = cv2.imread(path)  # BGR
+                if path.split(".")[-1].lower() == "heic":
+                    # Load HEIC image using Pillow with pillow-heif
+                    check_requirements("pillow-heif")
+
+                    from pillow_heif import register_heif_opener
+
+                    register_heif_opener()  # Register HEIF opener with Pillow
+                    with Image.open(path) as img:
+                        im0 = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)  # convert image to BGR nparray
+                else:
+                    im0 = imread(path)  # BGR
                 if im0 is None:
-                    raise FileNotFoundError(f"Image Not Found {path}")
-                paths.append(path)
-                imgs.append(im0)
-                info.append(f"image {self.count + 1}/{self.nf} {path}: ")
+                    LOGGER.warning(f"WARNING ⚠️ 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
@@ -370,7 +435,7 @@ class LoadImagesAndVideos:
         return paths, imgs, info
 
     def _new_video(self, path):
-        """Creates a new video capture object for the given path."""
+        """Creates a new video capture object for the given path and initializes video-related attributes."""
         self.frame = 0
         self.cap = cv2.VideoCapture(path)
         self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))
@@ -379,40 +444,50 @@ class LoadImagesAndVideos:
         self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT) / self.vid_stride)
 
     def __len__(self):
-        """Returns the number of batches in the object."""
-        return math.ceil(self.nf / self.bs)  # number of files
+        """Returns 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.
 
-    This class is designed to manage 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.
+    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): List of image paths or autogenerated filenames.
-        im0 (list): List of images stored as Numpy arrays.
-        mode (str): Type of data being processed, defaults to 'image'.
+        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(im): Validate and format a single image to a Numpy array.
+        _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):
-        """Initialize PIL and Numpy Dataloader."""
+        """Initializes a loader for PIL and Numpy images, converting inputs to a standardized format."""
         if not isinstance(im0, list):
             im0 = [im0]
-        self.paths = [getattr(im, "filename", f"image{i}.jpg") for i, im in enumerate(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)]
         self.im0 = [self._single_check(im) for im in im0]
         self.mode = "image"
         self.bs = len(self.im0)
 
     @staticmethod
     def _single_check(im):
-        """Validate and format an image to numpy array."""
+        """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):
             if im.mode != "RGB":
@@ -422,41 +497,48 @@ class LoadPilAndNumpy:
         return im
 
     def __len__(self):
-        """Returns the length of the 'im0' attribute."""
+        """Returns the length of the 'im0' attribute, representing the number of loaded images."""
         return len(self.im0)
 
     def __next__(self):
-        """Returns batch paths, images, processed images, None, ''."""
+        """Returns 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):
-        """Enables iteration for class LoadPilAndNumpy."""
+        """Iterates through PIL/numpy images, yielding paths, raw images, and metadata for processing."""
         self.count = 0
         return self
 
 
 class LoadTensor:
     """
-    Load images from torch.Tensor data.
+    A class for loading and processing tensor data for object detection tasks.
 
-    This class manages the loading and pre-processing of image data from PyTorch tensors for further processing.
+    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).
+        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 mode, set to 'image'.
-        paths (list): List of image paths or filenames.
-        count (int): Counter for iteration, initialized at 0 during `__iter__()`.
+        mode (str): Current processing mode, set to 'image'.
+        paths (List[str]): List of image paths or auto-generated filenames.
 
     Methods:
-        _single_check(im, stride): Validate and possibly modify the input tensor.
+        _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) -> None:
-        """Initialize Tensor Dataloader."""
+        """Initialize LoadTensor object for processing torch.Tensor image data."""
         self.im0 = self._single_check(im0)
         self.bs = self.im0.shape[0]
         self.mode = "image"
@@ -464,7 +546,7 @@ class LoadTensor:
 
     @staticmethod
     def _single_check(im, stride=32):
-        """Validate and format an image to torch.Tensor."""
+        """Validates and formats a single image tensor, ensuring correct shape and normalization."""
         s = (
             f"WARNING ⚠️ 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."
@@ -486,24 +568,24 @@ class LoadTensor:
         return im
 
     def __iter__(self):
-        """Returns an iterator object."""
+        """Yields an iterator object for iterating through tensor image data."""
         self.count = 0
         return self
 
     def __next__(self):
-        """Return next item in the iterator."""
+        """Yields 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."""
+        """Returns the batch size of the tensor input."""
         return self.bs
 
 
 def autocast_list(source):
-    """Merges a list of source of different types into a list of numpy arrays or PIL images."""
+    """Merges 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
@@ -519,26 +601,47 @@ def autocast_list(source):
     return files
 
 
-def get_best_youtube_url(url, use_pafy=True):
+def get_best_youtube_url(url, method="pytube"):
     """
     Retrieves the URL of the best quality MP4 video stream from a given YouTube video.
 
-    This function uses the pafy or yt_dlp library to extract the video info from YouTube. It then finds the highest
-    quality MP4 format that has video codec but no audio codec, and returns the URL of this video stream.
-
     Args:
         url (str): The URL of the YouTube video.
-        use_pafy (bool): Use the pafy package, default=True, otherwise use yt_dlp package.
+        method (str): The method to use for extracting video info. Options are "pytube", "pafy", and "yt-dlp".
+            Defaults to "pytube".
 
     Returns:
-        (str): The URL of the best quality MP4 video stream, or None if no suitable stream is found.
+        (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 use_pafy:
+    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  # noqa
 
         return pafy.new(url).getbestvideo(preftype="mp4").url
-    else:
+
+    elif method == "yt-dlp":
         check_requirements("yt-dlp")
         import yt_dlp