9.2 KiB
Models Cleanup & Truck Filter Implementation Plan
For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (
- [ ]) syntax for tracking.
Goal: Clean up the models directory, create TensorRT .engine files for all .pt models, add a "truck only" filter option, and fix the onnxruntime dependency issue.
Architecture: One script to export all .pt models to .engine, one cleanup step to remove duplicates, one template change to add truck filter, one fix to handle missing onnxruntime gracefully.
Tech Stack: Python, ultralytics YOLO, TensorRT (via ultralytics export), Flask templates
Spec: User request — April 2026
Context
Current State
- Models directory: 16 files (7
.pt, 7.onnx, 0.engine) - Duplicate file:
v4-best (1).pt(same size asv4-best.pt) — needs removal - No .engine files: All models run as .pt or .onnx; .engine (TensorRT) would be faster on Jetson
- Recent job failure:
job-e4e73c85failed withNo module named 'onnxruntime'when trying to run the second .onnx model
Filtering Explained
The class filter dropdown in the upload page controls which object classes the YOLO model detects:
| Filter Option | Value | Behavior |
|---|---|---|
| Use model defaults | "default" |
Uses model_config.known_classes from KNOWN_MODEL_CLASSES map (e.g., ["person", "sack"] for karung model) |
| sack only | "sack" |
Only detects sack objects |
| box only | "box" |
Only detects box objects |
| sack + box | "sack,box" |
Detects both sack and box |
| truck only | (missing) | Should detect only truck objects |
| all classes | "all" |
No filtering — model detects everything it was trained on |
Key difference: "Use model defaults" applies the known class filter from the registry. "All classes" passes None as the filter, letting the model detect all its trained classes. The pipeline code at src/pipeline.py:92-94 shows:
effective_filter = class_filter or (
model_config.known_classes if model_config.known_classes else None
)
Onnxruntime Issue
The error No module named 'onnxruntime' occurs when the pipeline tries to load .onnx models. Two options:
- Install onnxruntime (
pip install onnxruntime) - Convert .onnx models to .engine (TensorRT) which doesn't need onnxruntime
Since we're creating .engine files anyway, option 2 is preferred for Jetson.
Global Constraints
- Python >= 3.10
- Platform: Jetson (ARM64) with CUDA
- ultralytics already installed
- TensorRT available on Jetson
- .engine files are gitignored
- Models directory:
./models/
Task 1: Remove Duplicate Model File
Files:
- Delete:
models/v4-best (1).pt
Requirements:
-
Remove the duplicate
v4-best (1).ptfile (same content asv4-best.pt) -
Verify
v4-best.ptstill exists after deletion -
Step 1: Remove duplicate file
rm "/home/jetson/feedmill_semarang_project/feedmill_recounter/models/v4-best (1).pt"
- Step 2: Verify v4-best.pt still exists
ls -la /home/jetson/feedmill_semarang_project/feedmill_recounter/models/v4-best.pt
- Step 3: Commit
git add -A && git commit -m "chore: remove duplicate v4-best (1).pt model file"
Task 2: Create TensorRT .engine Files for All .pt Models
Files:
- Create:
scripts/export_engines.py
Requirements:
-
Script iterates over all
.ptfiles inmodels/ -
For each .pt file, export to .engine using
model.export(format='engine', device=0, half=True, imgsz=640) -
Skip if .engine already exists
-
Handle export failures gracefully (log warning, continue)
-
Print summary of successful/failed exports
-
Step 1: Create export script
#!/usr/bin/env python3
"""Export all .pt models to TensorRT .engine format for Jetson."""
from pathlib import Path
from ultralytics import YOLO
MODELS_DIR = Path(__file__).resolve().parent.parent / "models"
def main():
pt_files = sorted(MODELS_DIR.glob("*.pt"))
if not pt_files:
print("No .pt files found in", MODELS_DIR)
return
success = 0
failed = 0
skipped = 0
for pt_path in pt_files:
engine_path = pt_path.with_suffix(".engine")
if engine_path.exists():
print(f"[SKIP] {pt_path.name} — .engine already exists")
skipped += 1
continue
print(f"[INFO] Exporting {pt_path.name} to TensorRT engine...")
try:
model = YOLO(str(pt_path))
engine_path_str = model.export(format="engine", device=0, half=True, imgsz=640)
print(f"[OK] Exported: {engine_path_str}")
success += 1
except Exception as e:
print(f"[FAIL] {pt_path.name}: {e}")
failed += 1
print(f"\nSummary: {success} exported, {skipped} skipped, {failed} failed")
if __name__ == "__main__":
main()
- Step 2: Run the export script
cd /home/jetson/feedmill_semarang_project/feedmill_recounter
python scripts/export_engines.py
- Step 3: Verify .engine files created
ls -la models/*.engine
- Step 4: Commit
git add scripts/export_engines.py && git commit -m "feat: add TensorRT engine export script"
Task 3: Add "truck only" Filter Option
Files:
- Modify:
templates/index.html:255— add truck only option
Requirements:
-
Add
<option value="truck">truck only</option>after the "sack + box" option -
The value
"truck"will be split into["truck"]by the existing filter logic inapp.py:72 -
Step 1: Add truck only option to template
In templates/index.html, after line 254 (<option value="sack,box">sack + box</option>), add:
<option value="truck">truck only</option>
- Step 2: Verify template renders
curl -s http://192.168.192.93:9000/ | grep "truck only"
- Step 3: Commit
git add templates/index.html && git commit -m "feat: add truck only filter option to upload page"
Task 4: Handle Missing onnxruntime Gracefully
Files:
- Modify:
src/pipeline.py:95— catch import error and suggest .engine
Requirements:
-
When loading a .onnx model, if onnxruntime is not installed, raise a clear error message
-
Suggest the user either install onnxruntime or use the .engine version of the model
-
This prevents cryptic
ModuleNotFoundErrordeep in the stack -
Step 1: Add onnxruntime check in pipeline.py
In src/pipeline.py, before line 95 (shared_model = YOLO(model_config.path)), add:
# Check onnxruntime for .onnx models
if model_config.path.endswith(".onnx"):
try:
import onnxruntime # noqa: F401
except ImportError:
raise RuntimeError(
f"onnxruntime is not installed. Cannot load .onnx model '{model_config.filename}'. "
f"Either install it (pip install onnxruntime) or use the .engine version of this model."
)
- Step 2: Run tests
cd /home/jetson/feedmill_semarang_project/feedmill_recounter
python -m pytest tests/test_pipeline.py -v
- Step 3: Commit
git add src/pipeline.py && git commit -m "fix: graceful error when onnxruntime missing for .onnx models"
Task 5: Update Model Registry for .engine Files
Files:
- Modify:
src/model_registry.py:9-16— add .engine entries to KNOWN_MODEL_CLASSES
Requirements:
-
Add entries for .engine files so they get proper known_classes
-
Since .engine files have the same stem as .pt files, the existing mapping should work
-
But verify that .engine files are picked up by
scan_models()(they should be, sinceMODEL_EXTENSIONSincludes.engine) -
Step 1: Verify .engine files are scanned
cd /home/jetson/feedmill_semarang_project/feedmill_recounter
python3 -c "from src.model_registry import scan_models; models = scan_models('./models'); print([m.filename for m in models if m.filename.endswith('.engine')])"
- Step 2: If not scanned, check MODEL_EXTENSIONS includes .engine
The current code already has MODEL_EXTENSIONS = {".pt", ".onnx", ".engine"} — no change needed.
- Step 3: Verify known_classes are assigned to .engine models
python3 -c "from src.model_registry import scan_models; models = scan_models('./models'); [(print(m.filename, m.known_classes)) for m in models if m.filename.endswith('.engine')]"
- Step 4: No commit needed if everything works — this is verification only
Summary of Changes
| Task | File | Change |
|---|---|---|
| 1 | models/v4-best (1).pt |
DELETE |
| 2 | scripts/export_engines.py |
CREATE — export script |
| 3 | templates/index.html |
ADD — truck only filter option |
| 4 | src/pipeline.py |
ADD — onnxruntime check |
| 5 | src/model_registry.py |
VERIFY only — .engine already supported |
Execution Order
- Task 1 (quick cleanup)
- Task 3 (quick template fix)
- Task 4 (quick pipeline fix)
- Task 5 (verification)
- Task 2 (export engines — longest, can run in background)
Tasks 1, 3, 4 can be done in parallel. Task 2 should be last since it takes time.