nnspike.data.dataset

This module defines custom PyTorch Datasets for driving records, including image preprocessing and augmentation.

This module provides dataset classes for different machine learning tasks including regression, classification, and multi-task learning. All datasets support data augmentation through brightness/contrast adjustments, RGB shifts, and horizontal flipping based on course direction.

Modules:

• cv2: OpenCV library for image processing.

• torch: PyTorch library for tensor operations and neural networks.

• albumentations as A: Albumentations library for image augmentations.

• torchvision.transforms as transforms: PyTorch’s torchvision library for common image transformations.

• numpy as np: NumPy library for numerical operations.

• torch.utils.data.Dataset: Base class for all datasets in PyTorch.

• nnspike.utils.normalize_image: Custom function for image normalization.

• nnspike.constants.Mode: Enumeration for different operation modes.

Constants:

• transform_bright_shift (albumentations.ReplayCompose): Augmentation pipeline with random brightness/contrast adjustments and RGB shifts.

• transform_flip (albumentations.Compose): Augmentation pipeline for horizontal flipping of images.

Functions:

• _rand_relative_position(relative_position, position_variation): Adds random variation to relative position for data augmentation.

Classes:

• RegressionDataset(Dataset): Custom dataset class for regression tasks, predicting continuous values like steering angles.

• ClassificationDataset(Dataset): Custom dataset class for classification tasks, predicting discrete modes of operation.

• MultiTaskDataset(Dataset): Custom dataset class for multi-task learning, combining both regression and classification.

RegressionDataset Class:

Designed for predicting continuous values from images and relative position data.

Methods:

• __init__(self, inputs, outputs, roi, train_course, position_variation=0.00625):

  Initializes the dataset with input data, target values, region of interest, training course, and position variation.

• __len__(self):

  Returns the number of samples in the dataset.

• __getitem__(self, idx):

  Retrieves and processes the sample at the given index. Returns tuple of (roi_tensor, relative_position) and target_x.

ClassificationDataset Class:

Designed for predicting discrete modes/classes from images and relative position data.

Methods:

• __init__(self, inputs, outputs, roi, train_course, position_variation=0.00625, transform=None):

  Initializes the dataset with input data, class labels, region of interest, training course, position variation, and optional transforms.

• __len__(self):

  Returns the number of samples in the dataset.

• __getitem__(self, idx):

  Retrieves and processes the sample at the given index. Returns tuple of (roi_tensor, relative_position) and mode.

MultiTaskDataset Class:

Designed for simultaneous regression and classification tasks from the same input data.

Methods:

• __init__(self, inputs, outputs, roi, train_course, position_variation=0.00625):

  Initializes the dataset with input data, combined outputs (mode, target), region of interest, training course, and position variation.

• __len__(self):

  Returns the number of samples in the dataset.

• __getitem__(self, idx):

  Retrieves and processes the sample at the given index. Returns tuple of (roi_tensor, relative_position) and (mode, target_x).

Usage Examples:

# Regression dataset for predicting steering angles regression_dataset = RegressionDataset(

inputs=[(‘path/to/image.png’, 0.5, ‘course1’)], outputs=[100.0], roi=np.array([0, 0, 200, 200]), train_course=’course1’

)

# Classification dataset for predicting operation modes classification_dataset = ClassificationDataset(

inputs=[(‘path/to/image.png’, 0.5, ‘course1’)], outputs=[1], roi=np.array([0, 0, 200, 200]), train_course=’course1’

)

# Multi-task dataset for both regression and classification multitask_dataset = MultiTaskDataset(

inputs=[(‘path/to/image.png’, 0.5, ‘course1’)], outputs=[(1, 100.0)], roi=np.array([0, 0, 200, 200]), train_course=’course1’

)

# Using with DataLoader dataloader = torch.utils.data.DataLoader(regression_dataset, batch_size=4, shuffle=True) for (roi_area, relative_position), target in dataloader:

# Training loop here pass

Notes

• All datasets apply data augmentation including brightness/contrast adjustments and RGB shifts.

• Images are horizontally flipped when the course differs from the training course.

• Relative positions are augmented with random variation for improved generalization.

• Target values are normalized to [0, 1] range for regression tasks.

• The normalize_image function is used for image preprocessing.

• ROI (Region of Interest) defines the area of the image to extract for processing.

Classes

ClassificationDataset(inputs, outputs, roi, ...)	Dataset for classification tasks, particularly for predicting discrete modes of operation.
MultiTaskDataset(inputs, outputs, roi, ...)	Dataset for multi-task learning, combining both regression and classification tasks.
RegressionDataset(inputs, outputs, roi, ...)	Dataset for regression tasks, particularly for predicting continuous values like steering angles.

class nnspike.data.dataset.RegressionDataset(inputs, outputs, roi, train_course, position_variation=0.00625)

Dataset for regression tasks, particularly for predicting continuous values like steering angles.

This dataset handles image loading, preprocessing, and augmentation for regression tasks. It supports data augmentation through random brightness/contrast adjustments and RGB shifts.

preprocess

Transform to convert numpy arrays to PyTorch tensors.

Type:

transforms.ToTensor

__init__(inputs, outputs, roi, train_course, position_variation=0.00625)

class nnspike.data.dataset.ClassificationDataset(inputs, outputs, roi, train_course, position_variation=0.00625, transform=None)

Dataset for classification tasks, particularly for predicting discrete modes of operation.

This dataset handles image loading, preprocessing, and augmentation for classification tasks. It supports data augmentation through random brightness/contrast adjustments, RGB shifts, and optional custom transforms.

preprocess

Transform to convert numpy arrays to PyTorch tensors.

Type:

transforms.ToTensor

__init__(inputs, outputs, roi, train_course, position_variation=0.00625, transform=None)

class nnspike.data.dataset.MultiTaskDataset(inputs, outputs, roi, train_course, position_variation=0.00625)

Dataset for multi-task learning, combining both regression and classification tasks.

This dataset handles image loading, preprocessing, and augmentation for both regression and classification tasks simultaneously. It supports data augmentation through random brightness/contrast adjustments and RGB shifts.

preprocess

Transform to convert numpy arrays to PyTorch tensors.

Type:

transforms.ToTensor

__init__(inputs, outputs, roi, train_course, position_variation=0.00625)