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
This commit is contained in:
jetson committed 2026-09-21 13:37:01 +07:00
1 parent f9e411f407
commit 57afb62542
3 files changed
+39 -4

No files matched your search

+10 -1
View File
@@ -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,
+21
View File
@@ -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
+8 -3
View File
@@ -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",)
)