Azamatas/truck-fullness-detector

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)