Files
feedmill-recounter/docs/superpowers/plans/2026-09-18-models-cleanup.md
T

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 as v4-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-e4e73c85 failed with No 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:

  1. Install onnxruntime (pip install onnxruntime)
  2. 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).pt file (same content as v4-best.pt)

  • Verify v4-best.pt still 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 .pt files in models/

  • 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 in app.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 ModuleNotFoundError deep 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, since MODEL_EXTENSIONS includes .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

  1. Task 1 (quick cleanup)
  2. Task 3 (quick template fix)
  3. Task 4 (quick pipeline fix)
  4. Task 5 (verification)
  5. 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.