Reference for ultralytics/data/loaders.py#

@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 in-memory image (PIL/numpy) or list of images.
        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__: Return an iterator object for the class.
        __next__: Return source names, latest frames, and empty info strings for processing.
        __len__: Return the number of video streams.

    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
        try:
            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
                if not success or im is None:
                    raise ConnectionError(f"{st}Failed to read images from {s}")
                im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)[..., None] if self.cv2_flag == cv2.IMREAD_GRAYSCALE else im
                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()
        except Exception:
            self.close()  # release opened captures and stop started threads before re-raising
            raise
        LOGGER.info("")  # newline

def __iter__(self):
    """Return an iterator object and reset the frame counter."""
    self.count = -1
    return self

def __len__(self) -> int:
    """Return the number of video streams in the LoadStreams object."""
    return self.bs

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 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 not None and thread.is_alive():
            thread.join(timeout=5)  # Add timeout
    for cap in self.caps:  # Iterate through the stored VideoCapture objects
        if cap is None:
            continue
        try:
            cap.release()  # release video capture
        except Exception as e:
            LOGGER.warning(f"Could not release VideoCapture object: {e}")

def update(self, i: int, cap: cv2.VideoCapture, stream: str | int):
    """Read stream frames in daemon thread and update image buffer.

    Args:
        i (int): Index of the stream in the loader.
        cap (cv2.VideoCapture): Video capture object for the stream.
        stream (str | int): Stream source or webcam index, used to re-open the stream if the signal is lost.
    """
    n, f = 0, self.frames[i]  # frame number, total frames
    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()
                if not success or im is None:
                    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
                elif self.cv2_flag == cv2.IMREAD_GRAYSCALE:
                    im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)[..., None]
                if self.buffer:
                    self.imgs[i].append(im)
                else:
                    self.imgs[i] = [im]
        else:
            time.sleep(0.01)  # wait until the buffer has room

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__: Return an iterator object.
        __next__: Capture the next screenshot and return it.

    Examples:
        >>> loader = LoadScreenshots("screen 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 starting with "screen", optionally followed by a screen number
                and/or a "left top width height" capture region, e.g. "screen", "screen 1", or "screen 0 100 100
                640 480".
            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): Number of frames in the current video after applying vid_stride.
        fps (float): Frames per second of the current video as reported by OpenCV.
        count (int): Counter for iteration, initialized at 0 during __iter__().
        ni (int): Number of images.
        cv2_flag (int): OpenCV flag for image reading (grayscale, color/BGR, or unchanged for other channel counts).

    Methods:
        __init__: Initialize the LoadImagesAndVideos object.
        close: Release the current video capture object.
        __iter__: Return an iterator object for VideoStream or ImageFolder.
        __next__: Return the next batch of images or video frames along with their paths and metadata.
        _append_image: Append a decoded image to the batch lists, or warn and skip it if decoding failed.
        _new_video: Create a new video capture object for the given path.
        __len__: Return 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 *.txt or *.csv file containing paths to images and videos.
    """

    def __init__(self, path: str | Path | list | tuple, 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 | tuple): Path to an image/video file, directory, glob pattern, *.txt or *.csv file
                of source paths, 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).
        """
        source_path = path
        parent = None
        if isinstance(path, (str, Path)) and Path(path).suffix in {".txt", ".csv"}:  # txt/csv file with source paths
            parent, content = Path(path).parent, Path(path).read_text()
            if Path(path).suffix == ".txt":
                path = content.splitlines()
            else:
                rows = list(csv.reader(content.splitlines()))
                rows = rows[1:] if rows[:1] == [["source"]] else rows  # optional header row
                path = [p for row in rows for p in row]
            path = [p.strip() for p in path if p.strip()]
        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 parent and not (glob.glob(a, recursive=True) if "*" in a else os.path.exists(a)):
                a = str((parent / p).absolute())  # resolve unmatched manifest entries relative to the manifest
            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(glob.escape(a), "*.*"))))  # dir
            elif os.path.isfile(a):
                files.append(a)  # files (absolute or relative to CWD)
            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 = {1: cv2.IMREAD_GRAYSCALE, 3: cv2.IMREAD_COLOR}.get(channels, cv2.IMREAD_UNCHANGED)
        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 {source_path}. {FORMATS_HELP_MSG}")

def __iter__(self):
    """Iterate through image/video files, yielding source paths, images, and metadata."""
    self.count = 0
    return self

def __len__(self) -> int:
    """Return the number of batches in the dataset."""
    return math.ceil(self.nf / self.bs)  # number of batches

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()
                if success and im0 is not None:
                    if self.cv2_flag == cv2.IMREAD_GRAYSCALE:
                        im0 = cv2.cvtColor(im0, cv2.COLOR_BGR2GRAY)[..., None]
                    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, flush so a batch never spans two videos
                        self.count += 1
                        self.cap.release()
                        break
            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 imgs:  # flush so a batch never spans two videos, the next video opens on the next call
                    break
        else:
            # Handle image files
            self.mode = "image"
            if self.bs >= 8 and NUM_THREADS > 1 and (n := min(self.bs - len(imgs), self.ni - self.count)) >= 8:
                image_paths = self.files[self.count : self.count + n]
                # Keep fallback formats serial: their lazy PIL plugin registration is not thread-safe.
                if not any(Path(x).suffix.lower() in PIL_FALLBACK_SUFFIXES for x in image_paths):
                    with ThreadPool(min(n, NUM_THREADS)) as pool:
                        decoded = pool.map(partial(imread, flags=self.cv2_flag), image_paths)
                    for i, (image_path, im0) in enumerate(zip(image_paths, decoded)):
                        self._append_image(paths, imgs, info, image_path, im0, self.count + i + 1)
                    self.count += n  # move to the next batch of files
                    if self.count >= self.ni and imgs:  # flush images before starting videos
                        break
                    continue

            im0 = imread(path, flags=self.cv2_flag)  # BGR
            self._append_image(paths, imgs, info, path, im0, self.count + 1)
            self.count += 1  # move to the next file
            if self.count >= self.ni and imgs:  # end of image list, flush only a non-empty batch
                break

    return paths, imgs, info

def _append_image(self, paths: list, imgs: list, info: list, path: str, im0: np.ndarray | None, idx: int):
    """Append a decoded image to the batch lists, or warn and skip it if decoding failed."""
    if im0 is None:
        LOGGER.warning(f"Image Read Error {path}")
    else:
        paths.append(path)
        imgs.append(im0)
        info.append(f"image {idx}/{self.nf} {path}: ")

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 = self.cap.get(cv2.CAP_PROP_FPS)  # keep the float, int() turns 29.97 into 29 and stretches the video
    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 close(self):
    """Release the current video capture object, e.g. when inference stops before the video ends."""
    if self.cap:
        self.cap.release()

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]
        if not im0:  # an empty batch otherwise fails unnamed inside np.stack in Predictor.preprocess
            raise FileNotFoundError("No images found in source, predict requires at least one image.")
        # 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, channels) for im in im0]
        self.mode = "image"
        self.bs = len(self.im0)
        self.count = 0

def __iter__(self):
    """Iterate through PIL/numpy images, yielding paths, raw images, and metadata for processing."""
    self.count = 0
    return self

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

@staticmethod
def _single_check(im: Image.Image | np.ndarray, channels: int = 3) -> np.ndarray:
    """Validate an image and normalize its channel count.

    Notes:
        - PIL inputs are converted to NumPy and returned in OpenCV-compatible BGR order for color images.
        - NumPy color inputs are assumed to use OpenCV-compatible BGR order.
    """
    if not isinstance(im, (Image.Image, np.ndarray)):
        raise TypeError(f"Expected PIL/np.ndarray image type, but got {type(im)}")
    pil = isinstance(im, Image.Image)
    if pil:
        if im.mode.startswith("I;16"):  # convert() clips 16-bit values at 255, scale them like cv2 does
            im = Image.fromarray((np.asarray(im) >> 8).astype(np.uint8))
        flag = "L" if channels == 1 else "RGB"
        im = np.asarray(im if im.mode == flag else im.convert(flag))  # convert() copies even when mode matches
    im = np.atleast_3d(im)
    # Both routes validate here: a zero dimension divides by zero in LetterBox, and a batched array reads
    # shape[2] as a channel count it is not. Raised rather than asserted so `python -O` keeps the check, and
    # ahead of the cvtColor calls, which assert on an empty input instead of raising this message.
    if im.ndim != 3 or not all(im.shape):
        raise ValueError(f"Expected a single (H, W, C) image, but got array of shape {im.shape}")
    if pil:
        return im if channels == 1 else cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
    c = im.shape[2]
    if c == channels:
        return im
    if c == 2:  # gray + alpha
        im, c = im[..., :1], 1
    u8 = im.dtype == np.uint8  # cvtColor rejects dtypes NumPy indexing accepts, float64 among them
    if c == 1:
        if u8 and channels == 3:
            return cv2.cvtColor(im[..., 0], cv2.COLOR_GRAY2BGR)
        return np.repeat(im, channels, axis=2)
    if channels == 1:
        return cv2.cvtColor(im, cv2.COLOR_BGRA2GRAY if c == 4 else cv2.COLOR_BGR2GRAY)[..., None]
    return cv2.cvtColor(im, cv2.COLOR_BGRA2BGR) if u8 and c == 4 else np.ascontiguousarray(im[..., :3])

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: Validate and format 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 = [f"image{i}.jpg" for i in range(self.bs)]
        self.count = 0

def __iter__(self):
    """Return an iterator object for iterating through tensor image data."""
    self.count = 0
    return self

def __len__(self) -> int:
    """Return the batch size of the tensor input."""
    return self.bs

def __next__(self) -> tuple[list[str], torch.Tensor, list[str]]:
    """Return 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

@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 not all(im.shape) or im.shape[2] % stride or im.shape[3] % stride:
        raise ValueError(s)  # a zero dimension reaches im.max() below on an empty tensor
    if im.max() > 1.0 + (torch.finfo(im.dtype).eps if im.is_floating_point() else 0):
        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 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.

    Args:
        source (list[Any]): List of file paths, URLs, PIL images, or numpy arrays.

    Returns:
        (list[PIL.Image.Image | np.ndarray]): List of PIL images (paths and URLs are opened and EXIF-transposed) or
            numpy arrays.

    Raises:
        TypeError: If an element is not a supported source type.
    """
    files = []
    for im in source:
        if isinstance(im, (str, Path)):  # filename or uri
            if str(im).startswith("http"):  # requests follows HTTP 308 redirects that urllib lacks pre-3.11
                import requests  # scoped as slow import

                im = BytesIO(requests.get(im).content)
            im = Image.open(im)
            filename = im.filename
            im = ImageOps.exif_transpose(im)
            im.filename = filename
            files.append(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.
        - 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

        stream = (
            YouTube(url).streams.filter(file_extension="mp4", only_video=True).order_by("resolution").desc().first()
        )
        if stream:
            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
            if f["vcodec"] != "none" and f["acodec"] == "none" and f["ext"] == "mp4":
                return f.get("url")