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

# 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:
```python
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**

```bash
rm "/home/jetson/feedmill_semarang_project/feedmill_recounter/models/v4-best (1).pt"
```

- [ ] **Step 2: Verify v4-best.pt still exists**

```bash
ls -la /home/jetson/feedmill_semarang_project/feedmill_recounter/models/v4-best.pt
```

- [ ] **Step 3: Commit**

```bash
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**

```python
#!/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**

```bash
cd /home/jetson/feedmill_semarang_project/feedmill_recounter
python scripts/export_engines.py
```

- [ ] **Step 3: Verify .engine files created**

```bash
ls -la models/*.engine
```

- [ ] **Step 4: Commit**

```bash
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:

```html
<option value="truck">truck only</option>
```

- [ ] **Step 2: Verify template renders**

```bash
curl -s http://192.168.192.93:9000/ | grep "truck only"
```

- [ ] **Step 3: Commit**

```bash
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:

```python
    # 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**

```bash
cd /home/jetson/feedmill_semarang_project/feedmill_recounter
python -m pytest tests/test_pipeline.py -v
```

- [ ] **Step 3: Commit**

```bash
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**

```python
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**

```python
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.