diff --git a/chicken-detection-model-v26n-300e-best-2026-05-02-NEW.engine b/chicken-detection-model-v26n-300e-best-2026-05-02-NEW.engine index b189515..ca98117 100755 Binary files a/chicken-detection-model-v26n-300e-best-2026-05-02-NEW.engine and b/chicken-detection-model-v26n-300e-best-2026-05-02-NEW.engine differ diff --git a/chicken-detection-model-v26n-300e-best-2026-05-02-NEW.onnx b/chicken-detection-model-v26n-300e-best-2026-05-02-NEW.onnx index 55acf3a..a86fd41 100755 Binary files a/chicken-detection-model-v26n-300e-best-2026-05-02-NEW.onnx and b/chicken-detection-model-v26n-300e-best-2026-05-02-NEW.onnx differ diff --git a/export_engine.py b/export_engine.py new file mode 100644 index 0000000..b8081a6 --- /dev/null +++ b/export_engine.py @@ -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() diff --git a/src/chicken_counter/batch_runner.py b/src/chicken_counter/batch_runner.py index 0c97459..5a283d9 100755 --- a/src/chicken_counter/batch_runner.py +++ b/src/chicken_counter/batch_runner.py @@ -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.pipeline import run_pipeline from chicken_counter.report import build_batch_report, persist_batch_reports -from chicken_counter.tracking import DetectionTracker 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) 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] = [] 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, ) 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( CameraBatchResult( camera_id=camera_id, diff --git a/src/chicken_counter/tracking.py b/src/chicken_counter/tracking.py index ebc8750..99e2dfc 100755 --- a/src/chicken_counter/tracking.py +++ b/src/chicken_counter/tracking.py @@ -17,7 +17,7 @@ class DetectionTracker: self.config = config model_path = Path(config.detection.model_path) 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.verbose = config.performance.verbose self._infer_count = 0 @@ -47,7 +47,6 @@ class DetectionTracker: offset_x, offset_y = x1, y1 track_kwargs: dict = { - "task": "detect", "source": source, "persist": self.config.tracker.persist, "tracker": self.tracker_config_path,