Reference for ultralytics/data/dataset.py#

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

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

    Attributes:
        format_class (type[Format]): Formatter appended by build_transforms; subclasses override it per task.
        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.
        get_label_files: Return companion label files for the dataset's images.
        verify_args: Return the per-image verification function and its arguments.
        result_to_label: Convert one verification result into a label dict.
        verify_labels: Check box/segment consistency of the loaded labels.
        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()
    """

    format_class = Format

    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), **kwargs)

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

def _load_or_scan_cache(self, cache_path: Path, cache_hash: str) -> tuple[dict, bool]:
    """Load a dataset cache file if it matches the current version and hash, otherwise rescan and rebuild it.

    Args:
        cache_path (Path): Path of the cache file.
        cache_hash (str): Expected hash of the dataset files.

    Returns:
        (tuple): (cache dict, True if a valid existing cache file was loaded).
    """
    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"] == cache_hash  # identical hash
    except (FileNotFoundError, AssertionError, AttributeError, ModuleNotFoundError):
        cache, exists = self.cache_labels(cache_path), False  # run cache ops
    return cache, exists

def build_text_transforms(self, transforms: Compose, max_samples: int) -> Compose:
    """Insert text augmentation for text-based subclasses providing `category_freq`.

    Args:
        transforms (Compose): Transforms composed by build_transforms.
        max_samples (int): Maximum number of text samples per image.

    Returns:
        (Compose): Transforms with RandomLoadText inserted before Format when augmenting.
    """
    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(max_samples, 80),
            padding=True,
            padding_value=self._get_neg_texts(self.category_freq),
        )
        transforms.insert(-1, transform)
    return transforms

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(
        self.format_class(
            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 cache_labels(self, path: Path = Path("./labels.cache")) -> dict:
    """Cache dataset labels, check images and read shapes.

    This is the shared scanning skeleton for file-based datasets; subclasses customize it through the
    `get_label_files`, `get_cache_hash`, `verify_args`, `result_to_label` and `scan_summary` hooks instead of
    duplicating this method.

    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)
    with ThreadPool(NUM_THREADS) as pool:
        func, iterable = self.verify_args()
        results = pool.imap(func=func, iterable=iterable)
        pbar = TQDM(results, desc=desc, total=total)
        for result in pbar:
            label, nm_f, nf_f, ne_f, nc_f, msg = self.result_to_label(result)
            nm += nm_f
            nf += nf_f
            ne += ne_f
            nc += nc_f
            if label is not None:
                x["labels"].append(label)
            if msg:
                msgs.append(msg)
            pbar.desc = f"{desc} {self.scan_summary(nf, nm, ne, nc)}"
        pbar.close()

    if msgs:
        LOGGER.info("\n".join(msgs))
    if nf == 0:
        if self.augment:  # training requires labels; unlabeled val splits (e.g. COCO test-dev) only warn
            raise ValueError(f"{self.prefix}No labels found in {path}. {HELP_URL}")
        LOGGER.warning(f"{self.prefix}No labels found in {path}. {HELP_URL}")
    x["hash"] = self.get_cache_hash()
    x["results"] = nf, nm, ne, nc, total
    x["msgs"] = msgs  # warnings
    if x["labels"]:
        save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
    return x

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)

@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", "semantic_mask", "sem_masks", "depth"}:
            value = torch.stack(value, 0)
        elif k == "visuals":
            value = torch.nn.utils.rnn.pad_sequence(value, batch_first=True)
        if k in {"masks", "keypoints", "bboxes", "cls", "segments", "obb"}:
            value = torch.cat(value, 0)
        new_batch[k] = value
    if "batch_idx" in new_batch:
        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

def get_cache_hash(self) -> str:
    """Return the hash used to validate a label cache against the current dataset files.

    Returns:
        (str): Dataset cache hash.
    """
    return get_hash(self.label_files + self.im_files)

def get_label_files(self) -> list[str]:
    """Return the companion label files for the dataset's images, storing them on the instance.

    Returns:
        (list[str]): List of label file paths.
    """
    self.label_files = img2label_paths(self.im_files)
    return self.label_files

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.
    """
    label_files = self.get_label_files()
    cache_path = Path(label_files[0]).parent.with_suffix(".cache")
    cache, exists = self._load_or_scan_cache(cache_path, self.get_cache_hash())

    # 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}... {self.scan_summary(nf, nm, ne, nc)}"
        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
    labels = cache["labels"]
    if not labels:
        issues = "\n  ".join(sorted(set(cache["msgs"]))) or "no error details"
        raise RuntimeError(f"No valid images found in {cache_path}.\n  {issues}\n{HELP_URL}")
    [cache.pop(k) for k in ("hash", "version", "msgs")]  # remove items
    self.im_files = [lb["im_file"] for lb in labels]  # update im_files
    self.verify_labels(labels, cache_path)
    return labels

def result_to_label(self, result: list) -> tuple[dict | None, int, int, int, int, str]:
    """Convert one verification result into a label dict and scan counter increments.

    Args:
        result (list): One result from the verification function returned by `verify_args`.

    Returns:
        (tuple): (label dict or None, missing, found, empty, corrupt, message).
    """
    im_file, lb, shape, segments, keypoint, nm_f, nf_f, ne_f, nc_f, msg = result
    label = (
        {
            "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 im_file
        else None
    )
    return label, nm_f, nf_f, ne_f, nc_f, msg

def scan_summary(self, nf: int, nm: int, ne: int, nc: int) -> str:
    """Return a one-line summary of scan counters for progress bars and cache logs.

    Args:
        nf (int): Number of found images.
        nm (int): Number of missing labels.
        ne (int): Number of empty labels.
        nc (int): Number of corrupt images.

    Returns:
        (str): Scan summary message.
    """
    return f"{nf} images, {nm + ne} backgrounds, {nc} corrupt"

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

def verify_args(self) -> tuple:
    """Return the per-image verification function and its argument iterable used by `cache_labels`.

    Returns:
        (tuple): (verify function, zipped argument iterable) for ThreadPool.imap.
    """
    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]'"
        )
    return verify_image_label, 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),
    )

def verify_labels(self, labels: list[dict], cache_path: Path) -> None:
    """Check that the dataset is all boxes or all segments, removing mixed segments if necessary.

    Args:
        labels (list[dict]): List of label dictionaries.
        cache_path (Path): Path of the dataset cache file, used in warning messages.
    """
    # 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 self.use_segments and len_boxes != len_segments:
        raise ValueError(
            f"Segment dataset requires equal numbers of boxes and segments, but got len(segments) = "
            f"{len_segments}, len(boxes) = {len_boxes}. Please supply a segment dataset, not a detect dataset."
        )
    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}")

class DepthDataset(YOLODataset):
    """Dataset for monocular depth estimation with paired RGB + depth map loading.

    Extends YOLODataset to load depth ground truth maps alongside RGB images. Depth maps are stored as .npy files in a
    parallel directory structure (images/train/*.jpg → depth/train/*.npy).

    Examples:
        >>> dataset = DepthDataset(img_path="/data/nyu/images/train", data={"nc": 1})
    """

    format_class = DepthFormat

def _depth_path_for(self, im_file: str) -> str:
    """Map an image path to its companion depth .npy path (last 'images' path component → 'depth')."""
    parts = list(Path(im_file).parts)
    for i in range(len(parts) - 1, -1, -1):
        if parts[i] == "images":
            parts[i] = "depth"
            break
    return str(Path(*parts).with_suffix(".npy"))

def _load_depth(self, index):
    """Return the native-resolution depth map for an image, with non-finite values mapped to 0 (invalid)."""
    depth = np.load(self._depth_path_for(self.im_files[index])).astype(np.float32)
    return np.nan_to_num(depth, nan=0.0, posinf=0.0, neginf=0.0)

def build_transforms(self, hyp=None):
    """Build transforms for depth estimation.

    Args:
        hyp (dict): Hyperparameters.

    Returns:
        (Compose): Composed transforms.
    """
    # NOTE: For now following arguments are not supported
    hyp.mosaic = hyp.mixup = hyp.cutmix = hyp.copy_paste = 0.0
    transforms = super().build_transforms(hyp)
    if not self.augment:
        # stretch the image instead of padding
        transforms[-2] = LetterBox(new_shape=(self.imgsz, self.imgsz), scale_fill=True)
    return transforms

def get_cache_hash(self) -> str:
    """Return a hash over the paired depth and image files.

    Returns:
        (str): Dataset cache hash.
    """
    return get_hash(self.depth_files + self.im_files)

def get_image_and_label(self, index):
    """Load image, label, and depth map for the given index."""
    label = super().get_image_and_label(index)
    h, w = label["resized_shape"]
    depth = self._load_depth(index)
    if depth.shape[:2] != (h, w):
        depth = cv2.resize(depth, (w, h), interpolation=cv2.INTER_NEAREST)
    label["depth"] = depth
    return label

def get_label_files(self) -> list[str]:
    """Return the depth .npy paths paired with the dataset's images.

    Returns:
        (list[str]): List of depth file paths.
    """
    self.depth_files = [self._depth_path_for(f) for f in self.im_files]
    return self.depth_files

def result_to_label(self, result: tuple) -> tuple[dict | None, int, int, int, int, str]:
    """Convert one verify_image_depth result into a label dict and scan counter increments."""
    im_file, shape, nf_f, nm_f, nc_f, msg = result
    label = (
        {
            "im_file": im_file,
            "shape": shape,
            "cls": np.array([], dtype=np.float32),
            "bboxes": np.zeros((0, 4), dtype=np.float32),
            "segments": [],
            "normalized": True,
            "bbox_format": "xywh",
        }
        if im_file
        else None
    )
    return label, nm_f, nf_f, 0, nc_f, msg

def scan_summary(self, nf: int, nm: int, ne: int, nc: int) -> str:
    """Return a one-line summary of image-depth scan counters."""
    return f"{nf} images, {nm} missing depth, {nc} corrupt"

def verify_args(self) -> tuple:
    """Return the depth verification function and its argument iterable."""
    return verify_image_depth, zip(self.im_files, self.depth_files, repeat(self.prefix))

def verify_labels(self, labels: list[dict], cache_path: Path) -> None:
    """Skip box and segment checks; depth datasets carry no box or segment annotations."""

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)

@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
    # a background-only dataset sees no class, leaving every class an equally valid negative
    return category_freq or dict.fromkeys((t.strip() for text in texts for t in text), 0)

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

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

    Returns:
        (Compose): Composed transforms including text augmentation if applicable.
    """
    return self.build_text_transforms(super().build_transforms(hyp), self.data["nc"])

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

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_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)

@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

def _verify_instance_counts(self, labels: list[dict[str, Any]]) -> None:
    """Verify instance counts for known grounding datasets."""
    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 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.
    """
    return self.build_text_transforms(super().build_transforms(hyp), self.max_samples)

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)
    dropped = False
    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
        bboxes = []
        segments = []
        segmented = False
        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)
                raw_seg = ann.get("segmentation")
                segmented |= raw_seg is not None
                seg = raw_seg if isinstance(raw_seg, list) else []
                polygons = [
                    p
                    for p in seg
                    if isinstance(p, list)
                    and len(p) >= 6
                    and not len(p) % 2
                    and all(isinstance(c, (int, float)) for c in p)
                ]
                dropped |= bool(raw_seg) and (not isinstance(raw_seg, list) or len(polygons) < len(seg))
                if not polygons:  # keep one segment per box so an image mixing the two kinds stays aligned
                    cx, cy, bw, bh = box[1:]
                    x1, y1, x2, y2 = cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2
                    segments.append([cls, x1, y1, x2, y1, x2, y2, x1, y2])  # segments2boxes returns the box
                    continue
                elif len(polygons) > 1:
                    s = merge_multi_segment(polygons)
                    s = (np.concatenate(s, axis=0) / np.array([w, h], dtype=np.float32)).reshape(-1).tolist()
                else:
                    s = [j for i in polygons 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()
                    )
                segments.append([cls, *s])
        lb = np.array(bboxes, dtype=np.float32) if len(bboxes) else np.zeros((0, 5), dtype=np.float32)

        if segmented:
            segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in segments]  # (cls, xy1...)
            lb[:, 1:] = segments2boxes(segments)  # boxes follow the polygons
        else:
            segments = []  # no annotation carried a segmentation, so store no masks

        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,
            }
        )
    if dropped:
        LOGGER.warning(
            f"{self.json_file}: ignored segmentations that are not polygon point lists, such as RLE masks. "
            "Annotations left without a polygon use a segment shaped like their bounding box."
        )
    x["hash"] = self.get_cache_hash()
    save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
    return x

def get_cache_hash(self) -> str:
    """Return a hash over the annotation file and images scanned against it."""
    return get_hash([self.json_file, *self.scan_files])

def get_img_files(self, img_path: str) -> list[str]:
    """Return every image under `img_path`; the annotations, not `fraction`, decide which ones are used."""
    self.fraction = 1.0  # a truncated inventory would leave later images outside the cache key
    self.scan_files = super().get_img_files(img_path)
    return self.scan_files

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")
    cache, _ = self._load_or_scan_cache(cache_path, self.get_cache_hash())
    [cache.pop(k) for k in ("hash", "version")]  # remove items
    labels = cache["labels"]
    if not labels:
        raise RuntimeError(f"No images from {self.json_file} found in {self.img_path}. {HELP_URL}")
    if not any(label["texts"] for label in labels):  # category_freq is empty, so negative texts cannot be built
        raise RuntimeError(
            f"No annotations in {self.json_file} survived filtering. Every one is iscrowd, resolves to an empty "
            f"caption span or has a zero-size box. {HELP_URL}"
        )
    self._verify_instance_counts(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

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])
    """

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)

@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)

class SemanticDataset(YOLODataset):
    """Dataset for semantic segmentation with PNG mask labels.

    Expects a directory structure where each image has a corresponding PNG mask file with the same stem. Pixel values in
    masks represent class IDs, with 255 as the ignore label.

    The mask directory is specified in the dataset YAML via 'masks_dir' key, and mirrors the images/ directory structure
    (e.g., images/train/ -> masks/train/).

    Attributes:
        data (dict): Dataset configuration from YAML.
        mask_files (list[str]): List of mask file paths corresponding to images.
        include_class (np.ndarray | None): Class ids to keep per pixel (None keeps all).
    """

    format_class = SemanticFormat

    def __init__(self, *args, data: dict | None = None, **kwargs):
        """Initialize SemanticDataset.

        Args:
            *args (Any): Additional positional arguments for the parent class.
            data (dict): Dataset configuration dictionary.
            **kwargs (Any): Additional keyword arguments for the parent class.
        """
        self.data = data or {}
        self.label_mapping = self._parse_label_mapping(self.data.get("label_mapping"))
        self.mask_files = []
        self.include_class = None
        super().__init__(*args, data=data, **kwargs)

def _parse_label_mapping(self, mapping):
    """Normalize label_mapping entries from dataset YAML into integer-to-integer ids."""
    if mapping is None:
        return {}
    if not isinstance(mapping, dict):
        raise TypeError(f"Expected 'label_mapping' to be a dict in dataset YAML, but got {type(mapping).__name__}.")

    normalized = {}
    for src, dst in mapping.items():
        src = int(src)
        if isinstance(dst, str):
            dst = dst.strip()
            dst = 255 if dst == "ignore_label" else int(dst)
        elif dst is None:
            dst = 255
        else:
            dst = int(dst)
        normalized[src] = dst
    return normalized

def convert_label(self, label, inverse=False):
    """Convert label values using the dataset's label mapping.

    Args:
        label (np.ndarray): Segmentation label array to convert.
        inverse (bool): If True, apply inverse mapping (mapped -> original). Defaults to False.

    Returns:
        (np.ndarray): Label array with converted values.
    """
    temp = label.copy()
    if inverse:
        for v, k in self.label_mapping.items():
            label[temp == k] = v
    else:
        for k, v in self.label_mapping.items():
            label[temp == k] = v
    return label

def get_cache_hash(self) -> str:
    """Return a hash for semantic cache validation that also includes label_mapping changes.

    Returns:
        (str): Dataset cache hash.
    """
    mapping = json.dumps(self.label_mapping, sort_keys=True, separators=(",", ":"))
    return get_hash(self.im_files + self.mask_files + [f"label_mapping:{mapping}", "mask_bit_depth"])

def get_image_and_label(self, index):
    """Get image, label and semantic mask for the given index.

    Overrides parent to include semantic mask so that Mosaic/CopyPaste mix images
    also have their masks loaded.

    Args:
        index (int): Dataset index.

    Returns:
        (dict): Label dict with 'img', 'semantic_mask', and metadata.
    """
    label = super().get_image_and_label(index)
    h, w = label["img"].shape[:2]
    mask = self.load_mask(index, image_shape=(h, w))
    if self.include_class is not None:  # keep only selected classes; remap the rest to the ignore label
        mask[~np.isin(mask, self.include_class)] = 255
    # Resize mask to match the resized image dimensions
    if mask.shape[:2] != (h, w):
        mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
    label["semantic_mask"] = mask
    return label

def get_label_files(self) -> list[str]:
    """Return the mask PNG paths paired with the dataset's images.

    Returns:
        (list[str]): List of mask file paths.
    """
    self.mask_files = img2label_paths(self.im_files, label_dir=self.data.get("masks_dir", "masks"), suffix=".png")
    return self.mask_files

def get_labels(self) -> list[dict]:
    """Load semantic labels from cache or scan image-mask paths.

    Returns:
        (list[dict]): List of label dictionaries with mask file paths and image shapes.
    """
    labels = super().get_labels()
    self.mask_files = [lb["mask_file"] for lb in labels]
    return labels

def load_image(self, i, rect_mode=True):
    """Load an image for semantic segmentation, scaling the short side to imgsz when rect_mode=True."""
    return super().load_image(i, rect_mode=rect_mode, resize_short=self.augment)

def load_mask(self, index: int, image_shape: tuple[int, int] | None = None) -> np.ndarray:
    """Load a semantic mask and apply optional dataset label mapping."""
    mask_file = self.labels[index]["mask_file"]
    mask = cv2.imread(mask_file, cv2.IMREAD_GRAYSCALE)
    if mask is None:
        raise FileNotFoundError(f"Semantic mask not found or unreadable: {mask_file}")
    if mask.ndim == 3:
        mask = mask[..., 0]  # Windows patched cv2.imread expands grayscale reads to (H, W, 1)
    if int(self.data.get("nc", 0)) == 1 and self.labels[index]["is_1bit"]:
        mask[mask == 255] = 1  # cv2 expands 1-bit PNG foreground to 255.
    if self.label_mapping:
        mask = self.convert_label(mask, inverse=False)
    return mask.astype(np.uint8, copy=False)

def result_to_label(self, result: tuple) -> tuple[dict | None, int, int, int, int, str]:
    """Convert one verify_image_mask result into a label dict and scan counter increments."""
    im_file, mask_file, shape, is_1bit, nm_f, nf_f, nc_f, msg = result
    label = (
        {
            "im_file": im_file,
            "mask_file": mask_file,
            "shape": shape,
            "is_1bit": is_1bit,
            "cls": np.array([], dtype=np.float32),
            "bboxes": np.zeros((0, 4), dtype=np.float32),
            "segments": [],
            "normalized": True,
            "bbox_format": "xywh",
        }
        if im_file
        else None
    )
    return label, nm_f, nf_f, 0, nc_f, msg

def scan_summary(self, nf: int, nm: int, ne: int, nc: int) -> str:
    """Return a one-line summary of image-mask scan counters."""
    return f"{nf} images, {nm} missing masks, {nc} corrupt"

def update_labels(self, include_class: list[int] | None) -> None:
    """Update labels to include only specified classes.

    Args:
        include_class (list[int], optional): List of classes to include. If None, all classes are included.
    """
    if self.single_cls:
        raise NotImplementedError(
            "'single_cls=True' is not supported for semantic segmentation: it forces a single-channel "
            "model but cannot collapse multi-class masks. Use a dataset with 'nc: 1' for binary "
            "(foreground/background) segmentation instead."
        )
    self.include_class = None if include_class is None else np.asarray(include_class, dtype=np.int32).reshape(-1)
    if self.include_class is not None and int(self.data.get("nc", 0)) == 1:
        LOGGER.warning(
            "'classes' filtering is ignored for single-class (binary) semantic segmentation: keeping only "
            "the sole class would discard all background supervision."
        )
        self.include_class = None

def verify_args(self) -> tuple:
    """Return the mask verification function and its argument iterable."""
    return verify_image_mask, zip(
        self.im_files,
        self.mask_files,
        repeat(self.prefix),
        repeat(int(self.data.get("nc", 0)) == 1),
    )

def verify_labels(self, labels: list[dict], cache_path: Path) -> None:
    """Skip box and segment checks; semantic masks carry no box or segment annotations."""

class PolygonSemanticDataset(SemanticDataset, YOLODataset):
    """Semantic segmentation dataset that rasterizes YOLO polygon labels into masks on the fly.

    Used when the dataset YAML lacks 'masks_dir'. Pixels not covered by any polygon become a dedicated background class.
    Requires `add_polygon_background(data)` to be called first: for nc > 1 it bumps `data['nc']` to user_nc + 1 with
    background at `nc - 1`; for nc == 1 it keeps nc=1 and rasterizes a {0=bg, 1=fg} binary mask for use with
    BCEWithLogitsLoss.
    """

    def __init__(self, *args, data: dict | None = None, **kwargs):
        """Initialize PolygonSemanticDataset.

        Args:
            *args (Any): Additional positional arguments for the parent class.
            data (dict): Dataset configuration dictionary.
            **kwargs (Any): Additional keyword arguments for the parent class.
        """
        nc = (data or {}).get("nc") or len((data or {}).get("names", {}))
        self.bg_class_idx = data.get("bg_class_idx", max(int(nc) - 1, 0))
        super().__init__(*args, data=data, **kwargs)

def load_mask(self, index: int, image_shape: tuple[int, int] | None = None) -> np.ndarray:
    """Rasterize this image's polygons into a (H, W) uint8 semantic mask, bg = self.bg_class_idx."""
    h, w = image_shape
    label = self.labels[index]
    cls = label.get("cls")
    segments = label.get("segments") or []
    if cls is None or len(cls) == 0 or len(segments) == 0:
        return np.full((h, w), self.bg_class_idx, dtype=np.uint8)

    # Denormalize polygons (stored as normalized xy) to pixel coordinates at (h, w).
    scale = np.array([w, h], dtype=np.float32)
    polys = [np.asarray(s, dtype=np.float32).reshape(-1, 2) * scale for s in segments]
    # Returns (H, W) instance index map: 0 = no polygon, 1..N = sorted instance index.
    inst, sorted_idx = polygons2masks_overlap((h, w), polys, downsample_ratio=1)
    out = np.full((h, w), self.bg_class_idx, dtype=np.uint8)
    fg = inst > 0
    if int(self.data.get("nc", 0)) == 1:  # binary: fg=1 regardless of label cls value
        out[fg] = 1
    else:
        cls_arr = np.asarray(cls).reshape(-1).astype(np.int32)[sorted_idx]
        out[fg] = cls_arr[inst[fg] - 1].astype(np.uint8)
    return out

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.
        img_cache (np.ndarray): Contiguous uint8 buffer holding all cached images when caching in RAM.
        img_offsets (np.ndarray): Flat offset of each image within img_cache.
        img_shapes (list): (h, w, c) shape of each cached image.

    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.
        cache_images: Decode all images once into a single contiguous RAM buffer.
    """

    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)
        is_ndjson = (Path(root).parent / ".ndjson.yaml").is_file()
        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
        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
        if is_ndjson:
            self.samples = [(f, int(Path(f).parent.name)) for f, _ in self.samples]
        if args.single_cls:
            self.samples = [(f, 0) for f, _ in self.samples]
        self.samples = [[*list(x), Path(x[0]).with_suffix(".npy"), None] for x in self.samples]  # file, index, npy, im
        if self.cache_ram:
            self.cache_images()
        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:
        h, w, c = self.img_shapes[i]
        pos = self.img_offsets[i]
        im = self.img_cache[pos : pos + h * w * c].reshape(h, w, c)  # zero-copy view
    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 cache_images(self) -> None:
    """Decode all images once into a single contiguous uint8 buffer before DataLoader workers fork.

    A Python list of per-image arrays is duplicated into every forked worker by copy-on-write refcounting
    (https://github.com/ultralytics/ultralytics/issues/9824); one shared numpy buffer is read-only across
    workers instead, so RAM stays flat. Original image sizes are preserved for the transforms.
    """
    with ThreadPool(NUM_THREADS) as pool:
        ims = list(
            TQDM(
                pool.imap(lambda s: cv2.imread(s[0]), self.samples),
                total=len(self.samples),
                desc=f"{self.prefix}Caching images",
                disable=LOCAL_RANK > 0,
            )
        )
    self.img_shapes = [im.shape for im in ims]
    self.img_offsets = np.cumsum([0] + [im.size for im in ims[:-1]])
    self.img_cache = np.concatenate([im.reshape(-1) for im in ims])

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, corrupt, total, samples
        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

    # NOTE: ModuleNotFoundError to prevent numpy version conflicts when loading cache files created with different numpy versions
    except (FileNotFoundError, AssertionError, AttributeError, ModuleNotFoundError):
        # 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