# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com)
#
# MIT License
from __future__ import annotations
from typing import TYPE_CHECKING, ClassVar
from trtutils.image._detector import Detector
from trtutils.image._schema import InputSchema, OutputSchema
from trtutils.models._model import Model
if TYPE_CHECKING:
from pathlib import Path
from typing_extensions import Self
[docs]
class RTDETRv1(Detector, Model):
"""Alias of DETR with default args for RT-DETRv1."""
_model_type = "rtdetrv1"
_friendly_name = "RT-DETRv1"
_default_imgsz = 640
_valid_imgszs: ClassVar[list[int]] = [640]
_input_tensors: ClassVar[list[tuple[str, str]]] = [
("images", "image"),
("orig_target_sizes", "size"),
]
def __init__(
self: Self,
engine_path: Path | str,
warmup_iterations: int = 10,
input_range: tuple[float, float] = (0, 255),
preprocessor: str = "trt",
resize_method: str = "linear",
conf_thres: float = 0.1,
nms_iou_thres: float = 0.5,
mean: tuple[float, float, float] | None = None,
std: tuple[float, float, float] | None = None,
dla_core: int | None = None,
device: int | None = None,
backend: str = "auto",
*,
warmup: bool | None = None,
pagelocked_mem: bool | None = None,
unified_mem: bool | None = None,
cuda_graph: bool | None = None,
extra_nms: bool | None = None,
agnostic_nms: bool | None = None,
no_warn: bool | None = None,
verbose: bool | None = None,
) -> None:
Detector.__init__(
self,
engine_path=engine_path,
warmup_iterations=warmup_iterations,
input_range=input_range,
preprocessor=preprocessor,
resize_method=resize_method,
conf_thres=conf_thres,
nms_iou_thres=nms_iou_thres,
mean=mean,
std=std,
input_schema=InputSchema.RT_DETR,
output_schema=OutputSchema.DETR_LBS,
dla_core=dla_core,
device=device,
backend=backend,
warmup=warmup,
pagelocked_mem=pagelocked_mem,
unified_mem=unified_mem,
cuda_graph=cuda_graph,
extra_nms=extra_nms,
agnostic_nms=agnostic_nms,
no_warn=no_warn,
verbose=verbose,
)
[docs]
class RTDETRv2(Detector, Model):
"""Alias of DETR with default args for RT-DETRv2."""
_model_type = "rtdetrv2"
_friendly_name = "RT-DETRv2"
_default_imgsz = 640
_valid_imgszs: ClassVar[list[int]] = [640]
_input_tensors: ClassVar[list[tuple[str, str]]] = [
("image", "image"),
("orig_target_sizes", "size"),
]
def __init__(
self: Self,
engine_path: Path | str,
warmup_iterations: int = 10,
input_range: tuple[float, float] = (0, 255),
preprocessor: str = "trt",
resize_method: str = "linear",
conf_thres: float = 0.1,
nms_iou_thres: float = 0.5,
mean: tuple[float, float, float] | None = None,
std: tuple[float, float, float] | None = None,
dla_core: int | None = None,
device: int | None = None,
backend: str = "auto",
*,
warmup: bool | None = None,
pagelocked_mem: bool | None = None,
unified_mem: bool | None = None,
cuda_graph: bool | None = None,
extra_nms: bool | None = None,
agnostic_nms: bool | None = None,
no_warn: bool | None = None,
verbose: bool | None = None,
) -> None:
Detector.__init__(
self,
engine_path=engine_path,
warmup_iterations=warmup_iterations,
input_range=input_range,
preprocessor=preprocessor,
resize_method=resize_method,
conf_thres=conf_thres,
nms_iou_thres=nms_iou_thres,
mean=mean,
std=std,
input_schema=InputSchema.RT_DETR,
output_schema=OutputSchema.DETR_LBS,
dla_core=dla_core,
device=device,
backend=backend,
warmup=warmup,
pagelocked_mem=pagelocked_mem,
unified_mem=unified_mem,
cuda_graph=cuda_graph,
extra_nms=extra_nms,
agnostic_nms=agnostic_nms,
no_warn=no_warn,
verbose=verbose,
)
[docs]
class RTDETRv3(Detector, Model):
"""Alias of DETR with default args for RT-DETRv3."""
_model_type = "rtdetrv3"
_friendly_name = "RT-DETRv3"
_default_imgsz = 640
_valid_imgszs: ClassVar[list[int]] = [640]
_input_tensors: ClassVar[list[tuple[str, str]]] = [
("image", "image"),
("im_shape", "size"),
("scale_factor", "size"),
]
def __init__(
self: Self,
engine_path: Path | str,
warmup_iterations: int = 10,
input_range: tuple[float, float] = (0, 255),
preprocessor: str = "trt",
resize_method: str = "linear",
conf_thres: float = 0.1,
nms_iou_thres: float = 0.5,
mean: tuple[float, float, float] | None = None,
std: tuple[float, float, float] | None = None,
dla_core: int | None = None,
device: int | None = None,
backend: str = "auto",
*,
warmup: bool | None = None,
pagelocked_mem: bool | None = None,
unified_mem: bool | None = None,
cuda_graph: bool | None = None,
extra_nms: bool | None = None,
agnostic_nms: bool | None = None,
no_warn: bool | None = None,
verbose: bool | None = None,
) -> None:
Detector.__init__(
self,
engine_path=engine_path,
warmup_iterations=warmup_iterations,
input_range=input_range,
preprocessor=preprocessor,
resize_method=resize_method,
conf_thres=conf_thres,
nms_iou_thres=nms_iou_thres,
mean=mean,
std=std,
input_schema=InputSchema.RT_DETR_V3,
output_schema=OutputSchema.RT_DETR_V3,
dla_core=dla_core,
device=device,
backend=backend,
warmup=warmup,
pagelocked_mem=pagelocked_mem,
unified_mem=unified_mem,
cuda_graph=cuda_graph,
extra_nms=extra_nms,
agnostic_nms=agnostic_nms,
no_warn=no_warn,
verbose=verbose,
)