# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

from __future__ import annotations

import hashlib
import glob
import json
import math
import os
import random
from collections import defaultdict
from copy import deepcopy
from itertools import repeat
from multiprocessing.pool import ThreadPool
from pathlib import Path
from time import perf_counter
from typing import Any

import cv2
import numpy as np
import torch
from PIL import Image
from torch.utils.data import ConcatDataset

from ultralytics.utils import DEFAULT_CFG, LOCAL_RANK, LOGGER, NUM_THREADS, TQDM, colorstr
from ultralytics.utils.instance import Instances
from ultralytics.utils.ops import resample_segments, segments2boxes
from ultralytics.utils.plotting_3d import decode_cut_partial_side_edge_from_gt, decode_visible_face_edge_from_gt
from ultralytics.utils.torch_utils import TORCHVISION_0_18

from .augment import (
    Compose,
    Format,
    LetterBox,
    RandomHSV,
    RandomLoadText,
    classify_augmentations,
    classify_transforms,
    v8_transforms,
)
from .base import BaseDataset
from .converter import merge_multi_segment
from .ground3d_augment import (
    adjust_calib_for_roi_crop,
    apply_simul_transform,
    build_final_resized_calib,
    compute_centered_roi_bounds,
    compute_simul_calib,
    compute_vanishing_point_x,
    compute_vanishing_point_y,
    normalize_roi_depth,
    pack_labels_to_48,
    parse_ground_3d_label_file,
    read_calib_from_path,
    remap_labels_to_roi,
    unpack_labels_from_48,
)
from .utils import (
    FORMATS_HELP_MSG,
    HELP_URL,
    IMG_FORMATS,
    check_file_speeds,
    get_hash,
    img2label_paths,
    load_dataset_cache_file,
    save_dataset_cache_file,
    verify_image,
    verify_image_label,
)

# Ultralytics dataset *.cache version, >= 1.0.0 for Ultralytics YOLO models
DATASET_CACHE_VERSION = "1.0.3"
GROUND3D_REQUIRED_DATA_KEYS = ("path", "ori_img_size", "roi", "virtual_fx", "virtual_camera_prob", "crop_center_mode")
GROUND3D_VALID_CROP_CENTER_MODES = {"cxvy", "vxvy"}


class Ground3DImageReadError(RuntimeError):
    """Raised when a Ground3D sample image is unreadable or fails validation."""


class Ground3DCalibrationError(RuntimeError):
    """Raised when a Ground3D sample is missing required calibration."""


def _resize_ground3d_image_in_steps(
    img: np.ndarray, target_size: tuple[int, int], interpolation: int = cv2.INTER_LINEAR
) -> np.ndarray:
    """Resize a Ground3D image with repeated 0.5x downsampling before the final resize."""
    target_w, target_h = int(target_size[0]), int(target_size[1])
    current_h, current_w = img.shape[:2]

    if (current_w, current_h) == (target_w, target_h):
        return img

    if target_w < current_w and target_h < current_h:
        while True:
            next_w = math.ceil(current_w * 0.5)
            next_h = math.ceil(current_h * 0.5)
            if next_w < target_w or next_h < target_h:
                break

            img = cv2.resize(img, (next_w, next_h), interpolation=interpolation)
            current_h, current_w = img.shape[:2]
            if (current_w, current_h) == (target_w, target_h):
                return img

    return cv2.resize(img, (target_w, target_h), interpolation=interpolation)


def _require_ground3d_data_value(data: dict[str, Any], key: str) -> Any:
    """Return a required Ground3D config value or raise a clear error."""
    if key not in data or data[key] is None:
        raise ValueError(f"Ground3D dataset config must define '{key}' explicitly; defaults are not allowed.")
    return data[key]


def _validate_ground3d_pair(value: Any, key: str) -> tuple[int, int]:
    """Validate a Ground3D size-like config field and return it as an integer pair."""
    if not isinstance(value, (list, tuple)) or len(value) != 2:
        raise ValueError(f"Ground3D dataset config field '{key}' must be a 2-item list/tuple, but got: {value!r}")
    pair = (int(value[0]), int(value[1]))
    if pair[0] <= 0 or pair[1] <= 0:
        raise ValueError(f"Ground3D dataset config field '{key}' must contain positive values, but got: {value!r}")
    return pair


def _validate_ground3d_crop_center_mode(value: Any) -> str:
    """Validate the Ground3D crop-center mode."""
    if value not in GROUND3D_VALID_CROP_CENTER_MODES:
        valid_modes = ", ".join(sorted(GROUND3D_VALID_CROP_CENTER_MODES))
        raise ValueError(
            f"Ground3D dataset config field 'crop_center_mode' must be one of {{{valid_modes}}}, but got: {value!r}"
        )
    return str(value)


def _validate_ground3d_positive_float(value: Any, key: str) -> float:
    """Validate a positive numeric Ground3D config field."""
    value = float(value)
    if value <= 0:
        raise ValueError(f"Ground3D dataset config field '{key}' must be positive, but got: {value!r}")
    return value


def _validate_ground3d_float(value: Any, key: str) -> float:
    """Validate a numeric Ground3D config field."""
    try:
        return float(value)
    except (TypeError, ValueError) as e:
        raise ValueError(f"Ground3D dataset config field '{key}' must be numeric, but got: {value!r}") from e


def _ground3d_has_finite_targets(labels_3d: np.ndarray | None) -> bool:
    """Return whether a parsed 3D label array contains any real 3D supervision."""
    if labels_3d is None:
        return False
    labels_3d = np.asarray(labels_3d)
    return labels_3d.size > 0 and np.isfinite(labels_3d).any()


class YOLODataset(BaseDataset):
    """Dataset class for loading object detection and/or segmentation labels in YOLO format.

    This class supports loading data for object detection, segmentation, pose estimation, and oriented bounding box
    (OBB) tasks using the YOLO format.

    Attributes:
        use_segments (bool): Indicates if segmentation masks should be used.
        use_keypoints (bool): Indicates if keypoints should be used for pose estimation.
        use_obb (bool): Indicates if oriented bounding boxes should be used.
        data (dict): Dataset configuration dictionary.

    Methods:
        cache_labels: Cache dataset labels, check images and read shapes.
        get_labels: Return list of label dictionaries for YOLO training.
        build_transforms: Build and append transforms to the list.
        close_mosaic: Disable mosaic, copy_paste, mixup and cutmix augmentations and build transformations.
        update_labels_info: Update label format for different tasks.
        collate_fn: Collate data samples into batches.

    Examples:
        >>> dataset = YOLODataset(img_path="path/to/images", data={"names": {0: "person"}}, task="detect")
        >>> dataset.get_labels()
    """

    def __init__(self, *args, data: dict | None = None, task: str = "detect", **kwargs):
        """Initialize the YOLODataset.

        Args:
            data (dict, optional): Dataset configuration dictionary.
            task (str): Task type, one of 'detect', 'segment', 'pose', or 'obb'.
            *args (Any): Additional positional arguments for the parent class.
            **kwargs (Any): Additional keyword arguments for the parent class.
        """
        self.use_segments = task == "segment"
        self.use_keypoints = task == "pose"
        self.use_obb = task == "obb"
        self.data = data
        assert not (self.use_segments and self.use_keypoints), "Can not use both segments and keypoints."
        super().__init__(*args, channels=self.data.get("channels", 3), use_yuv444=self.data.get("use_yuv444", False), **kwargs)

    def cache_labels(self, path: Path = Path("./labels.cache")) -> dict:
        """Cache dataset labels, check images and read shapes.

        Args:
            path (Path): Path where to save the cache file.

        Returns:
            (dict): Dictionary containing cached labels and related information.
        """
        x = {"labels": []}
        nm, nf, ne, nc, msgs = 0, 0, 0, 0, []  # number missing, found, empty, corrupt, messages
        desc = f"{self.prefix}Scanning {path.parent / path.stem}..."
        total = len(self.im_files)
        nkpt, ndim = self.data.get("kpt_shape", (0, 0))
        if self.use_keypoints and (nkpt <= 0 or ndim not in {2, 3}):
            raise ValueError(
                "'kpt_shape' in data.yaml missing or incorrect. Should be a list with [number of "
                "keypoints, number of dims (2 for x,y or 3 for x,y,visible)], i.e. 'kpt_shape: [17, 3]'"
            )
        with ThreadPool(NUM_THREADS) as pool:
            results = pool.imap(
                func=verify_image_label,
                iterable=zip(
                    self.im_files,
                    self.label_files,
                    repeat(self.prefix),
                    repeat(self.use_keypoints),
                    repeat(len(self.data["names"])),
                    repeat(nkpt),
                    repeat(ndim),
                    repeat(self.single_cls),
                ),
            )
            pbar = TQDM(results, desc=desc, total=total)
            for im_file, lb, shape, segments, keypoint, nm_f, nf_f, ne_f, nc_f, msg in pbar:
                nm += nm_f
                nf += nf_f
                ne += ne_f
                nc += nc_f
                if im_file:
                    x["labels"].append(
                        {
                            "im_file": im_file,
                            "shape": shape,
                            "cls": lb[:, 0:1],  # n, 1
                            "bboxes": lb[:, 1:],  # n, 4
                            "segments": segments,
                            "keypoints": keypoint,
                            "normalized": True,
                            "bbox_format": "xywh",
                        }
                    )
                if msg:
                    msgs.append(msg)
                pbar.desc = f"{desc} {nf} images, {nm + ne} backgrounds, {nc} corrupt"
            pbar.close()

        if msgs:
            LOGGER.info("\n".join(msgs))
        if nf == 0:
            LOGGER.warning(f"{self.prefix}No labels found in {path}. {HELP_URL}")
        x["hash"] = get_hash(self.label_files + self.im_files)
        x["results"] = nf, nm, ne, nc, len(self.im_files)
        x["msgs"] = msgs  # warnings
        save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
        return x

    def get_labels(self) -> list[dict]:
        """Return list of label dictionaries for YOLO training.

        This method loads labels from disk or cache, verifies their integrity, and prepares them for training.

        Returns:
            (list[dict]): List of label dictionaries, each containing information about an image and its annotations.
        """
        self.label_files = img2label_paths(self.im_files)
        cache_path = Path(self.label_files[0]).parent.with_suffix(".cache")
        try:
            cache, exists = load_dataset_cache_file(cache_path), True  # attempt to load a *.cache file
            assert cache["version"] == DATASET_CACHE_VERSION  # matches current version
            assert cache["hash"] == get_hash(self.label_files + self.im_files)  # identical hash
        except (FileNotFoundError, AssertionError, AttributeError, ModuleNotFoundError):
            cache, exists = self.cache_labels(cache_path), False  # run cache ops

        # Display cache
        nf, nm, ne, nc, n = cache.pop("results")  # found, missing, empty, corrupt, total
        if exists and LOCAL_RANK in {-1, 0}:
            d = f"Scanning {cache_path}... {nf} images, {nm + ne} backgrounds, {nc} corrupt"
            TQDM(None, desc=self.prefix + d, total=n, initial=n)  # display results
            if cache["msgs"]:
                LOGGER.info("\n".join(cache["msgs"]))  # display warnings

        # Read cache
        [cache.pop(k) for k in ("hash", "version", "msgs")]  # remove items
        labels = cache["labels"]
        if not labels:
            raise RuntimeError(
                f"No valid images found in {cache_path}. Images with incorrectly formatted labels are ignored. {HELP_URL}"
            )
        self.im_files = [lb["im_file"] for lb in labels]  # update im_files

        # Check if the dataset is all boxes or all segments
        lengths = ((len(lb["cls"]), len(lb["bboxes"]), len(lb["segments"])) for lb in labels)
        len_cls, len_boxes, len_segments = (sum(x) for x in zip(*lengths))
        if len_segments and len_boxes != len_segments:
            LOGGER.warning(
                f"Box and segment counts should be equal, but got len(segments) = {len_segments}, "
                f"len(boxes) = {len_boxes}. To resolve this only boxes will be used and all segments will be removed. "
                "To avoid this please supply either a detect or segment dataset, not a detect-segment mixed dataset."
            )
            for lb in labels:
                lb["segments"] = []
        if len_cls == 0:
            LOGGER.warning(f"Labels are missing or empty in {cache_path}, training may not work correctly. {HELP_URL}")
        return labels

    def build_transforms(self, hyp: dict | None = None) -> Compose:
        """Build and append transforms to the list.

        Args:
            hyp (dict, optional): Hyperparameters for transforms.

        Returns:
            (Compose): Composed transforms.
        """
        if self.augment:
            hyp.mosaic = hyp.mosaic if self.augment and not self.rect else 0.0
            hyp.mixup = hyp.mixup if self.augment and not self.rect else 0.0
            hyp.cutmix = hyp.cutmix if self.augment and not self.rect else 0.0
            transforms = v8_transforms(self, self.imgsz, hyp)
        else:
            transforms = Compose([LetterBox(new_shape=(self.imgsz, self.imgsz), scaleup=False)])
        transforms.append(
            Format(
                bbox_format="xywh",
                normalize=True,
                return_mask=self.use_segments,
                return_keypoint=self.use_keypoints,
                return_obb=self.use_obb,
                batch_idx=True,
                mask_ratio=hyp.mask_ratio,
                mask_overlap=hyp.overlap_mask,
                bgr=hyp.bgr if self.augment else 0.0,  # only affect training.
            )
        )
        return transforms

    def close_mosaic(self, hyp: dict) -> None:
        """Disable mosaic, copy_paste, mixup and cutmix augmentations by setting their probabilities to 0.0.

        Args:
            hyp (dict): Hyperparameters for transforms.
        """
        hyp.mosaic = 0.0
        hyp.copy_paste = 0.0
        hyp.mixup = 0.0
        hyp.cutmix = 0.0
        self.transforms = self.build_transforms(hyp)

    def update_labels_info(self, label: dict) -> dict:
        """Update label format for different tasks.

        Args:
            label (dict): Label dictionary containing bboxes, segments, keypoints, etc.

        Returns:
            (dict): Updated label dictionary with instances.

        Notes:
            cls is not with bboxes now, classification and semantic segmentation need an independent cls label
            Can also support classification and semantic segmentation by adding or removing dict keys there.
        """
        bboxes = label.pop("bboxes")
        segments = label.pop("segments", [])
        keypoints = label.pop("keypoints", None)
        bbox_format = label.pop("bbox_format")
        normalized = label.pop("normalized")

        # NOTE: do NOT resample oriented boxes
        segment_resamples = 100 if self.use_obb else 1000
        if len(segments) > 0:
            # make sure segments interpolate correctly if original length is greater than segment_resamples
            max_len = max(len(s) for s in segments)
            segment_resamples = (max_len + 1) if segment_resamples < max_len else segment_resamples
            # list[np.array(segment_resamples, 2)] * num_samples
            segments = np.stack(resample_segments(segments, n=segment_resamples), axis=0)
        else:
            segments = np.zeros((0, segment_resamples, 2), dtype=np.float32)
        label["instances"] = Instances(bboxes, segments, keypoints, bbox_format=bbox_format, normalized=normalized)
        return label

    @staticmethod
    def collate_fn(batch: list[dict]) -> dict:
        """Collate data samples into batches.

        Args:
            batch (list[dict]): List of dictionaries containing sample data.

        Returns:
            (dict): Collated batch with stacked tensors.
        """
        new_batch = {}
        batch = [dict(sorted(b.items())) for b in batch]  # make sure the keys are in the same order
        keys = batch[0].keys()
        values = list(zip(*[list(b.values()) for b in batch]))
        for i, k in enumerate(keys):
            value = values[i]
            if k in {"img", "text_feats", "sem_masks"}:
                value = torch.stack(value, 0)
            elif k == "visuals":
                value = torch.nn.utils.rnn.pad_sequence(value, batch_first=True)
            elif k == "camera_mode":
                value = tuple(value)
            if k in {
                "masks",
                "keypoints",
                "bboxes",
                "cls",
                "segments",
                "obb",
                "difficulties",
                "difficulty_levels",
                "labels_3d",
                "edge_faces_points_2d",
                "edge_faces_depths",
                "edge_faces_valid",
                "edge_partial_points_2d",
                "edge_partial_depths",
                "edge_partial_face_type",
                "edge_partial_valid",
            }:
                value = torch.cat(value, 0)
            new_batch[k] = value
        new_batch["batch_idx"] = list(new_batch["batch_idx"])
        for i in range(len(new_batch["batch_idx"])):
            new_batch["batch_idx"][i] += i  # add target image index for build_targets()
        new_batch["batch_idx"] = torch.cat(new_batch["batch_idx"], 0)
        return new_batch


class YOLOGroundDataset(YOLODataset):
    """Dataset class for loading ground 2D detection labels with difficulty scores.

    This class extends YOLODataset to support custom annotation format with:
    - String class names mapped to numeric IDs via class_map
    - Difficulty scores for each bounding box
    - Optional minimum box size filtering
    - Optional YUV444 color space support

    Attributes:
        class_map (dict): Mapping from class names to class IDs.
        min_wh (float): Minimum box width/height in pixels for filtering.
        use_yuv444 (bool): Whether to use YUV444 color space.

    Examples:
        >>> data = {"class_map": {"car": 0, "pedestrian": 1}, "min_wh": 2.0}
        >>> dataset = YOLOGroundDataset(img_path="path/to/images", data=data, task="detect")
    """

    def __init__(self, *args, data: dict | None = None, task: str = "detect", **kwargs):
        """Initialize the YOLOGroundDataset with class mapping and difficulty support.

        Args:
            data (dict, optional): Dataset configuration with 'class_map' and optional 'min_wh', 'use_yuv444'.
            task (str): Task type, should be 'detect' for ground dataset.
            *args (Any): Additional positional arguments for the parent class.
            **kwargs (Any): Additional keyword arguments for the parent class.
        """
        self.class_map = data.get("class_map", {})
        self.min_wh = data.get("min_wh", 2.0)
        self.use_yuv444 = data.get("use_yuv444", False)
        self.difficulty_weights = data.get("difficulty_weights", [1.0])
        super().__init__(*args, data=data, task=task, **kwargs)

    def get_img_files(self, img_path: str | list[str]) -> list[str]:
        """Read image files while preserving source paths for label lookup.

        Ground 2D train/val list files may live under one ``Detection2D_*`` root while the corresponding images should
        be loaded from the YAML ``path`` root. The original list paths are retained for label lookup.
        """
        try:
            files = []
            for p in img_path if isinstance(img_path, list) else [img_path]:
                p = Path(p)
                if p.is_dir():
                    files += glob.glob(str(p / "**" / "*.*"), recursive=True)
                elif p.is_file():
                    with open(p, encoding="utf-8") as t:
                        lines = t.read().strip().splitlines()
                        parent = str(p.parent) + os.sep
                        files += [x.replace("./", parent) if x.startswith("./") else x for x in lines]
                else:
                    raise FileNotFoundError(f"{self.prefix}{p} does not exist")
            entries = sorted(x.replace("/", os.sep) for x in files if x.strip())
            assert entries, f"{self.prefix}No entries found in {img_path}."
        except Exception as e:
            raise FileNotFoundError(f"{self.prefix}Error loading data from {img_path}\n{HELP_URL}") from e
        if self.fraction < 1:
            entries = entries[: round(len(entries) * self.fraction)]

        self.ground_label_files = []
        im_files = []
        for f in entries:
            suffix = Path(f).suffix.lower().lstrip(".")
            if suffix in IMG_FORMATS:
                self.ground_label_files.append(img2label_paths([f])[0])
                im_files.append(self._ground_image_path_from_label_path(f))
            elif suffix == "txt":
                self.ground_label_files.append(f)
                im_files.append(self._ground_image_path_from_label_path(f, label_path=True))
            else:
                raise FileNotFoundError(f"{self.prefix}Unsupported ground list entry: {f}. {FORMATS_HELP_MSG}")
        assert im_files, f"{self.prefix}No images found in {img_path}. {FORMATS_HELP_MSG}"
        check_file_speeds(im_files, prefix=self.prefix)
        return im_files

    def _ground_image_path_from_label_path(self, im_file: str, label_path: bool = False) -> str:
        """Map a list-file path from a ``Detection2D_*``/``2Ddetection_*`` root onto the configured image root."""
        image_root = Path(self.data["path"])
        parts = Path(im_file).parts
        detection_root_index = next(
            (i for i, part in enumerate(parts) if part.lower().startswith(("detection2d", "2ddetection"))),
            None,
        )
        if detection_root_index is not None and detection_root_index + 1 < len(parts):
            rel_parts = list(parts[detection_root_index + 1 :])
        else:
            rel_path = Path(im_file) if not Path(im_file).is_absolute() else Path(*parts[1:])
            rel_parts = list(rel_path.parts)

        if label_path:
            image_path = image_root.joinpath(*rel_parts).with_suffix(".jpg")
            if image_path.exists():
                return str(image_path)
            image_rel_parts = ["images" if part == "labels" else part for part in rel_parts]
            image_path = image_root.joinpath(*image_rel_parts).with_suffix(".jpg")
            if image_path.exists():
                return str(image_path)
            for suffix in ("png", *sorted(IMG_FORMATS - {"jpg", "jpeg", "png"}), "jpeg"):
                candidate = image_path.with_suffix(f".{suffix}")
                if candidate.exists():
                    return str(candidate)
            return str(image_path)

        return str(image_root.joinpath(*rel_parts))

    def cache_labels(self, path: Path = Path("./labels.cache")) -> dict:
        """Cache dataset labels with difficulty scores, check images and read shapes.

        Args:
            path (Path): Path where to save the cache file.

        Returns:
            (dict): Dictionary containing cached labels and related information.
        """
        from ultralytics.data.utils import verify_image_label_ground

        x = {"labels": []}
        nm, nf, ne, nc, msgs = 0, 0, 0, 0, []  # number missing, found, empty, corrupt, messages
        desc = f"{self.prefix}Scanning {path.parent / path.stem}..."
        total = len(self.im_files)
        with ThreadPool(NUM_THREADS) as pool:
            results = pool.imap(
                func=verify_image_label_ground,
                iterable=zip(
                    self.im_files,
                    self.label_files,
                    repeat(self.prefix),
                    repeat(self.class_map),
                ),
            )
            pbar = TQDM(results, desc=desc, total=total)
            for im_file, lb, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar:
                nm += nm_f
                nf += nf_f
                ne += ne_f
                nc += nc_f
                if im_file:
                    # Map raw difficulty values to loss weights using difficulty_weights lookup
                    dw = self.difficulty_weights
                    if len(lb) > 0:
                        raw_diff = lb[:, 5].astype(int).clip(0, len(dw) - 1)
                        loss_weights = np.array([dw[d] for d in raw_diff], dtype=np.float32).reshape(-1, 1)
                        difficulty_levels = raw_diff.astype(np.int64).reshape(-1, 1)
                    else:
                        loss_weights = np.zeros((0, 1), dtype=np.float32)
                        difficulty_levels = np.zeros((0, 1), dtype=np.int64)

                    x["labels"].append(
                        {
                            "im_file": im_file,
                            "shape": shape,
                            "cls": lb[:, 0:1],  # n, 1
                            "bboxes": lb[:, 1:5],  # n, 4
                            "difficulties": loss_weights,  # n, 1 (pre-computed loss weights)
                            "difficulty_levels": difficulty_levels,  # n, 1 raw class target in [0, 3]
                            "segments": segments,
                            "keypoints": None,
                            "normalized": True,
                            "bbox_format": "xywh",
                        }
                    )
                if msg:
                    msgs.append(msg)
                pbar.desc = f"{desc} {nf} images, {nm + ne} backgrounds, {nc} corrupt"
            pbar.close()

        if msgs:
            LOGGER.info("\n".join(msgs))
        if nf == 0:
            LOGGER.warning(f"{self.prefix}No labels found in {path}. {HELP_URL}")
        x["hash"] = get_hash(self.label_files + self.im_files)
        x["results"] = nf, nm, ne, nc, len(self.im_files)
        x["msgs"] = msgs  # warnings
        save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
        return x

    def get_labels(self) -> list[dict]:
        """Return ground labels derived from the original train/val list paths."""
        self.label_files = getattr(self, "ground_label_files", img2label_paths(self.im_files))
        cache_path = Path(self.label_files[0]).parent.with_suffix(".cache")
        try:
            cache, exists = load_dataset_cache_file(cache_path), True
            assert cache["version"] == DATASET_CACHE_VERSION
            assert cache["hash"] == get_hash(self.label_files + self.im_files)
        except (FileNotFoundError, AssertionError, AttributeError, ModuleNotFoundError):
            cache, exists = self.cache_labels(cache_path), False

        nf, nm, ne, nc, n = cache.pop("results")
        if exists and LOCAL_RANK in {-1, 0}:
            d = f"Scanning {cache_path}... {nf} images, {nm + ne} backgrounds, {nc} corrupt"
            TQDM(None, desc=self.prefix + d, total=n, initial=n)
            if cache["msgs"]:
                LOGGER.info("\n".join(cache["msgs"]))

        [cache.pop(k) for k in ("hash", "version", "msgs")]
        labels = cache["labels"]
        if not labels:
            raise RuntimeError(
                f"No valid images found in {cache_path}. Images with incorrectly formatted labels are ignored. {HELP_URL}"
            )
        self.im_files = [lb["im_file"] for lb in labels]
        if sum(len(lb["cls"]) for lb in labels) == 0:
            LOGGER.warning(f"Labels are missing or empty in {cache_path}, training may not work correctly. {HELP_URL}")
        return labels

    def update_labels_info(self, label: dict) -> dict:
        """Update label format with difficulty scores.

        Args:
            label (dict): Label dictionary containing bboxes, difficulties, etc.

        Returns:
            (dict): Updated label dictionary with instances including difficulties.
        """
        bboxes = label.pop("bboxes")
        difficulties = label.pop("difficulties", None)
        difficulty_levels = label.pop("difficulty_levels", None)
        segments = label.pop("segments", [])
        keypoints = label.pop("keypoints", None)
        bbox_format = label.pop("bbox_format")
        normalized = label.pop("normalized")

        # Convert empty segments list to numpy array (required by Instances.denormalize)
        if len(segments) == 0:
            segments = np.zeros((0, 1000, 2), dtype=np.float32)

        # Create Instances with difficulties
        label["instances"] = Instances(
            bboxes,
            segments,
            keypoints,
            difficulties=difficulties,
            difficulty_levels=difficulty_levels,
            bbox_format=bbox_format,
            normalized=normalized,
        )
        return label

    def __getitem__(self, index: int) -> dict:
        """Return transformed label with post-augmentation min_wh filtering.

        Filters out boxes whose width and height are both smaller than min_wh pixels after all augmentations
        (resize, crop, mosaic, etc.) have been applied, ensuring the filter operates on final pixel sizes.
        """
        labels = self.transforms(self.get_image_and_label(index))
        if self.min_wh > 0 and "bboxes" in labels and len(labels["bboxes"]) > 0:
            bboxes = labels["bboxes"]  # normalized xywh after Format
            _, h, w = labels["img"].shape  # CHW
            w_px = bboxes[:, 2] * w
            h_px = bboxes[:, 3] * h
            valid = (w_px >= self.min_wh) | (h_px >= self.min_wh)
            if not valid.all():
                labels["bboxes"] = bboxes[valid]
                labels["cls"] = labels["cls"][valid]
                if "batch_idx" in labels:
                    labels["batch_idx"] = labels["batch_idx"][valid]
                if "difficulties" in labels:
                    labels["difficulties"] = labels["difficulties"][valid]
                if "difficulty_levels" in labels:
                    labels["difficulty_levels"] = labels["difficulty_levels"][valid]
        return labels


class YOLOGround3DDataset(YOLOGroundDataset):
    """Dataset for ground 3D detection backed by GT list files and on-the-fly label loading.

    Extends YOLOGroundDataset to support joint 2D+3D detection with:
    - On-the-fly parsing of 19-col (complete_3d) and 51-col (face_3d) label files
    - A simple GT-manifest list used for dataset length and sample lookup
    - Only color-space augmentations (3D labels are in camera space)
    - Virtual camera / fisheye augmentation support

    Attributes:
        face_3d_classes (set): Class IDs with 4-face annotations.
        complete_3d_classes (set): Class IDs with whole-box 3D only.
        norm_scales_3d (dict): Normalization scales for 3D loss computation.
    """

    def __init__(
        self,
        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,
        data: dict | None = None,
        task: str = "detect",
    ):
        init_start = perf_counter()
        data = data or {}
        self.use_segments = task == "segment"
        self.use_keypoints = task == "pose"
        self.use_obb = task == "obb"
        assert not (self.use_segments and self.use_keypoints), "Can not use both segments and keypoints."
        missing_keys = [key for key in GROUND3D_REQUIRED_DATA_KEYS if key not in data or data[key] is None]
        if missing_keys:
            missing = ", ".join(missing_keys)
            raise ValueError(f"Ground3D dataset config must define required field(s) explicitly: {missing}.")

        path_value = str(_require_ground3d_data_value(data, "path")).strip()    # 检查 3D dataset 必需配置是否缺失。依赖文件顶部 GROUND3D_REQUIRED_DATA_KEYS = ("path", "ori_img_size", "roi", "virtual_fx", "virtual_camera_prob", "crop_center_mode")。
        if not path_value:
            raise ValueError("Ground3D dataset config field 'path' must be a non-empty string.")

        normalized_ground3d_cfg = {
            "path": str(Path(path_value).expanduser()),                                                                               # 读取 data["path"]，也就是数据根目录。
            "ori_img_size": list(_validate_ground3d_pair(_require_ground3d_data_value(data, "ori_img_size"), "ori_img_size")),        # 
            "roi": list(_validate_ground3d_pair(_require_ground3d_data_value(data, "roi"), "roi")),
            "virtual_fx": _validate_ground3d_positive_float(_require_ground3d_data_value(data, "virtual_fx"), "virtual_fx"),
            "virtual_camera_prob": _validate_ground3d_float(
                _require_ground3d_data_value(data, "virtual_camera_prob"), "virtual_camera_prob"
            ),
            "crop_center_mode": _validate_ground3d_crop_center_mode(_require_ground3d_data_value(data, "crop_center_mode")),
        }

        self.data = {**data, **normalized_ground3d_cfg}
        self.class_map = data.get("class_map", {})                           # class_map：字符串类别名到 class id 的映射。3D label 里类别可能是名字，需要映射成数字类别
        self.min_wh = data.get("min_wh", 2.0)                                # min_wh：过滤过小 2D 框的阈值，默认 2 像素
        self.use_yuv444 = data.get("use_yuv444", False)
        self.difficulty_weights = data.get("difficulty_weights", [1.0])
        self.face_3d_classes = set(data.get("face_3d_classes", []))
        self.complete_3d_classes = set(data.get("complete_3d_classes", []))
        self.norm_scales_3d = data.get("norm_scales_3d", {})                 # 3D loss 归一化尺度，例如 depth、尺寸、offset 等归一化参数
        self.image_root = Path(self.data["path"])                            # 数据根目录 Path 对象，后续从 label 相对路径推导 image/calib 路径。
        self.ori_img_size = self.data["ori_img_size"]
        self.roi = self.data["roi"]
        self.virtual_fx = self.data["virtual_fx"]
        self.virtual_camera_prob = self.data["virtual_camera_prob"]
        self.crop_center_mode = self.data["crop_center_mode"]
        self.virtual_camera_val_zoom = data.get("virtual_camera_val_zoom", False)           # 验证阶段是否也启用 virtual camera zoom 相关逻辑。默认关闭
        self._ground3d_shape = (int(self.ori_img_size[1]), int(self.ori_img_size[0]))
        self._include_class = None
        self.face_visibility_score_thresh = float(
            getattr(hyp, "face_visibility_score_thresh", DEFAULT_CFG.face_visibility_score_thresh)       # 可见面分数阈值，从 hyp 里取，取不到就用 DEFAULT_CFG。后续构建 edge GT / face GT 时判断某个 face 是否有效。
        )
        self.precompute_edge_gt = float(getattr(hyp, "edge_aux_loss_gain", 1.0)) > 0
        self.profile_timing = bool(getattr(hyp, "batch_timing", False))
        self.skip_bad_images = bool(data.get("skip_bad_images", True))
        self.strict_image_shape = bool(data.get("strict_image_shape", "ori_img_size" in data))
        self.min_image_side = max(int(data.get("min_image_side", 10)), 1)

        self.img_path = img_path
        self.imgsz = imgsz
        self.augment = augment
        self.single_cls = single_cls
        self.prefix = prefix
        self.fraction = fraction
        self.channels = self.data.get("channels", 3)
        self.cv2_flag = cv2.IMREAD_GRAYSCALE if self.channels == 1 else cv2.IMREAD_COLOR         # 根据通道数决定 OpenCV 按灰度还是彩色读取。Ground 3D 通常是 3 通道，所以一般是 cv2.IMREAD_COLOR

        self.label_entries = self._load_ground3d_label_entries(self.img_path)
        self.labels = self.label_entries
        self.update_labels(include_class=classes)
        self.ni = len(self.labels)
        self._bad_image_mask = np.zeros(self.ni, dtype=bool)
        self._warned_bad_image_indices: set[int] = set()

        self.rect = rect
        self.batch_size = batch_size
        self.stride = stride
        self.pad = pad
        if self.rect:
            assert self.batch_size is not None
            self.set_rectangle()

        self.buffer = [] 
        self.max_buffer_length = min((self.ni, self.batch_size * 8, 1000)) if self.augment else 0          # 训练增强时最多缓存 min(样本数, batch*8, 1000) 张；验证或不增强时为 0
        self.ims, self.im_hw0, self.im_hw = [None] * self.ni, [None] * self.ni, [None] * self.ni
        self.npy_files = []
        self.cache = None
        if cache:
            LOGGER.warning(f"{self.prefix}Ground3D dataset ignores cache={cache} and loads samples on-the-fly")

        self.transforms = self.build_transforms(hyp=hyp)
        if LOCAL_RANK in {-1, 0}:
            LOGGER.info(
                f"{self.prefix}Ground3D dataset ready with {self.ni:,} samples in {perf_counter() - init_start:.1f}s"
            )

    def _load_ground3d_label_entries(self, img_path: str | list[str]) -> list[tuple[str, str]]:
        """Load Ground3D GT list files and return `(list_root, rel_label)` entries. 说明返回二元组：list_root 是 list 文件所在目录，rel_label 是 list 里写的 label 相对路径"""
        items = img_path if isinstance(img_path, list) else [img_path]
        label_entries = []
        log_progress = LOCAL_RANK in {-1, 0}
        scan_start = perf_counter()
        total_entries = 0

        if log_progress:
            LOGGER.info(f"{self.prefix}Initializing Ground3D dataset from {len(items)} GT list file(s)...")

        try:
            for item in items:
                item = Path(item)
                if not item.is_file():
                    raise FileNotFoundError(f"{self.prefix}Ground3D dataset expects GT list files, but got: {item}")
                if log_progress:
                    LOGGER.info(f"{self.prefix}Reading GT list: {item}")

                with open(item, encoding="utf-8") as f:
                    file_has_entries = False
                    for line in f:
                        entry = line.strip()
                        if not entry or entry.lstrip().startswith("#"):
                            continue
                        if Path(entry).suffix.lower() != ".txt":
                            raise ValueError(f"{self.prefix}Ground3D GT list entries must be label .txt paths, but got: {entry}")

                        file_has_entries = True
                        label_entries.append((str(item.parent.resolve()), entry))
                        total_entries += 1
                        if log_progress and total_entries % 100_000 == 0:
                            LOGGER.info(
                                f"{self.prefix}Parsed {total_entries:,} GT entries in {perf_counter() - scan_start:.1f}s"
                            )

                if not file_has_entries:
                    raise FileNotFoundError(f"{self.prefix}No GT entries found in {item}.")

            assert label_entries, f"{self.prefix}No GT entries found in {img_path}."
        except Exception as e:
            raise FileNotFoundError(f"{self.prefix}Error loading Ground3D GT list from {img_path}\n{HELP_URL}") from e

        if self.fraction < 1:
            n = round(len(label_entries) * self.fraction)
            label_entries = label_entries[:n]
            if log_progress:
                LOGGER.info(
                    f"{self.prefix}Applied dataset fraction={self.fraction:g}: keeping {len(label_entries):,}/{total_entries:,} entries"
                )

        if log_progress:
            LOGGER.info(
                f"{self.prefix}Loaded {len(label_entries):,} Ground3D GT entries from {len(items)} list file(s) "
                f"in {perf_counter() - scan_start:.1f}s"
            )
            LOGGER.info(f"{self.prefix}Checking sampled image access speed...")
        sample_entries = random.sample(label_entries, min(5, len(label_entries)))
        check_file_speeds([self._entry_to_image_file(entry) for entry in sample_entries], prefix=self.prefix)
        return label_entries                    # label 路径里的 labels 会被换成 images，后缀优先 .png，找不到再尝试 .jpg。

    @staticmethod
    def _label_rel_to_image_rel(rel_label: Path) -> Path:
        """Convert a relative Ground3D label path to its matching image path."""
        parts = list(rel_label.parts)
        if "labels" in parts:
            parts[len(parts) - 1 - parts[::-1].index("labels")] = "images"
            rel_label = Path(*parts)
        return rel_label.with_suffix(".png")

    @staticmethod
    def _entry_to_rel_label(label_entry: tuple[str, str]) -> Path:
        """Return the relative label path stored in a Ground3D GT entry."""
        return Path(label_entry[1])

    def _entry_to_label_file(self, label_entry: tuple[str, str]) -> str:
        """Resolve a Ground3D GT entry to its absolute label file path."""
        list_root, _ = label_entry
        return str((Path(list_root) / self._entry_to_rel_label(label_entry)).resolve())

    @staticmethod
    def _label_rel_to_calib_rel(rel_label: Path) -> Path:
        """Convert a relative Ground3D label path to its matching calibration path."""
        parts = list(rel_label.parts)
        if "labels" in parts:
            parts[len(parts) - 1 - parts[::-1].index("labels")] = "calib"
            rel_label = Path(*parts)
        return rel_label.with_suffix(".json")

    def _entry_to_calib_file(self, label_entry: tuple[str, str]) -> str:
        """Resolve a Ground3D GT entry to its absolute calibration file path under the label root."""
        list_root, _ = label_entry
        return str((Path(list_root) / self._label_rel_to_calib_rel(self._entry_to_rel_label(label_entry))).resolve())

    def _entry_to_image_file(self, label_entry: tuple[str, str]) -> str:
        """Resolve a Ground3D GT entry to its absolute image file path."""
        rel_image = self._label_rel_to_image_rel(self._entry_to_rel_label(label_entry))
        image_file = (self.image_root / rel_image).resolve()
        if image_file.exists():
            return str(image_file)

        jpg_image_file = image_file.with_suffix(".jpg")
        if jpg_image_file.exists():
            return str(jpg_image_file)

        return str(image_file)

    def update_labels(self, include_class: list[int] | None) -> None:
        """Record class filtering rules and apply them lazily on parsed Ground3D labels."""
        self._include_class = np.array(include_class, dtype=np.int64).reshape(1, -1) if include_class is not None else None

    def set_rectangle(self) -> None:
        """Configure rectangular batching from the fixed Ground3D source shape."""
        if self.ni == 0:
            self.batch = np.zeros(0, dtype=int)
            self.batch_shapes = np.zeros((0, 2), dtype=int)
            return

        bi = np.floor(np.arange(self.ni) / self.batch_size).astype(int)
        nb = bi[-1] + 1
        ori_h, ori_w = self._ground3d_shape
        ar = ori_h / ori_w if ori_w else 1.0

        if ar < 1:
            base_shape = [ar, 1.0]
        elif ar > 1:
            base_shape = [1.0, 1.0 / ar]
        else:
            base_shape = [1.0, 1.0]

        imgsz = np.array(self.imgsz if isinstance(self.imgsz, (list, tuple)) else [self.imgsz, self.imgsz], dtype=np.float32)
        shapes = np.tile(np.array(base_shape, dtype=np.float32), (nb, 1))
        self.batch_shapes = np.ceil(shapes * imgsz / self.stride + self.pad).astype(int) * self.stride
        self.batch = bi

    def _filter_ground3d_labels(self, lb_2d: dict[str, Any], lb_3d: np.ndarray | None) -> tuple[dict[str, Any], np.ndarray | None]:
        """Apply class selection lazily to on-the-fly Ground3D labels."""
        if self._include_class is not None and len(lb_2d["cls"]):
            keep = (lb_2d["cls"] == self._include_class).any(1)
            lb_2d = {
                **lb_2d,
                "cls": lb_2d["cls"][keep],
                "bboxes": lb_2d["bboxes"][keep],
                "difficulties": lb_2d["difficulties"][keep],
                "difficulty_levels": lb_2d["difficulty_levels"][keep],
                "segments": [lb_2d["segments"][si] for si, idx in enumerate(keep.tolist()) if idx] if lb_2d["segments"] else [],
                "keypoints": lb_2d["keypoints"][keep] if lb_2d["keypoints"] is not None else None,
            }
            if lb_3d is not None:
                lb_3d = lb_3d[keep]

        if self.single_cls and len(lb_2d["cls"]):
            lb_2d = {**lb_2d, "cls": lb_2d["cls"].copy()}
            lb_2d["cls"][:, 0] = 0

        return lb_2d, lb_3d

    def _mark_bad_image(self, index: int, im_file: str, reason: str) -> None:
        """Remember bad images so later epochs skip them immediately."""
        self._bad_image_mask[index] = True
        if index not in self._warned_bad_image_indices:
            LOGGER.warning(f"{self.prefix}Skipping bad Ground3D image [{index}] {im_file}: {reason}")
            self._warned_bad_image_indices.add(index)

    def _validate_ground3d_image(self, index: int, im_file: str, img: np.ndarray | None) -> np.ndarray:
        """Validate a decoded Ground3D image before calibration and label work."""
        if self._bad_image_mask[index]:
            raise Ground3DImageReadError(f"{im_file}: previously marked unreadable or invalid")
        if img is None:
            self._mark_bad_image(index, im_file, "cv2.imread() returned None")
            raise Ground3DImageReadError(f"{im_file}: cv2.imread() returned None")
        if not isinstance(img, np.ndarray):
            self._mark_bad_image(index, im_file, f"decoder returned {type(img).__name__}, expected ndarray")
            raise Ground3DImageReadError(f"{im_file}: decoder returned non-array image data")

        shape = tuple(int(v) for v in img.shape[:2])
        if self.channels == 1:
            valid_channels = img.ndim == 2
            expected_shape = "HxW"
        else:
            valid_channels = img.ndim == 3 and img.shape[2] == self.channels
            expected_shape = f"HxWx{self.channels}"
        if not valid_channels:
            self._mark_bad_image(index, im_file, f"unexpected decoded shape {img.shape}, expected {expected_shape}")
            raise Ground3DImageReadError(f"{im_file}: unexpected decoded shape {img.shape}")
        if min(shape) < self.min_image_side:
            self._mark_bad_image(index, im_file, f"decoded image is too small: {shape}")
            raise Ground3DImageReadError(f"{im_file}: decoded image is too small: {shape}")
        if self.strict_image_shape and shape != self._ground3d_shape:
            self._mark_bad_image(
                index,
                im_file,
                f"decoded shape {shape} does not match expected {self._ground3d_shape}",
            )
            raise Ground3DImageReadError(f"{im_file}: decoded shape {shape} does not match expected {self._ground3d_shape}")
        return img

    def _read_ground3d_image(self, index: int, im_file: str) -> np.ndarray:
        """Read and validate a Ground3D image, failing with a skip-safe exception."""
        if self._bad_image_mask[index]:
            raise Ground3DImageReadError(f"{im_file}: previously marked unreadable or invalid")
        img = cv2.imread(im_file, self.cv2_flag)
        return self._validate_ground3d_image(index, im_file, img)

    def _load_transformed_ground3d_sample(self, index: int) -> dict[str, Any]:
        """Load, transform, and finalize one Ground3D sample."""
        labels = self.transforms(self.get_image_and_label(index))

        # Convert labels_3d: numpy/None -> tensor
        if labels.get("labels_3d") is not None:
            if isinstance(labels["labels_3d"], np.ndarray):
                labels["labels_3d"] = torch.from_numpy(labels["labels_3d"].copy())
        else:
            nl = len(labels["bboxes"]) if "bboxes" in labels else 0
            # Keep absent 3D supervision as NaN so downstream 3D loss/metrics skip pure 2D objects.
            labels["labels_3d"] = torch.full((nl, 42), float("nan"), dtype=torch.float32)

        if self.min_wh > 0 and "bboxes" in labels and len(labels["bboxes"]) > 0:
            bboxes = labels["bboxes"]
            _, h, w = labels["img"].shape
            w_px = bboxes[:, 2] * w
            h_px = bboxes[:, 3] * h
            valid = (w_px >= self.min_wh) | (h_px >= self.min_wh)
            if not valid.all():
                labels["bboxes"] = bboxes[valid]
                labels["cls"] = labels["cls"][valid]
                if "batch_idx" in labels:
                    labels["batch_idx"] = labels["batch_idx"][valid]
                if "difficulties" in labels:
                    labels["difficulties"] = labels["difficulties"][valid]
                if "difficulty_levels" in labels:
                    labels["difficulty_levels"] = labels["difficulty_levels"][valid]
                labels["labels_3d"] = labels["labels_3d"][valid]
        edge_start = perf_counter() if self.profile_timing else None
        self._build_edge_gt(labels)
        if edge_start is not None and isinstance(labels.get("_profile_timing"), dict):
            labels["_profile_timing"]["precompute_edge_gt"] = perf_counter() - edge_start
        return labels

    def get_image_and_label(self, index):
        """Load a mono3D sample through an explicit ROI or virtual-camera preprocessing pipeline."""
        label_entry = self.label_entries[index]
        label_file = self._entry_to_label_file(label_entry)
        im_file = self._entry_to_image_file(label_entry)
        calib_file = self._entry_to_calib_file(label_entry)
        label = {"im_file": im_file}
        timings = {} if self.profile_timing else None

        t0 = perf_counter()
        img = self._read_ground3d_image(index, im_file)
        if timings is not None:
            timings["read_image"] = perf_counter() - t0

        ori_h, ori_w = img.shape[:2]
        t0 = perf_counter()
        lb_2d, lb_3d = parse_ground_3d_label_file(
            label_file,
            self.class_map,
            self.difficulty_weights,
            self.face_3d_classes,
            self.complete_3d_classes,
            self.min_wh,
        )
        lb_2d, lb_3d = self._filter_ground3d_labels(lb_2d, lb_3d)
        if timings is not None:
            timings["parse_label"] = perf_counter() - t0
        has_3d_targets = _ground3d_has_finite_targets(lb_3d)
        t0 = perf_counter()
        raw_calib = read_calib_from_path(im_file, image_root=self.image_root, extra_calib_candidates=[calib_file])
        if timings is not None:
            timings["read_calib"] = perf_counter() - t0
        if raw_calib is None and has_3d_targets:
            expected_calib = Path(calib_file)
            expected_calib = (
                expected_calib if expected_calib.name == "camera4.json" else expected_calib.parent / "L2_calib" / "camera4.json"
            )
            LOGGER.error(f"{self.prefix}Missing Ground3D calibration [{index}] {im_file} (expected {expected_calib})")
            raise Ground3DCalibrationError(
                f"{im_file}: calibration file not found (expected clip-level camera file at {expected_calib})"
            )

        target_w, target_h = self.imgsz if isinstance(self.imgsz, (list, tuple)) else (self.imgsz, self.imgsz)
        roi_w, roi_h = self.roi
        use_virtual_camera = (
            raw_calib is not None and 0 < self.virtual_camera_prob <= 1 and random.random() < self.virtual_camera_prob
        )
        use_virtual_camera_zoom = self.augment or self.virtual_camera_val_zoom

        roi_active = (raw_calib is not None or not has_3d_targets) and (roi_h < ori_h or roi_w < ori_w)
        sample_ori_shape = (ori_h, ori_w)

        vp_x = compute_vanishing_point_x(raw_calib, ori_w)
        vp_y = compute_vanishing_point_y(raw_calib, ori_h)
        crop_center_x = vp_x if self.crop_center_mode == "vxvy" else ori_w / 2

        label["camera_mode"] = "virtual" if use_virtual_camera else "roi"

        if use_virtual_camera:
            t0 = perf_counter()
            simul_calib = compute_simul_calib(
                raw_calib,
                (ori_w, ori_h),
                (target_w, target_h),
                crop_center_x,
                vp_y,
                target_fx=self.virtual_fx,
                augment=use_virtual_camera_zoom,
            )
            if timings is not None:
                timings["compute_simul_calib"] = perf_counter() - t0
            labels_48 = pack_labels_to_48(lb_2d, lb_3d)
            t0 = perf_counter()
            img, labels_48 = apply_simul_transform(
                img, labels_48, simul_calib, raw_calib, (target_w, target_h), augment=self.augment
            )
            if timings is not None:
                timings["apply_simul_transform"] = perf_counter() - t0
            lb_2d, lb_3d = unpack_labels_from_48(labels_48)
            final_calib = {
                "fx": simul_calib["fx"],
                "fy": simul_calib["fy"],
                "cx": simul_calib["cx"],
                "cy": simul_calib["cy"],
                "depth_scale": simul_calib["depth_scale"],
            }
            sample_ori_shape = (target_h, target_w)
        else:
            crop_bounds = (0, 0, ori_w, ori_h)
            roi_crop_time = 0.0
            if roi_active:
                t0 = perf_counter()
                crop_bounds = compute_centered_roi_bounds(ori_w, ori_h, roi_w, roi_h, crop_center_x, vp_y)
                crop_x1, crop_y1, crop_x2, crop_y2 = crop_bounds
                img = img[crop_y1:crop_y2, crop_x1:crop_x2]
                sample_ori_shape = (crop_y2 - crop_y1, crop_x2 - crop_x1)
                roi_crop_time = perf_counter() - t0

            t0 = perf_counter()
            calib_for_resize = adjust_calib_for_roi_crop(raw_calib, ori_w, ori_h, crop_bounds)
            img = _resize_ground3d_image_in_steps(img, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
            final_calib = build_final_resized_calib(
                calib_for_resize["focal_u"],
                calib_for_resize["focal_v"],
                calib_for_resize["cu"],
                calib_for_resize["cv"],
                calib_for_resize["src_w"],
                calib_for_resize["src_h"],
                target_w,
                target_h,
                self.virtual_fx,
                distort_coeffs=calib_for_resize["distort_coeffs"],
            )
            if timings is not None:
                timings["roi_resize"] = perf_counter() - t0
                if roi_crop_time > 0:
                    timings["roi_crop"] = roi_crop_time
            if roi_active:
                t0 = perf_counter()
                lb_2d, lb_3d = remap_labels_to_roi(lb_2d, lb_3d, ori_w, ori_h, crop_bounds)
                if timings is not None:
                    timings["remap_labels_to_roi"] = perf_counter() - t0
            t0 = perf_counter()
            lb_3d = normalize_roi_depth(lb_3d, final_calib["fx"], self.virtual_fx)
            if timings is not None:
                timings["normalize_roi_depth"] = perf_counter() - t0

        label["ori_shape"] = sample_ori_shape
        label["resized_shape"] = (target_h, target_w)
        label["ratio_pad"] = (target_h / sample_ori_shape[0], target_w / sample_ori_shape[1])
        label["img"] = img
        label["calib"] = final_calib
        if timings is not None:
            timings["total_sample_build"] = sum(timings.values())
            label["_profile_timing"] = timings
        if self.rect:
            label["rect_shape"] = self.batch_shapes[self.batch[index]]

        label.update(lb_2d)
        label["labels_3d"] = lb_3d
        return self.update_labels_info(label)

    def build_transforms(self, hyp=None):
        """Build transforms: only color-space augmentations for 3D training.

        No geometric augmentations (Mosaic, RandomPerspective, MixUp, CutMix, RandomFlip)
        since 3D labels are in camera space, not image space.
        No LetterBox since image is already resized to target in get_image_and_label.
        """
        transforms = []
        if self.augment and hyp:
            transforms.append(RandomHSV(hgain=hyp.hsv_h, sgain=hyp.hsv_s, vgain=hyp.hsv_v))
        transforms.append(
            Format(
                bbox_format="xywh",
                normalize=True,
                batch_idx=True,
                bgr=hyp.bgr if self.augment and hyp else 0.0,
            )
        )
        return Compose(transforms)

    def update_labels_info(self, label):
        """Update label format, preserving labels_3d, calib, and camera mode in the dict."""
        labels_3d = label.pop("labels_3d", None)
        calib = label.pop("calib", None)
        camera_mode = label.pop("camera_mode", None)
        label = super().update_labels_info(label)
        label["labels_3d"] = labels_3d
        label["calib"] = calib
        label["camera_mode"] = camera_mode
        return label

    @staticmethod
    def _empty_edge_gt(nl: int) -> dict[str, torch.Tensor]:
        """Return empty precomputed edge GT tensors aligned with the object dimension."""
        return {
            "edge_faces_points_2d": torch.zeros((nl, 4, 5, 2), dtype=torch.float32),
            "edge_faces_depths": torch.zeros((nl, 4, 5), dtype=torch.float32),
            "edge_faces_valid": torch.zeros((nl, 4), dtype=torch.bool),
            "edge_partial_points_2d": torch.zeros((nl, 5, 2), dtype=torch.float32),
            "edge_partial_depths": torch.zeros((nl, 5), dtype=torch.float32),
            "edge_partial_face_type": torch.full((nl,), -1, dtype=torch.long),
            "edge_partial_valid": torch.zeros((nl,), dtype=torch.bool),
        }

    def _build_edge_gt(self, labels: dict[str, Any]) -> None:
        """Precompute GT edge supervision on CPU inside dataloader workers."""
        nl = int(len(labels.get("labels_3d", [])))
        labels.update(self._empty_edge_gt(nl))
        if not self.precompute_edge_gt or nl == 0:
            return

        calib = labels.get("calib")
        if not isinstance(calib, dict):
            return

        labels_3d = labels["labels_3d"]
        cls = labels["cls"]
        bboxes = labels["bboxes"]
        _, img_h, img_w = labels["img"].shape

        labels_3d_np = labels_3d.cpu().numpy() if isinstance(labels_3d, torch.Tensor) else np.asarray(labels_3d)
        cls_np = cls.view(-1).cpu().numpy().astype(int) if isinstance(cls, torch.Tensor) else np.asarray(cls).reshape(-1).astype(int)
        bboxes_np = bboxes.cpu().numpy() if isinstance(bboxes, torch.Tensor) else np.asarray(bboxes)

        xyxy = np.zeros((nl, 4), dtype=np.float32)
        xyxy[:, 0] = (bboxes_np[:, 0] - bboxes_np[:, 2] / 2) * img_w
        xyxy[:, 1] = (bboxes_np[:, 1] - bboxes_np[:, 3] / 2) * img_h
        xyxy[:, 2] = (bboxes_np[:, 0] + bboxes_np[:, 2] / 2) * img_w
        xyxy[:, 3] = (bboxes_np[:, 1] + bboxes_np[:, 3] / 2) * img_h

        for obj_idx in range(nl):
            target_42 = labels_3d_np[obj_idx]
            cls_id = int(cls_np[obj_idx])
            bbox_xyxy = xyxy[obj_idx]

            partial_edge = decode_cut_partial_side_edge_from_gt(
                target_42,
                cls_id,
                calib,
                int(img_w),
                int(img_h),
                self.face_3d_classes,
                self.complete_3d_classes,
                bbox_xyxy=bbox_xyxy,
            )
            if partial_edge is not None:
                labels["edge_partial_points_2d"][obj_idx] = torch.from_numpy(partial_edge["points_2d"])
                labels["edge_partial_depths"][obj_idx] = torch.from_numpy(partial_edge["depths"])
                labels["edge_partial_face_type"][obj_idx] = int(partial_edge["face_type"])
                labels["edge_partial_valid"][obj_idx] = True

            for face_type in range(4):
                decoded = partial_edge if partial_edge is not None and int(partial_edge["face_type"]) == face_type else None
                if decoded is None:
                    decoded = decode_visible_face_edge_from_gt(
                        target_42,
                        cls_id,
                        calib,
                        int(img_w),
                        int(img_h),
                        self.face_3d_classes,
                        self.complete_3d_classes,
                        face_type=face_type,
                        score_thr=self.face_visibility_score_thresh,
                        bbox_xyxy=bbox_xyxy,
                    )
                if decoded is None:
                    continue
                labels["edge_faces_points_2d"][obj_idx, face_type] = torch.from_numpy(decoded["points_2d"])
                labels["edge_faces_depths"][obj_idx, face_type] = torch.from_numpy(decoded["depths"])
                labels["edge_faces_valid"][obj_idx, face_type] = True

    def __getitem__(self, index):
        """Return transformed label with post-augmentation min_wh filtering for both 2D and 3D."""
        if not self.skip_bad_images:
            return self._load_transformed_ground3d_sample(int(index))

        last_error = None
        start_index = int(index)
        for offset in range(self.ni):
            sample_index = (start_index + offset) % self.ni
            try:
                return self._load_transformed_ground3d_sample(sample_index)
            except Ground3DImageReadError as exc:
                last_error = exc
                continue

        raise RuntimeError(
            f"{self.prefix}Failed to load a valid Ground3D image after checking all {self.ni} samples."
        ) from last_error


class YOLOMultiModalDataset(YOLODataset):
    """Dataset class for loading object detection and/or segmentation labels in YOLO format with multi-modal support.

    This class extends YOLODataset to add text information for multi-modal model training, enabling models to process
    both image and text data.

    Methods:
        update_labels_info: Add text information for multi-modal model training.
        build_transforms: Enhance data transformations with text augmentation.

    Examples:
        >>> dataset = YOLOMultiModalDataset(img_path="path/to/images", data={"names": {0: "person"}}, task="detect")
        >>> batch = next(iter(dataset))
        >>> print(batch.keys())  # Should include 'texts'
    """

    def __init__(self, *args, data: dict | None = None, task: str = "detect", **kwargs):
        """Initialize a YOLOMultiModalDataset.

        Args:
            data (dict, optional): Dataset configuration dictionary.
            task (str): Task type, one of 'detect', 'segment', 'pose', or 'obb'.
            *args (Any): Additional positional arguments for the parent class.
            **kwargs (Any): Additional keyword arguments for the parent class.
        """
        super().__init__(*args, data=data, task=task, **kwargs)

    def update_labels_info(self, label: dict) -> dict:
        """Add text information for multi-modal model training.

        Args:
            label (dict): Label dictionary containing bboxes, segments, keypoints, etc.

        Returns:
            (dict): Updated label dictionary with instances and texts.
        """
        labels = super().update_labels_info(label)
        # NOTE: some categories are concatenated with its synonyms by `/`.
        # NOTE: and `RandomLoadText` would randomly select one of them if there are multiple words.
        labels["texts"] = [v.split("/") for _, v in self.data["names"].items()]

        return labels

    def build_transforms(self, hyp: dict | None = None) -> Compose:
        """Enhance data transformations with optional text augmentation for multi-modal training.

        Args:
            hyp (dict, optional): Hyperparameters for transforms.

        Returns:
            (Compose): Composed transforms including text augmentation if applicable.
        """
        transforms = super().build_transforms(hyp)
        if self.augment:
            # NOTE: hard-coded the args for now.
            # NOTE: this implementation is different from official yoloe,
            # the strategy of selecting negative is restricted in one dataset,
            # while official pre-saved neg embeddings from all datasets at once.
            transform = RandomLoadText(
                max_samples=min(self.data["nc"], 80),
                padding=True,
                padding_value=self._get_neg_texts(self.category_freq),
            )
            transforms.insert(-1, transform)
        return transforms

    @property
    def category_names(self):
        """Return category names for the dataset.

        Returns:
            (set[str]): Set of class names.
        """
        names = self.data["names"].values()
        return {n.strip() for name in names for n in name.split("/")}  # category names

    @property
    def category_freq(self):
        """Return frequency of each category in the dataset."""
        texts = [v.split("/") for v in self.data["names"].values()]
        category_freq = defaultdict(int)
        for label in self.labels:
            for c in label["cls"].squeeze(-1):  # to check
                text = texts[int(c)]
                for t in text:
                    t = t.strip()
                    category_freq[t] += 1
        return category_freq

    @staticmethod
    def _get_neg_texts(category_freq: dict, threshold: int = 100) -> list[str]:
        """Get negative text samples based on frequency threshold."""
        threshold = min(max(category_freq.values()), 100)
        return [k for k, v in category_freq.items() if v >= threshold]


class GroundingDataset(YOLODataset):
    """Dataset class for object detection tasks using annotations from a JSON file in grounding format.

    This dataset is designed for grounding tasks where annotations are provided in a JSON file rather than the standard
    YOLO format text files.

    Attributes:
        json_file (str): Path to the JSON file containing annotations.

    Methods:
        get_img_files: Return empty list as image files are read in get_labels.
        get_labels: Load annotations from a JSON file and prepare them for training.
        build_transforms: Configure augmentations for training with optional text loading.

    Examples:
        >>> dataset = GroundingDataset(img_path="path/to/images", json_file="annotations.json", task="detect")
        >>> len(dataset)  # Number of valid images with annotations
    """

    def __init__(self, *args, task: str = "detect", json_file: str = "", max_samples: int = 80, **kwargs):
        """Initialize a GroundingDataset for object detection.

        Args:
            json_file (str): Path to the JSON file containing annotations.
            task (str): Must be 'detect' or 'segment' for GroundingDataset.
            max_samples (int): Maximum number of samples to load for text augmentation.
            *args (Any): Additional positional arguments for the parent class.
            **kwargs (Any): Additional keyword arguments for the parent class.
        """
        assert task in {"detect", "segment"}, "GroundingDataset currently only supports `detect` and `segment` tasks"
        self.json_file = json_file
        self.max_samples = max_samples
        super().__init__(*args, task=task, data={"channels": 3}, **kwargs)

    def get_img_files(self, img_path: str) -> list:
        """The image files would be read in `get_labels` function, return empty list here.

        Args:
            img_path (str): Path to the directory containing images.

        Returns:
            (list): Empty list as image files are read in get_labels.
        """
        return []

    def verify_labels(self, labels: list[dict[str, Any]]) -> None:
        """Verify the number of instances in the dataset matches expected counts.

        This method checks if the total number of bounding box instances in the provided labels matches the expected
        count for known datasets. It performs validation against a predefined set of datasets with known instance
        counts.

        Args:
            labels (list[dict[str, Any]]): List of label dictionaries, where each dictionary contains dataset
                annotations. Each label dict must have a 'bboxes' key with a numpy array or tensor containing bounding
                box coordinates.

        Raises:
            AssertionError: If the actual instance count doesn't match the expected count for a recognized dataset.

        Notes:
            For unrecognized datasets (those not in the predefined expected_counts),
            a warning is logged and verification is skipped.
        """
        expected_counts = {
            "final_mixed_train_no_coco_segm": 3662412,
            "final_mixed_train_no_coco": 3681235,
            "final_flickr_separateGT_train_segm": 638214,
            "final_flickr_separateGT_train": 640704,
        }

        instance_count = sum(label["bboxes"].shape[0] for label in labels)
        for data_name, count in expected_counts.items():
            if data_name in self.json_file:
                assert instance_count == count, f"'{self.json_file}' has {instance_count} instances, expected {count}."
                return
        LOGGER.warning(f"Skipping instance count verification for unrecognized dataset '{self.json_file}'")

    def cache_labels(self, path: Path = Path("./labels.cache")) -> dict[str, Any]:
        """Load annotations from a JSON file, filter, and normalize bounding boxes for each image.

        Args:
            path (Path): Path where to save the cache file.

        Returns:
            (dict[str, Any]): Dictionary containing cached labels and related information.
        """
        x = {"labels": []}
        LOGGER.info("Loading annotation file...")
        with open(self.json_file) as f:
            annotations = json.load(f)
        images = {f"{x['id']:d}": x for x in annotations["images"]}
        img_to_anns = defaultdict(list)
        for ann in annotations["annotations"]:
            img_to_anns[ann["image_id"]].append(ann)
        for img_id, anns in TQDM(img_to_anns.items(), desc=f"Reading annotations {self.json_file}"):
            img = images[f"{img_id:d}"]
            h, w, f = img["height"], img["width"], img["file_name"]
            im_file = Path(self.img_path) / f
            if not im_file.exists():
                continue
            self.im_files.append(str(im_file))
            bboxes = []
            segments = []
            cat2id = {}
            texts = []
            for ann in anns:
                if ann["iscrowd"]:
                    continue
                box = np.array(ann["bbox"], dtype=np.float32)
                box[:2] += box[2:] / 2
                box[[0, 2]] /= float(w)
                box[[1, 3]] /= float(h)
                if box[2] <= 0 or box[3] <= 0:
                    continue

                caption = img["caption"]
                cat_name = " ".join([caption[t[0] : t[1]] for t in ann["tokens_positive"]]).lower().strip()
                if not cat_name:
                    continue

                if cat_name not in cat2id:
                    cat2id[cat_name] = len(cat2id)
                    texts.append([cat_name])
                cls = cat2id[cat_name]  # class
                box = [cls, *box.tolist()]
                if box not in bboxes:
                    bboxes.append(box)
                    if ann.get("segmentation") is not None:
                        if len(ann["segmentation"]) == 0:
                            segments.append(box)
                            continue
                        elif len(ann["segmentation"]) > 1:
                            s = merge_multi_segment(ann["segmentation"])
                            s = (np.concatenate(s, axis=0) / np.array([w, h], dtype=np.float32)).reshape(-1).tolist()
                        else:
                            s = [j for i in ann["segmentation"] for j in i]  # all segments concatenated
                            s = (
                                (np.array(s, dtype=np.float32).reshape(-1, 2) / np.array([w, h], dtype=np.float32))
                                .reshape(-1)
                                .tolist()
                            )
                        s = [cls, *s]
                        segments.append(s)
            lb = np.array(bboxes, dtype=np.float32) if len(bboxes) else np.zeros((0, 5), dtype=np.float32)

            if segments:
                classes = np.array([x[0] for x in segments], dtype=np.float32)
                segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in segments]  # (cls, xy1...)
                lb = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1)  # (cls, xywh)
            lb = np.array(lb, dtype=np.float32)

            x["labels"].append(
                {
                    "im_file": im_file,
                    "shape": (h, w),
                    "cls": lb[:, 0:1],  # n, 1
                    "bboxes": lb[:, 1:],  # n, 4
                    "segments": segments,
                    "normalized": True,
                    "bbox_format": "xywh",
                    "texts": texts,
                }
            )
        x["hash"] = get_hash(self.json_file)
        save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
        return x

    def get_labels(self) -> list[dict]:
        """Load labels from cache or generate them from JSON file.

        Returns:
            (list[dict]): List of label dictionaries, each containing information about an image and its annotations.
        """
        cache_path = Path(self.json_file).with_suffix(".cache")
        try:
            cache, _ = load_dataset_cache_file(cache_path), True  # attempt to load a *.cache file
            assert cache["version"] == DATASET_CACHE_VERSION  # matches current version
            assert cache["hash"] == get_hash(self.json_file)  # identical hash
        except (FileNotFoundError, AssertionError, AttributeError, ModuleNotFoundError):
            cache, _ = self.cache_labels(cache_path), False  # run cache ops
        [cache.pop(k) for k in ("hash", "version")]  # remove items
        labels = cache["labels"]
        self.verify_labels(labels)
        self.im_files = [str(label["im_file"]) for label in labels]
        if LOCAL_RANK in {-1, 0}:
            LOGGER.info(f"Load {self.json_file} from cache file {cache_path}")
        return labels

    def build_transforms(self, hyp: dict | None = None) -> Compose:
        """Configure augmentations for training with optional text loading.

        Args:
            hyp (dict, optional): Hyperparameters for transforms.

        Returns:
            (Compose): Composed transforms including text augmentation if applicable.
        """
        transforms = super().build_transforms(hyp)
        if self.augment:
            # NOTE: hard-coded the args for now.
            # NOTE: this implementation is different from official yoloe,
            # the strategy of selecting negative is restricted in one dataset,
            # while official pre-saved neg embeddings from all datasets at once.
            transform = RandomLoadText(
                max_samples=min(self.max_samples, 80),
                padding=True,
                padding_value=self._get_neg_texts(self.category_freq),
            )
            transforms.insert(-1, transform)
        return transforms

    @property
    def category_names(self):
        """Return unique category names from the dataset."""
        return {t.strip() for label in self.labels for text in label["texts"] for t in text}

    @property
    def category_freq(self):
        """Return frequency of each category in the dataset."""
        category_freq = defaultdict(int)
        for label in self.labels:
            for text in label["texts"]:
                for t in text:
                    t = t.strip()
                    category_freq[t] += 1
        return category_freq

    @staticmethod
    def _get_neg_texts(category_freq: dict, threshold: int = 100) -> list[str]:
        """Get negative text samples based on frequency threshold."""
        threshold = min(max(category_freq.values()), 100)
        return [k for k, v in category_freq.items() if v >= threshold]


class YOLOConcatDataset(ConcatDataset):
    """Dataset as a concatenation of multiple datasets.

    This class is useful to assemble different existing datasets for YOLO training, ensuring they use the same collation
    function.

    Methods:
        collate_fn: Static method that collates data samples into batches using YOLODataset's collation function.

    Examples:
        >>> dataset1 = YOLODataset(...)
        >>> dataset2 = YOLODataset(...)
        >>> combined_dataset = YOLOConcatDataset([dataset1, dataset2])
    """

    @staticmethod
    def collate_fn(batch: list[dict]) -> dict:
        """Collate data samples into batches.

        Args:
            batch (list[dict]): List of dictionaries containing sample data.

        Returns:
            (dict): Collated batch with stacked tensors.
        """
        return YOLODataset.collate_fn(batch)

    def close_mosaic(self, hyp: dict) -> None:
        """Disable mosaic, copy_paste, mixup and cutmix augmentations by setting their probabilities to 0.0.

        Args:
            hyp (dict): Hyperparameters for transforms.
        """
        for dataset in self.datasets:
            if not hasattr(dataset, "close_mosaic"):
                continue
            dataset.close_mosaic(hyp)


# TODO: support semantic segmentation
class SemanticDataset(BaseDataset):
    """Semantic Segmentation Dataset."""

    def __init__(self):
        """Initialize a SemanticDataset object."""
        super().__init__()


class ClassificationDataset:
    """Dataset class for image classification tasks wrapping torchvision ImageFolder functionality.

    This class offers functionalities like image augmentation, caching, and verification. It's designed to efficiently
    handle large datasets for training deep learning models, with optional image transformations and caching mechanisms
    to speed up training.

    Attributes:
        cache_ram (bool): Indicates if caching in RAM is enabled.
        cache_disk (bool): Indicates if caching on disk is enabled.
        samples (list): A list of lists, each containing the path to an image, its class index, path to its .npy cache
            file (if caching on disk), and optionally the loaded image array (if caching in RAM).
        torch_transforms (callable): PyTorch transforms to be applied to the images.
        root (str): Root directory of the dataset.
        prefix (str): Prefix for logging and cache filenames.

    Methods:
        __getitem__: Return transformed image and class index for the given sample index.
        __len__: Return the total number of samples in the dataset.
        verify_images: Verify all images in dataset.
    """

    def __init__(self, root: str, args, augment: bool = False, prefix: str = ""):
        """Initialize YOLO classification dataset with root directory, arguments, augmentations, and cache settings.

        Args:
            root (str): Path to the dataset directory where images are stored in a class-specific folder structure.
            args (Namespace): Configuration containing dataset-related settings such as image size, augmentation
                parameters, and cache settings.
            augment (bool, optional): Whether to apply augmentations to the dataset.
            prefix (str, optional): Prefix for logging and cache filenames, aiding in dataset identification.
        """
        import torchvision  # scope for faster 'import ultralytics'

        # Base class assigned as attribute rather than used as base class to allow for scoping slow torchvision import
        if TORCHVISION_0_18:  # 'allow_empty' argument first introduced in torchvision 0.18
            self.base = torchvision.datasets.ImageFolder(root=root, allow_empty=True)
        else:
            self.base = torchvision.datasets.ImageFolder(root=root)
        self.samples = self.base.samples
        self.root = self.base.root

        # Initialize attributes
        if augment and args.fraction < 1.0:  # reduce training fraction
            self.samples = self.samples[: round(len(self.samples) * args.fraction)]
        self.prefix = colorstr(f"{prefix}: ") if prefix else ""
        self.cache_ram = args.cache is True or str(args.cache).lower() == "ram"  # cache images into RAM
        if self.cache_ram:
            LOGGER.warning(
                "Classification `cache_ram` training has known memory leak in "
                "https://github.com/ultralytics/ultralytics/issues/9824, setting `cache_ram=False`."
            )
            self.cache_ram = False
        self.cache_disk = str(args.cache).lower() == "disk"  # cache images on hard drive as uncompressed *.npy files
        self.samples = self.verify_images()  # filter out bad images
        self.samples = [[*list(x), Path(x[0]).with_suffix(".npy"), None] for x in self.samples]  # file, index, npy, im
        scale = (1.0 - args.scale, 1.0)  # (0.08, 1.0)
        self.torch_transforms = (
            classify_augmentations(
                size=args.imgsz,
                scale=scale,
                hflip=args.fliplr,
                vflip=args.flipud,
                erasing=args.erasing,
                auto_augment=args.auto_augment,
                hsv_h=args.hsv_h,
                hsv_s=args.hsv_s,
                hsv_v=args.hsv_v,
            )
            if augment
            else classify_transforms(size=args.imgsz)
        )

    def __getitem__(self, i: int) -> dict:
        """Return transformed image and class index for the given sample index.

        Args:
            i (int): Index of the sample to retrieve.

        Returns:
            (dict): Dictionary containing the image and its class index.
        """
        f, j, fn, im = self.samples[i]  # filename, index, filename.with_suffix('.npy'), image
        if self.cache_ram:
            if im is None:  # Warning: two separate if statements required here, do not combine this with previous line
                im = self.samples[i][3] = cv2.imread(f)
        elif self.cache_disk:
            if not fn.exists():  # load npy
                np.save(fn.as_posix(), cv2.imread(f), allow_pickle=False)
            im = np.load(fn)
        else:  # read image
            im = cv2.imread(f)  # BGR
        # Convert NumPy array to PIL image
        im = Image.fromarray(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))
        sample = self.torch_transforms(im)
        return {"img": sample, "cls": j}

    def __len__(self) -> int:
        """Return the total number of samples in the dataset."""
        return len(self.samples)

    def verify_images(self) -> list[tuple]:
        """Verify all images in dataset.

        Returns:
            (list[tuple]): List of valid samples after verification.
        """
        desc = f"{self.prefix}Scanning {self.root}..."
        path = Path(self.root).with_suffix(".cache")  # *.cache file path

        try:
            check_file_speeds([file for (file, _) in self.samples[:5]], prefix=self.prefix)  # check image read speeds
            cache = load_dataset_cache_file(path)  # attempt to load a *.cache file
            assert cache["version"] == DATASET_CACHE_VERSION  # matches current version
            assert cache["hash"] == get_hash([x[0] for x in self.samples])  # identical hash
            nf, nc, n, samples = cache.pop("results")  # found, missing, empty, corrupt, total
            if LOCAL_RANK in {-1, 0}:
                d = f"{desc} {nf} images, {nc} corrupt"
                TQDM(None, desc=d, total=n, initial=n)
                if cache["msgs"]:
                    LOGGER.info("\n".join(cache["msgs"]))  # display warnings
            return samples

        except (FileNotFoundError, AssertionError, AttributeError):
            # Run scan if *.cache retrieval failed
            nf, nc, msgs, samples, x = 0, 0, [], [], {}
            with ThreadPool(NUM_THREADS) as pool:
                results = pool.imap(func=verify_image, iterable=zip(self.samples, repeat(self.prefix)))
                pbar = TQDM(results, desc=desc, total=len(self.samples))
                for sample, nf_f, nc_f, msg in pbar:
                    if nf_f:
                        samples.append(sample)
                    if msg:
                        msgs.append(msg)
                    nf += nf_f
                    nc += nc_f
                    pbar.desc = f"{desc} {nf} images, {nc} corrupt"
                pbar.close()
            if msgs:
                LOGGER.info("\n".join(msgs))
            x["hash"] = get_hash([x[0] for x in self.samples])
            x["results"] = nf, nc, len(samples), samples
            x["msgs"] = msgs  # warnings
            save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
            return samples
