object_detection/pt/src/data/yolod/datasets/datasets_wrapper.py

# /*---------------------------------------------------------------------------------------------
#  * Copyright (c) 2025 STMicroelectronics.
#  * All rights reserved.
#  *
#  * This software is licensed under terms that can be found in the LICENSE file in
#  * the root directory of this software component.
#  * If no LICENSE file comes with this software, it is provided AS-IS.
#  *--------------------------------------------------------------------------------------------*/
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

import bisect
import copy
import os
import random
from abc import ABCMeta, abstractmethod
from functools import partial, wraps
from multiprocessing.pool import ThreadPool
import psutil
from loguru import logger
from tqdm import tqdm

import numpy as np

from torch.utils.data.dataset import ConcatDataset as torchConcatDataset
from torch.utils.data.dataset import Dataset as torchDataset


class ConcatDataset(torchConcatDataset):
    def __init__(self, datasets):
        super(ConcatDataset, self).__init__(datasets)
        if hasattr(self.datasets[0], "input_dim"):
            self._input_dim = self.datasets[0].input_dim
            self.input_dim = self.datasets[0].input_dim

    def pull_item(self, idx):
        if idx < 0:
            if -idx > len(self):
                raise ValueError(
                    "absolute value of index should not exceed dataset length"
                )
            idx = len(self) + idx
        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
        if dataset_idx == 0:
            sample_idx = idx
        else:
            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
        return self.datasets[dataset_idx].pull_item(sample_idx)


class MixConcatDataset(torchConcatDataset):
    def __init__(self, datasets):
        super(MixConcatDataset, self).__init__(datasets)
        if hasattr(self.datasets[0], "input_dim"):
            self._input_dim = self.datasets[0].input_dim
            self.input_dim = self.datasets[0].input_dim

    def __getitem__(self, index):

        if not isinstance(index, int):
            idx = index[1]
        if idx < 0:
            if -idx > len(self):
                raise ValueError(
                    "absolute value of index should not exceed dataset length"
                )
            idx = len(self) + idx
        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
        if dataset_idx == 0:
            sample_idx = idx
        else:
            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
        if not isinstance(index, int):
            index = (index[0], sample_idx, index[2])

        return self.datasets[dataset_idx][index]


class Dataset(torchDataset):
    """ This class is a subclass of the base :class:`torch.utils.data.Dataset`,
    that enables on the fly resizing of the ``input_dim``.

    Args:
        input_dimension (tuple): (width,height) tuple with default dimensions of the network
    """

    def __init__(self, input_dimension, mosaic=True):
        super().__init__()
        self.__input_dim = input_dimension[:2]
        self.enable_mosaic = mosaic

    @property
    def input_dim(self):
        """
        Dimension that can be used by transforms to set the correct image size, etc.
        This allows transforms to have a single source of truth
        for the input dimension of the network.

        Return:
            list: Tuple containing the current width,height
        """
        if hasattr(self, "_input_dim"):
            return self._input_dim
        return self.__input_dim

    @staticmethod
    def mosaic_getitem(getitem_fn):
        """
        Decorator method that needs to be used around the ``__getitem__`` method. |br|
        This decorator enables the closing mosaic

        Example:
            >>> class CustomSet(ln.data.Dataset):
            ...     def __len__(self):
            ...         return 10
            ...     @ln.data.Dataset.mosaic_getitem
            ...     def __getitem__(self, index):
            ...         return self.enable_mosaic
        """

        @wraps(getitem_fn)
        def wrapper(self, index):
            if not isinstance(index, int):
                self.enable_mosaic = index[0]
                index = index[1]

            ret_val = getitem_fn(self, index)

            return ret_val

        return wrapper


class CacheDataset(Dataset, metaclass=ABCMeta):
    """ This class is a subclass of the base :class:`yolox.data.datasets.Dataset`,
    that enables cache images to ram or disk.

    Args:
        input_dimension (tuple): (width,height) tuple with default dimensions of the network
        num_imgs (int): datset size
        data_dir (str): the root directory of the dataset, e.g. `/path/to/COCO`.
        cache_dir_name (str): the name of the directory to cache to disk,
            e.g. `"custom_cache"`. The files cached to disk will be saved
            under `/path/to/COCO/custom_cache`.
        path_filename (str): a list of paths to the data relative to the `data_dir`,
            e.g. if you have data `/path/to/COCO/train/1.jpg`, `/path/to/COCO/train/2.jpg`,
            then `path_filename = ['train/1.jpg', ' train/2.jpg']`.
        cache (bool): whether to cache the images to ram or disk.
        cache_type (str): the type of cache,
            "ram" : Caching imgs to ram for fast training.
            "disk": Caching imgs to disk for fast training.
    """

    def __init__(
        self,
        input_dimension,
        num_imgs=None,
        data_dir=None,
        cache_dir_name=None,
        path_filename=None,
        cache=False,
        cache_type="ram",
    ):
        super().__init__(input_dimension)
        self.cache = cache
        self.cache_type = cache_type

        if self.cache and self.cache_type == "disk":
            self.cache_dir = os.path.join(data_dir, cache_dir_name)
            self.path_filename = path_filename

        if self.cache and self.cache_type == "ram":
            self.imgs = None

        if self.cache:
            self.cache_images(
                num_imgs=num_imgs,
                data_dir=data_dir,
                cache_dir_name=cache_dir_name,
                path_filename=path_filename,
            )

    def __del__(self):
        if self.cache and self.cache_type == "ram":
            del self.imgs

    @abstractmethod
    def read_img(self, index):
        """
        Given index, return the corresponding image

        Args:
            index (int): image index
        """
        raise NotImplementedError

    def cache_images(
        self,
        num_imgs=None,
        data_dir=None,
        cache_dir_name=None,
        path_filename=None,
    ):
        assert num_imgs is not None, "num_imgs must be specified as the size of the dataset"
        if self.cache_type == "disk":
            assert (data_dir and cache_dir_name and path_filename) is not None, \
                "data_dir, cache_name and path_filename must be specified if cache_type is disk"
            self.path_filename = path_filename

        mem = psutil.virtual_memory()
        mem_required = self.cal_cache_occupy(num_imgs)
        gb = 1 << 30

        if self.cache_type == "ram":
            if mem_required > mem.available:
                self.cache = False
            else:
                logger.info(
                    f"{mem_required / gb:.1f}GB RAM required, "
                    f"{mem.available / gb:.1f}/{mem.total / gb:.1f}GB RAM available, "
                    f"Since the first thing we do is cache, "
                    f"there is no guarantee that the remaining memory space is sufficient"
                )

        if self.cache and self.imgs is None:
            if self.cache_type == 'ram':
                self.imgs = [None] * num_imgs
                logger.info("You are using cached images in RAM to accelerate training!")
            else:   # 'disk'
                if not os.path.exists(self.cache_dir):
                    os.mkdir(self.cache_dir)
                    logger.warning(
                        f"\n*******************************************************************\n"
                        f"You are using cached images in DISK to accelerate training.\n"
                        f"This requires large DISK space.\n"
                        f"Make sure you have {mem_required / gb:.1f} "
                        f"available DISK space for training your dataset.\n"
                        f"*******************************************************************\\n"
                    )
                else:
                    logger.info(f"Found disk cache at {self.cache_dir}")
                    return

            logger.info(
                "Caching images...\n"
                "This might take some time for your dataset"
            )

            num_threads = min(8, max(1, os.cpu_count() - 1))
            b = 0
            load_imgs = ThreadPool(num_threads).imap(
                partial(self.read_img, use_cache=False),
                range(num_imgs)
            )
            pbar = tqdm(enumerate(load_imgs), total=num_imgs)
            for i, x in pbar:   # x = self.read_img(self, i, use_cache=False)
                if self.cache_type == 'ram':
                    self.imgs[i] = x
                else:   # 'disk'
                    cache_filename = f'{self.path_filename[i].split(".")[0]}.npy'
                    cache_path_filename = os.path.join(self.cache_dir, cache_filename)
                    os.makedirs(os.path.dirname(cache_path_filename), exist_ok=True)
                    np.save(cache_path_filename, x)
                b += x.nbytes
                pbar.desc = \
                    f'Caching images ({b / gb:.1f}/{mem_required / gb:.1f}GB {self.cache_type})'
            pbar.close()

    def cal_cache_occupy(self, num_imgs):
        cache_bytes = 0
        num_samples = min(num_imgs, 32)
        for _ in range(num_samples):
            img = self.read_img(index=random.randint(0, num_imgs - 1), use_cache=False)
            cache_bytes += img.nbytes
        mem_required = cache_bytes * num_imgs / num_samples
        return mem_required


def cache_read_img(use_cache=True):
    def decorator(read_img_fn):
        """
        Decorate the read_img function to cache the image

        Args:
            read_img_fn: read_img function
            use_cache (bool, optional): For the decorated read_img function,
                whether to read the image from cache.
                Defaults to True.
        """
        @wraps(read_img_fn)
        def wrapper(self, index, use_cache=use_cache):
            cache = self.cache and use_cache
            if cache:
                if self.cache_type == "ram":
                    img = self.imgs[index]
                    img = copy.deepcopy(img)
                elif self.cache_type == "disk":
                    img = np.load(
                        os.path.join(
                            self.cache_dir, f"{self.path_filename[index].split('.')[0]}.npy"))
                else:
                    raise ValueError(f"Unknown cache type: {self.cache_type}")
            else:
                img = read_img_fn(self, index)
            return img
        return wrapper
    return decorator