Reference for ultralytics/data/base.py

def __getitem__(self, index):
    """Return transformed label information for given index."""
    return self.transforms(self.get_image_and_label(index))

def __len__(self):
    """Return the length of the labels list for the dataset."""
    return len(self.labels)

def build_transforms(self, hyp=None):
    """
    Users can customize augmentations here.

    Examples:
        >>> if self.augment:
        ...     # Training transforms
        ...     return Compose([])
        >>> else:
        ...    # Val transforms
        ...    return Compose([])
    """
    raise NotImplementedError

def cache_images(self):
    """Cache images to memory or disk for faster training."""
    b, gb = 0, 1 << 30  # bytes of cached images, bytes per gigabytes
    fcn, storage = (self.cache_images_to_disk, "Disk") if self.cache == "disk" else (self.load_image, "RAM")
    with ThreadPool(NUM_THREADS) as pool:
        results = pool.imap(fcn, range(self.ni))
        pbar = TQDM(enumerate(results), total=self.ni, disable=LOCAL_RANK > 0)
        for i, x in pbar:
            if self.cache == "disk":
                b += self.npy_files[i].stat().st_size
            else:  # 'ram'
                self.ims[i], self.im_hw0[i], self.im_hw[i] = x  # im, hw_orig, hw_resized = load_image(self, i)
                b += self.ims[i].nbytes
            pbar.desc = f"{self.prefix}Caching images ({b / gb:.1f}GB {storage})"
        pbar.close()

def cache_images_to_disk(self, i):
    """Save an image as an *.npy file for faster loading."""
    f = self.npy_files[i]
    if not f.exists():
        np.save(f.as_posix(), cv2.imread(self.im_files[i]), allow_pickle=False)

def check_cache_disk(self, safety_margin=0.5):
    """
    Check if there's enough disk space for caching images.

    Args:
        safety_margin (float, optional): Safety margin factor for disk space calculation.

    Returns:
        (bool): True if there's enough disk space, False otherwise.
    """
    import shutil

    b, gb = 0, 1 << 30  # bytes of cached images, bytes per gigabytes
    n = min(self.ni, 30)  # extrapolate from 30 random images
    for _ in range(n):
        im_file = random.choice(self.im_files)
        im = cv2.imread(im_file)
        if im is None:
            continue
        b += im.nbytes
        if not os.access(Path(im_file).parent, os.W_OK):
            self.cache = None
            LOGGER.info(f"{self.prefix}Skipping caching images to disk, directory not writeable ⚠️")
            return False
    disk_required = b * self.ni / n * (1 + safety_margin)  # bytes required to cache dataset to disk
    total, used, free = shutil.disk_usage(Path(self.im_files[0]).parent)
    if disk_required > free:
        self.cache = None
        LOGGER.info(
            f"{self.prefix}{disk_required / gb:.1f}GB disk space required, "
            f"with {int(safety_margin * 100)}% safety margin but only "
            f"{free / gb:.1f}/{total / gb:.1f}GB free, not caching images to disk ⚠️"
        )
        return False
    return True

def check_cache_ram(self, safety_margin=0.5):
    """
    Check if there's enough RAM for caching images.

    Args:
        safety_margin (float, optional): Safety margin factor for RAM calculation.

    Returns:
        (bool): True if there's enough RAM, False otherwise.
    """
    b, gb = 0, 1 << 30  # bytes of cached images, bytes per gigabytes
    n = min(self.ni, 30)  # extrapolate from 30 random images
    for _ in range(n):
        im = cv2.imread(random.choice(self.im_files))  # sample image
        if im is None:
            continue
        ratio = self.imgsz / max(im.shape[0], im.shape[1])  # max(h, w)  # ratio
        b += im.nbytes * ratio**2
    mem_required = b * self.ni / n * (1 + safety_margin)  # GB required to cache dataset into RAM
    mem = psutil.virtual_memory()
    if mem_required > mem.available:
        self.cache = None
        LOGGER.info(
            f"{self.prefix}{mem_required / gb:.1f}GB RAM required to cache images "
            f"with {int(safety_margin * 100)}% safety margin but only "
            f"{mem.available / gb:.1f}/{mem.total / gb:.1f}GB available, not caching images ⚠️"
        )
        return False
    return True

def get_image_and_label(self, index):
    """
    Get and return label information from the dataset.

    Args:
        index (int): Index of the image to retrieve.

    Returns:
        (dict): Label dictionary with image and metadata.
    """
    label = deepcopy(self.labels[index])  # requires deepcopy() https://github.com/ultralytics/ultralytics/pull/1948
    label.pop("shape", None)  # shape is for rect, remove it
    label["img"], label["ori_shape"], label["resized_shape"] = self.load_image(index)
    label["ratio_pad"] = (
        label["resized_shape"][0] / label["ori_shape"][0],
        label["resized_shape"][1] / label["ori_shape"][1],
    )  # for evaluation
    if self.rect:
        label["rect_shape"] = self.batch_shapes[self.batch[index]]
    return self.update_labels_info(label)

def get_img_files(self, img_path):
    """
    Read image files from the specified path.

    Args:
        img_path (str | List[str]): Path or list of paths to image directories or files.

    Returns:
        (List[str]): List of image file paths.

    Raises:
        FileNotFoundError: If no images are found or the path doesn't exist.
    """
    try:
        f = []  # image files
        for p in img_path if isinstance(img_path, list) else [img_path]:
            p = Path(p)  # os-agnostic
            if p.is_dir():  # dir
                f += glob.glob(str(p / "**" / "*.*"), recursive=True)
                # F = list(p.rglob('*.*'))  # pathlib
            elif p.is_file():  # file
                with open(p, encoding="utf-8") as t:
                    t = t.read().strip().splitlines()
                    parent = str(p.parent) + os.sep
                    f += [x.replace("./", parent) if x.startswith("./") else x for x in t]  # local to global path
                    # F += [p.parent / x.lstrip(os.sep) for x in t]  # local to global path (pathlib)
            else:
                raise FileNotFoundError(f"{self.prefix}{p} does not exist")
        im_files = sorted(x.replace("/", os.sep) for x in f if x.split(".")[-1].lower() in IMG_FORMATS)
        # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS])  # pathlib
        assert im_files, f"{self.prefix}No images found in {img_path}. {FORMATS_HELP_MSG}"
    except Exception as e:
        raise FileNotFoundError(f"{self.prefix}Error loading data from {img_path}\n{HELP_URL}") from e
    if self.fraction < 1:
        im_files = im_files[: round(len(im_files) * self.fraction)]  # retain a fraction of the dataset
    return im_files

def get_labels(self):
    """
    Users can customize their own format here.

    Note:
        Ensure output is a dictionary with the following keys:
        ```python
        dict(
            im_file=im_file,
            shape=shape,  # format: (height, width)
            cls=cls,
            bboxes=bboxes,  # xywh
            segments=segments,  # xy
            keypoints=keypoints,  # xy
            normalized=True,  # or False
            bbox_format="xyxy",  # or xywh, ltwh
        )
        ```
    """
    raise NotImplementedError

def load_image(self, i, rect_mode=True):
    """
    Load an image from dataset index 'i'.

    Args:
        i (int): Index of the image to load.
        rect_mode (bool, optional): Whether to use rectangular resizing.

    Returns:
        (np.ndarray): Loaded image.
        (tuple): Original image dimensions (h, w).
        (tuple): Resized image dimensions (h, w).

    Raises:
        FileNotFoundError: If the image file is not found.
    """
    im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i]
    if im is None:  # not cached in RAM
        if fn.exists():  # load npy
            try:
                im = np.load(fn)
            except Exception as e:
                LOGGER.warning(f"{self.prefix}WARNING ⚠️ Removing corrupt *.npy image file {fn} due to: {e}")
                Path(fn).unlink(missing_ok=True)
                im = cv2.imread(f)  # BGR
        else:  # read image
            im = cv2.imread(f)  # BGR
        if im is None:
            raise FileNotFoundError(f"Image Not Found {f}")

        h0, w0 = im.shape[:2]  # orig hw
        if rect_mode:  # resize long side to imgsz while maintaining aspect ratio
            r = self.imgsz / max(h0, w0)  # ratio
            if r != 1:  # if sizes are not equal
                w, h = (min(math.ceil(w0 * r), self.imgsz), min(math.ceil(h0 * r), self.imgsz))
                im = cv2.resize(im, (w, h), interpolation=cv2.INTER_LINEAR)
        elif not (h0 == w0 == self.imgsz):  # resize by stretching image to square imgsz
            im = cv2.resize(im, (self.imgsz, self.imgsz), interpolation=cv2.INTER_LINEAR)

        # Add to buffer if training with augmentations
        if self.augment:
            self.ims[i], self.im_hw0[i], self.im_hw[i] = im, (h0, w0), im.shape[:2]  # im, hw_original, hw_resized
            self.buffer.append(i)
            if 1 < len(self.buffer) >= self.max_buffer_length:  # prevent empty buffer
                j = self.buffer.pop(0)
                if self.cache != "ram":
                    self.ims[j], self.im_hw0[j], self.im_hw[j] = None, None, None

        return im, (h0, w0), im.shape[:2]

    return self.ims[i], self.im_hw0[i], self.im_hw[i]

def set_rectangle(self):
    """Set the shape of bounding boxes for YOLO detections as rectangles."""
    bi = np.floor(np.arange(self.ni) / self.batch_size).astype(int)  # batch index
    nb = bi[-1] + 1  # number of batches

    s = np.array([x.pop("shape") for x in self.labels])  # hw
    ar = s[:, 0] / s[:, 1]  # aspect ratio
    irect = ar.argsort()
    self.im_files = [self.im_files[i] for i in irect]
    self.labels = [self.labels[i] for i in irect]
    ar = ar[irect]

    # Set training image shapes
    shapes = [[1, 1]] * nb
    for i in range(nb):
        ari = ar[bi == i]
        mini, maxi = ari.min(), ari.max()
        if maxi < 1:
            shapes[i] = [maxi, 1]
        elif mini > 1:
            shapes[i] = [1, 1 / mini]

    self.batch_shapes = np.ceil(np.array(shapes) * self.imgsz / self.stride + self.pad).astype(int) * self.stride
    self.batch = bi  # batch index of image

def update_labels(self, include_class: Optional[list]):
    """
    Update labels to include only specified classes.

    Args:
        include_class (list, optional): List of classes to include. If None, all classes are included.
    """
    include_class_array = np.array(include_class).reshape(1, -1)
    for i in range(len(self.labels)):
        if include_class is not None:
            cls = self.labels[i]["cls"]
            bboxes = self.labels[i]["bboxes"]
            segments = self.labels[i]["segments"]
            keypoints = self.labels[i]["keypoints"]
            j = (cls == include_class_array).any(1)
            self.labels[i]["cls"] = cls[j]
            self.labels[i]["bboxes"] = bboxes[j]
            if segments:
                self.labels[i]["segments"] = [segments[si] for si, idx in enumerate(j) if idx]
            if keypoints is not None:
                self.labels[i]["keypoints"] = keypoints[j]
        if self.single_cls:
            self.labels[i]["cls"][:, 0] = 0

def update_labels_info(self, label):
    """Custom your label format here."""
    return label