From 57afb6254295d85743e306ab98f59699d3c24129 Mon Sep 17 00:00:00 2001
From: jetson <jetson@zenai-jetson-kantor.lan>
Date: Mon, 21 Sep 2026 13:37:01 +0700
Subject: [PATCH] fix: auto-use truck-detector when model lacks truck class

When a model without 'truck' class is selected (best, karung-dimuat,
yolo11n-bbox), the pipeline now automatically loads the truck-detector
model for truck detection. This ensures the batch manager transitions
to COUNTING_SACKS state and draws bounding boxes on output video.

Changes:
- Added find_truck_detector() helper in model_registry.py
- Added truck_model_config parameter to run_pipeline()
- job.py auto-selects truck-detector when model lacks truck class
- Verified with best.engine + batch-3.mp4: bounding boxes drawn
---
 src/job.py            | 11 ++++++++++-
 src/model_registry.py | 21 +++++++++++++++++++++
 src/pipeline.py       | 11 ++++++++---
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/src/job.py b/src/job.py
index 5222beb..3a3f2bf 100644
--- a/src/job.py
+++ b/src/job.py
@@ -13,7 +13,7 @@ from dataclasses import dataclass, field
 from enum import Enum, auto
 from pathlib import Path
 
-from src.model_registry import ModelConfig
+from src.model_registry import ModelConfig, find_truck_detector
 from src.pipeline import run_pipeline, PipelineResult
 
 
@@ -176,6 +176,14 @@ class JobQueue:
                     job.progress = i / total_models
                     class_filter = job.class_filters.get(model_cfg.filename)
 
+                # Auto-select truck detector if model lacks "truck" class
+                truck_model_config = None
+                if "truck" not in (model_cfg.known_classes or []):
+                    truck_model_config = find_truck_detector(os.path.join(os.path.dirname(model_cfg.path), "..", "models"))
+                    # If not found relative, try absolute models dir
+                    if truck_model_config is None:
+                        truck_model_config = find_truck_detector("./models")
+
                 output_path = os.path.join(
                     job.output_dir,
                     f"{model_cfg.stem}_annotated.mp4",
@@ -186,6 +194,7 @@ class JobQueue:
                     model_config=model_cfg,
                     output_path=output_path,
                     class_filter=class_filter,
+                    truck_model_config=truck_model_config,
                     preview_queue=job.preview_queue,
                     preview_every_n=2,
                     cancel_check=lambda: job.status == JobStatus.CANCELLED,
diff --git a/src/model_registry.py b/src/model_registry.py
index 5b75965..a1438aa 100644
--- a/src/model_registry.py
+++ b/src/model_registry.py
@@ -104,3 +104,24 @@ def scan_model_groups(models_dir: str) -> list[ModelGroup]:
             )
         )
     return groups
+
+
+def find_truck_detector(models_dir: str) -> ModelConfig | None:
+    """Find the truck-detector model in models directory.
+
+    Returns ModelConfig for the truck-detector if available, None otherwise.
+    Prefers .engine > .pt > .onnx.
+    """
+    p = Path(models_dir)
+    if not p.is_dir():
+        return None
+    for ext in FORMAT_PREFERENCE:
+        path = p / f"truck-detector{ext}"
+        if path.exists():
+            return ModelConfig(
+                filename=path.name,
+                path=str(path.resolve()),
+                stem="truck-detector",
+                known_classes=KNOWN_MODEL_CLASSES.get("truck-detector", ["truck"]),
+            )
+    return None
diff --git a/src/pipeline.py b/src/pipeline.py
index 0db4913..b9212b6 100644
--- a/src/pipeline.py
+++ b/src/pipeline.py
@@ -58,6 +58,7 @@ def run_pipeline(
     model_config: ModelConfig,
     output_path: str,
     class_filter: list[str] | None = None,
+    truck_model_config: ModelConfig | None = None,
     sack_conf: float = 0.4,
     truck_conf: float = 0.5,
     truck_det_interval: int = 15,
@@ -119,10 +120,14 @@ def run_pipeline(
         shared_model, conf=sack_conf, class_filter=effective_filter
     )
 
-    # Truck detector: if model has "truck" class, use same model
-    truck_has_truck = "truck" in (model_config.known_classes or [])
+    # Truck detector: use explicit truck_model_config, or same model if it has "truck" class
     truck_detector = None
-    if truck_has_truck:
+    if truck_model_config is not None:
+        # Load a separate truck detection model
+        truck_shared = YOLO(truck_model_config.path)
+        truck_detector = BaseDetector(truck_shared, conf=truck_conf, class_filter=("truck",))
+    elif "truck" in (model_config.known_classes or []):
+        # Use the same model for truck detection
         truck_detector = BaseDetector(
             shared_model, conf=truck_conf, class_filter=("truck",)
         )