nnspike.data

Data processing and dataset management for nnspike.

This module provides comprehensive data handling capabilities for neural network training with the LEGO SPIKE robot, including dataset classes, data augmentation, and preprocessing utilities.

Modules:: aug: Data augmentation functions for improving model robustness dataset: PyTorch dataset classes for different learning tasks preprocess: Data preprocessing and label management utilities

Example

Basic usage for creating and augmenting a dataset:

from nnspike.data import RegressionDataset, augment_dataset

# Create a dataset dataset = RegressionDataset(data_dir=”path/to/data”)

# Apply augmentation augmented_data = augment_dataset(dataset, factor=2)

nnspike.data.random_shift_scale_rotate(image, shift_limit=0.0625, scale_limit=0.1, rotate_limit=15)[source]

Apply a random shift, scale, and rotation transformation to an input image.

Parameters:

image (np.ndarray) – The input image to be transformed.
shift_limit (float, optional) – Maximum fraction of total height/width to shift the image. Default is 0.0625.
scale_limit (float, optional) – Maximum scaling factor. Default is 0.1.
rotate_limit (int, optional) – Maximum rotation angle in degrees. Default is 15.

Returns:

A tuple containing:

transformed_image (np.ndarray): The transformed image.
params (dict): The parameters used for the transformation.

Return type:

tuple

Example

>>> import numpy as np
>>> image = np.random.rand(100, 100, 3)
>>> transformed_image, params = random_shift_scale_rotate(image)

nnspike.data.augment_dataset(df, p, export_path)[source]

Augment images in a dataset based on specified conditions and save them to a new location.

This function filters a DataFrame to include only rows where ‘use’ is True, then randomly selects rows based on probability p. For each selected row, it loads the image, applies sequential transformations (shift-scale-rotate followed by perspective transform), and saves the augmented image to the export directory. The function returns a new DataFrame containing metadata for all augmented images.

Parameters:

df (pd.DataFrame) – DataFrame containing image paths and metadata. Must include columns: ‘image_path’, ‘use’, ‘frame_number’, ‘mode’, ‘course’, ‘motor_a_relative_position’, ‘motor_b_relative_position’, ‘data_type’.
p (float) – Probability threshold (0.0-1.0) for applying augmentation to each row where ‘use’ is True. Higher values result in more augmented images.
export_path (str) – Directory path where augmented images will be saved. Must not already exist - the function will create it.

Returns:

DataFrame containing metadata of augmented images with the same: structure as the input DataFrame. The ‘image_path’ column will contain paths to the newly created augmented images, and ‘target_x’, ‘left_x’, ‘right_x’ will be set to None.

Return type:

pd.DataFrame

Raises:

FileExistsError – If the export_path directory already exists.

Example

>>> import pandas as pd
>>> df = pd.DataFrame(
...     {
...         "image_path": ["/path/to/image1.jpg", "/path/to/image2.jpg"],
...         "use": [True, True],
...         "frame_number": [1, 2],
...         "mode": ["A", "B"],
...         "course": ["course1", "course1"],
...         "motor_a_relative_position": [0, 10],
...         "motor_b_relative_position": [0, -5],
...         "data_type": ["train", "train"],
...     }
... )
>>> augmented_df = augment_dataset(df, p=0.5, export_path="/path/to/augmented")

Note

The function applies two sequential transformations: 1. Random shift, scale, rotate with limits: shift_limit=0.05, scale_limit=0.05, rotate_limit=5 2. Perspective transform with scale=(0.01, 0.05)

class nnspike.data.RegressionDataset(inputs, outputs, roi, train_course, position_variation=0.00625)[source]

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)[source]

class nnspike.data.ClassificationDataset(inputs, outputs, roi, train_course, position_variation=0.00625, transform=None)[source]

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)[source]

class nnspike.data.MultiTaskDataset(inputs, outputs, roi, train_course, position_variation=0.00625)[source]

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)[source]

nnspike.data.balance_dataset(df, col_name, max_samples, num_bins)[source]

Balances the dataset by limiting the number of samples in each bin of a specified column.

This function creates a histogram of the specified column and ensures that no bin has more than max_samples samples. If a bin exceeds this limit, excess samples are randomly removed to balance the dataset.

Parameters:

df (pd.DataFrame) – The input DataFrame containing the data to be balanced.
col_name (str) – The name of the column to be used for creating bins.
max_samples (int) – The maximum number of samples allowed per bin.
num_bins (int) – The number of bins to divide the column into.

Returns:

A DataFrame with the dataset balanced according to the specified column and bin limits.

Return type:

pd.DataFrame

Note

Make sure the column does not have

None/Nan
empty string

Otherwise, ValueError: autodetected range of [nan, nan] is not finite may raise

nnspike.data.sort_by_frames_number(df)[source]

Sorts a DataFrame by the frame number extracted from the ‘image_path’ column.

This function extracts the frame number from the ‘image_path’ column of the DataFrame, sorts the DataFrame based on these frame numbers, and keeps the ‘frame_number’ column as the 2nd column in the DataFrame.

Parameters:: df (pd.DataFrame) – The input DataFrame containing an ‘image_path’ column with file paths that include frame numbers in the format ‘frame_<number>’.
Returns:: The sorted DataFrame with rows ordered by the extracted frame numbers and the frame_number column as the 2nd column.
Return type:: pd.DataFrame

nnspike.data.create_label_dataframe(path_pattern, course)[source]

Creates a comprehensive DataFrame with image paths and associated metadata for labeling tasks.

This function searches for image files matching the given path pattern and constructs a DataFrame containing the paths to these images along with multiple columns for sensor data and robot control information. The DataFrame includes columns for: - Image metadata: image_path, course, data_type, use flag - Target and mode information: mode, target_x - Motor data: motor_a_speed, motor_b_speed, motor_a_relative_position, motor_b_relative_position - Sensor readings: distance_sensor, color_reflected, color_ambient, color_value

All sensor and motor columns are initialized with default values (0 for numeric fields, NaN for mode and target_x, True for use flag) to be populated later by other functions.

Parameters:

path_pattern (str) – A glob pattern to match image file paths.
course (str) – The name of the course associated with the images.

Returns:

A DataFrame containing the image paths and comprehensive metadata columns with default values for subsequent data population.

Return type:

pd.DataFrame

nnspike.data.set_spike_status(label_df, status_df)[source]

Merges comprehensive sensor and motor data from status_df into label_df based on matching frame numbers.

This function performs a comprehensive merge of robot sensor and control data from status_df into label_df by matching frame_number values. It updates multiple columns including: - Motor control: motor_a_speed, motor_b_speed, motor_a_relative_position, motor_b_relative_position - Sensor readings: distance_sensor, color_reflected, color_ambient, color_value - Robot mode: mode (preserving manual labels from label_df when available)

The merge is performed as a left join, preserving all rows in label_df and only updating sensor/motor data where matching frame numbers exist in status_df. If label_df doesn’t have a frame_number column, it will be automatically added by calling sort_by_frames_number().

Parameters:

label_df (pd.DataFrame) – The label DataFrame containing image paths and metadata. Must have or be able to generate a frame_number column.
status_df (pd.DataFrame) – The status DataFrame containing comprehensive sensor and motor data with frame_number column for matching.

Returns:

The updated label DataFrame with sensor and motor data merged from status_df. All original columns are preserved, and sensor data is filled where frame matches exist.

Return type:

pd.DataFrame

Modules

`aug`	Image augmentation utilities for neural network training data.
`dataset`	This module defines custom PyTorch Datasets for driving records, including image preprocessing and augmentation.
`preprocess`	This module provides functions for preprocessing and preparing datasets of image frames with sensor data.