Piece / piece-with-pin classifier (two cameras, EfficientNet-B0)

Classifies the piece held by the SO-101 as nopin or pin (perno / tornillo) from the front and the wrist camera. About 85 % accurate on recording sessions it was not trained on, in the same room, with the same operator and lighting as the training sessions. Not tested on another day, lighting or batch of pieces.

Model

One ImageNet-pretrained EfficientNet-B0 (torchvision weights, 4.0 M parameters) shared by both cameras, fine-tuned (backbone at a tenth of the head learning rate). Each camera's feature map is pooled (mean + max), the two are concatenated and a linear layer gives the two logits. Input: two 448x336 RGB images (centre-cropped to 4:3, never stretched). Training: 25 epochs, class-weighted cross-entropy, AdamW + OneCycle, strong augmentation (rotation, zoom, shift, flip, saturation / white balance, blur, noise).

Data

1004 front+wrist pairs from the MikelRubi/so101-pick-cube_* recordings (2026-10-23 and 2026-10-24):

  • 654 from the last 3 frames of 218 pick-cube episodes (112 pin, 106 nopin), labelled by eye from the last frames (pin_labels.json); 16 of the pin and 72 of the nopin labels are "moderate" (blurred or partial view). 39 further episodes were too ambiguous and 5 had no piece in the gripper or no view of the tray, so they were left out.
  • 350 snapshot pairs (186 pin, 164 nopin) taken on purpose with 17b_record_real_so101_live.py (keys p / s), stored in each repo's pin_snapshots/ folder.

Evaluation

5-fold cross-validation where every fold holds out whole recording sessions (frames of one episode are near-duplicates, so a random split would look better than it is):

subset n accuracy majority-class baseline
all pairs 1004 0.85 0.52
episode last frames 654 0.87 0.51
snapshots 350 0.83 0.53
per episode 568 0.84 0.52
per episode, clear labels only 480 0.84 0.59

Per fold: 0.76, 0.85, 1.00, 0.88, 0.88. The final model in this repo is trained on all the data and was not itself evaluated; the numbers above are from the five cross-validation models. Labels are a visual judgement, not ground truth.

Same cross-validation, other runs: ResNet18 (11.2 M parameters) at 448x336 gave 0.85 on all pairs, 0.87 on episode frames, 0.79 on snapshots, 0.82 per episode; ResNet18 trained and tested on the episode frames only gave 0.84 (0.87 when the snapshots were added to the training set). A small CNN from scratch stayed at chance level (0.55). Earlier ResNet18 runs on fewer pairs: 0.66 (627 pairs), 0.81 (905 pairs). More data, pretrained weights and resolution all mattered.

Files

model.pt is a dict with state_dict, channels, height, width, classes, backbone (efficientnet_b0), cameras (both). Load it with load_model / predict from 22_train_pin_classifier.py in demos/06_imitation_learning/ of the LeRobot_ workspace (needs torch and torchvision). metrics.json is the training-run report of the final model; the cross-validation report is in the table above.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support