Gear Manipulation Dataset — PyTorch Loaders

A pair of PyTorch Dataset classes for loading robot manipulation data that combines RGB video frames with proprioceptive sensor streams. The data covers a robotic arm picking, inserting, removing, and placing gears of three sizes.

Action Classes

Both datasets share the same 13-class action vocabulary:

Datasets

SegDataset — Frame-level segmentation

Loads a full recording and returns a per-frame action label for every frame. Use this when you need dense, frame-by-frame supervision.

Expected directory layout

dataset_path/
 train/
 annotations/
 recording_001/
 action_labels.txt ← one action string per line, one per frame
 recording_002/
 action_labels.txt
 ...
 data/
 recording_001/
 frames/ ← 00000.jpg, 00001.jpg, ...
 proprioception/ ← 00000.npz, 00001.npz, ...
 recording_002/
 ...
 val/ (same structure)
 test/ (same structure)

action_labels.txt format

pick_small_gear
pick_small_gear
insert_small_gear
no_action
...

One action string per line, matching the keys in action_dict. The number of lines must equal the number of image files.

What __getitem__ returns

Quick start

from datamodule.SegDataset import SegDataset
dataset = SegDataset(dataset_path="/path/to/dataset", type="train")
print(len(dataset)) # number of recordings
sample = dataset[0]
print(sample["frames"].shape) # (N, H, W, 3)
print(sample["labels"].shape) # (N,)

ActionRecognitionDataset — Segment-level recognition

Splits each recording into action segments defined by a segments.txt file and returns one sample per segment along with a success flag. Use this for clip-level classification or success prediction.

Additional annotation file — segments.txt

Placed alongside action_labels.txt inside each recording's annotation folder:

frame_start	frame_end	success	text
0	24	1	pick small gear
25	61	1	insert small gear
62	80	0	

• Tab-separated with a header row.
• text is the human-readable action label (spaces, not underscores). An empty text field maps to no_action.
• success is 1 (succeeded) or 0 (failed).

What __getitem__ returns

Same proprioceptive keys as SegDataset, plus:

Note: frames and all proprioceptive arrays cover only the frames in [frame_start, frame_end] (inclusive).

Quick start

from datamodule.ActionRecognitionDataset import ActionRecognitionDataset
dataset = ActionRecognitionDataset(dataset_path="/path/to/dataset", type="train")
print(len(dataset)) # number of segments (across all recordings)
sample = dataset[0]
print(sample["frames"].shape) # (N, H, W, 3) — N = segment length
print(sample["label"]) # e.g. 3
print(sample["success"]) # 0 or 1

Proprioception .npz Format

Each .npz file stores one or more time-steps of sensor data (the sensor runs at a higher frequency than the camera). The dataset loader always takes the last sample in each file so that the proprioception timestamp aligns with the corresponding image frame:

proprioception[key][-1] # last reading in the file → aligned with the image

The keys present in the .npz files determine which proprioceptive signals appear in the returned dict — no hard-coding is needed.

Visualization (visualize.py)

visualize.py provides animated GIF generators for quick dataset inspection.

from visualize import animate_seg_sample, animate_action_sample

Both functions render the RGB video alongside scrolling time-series plots of every available proprioceptive signal. A red vertical line tracks the current frame.

Running the built-in demo

python visualize.py
# Produces:
# action_successful.gif — first successful segment found in train split
# action_failed.gif — first failed segment found in train split

Evaluation (metrics.py)

metrics.py provides three evaluation functions, one per task. All values are returned as percentages in a plain dict so they are easy to log (e.g. to wandb or a CSV).

recognition_metrics(y_true, y_pred)

Clip-level metrics for ActionRecognitionDataset. Expects one integer class ID per clip.

from metrics import recognition_metrics
from datamodule.ActionRecognitionDataset import ActionRecognitionDataset
dataset = ActionRecognitionDataset(dataset_path=..., type="test")
all_labels, all_preds = [], []
for sample in dataset:
 label = sample["label"]
 pred = model(sample["frames"]) # → int class id
 all_labels.append(label)
 all_preds.append(pred)
results = recognition_metrics(all_labels, all_preds)
# {'accuracy': 82.5, 'precision': 82.5, 'recall': 82.5, 'f1': 82.5}

For micro-averaging over a single-label classification task, accuracy = precision = recall = F1. They are all returned for API consistency with the other two functions.

segmentation_metrics(recognized_seqs, ground_truth_seqs, ...)

Frame-wise metrics for SegDataset. Both arguments are lists of lists of label strings (one list per recording).

from metrics import segmentation_metrics
from datamodule.SegDataset import SegDataset
from datamodule.ActionRecognitionDataset import action_dict
inv = {v: k for k, v in action_dict.items()} # int → string
dataset = SegDataset(dataset_path=..., type="test")
gt_seqs, pred_seqs = [], []
for sample in dataset:
 gt_seqs.append([inv[l] for l in sample["labels"]])
 preds = model(sample["frames"]) # → (N,) int array
 pred_seqs.append([inv[p] for p in preds])
results = segmentation_metrics(
 pred_seqs,
 gt_seqs,
 overlap_thresholds=(0.10, 0.25, 0.50),
 bg_class=("no_action",), # frames with this label are treated as background
)
# {'accuracy': 78.3, 'edit': 85.1, 'F1@0.10': 91.2, 'F1@0.25': 87.4, 'F1@0.50': 72.6}

The IoU thresholds can be changed via the overlap_thresholds argument. The bg_class argument controls which label strings are ignored when building segments (defaults to ("background",); set to ("no_action",) for this dataset).

anomaly_detection_metrics(y_true, y_pred)

Binary metrics for success/failure detection. Uses the success flag from ActionRecognitionDataset (inverted: success=0 → anomaly label 1).

from metrics import anomaly_detection_metrics
from datamodule.ActionRecognitionDataset import ActionRecognitionDataset
dataset = ActionRecognitionDataset(dataset_path=..., type="test")
all_labels, all_preds = [], []
for sample in dataset:
 gt_anomaly = 1 - sample["success"] # 0 = normal, 1 = anomaly
 pred_anomaly = anomaly_model(sample) # → 0 or 1
 all_labels.append(gt_anomaly)
 all_preds.append(pred_anomaly)
results = anomaly_detection_metrics(all_labels, all_preds)
# {'accuracy': 91.0, 'precision': 91.0, 'recall': 91.0, 'f1': 91.0}

Running the built-in demo

python metrics.py

Runs all three functions on synthetic data with realistic noise levels and prints a formatted table for each task.

Dependencies

torch
numpy
Pillow
matplotlib
scikit-learn

Install with:

pip install torch numpy Pillow matplotlib scikit-learn