forked from zakaria/chicken-counting-sukawarna-det
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
@@ -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()
|
||||
Reference in new issue
Block a user