fruit_project.utils.datasets.det_dataset

Classes

DET_DS

A custom dataset class for object detection tasks.

Functions

format_for_hf_processor(boxes, labels, idx)

Convert back to HF format

Module Contents

class fruit_project.utils.datasets.det_dataset.DET_DS(root_dir: str | None, split: str, config_file: str, transforms: albumentations.Compose | None = None, processor=None, normalize: bool = False)[source]

Bases: torch.utils.data.Dataset

A custom dataset class for object detection tasks.

This dataset class loads images and their corresponding labels from specified directories, applies transformations if provided, and returns the processed image along with target annotations.

root_dir[source]

The root directory containing the dataset.

Type:

str

split[source]

The dataset split (e.g., ‘train’, ‘val’, ‘test’).

Type:

str

config_file

The path to the configuration file containing class names and folder structure.

Type:

str

transforms[source]

A function or object to apply transformations to the images and annotations.

Type:

Albumentations Compose, optional

image_paths[source]

A list of valid image file paths.

Type:

list

labels[source]

A list of class names.

Type:

list

id2lbl[source]

A mapping from class IDs to class names.

Type:

dict

lbl2id[source]

A mapping from class names to class IDs.

Type:

dict

The configuration file (YAML) should contain:

  • names: List of class names

  • folders (optional): Dictionary with keys ‘images’, ‘labels’, ‘train’, ‘val’, ‘test’ specifying the folder names. Defaults to standard names if not provided.

  • folders.structure (optional): Either ‘type_first’ (default) for images/train structure or ‘split_first’ for train/images structure.

__len__()[source]

Returns the number of valid images in the dataset.

__getitem__(idx)[source]

Returns the processed image and target annotations for the given index.

Parameters:

  • root_dir (str) – The root directory containing the dataset.

  • split (str) – The dataset split (e.g., ‘train’, ‘val’, ‘test’).

  • config_file (str) – The path to the configuration file containing class names and folder structure.

  • transforms (Albumentations Compose, optional) – A function or object to apply transformations to the images and annotations.

Raises:

  • FileNotFoundError – If the configuration file or label files are not found.

  • ValueError – If an image cannot be loaded or is invalid.

root_dir: pathlib.Path[source]
split[source]
transforms = None[source]
config_dir[source]
processor = None[source]
normalize = False[source]
labels[source]
id2lbl[source]
lbl2id[source]
image_paths = [][source]
label_paths = [][source]
__len__()[source]
Returns:

The number of valid images in the dataset.

Return type:

int

__getitem__(idx)[source]

Retrieves the processed image and target annotations for the given index.

Parameters:

idx (int) – The index of the image to retrieve.

Returns:

A tuple containing:

  • img (numpy.ndarray): The processed image.

  • target (dict): A dictionary containing target annotations, including:

    • image_id (int): The index of the image.

    • annotations (list): A list of dictionaries with bounding box, category ID, area, and iscrowd flag.

    • orig_size (torch.Tensor): The original size of the image (height, width).

Return type:

tuple

get_raw_item(idx: int)[source]

Fetches a raw, untransformed image and its annotations. This is a helper method for multi-sample augmentations like Mosaic.

fruit_project.utils.datasets.det_dataset.format_for_hf_processor(boxes, labels, idx)[source]

Convert back to HF format