Model Checkpoints
Here are every model checkpoint and what each one does, in the order it was created.
1. Corner Detector — ml/checkpoints/basket_edges_resnet18.pth
Purpose: Predict the 4 corners of the basket rim from any camera frame.
- Architecture: ResNet18, ImageNet-pretrained, classifier replaced with
Linear(512 → 8) + Sigmoid
- Input:
224×224 normalized RGB frame (full camera image)
- Output: 8 floats in
[0,1] — normalized (x, y) for each of 4 corners
- Training script:
ml/train_basket_edges.py
- Trained on: 427 hand-labeled rows in
basket_edges.csv
- Loss: SmoothL1
- Validation MAE:
1.74% normalized (≈ 1–2% of image dimension per coord)
- Use: Stage 1 of the warped-fill pipeline; gives corners → perspective warp
- Status: Production-ready. This is your active corner detector.
2. Sand UNet v1 — ml/checkpoints/sand_unet.pth
Purpose: Predict per-pixel sand probability mask on a warped basket crop.
- Architecture: Custom TinyUNet — 840k params, 4 down/up steps, sigmoid output
- Input:
256×128 normalized RGB perspective-warped basket crop
- Output:
1×128×256 sand probability mask
- Training script:
ml/train_sand_unet.py
- Trained on: 248 rows from
basket_edges.csv
- Loss:
L1(mean(mask), fill_pct/100) + 0.5·TV + 0.3·saturation
- Validation MAE:
8.47% (25-frame val)
- Special: Weakly supervised — never sees pixel labels, learns from scalar
fill_pct
- Status: Superseded by v2. Kept for reference.
3. Sand UNet v2 — ml/checkpoints/sand_unet_v2.pth ★
Purpose: Same as v1 but on the extended dataset with color augmentation.
- Architecture: Same TinyUNet (840k params)
- Input: Same —
256×128 warped crop
- Output: Same —
1×128×256 mask
- Training script:
ml/train_sand_unet.py --color_aug
- Trained on: 577 rows from
basket_edges_extended.csv (hand + predicted corners)
- Augmentation:
ColorJitter(brightness/contrast/saturation/hue) — fixes light-beige sand
- Validation MAE:
5.28% (115-frame val)
- Status: Current best UNet. Use this for fill % + spatial mask.
4. Warped EffNet v1 — ml/checkpoints/fill_warped_effnet.pth
Purpose: Direct fill regressor on warped crops.
- Architecture: EfficientNet-B0, ImageNet-pretrained, classifier
Linear(1280 → 1) + Sigmoid
- Input:
224×224 resized warped basket crop
- Output: 1 float — fill in
[0,1]
- Training script:
ml/train_fill_warped_effnet.py
- Trained on: 248 rows
- Augmentation:
RandomHFlip, RandomPerspective, ColorJitter
- Training phases: Head-only (8 ep) → fine-tune last 3 blocks (20 ep)
- Validation MAE:
7.34%
- Status: Superseded by v2.
5. Warped EffNet v2 — ml/checkpoints/fill_warped_effnet_v2.pth ★
Purpose: Same as v1 but on extended dataset.
- Architecture: Same EfficientNet-B0
- Input: Same
- Output: Same
- Trained on: 577 rows
- Training phases: Head-only (8 ep) → fine-tune last 3 blocks (25 ep)
- Validation MAE:
3.77% (115-frame val)
- Status: Current best fill regressor when you don't need a mask.
6. Warped ResNet18 — ml/checkpoints/fill_warped_resnet18.pth
Purpose: Earlier baseline, dominated by EffNet.
- Architecture: ResNet18, ImageNet-pretrained
- Input:
224×224 warped crop
- Trained on: 130 rows
- Augmentation: None
- Validation MAE:
9.79%
- Status: Obsolete. Kept only as a sanity-check baseline; don't use.
7. Multi-task EffNet — ml/checkpoints/multitask_effnet.pth
Purpose: Single model that predicts BOTH corners and fill percent in one forward pass.
- Architecture: EfficientNet-B0 backbone → split into corner head (8 floats) + fill head (1 float)
- Input:
224×224 full camera frame
- Output:
(corners, fill) — both from one inference
- Training script:
ml/train_multitask.py
- Trained on: 248 rows
- Augmentation:
albumentations: HFlip + Affine + brightness/contrast/saturation/hue + RandomShadow
- Loss:
SmoothL1(corners) + 0.5·MSE(fill) (masked when fill missing)
- Validation MAE: corner
7.04% / fill 9.65%
- Status: Worse than two specialist models. Useful only for single-pass-inference latency requirements. Negative-result reference.
8. HF EfficientNet — ml/checkpoints/best_model.pth
Purpose: Production fill model your team trained externally, hosted on HuggingFace.
- Architecture: EfficientNet-B0, classifier
Linear(1280 → 1) + Sigmoid (same as v2)
- Input:
224×224 full camera frame
- Output: 1 float — fill in
[0,1]
- Training: Done before this session, possibly on the same training set we're evaluating on
- Repository:
Azamatas/truck-fullness-detector
- Auto-download: Via
ml/predict.py
- Validation MAE:
3.85% on the 115-frame val (suspect: probably saw most of these in training)
- Status: Currently wired into
server_side/myapp/fullness.py. Production model. Likely overfit; real-world accuracy unknown.
Quick Reference
-
Predict basket corners on a new frame
- Model:
basket_edges_resnet18.pth
- Validation MAE:
1.74%
-
Get a fill number, no debugging needed
- Model:
fill_warped_effnet_v2.pth
- Validation MAE:
3.77%
-
Get fill + a spatial sand mask you can inspect
- Model:
sand_unet_v2.pth
- Validation MAE:
5.28%
-
End-to-end fill (corners → warp → fill)
- Model: EffNet v2 + corner detector
- Validation MAE:
3.81% mean / ≈2% median
-
Best blend (no data leak)
- Model:
0.7·EffNet_v2 + 0.3·UNet_v2
- Validation MAE:
3.51% / 1.98% median
-
Single-pass corners + fill (latency-bound)
- Model:
multitask_effnet.pth
- Validation MAE:
7% / 10%
-
Production today (with caveat)
- Model:
HF best_model.pth
- Validation MAE:
~3.9% (suspect leak)
