Remove shared tracker between cameras, fix task warning, add export script
- Remove shared DetectionTracker across cameras to prevent state leakage - Add task="detect" to YOLO constructor to suppress warning - Add export_engine.py script for .pt to .engine conversion - Regenerate ONNX and TensorRT engine with latest settings
This commit is contained in:
1 parent
af4e514357
commit
d31ad05f0a
5 files changed
+80
-18
No files matched your search
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,78 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Export a YOLO .pt model to TensorRT .engine.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python3 export_engine.py chicken-detection-model-v26n-300e-best-2026-05-02-NEW.pt
|
||||||
|
python3 export_engine.py model.pt --imgsz 640 --half --workspace 4
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def export_engine(
|
||||||
|
model_path: str | Path,
|
||||||
|
*,
|
||||||
|
imgsz: int = 640,
|
||||||
|
half: bool = True,
|
||||||
|
int8: bool = False,
|
||||||
|
batch: int = 1,
|
||||||
|
workspace: int = 4, # GB
|
||||||
|
simplify: bool = True,
|
||||||
|
opset: int = 17,
|
||||||
|
verbose: bool = True,
|
||||||
|
) -> str:
|
||||||
|
from ultralytics import YOLO
|
||||||
|
|
||||||
|
model = YOLO(model_path, task="detect")
|
||||||
|
|
||||||
|
output = model.export(
|
||||||
|
format="engine",
|
||||||
|
imgsz=imgsz,
|
||||||
|
half=half,
|
||||||
|
int8=int8,
|
||||||
|
batch=batch,
|
||||||
|
workspace=workspace,
|
||||||
|
simplify=simplify,
|
||||||
|
opset=opset,
|
||||||
|
verbose=verbose,
|
||||||
|
)
|
||||||
|
|
||||||
|
print(f"\nExported to: {output}")
|
||||||
|
return str(output)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="Export YOLO .pt → TensorRT .engine")
|
||||||
|
parser.add_argument("model", help="Path to .pt model file")
|
||||||
|
parser.add_argument("--imgsz", type=int, default=640, help="Input image size (default: 640)")
|
||||||
|
parser.add_argument("--half", action="store_true", default=True, help="FP16 precision (default: on)")
|
||||||
|
parser.add_argument("--no-half", dest="half", action="store_false", help="FP32 precision")
|
||||||
|
parser.add_argument("--int8", action="store_true", help="INT8 quantization (needs calibration)")
|
||||||
|
parser.add_argument("--batch", type=int, default=1, help="Batch size (default: 1)")
|
||||||
|
parser.add_argument("--workspace", type=int, default=4, help="GPU workspace in GB (default: 4)")
|
||||||
|
parser.add_argument("--opset", type=int, default=17, help="ONNX opset version (default: 17)")
|
||||||
|
parser.add_argument("--quiet", action="store_true", help="Suppress verbose output")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
if not Path(args.model).exists():
|
||||||
|
print(f"error: model file not found: {args.model}", file=sys.stderr)
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
export_engine(
|
||||||
|
args.model,
|
||||||
|
imgsz=args.imgsz,
|
||||||
|
half=args.half,
|
||||||
|
int8=args.int8,
|
||||||
|
batch=args.batch,
|
||||||
|
workspace=args.workspace,
|
||||||
|
opset=args.opset,
|
||||||
|
verbose=not args.quiet,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -10,7 +10,6 @@ from chicken_counter.compress import compress_video_to_target
|
|||||||
from chicken_counter.config import BatchSettings, build_camera_config_from_batch
|
from chicken_counter.config import BatchSettings, build_camera_config_from_batch
|
||||||
from chicken_counter.pipeline import run_pipeline
|
from chicken_counter.pipeline import run_pipeline
|
||||||
from chicken_counter.report import build_batch_report, persist_batch_reports
|
from chicken_counter.report import build_batch_report, persist_batch_reports
|
||||||
from chicken_counter.tracking import DetectionTracker
|
|
||||||
from chicken_counter.types import CameraBatchResult
|
from chicken_counter.types import CameraBatchResult
|
||||||
|
|
||||||
|
|
||||||
@@ -29,20 +28,6 @@ def run_daily_batch(settings: BatchSettings, date: str | None = None, *, verbose
|
|||||||
discovery = discover_camera_videos(day_dir, settings)
|
discovery = discover_camera_videos(day_dir, settings)
|
||||||
camera_order = sorted(settings.cameras.items(), key=lambda item: item[1].camera_num)
|
camera_order = sorted(settings.cameras.items(), key=lambda item: item[1].camera_num)
|
||||||
|
|
||||||
first_camera_id = next(
|
|
||||||
camera_id for camera_id, _preset in camera_order if camera_id in discovery.found
|
|
||||||
)
|
|
||||||
first_source = discovery.found[first_camera_id]
|
|
||||||
init_output_path = output_dir / f"{first_camera_id}_vis.mp4" if not no_video else None
|
|
||||||
init_config = build_camera_config_from_batch(
|
|
||||||
settings,
|
|
||||||
first_camera_id,
|
|
||||||
source=first_source,
|
|
||||||
output_path=init_output_path,
|
|
||||||
checkpoint_dir=output_dir / "checkpoints" / first_camera_id,
|
|
||||||
)
|
|
||||||
shared_tracker = DetectionTracker(init_config)
|
|
||||||
|
|
||||||
camera_results: list[CameraBatchResult] = []
|
camera_results: list[CameraBatchResult] = []
|
||||||
report_path = output_dir / f"counts_{run_date}.json"
|
report_path = output_dir / f"counts_{run_date}.json"
|
||||||
|
|
||||||
@@ -73,7 +58,7 @@ def run_daily_batch(settings: BatchSettings, date: str | None = None, *, verbose
|
|||||||
checkpoint_dir=checkpoint_dir,
|
checkpoint_dir=checkpoint_dir,
|
||||||
)
|
)
|
||||||
camera_config.performance.verbose = verbose
|
camera_config.performance.verbose = verbose
|
||||||
pipeline_result = run_pipeline(camera_config, tracker=shared_tracker, show_progress=show_progress)
|
pipeline_result = run_pipeline(camera_config, show_progress=show_progress)
|
||||||
camera_results.append(
|
camera_results.append(
|
||||||
CameraBatchResult(
|
CameraBatchResult(
|
||||||
camera_id=camera_id,
|
camera_id=camera_id,
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ class DetectionTracker:
|
|||||||
self.config = config
|
self.config = config
|
||||||
model_path = Path(config.detection.model_path)
|
model_path = Path(config.detection.model_path)
|
||||||
self.model_kind = model_path.suffix.lower().lstrip(".") or "unknown"
|
self.model_kind = model_path.suffix.lower().lstrip(".") or "unknown"
|
||||||
self.model = YOLO(config.detection.model_path)
|
self.model = YOLO(config.detection.model_path, task="detect")
|
||||||
self.tracker_config_path = str(Path(config.tracker.tracker_config_path))
|
self.tracker_config_path = str(Path(config.tracker.tracker_config_path))
|
||||||
self.verbose = config.performance.verbose
|
self.verbose = config.performance.verbose
|
||||||
self._infer_count = 0
|
self._infer_count = 0
|
||||||
@@ -47,7 +47,6 @@ class DetectionTracker:
|
|||||||
offset_x, offset_y = x1, y1
|
offset_x, offset_y = x1, y1
|
||||||
|
|
||||||
track_kwargs: dict = {
|
track_kwargs: dict = {
|
||||||
"task": "detect",
|
|
||||||
"source": source,
|
"source": source,
|
||||||
"persist": self.config.tracker.persist,
|
"persist": self.config.tracker.persist,
|
||||||
"tracker": self.tracker_config_path,
|
"tracker": self.tracker_config_path,
|
||||||
|
|||||||
Reference in new issue
Block a user