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:
1 parent
f9e411f407
commit
57afb62542
3 files changed
+39
-4
No files matched your search
+10
-1
@@ -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,
|
||||
|
||||
@@ -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
@@ -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",)
|
||||
)
|
||||
|
||||
Reference in new issue
Block a user