Example: detector.py

# Copyright (c) 2024 Justin Davis (davisjustin302@gmail.com)
#
# MIT License
"""File showcasing the Detector class."""

from __future__ import annotations

import time
from pathlib import Path

import cv2

from trtutils import set_log_level
from trtutils.image import Detector


def main() -> None:
    engine_dir = Path(__file__).resolve().parent.parent.parent / "data" / "engines"
    engines = [
        engine_dir / "trt_yolov7t.engine",
        engine_dir / "trt_yolov8n.engine",
        engine_dir / "trt_yolov9t.engine",
        engine_dir / "trt_yolov10n.engine",
        engine_dir / "trt_yolov7t_dla.engine",
        engine_dir / "trt_yolov8n_dla.engine",
        engine_dir / "trt_yolov9t_dla.engine",
        engine_dir / "trt_yolov10n_dla.engine",
    ]

    img_path = str(Path(__file__).resolve().parent.parent.parent / "data" / "horse.jpg")
    img = cv2.imread(img_path)
    if img is None:
        err_msg = f"Failed to load image from {img_path}"
        raise FileNotFoundError(err_msg)

    for engine in engines:
        detector = Detector(engine, warmup=True, preprocessor="cuda")
        print(detector.name)

        t0 = time.perf_counter()
        output = detector.run([img])
        if not isinstance(output[0], list):
            err_msg = "Expected postprocessed output"
            raise TypeError(err_msg)
        bboxes = detector.get_detections(output)[0]
        t1 = time.perf_counter()

        print(f"RUN, bboxes: {len(bboxes)}, in {round((t1 - t0) * 1000.0, 2)}")

        # OR

        # end2end makes a few memory optimzations by avoiding extra GPU
        # memory transfers
        t0 = time.perf_counter()
        bboxes = detector.end2end([img])[0]
        t1 = time.perf_counter()

        print(f"END2END: bboxes: {len(bboxes)}, in {round((t1 - t0) * 1000.0, 2)}")

        del detector


if __name__ == "__main__":
    set_log_level("ERROR")
    main()