refactor: Refactor tensorrt_yolo.infer module and remove deprecated t…

…yping types (Tuple, List, Set, Dict)

tensorrt_yolo/c_lib_wrap.py.in

@@ -25,8 +25,9 @@
 import os
 import re
 import sys
+from collections.abc import Mapping
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Optional
 
 from loguru import logger
 
@@ -67,12 +68,12 @@ def find_library(base_path: Path, lib_name: str) -> Optional[Path]:
     return None
 
 
-def check_libs(libs_dir: Dict[str, Path]) -> bool:
+def check_libs(libs_dir: Mapping[str, Path]) -> bool:
     """
     Check if all necessary libraries are found.
 
     Args:
-        libs_dir (Dict[str, Path]): A dictionary containing library paths.
+        libs_dir (Mapping[str, Path]): A dictionary containing library paths.
 
     Returns:
         bool: True if all libraries are found, False otherwise.
@@ -88,13 +89,13 @@ def check_libs(libs_dir: Dict[str, Path]) -> bool:
     return False
 
 
-def update_libs(user_specified_dir: Path, libs_dir: Dict[str, Path]) -> bool:
+def update_libs(user_specified_dir: Path, libs_dir: Mapping[str, Path]) -> bool:
     """
     Update library directories and check if all libraries are found.
 
     Args:
         user_specified_dir (Path): The directory to check.
-        libs_dir (Dict[str, Path]): A dictionary to store library paths.
+        libs_dir (Mapping[str, Path]): A dictionary to store library paths.
 
     Returns:
         bool: True if all libraries are found, False otherwise.
@@ -122,7 +123,7 @@ def add_dll_search_dir(dir_path: Path) -> None:
 
 
 if os.name == "nt":
-    libs_dir: Dict[str, Path] = {}
+    libs_dir: dict[str, Path] = {}
     need_lib_dirs = [Path(r"${CUDA_PATH}"), Path(r"${TENSORRT_PATH}")]
     sys_paths = [Path(p) for p in os.environ["PATH"].strip().split(";")]
 

tensorrt_yolo/cli.py

@@ -129,11 +129,11 @@ def infer(engine, mode, input, output, labels, cudagraph):
         sys.exit(1)
 
     if output:
-        from .infer import generate_labels_with_colors
+        from .infer import generate_labels
 
         output_dir = Path(output)
         output_dir.mkdir(parents=True, exist_ok=True)
-        labels = generate_labels_with_colors(labels)
+        labels = generate_labels(labels)
 
     import cv2
     from rich.progress import track

tensorrt_yolo/export/__init__.py

@@ -2,7 +2,7 @@
 import warnings
 from copy import deepcopy
 from pathlib import Path
-from typing import Optional, Tuple
+from typing import Optional
 
 import onnx
 from loguru import logger
@@ -145,7 +145,7 @@ def load_model(version: str, weights: str, repo_dir: Optional[str] = None) -> Op
 
 def update_model(
     model: torch.nn.Module, version: str, dynamic: bool, max_boxes: int, iou_thres: float, conf_thres: float
-) -> Tuple[Optional[torch.nn.Module], str]:
+) -> tuple[Optional[torch.nn.Module], str]:
     """
     Update YOLO model with dynamic settings.
 
@@ -158,7 +158,7 @@ def update_model(
         conf_thres (float): Confidence threshold for object detection.
 
     Returns:
-        Tuple[Optional[torch.nn.Module], str]:
+        tuple[Optional[torch.nn.Module], str]:
             - Updated YOLO model or None if the version is not supported.
             - The name of the detection head.
     """

tensorrt_yolo/export/head.py

@@ -24,7 +24,6 @@
 # ==============================================================================
 import copy
 import math
-from typing import Tuple
 
 import torch
 import torch.nn.functional as F
@@ -63,7 +62,7 @@ def forward(
         score_activation: int = 0,
         class_agnostic: int = 1,
         plugin_version: str = '1',
-    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         batch_size, num_boxes, num_classes = scores.shape
         num_dets = torch.randint(0, max_output_boxes, (batch_size, 1), dtype=torch.int32)
         det_boxes = torch.randn(batch_size, max_output_boxes, 4, dtype=torch.float32)
@@ -85,7 +84,7 @@ def symbolic(
         score_activation: int = 0,
         class_agnostic: int = 1,
         plugin_version: str = '1',
-    ) -> Tuple[Value, Value, Value, Value]:
+    ) -> tuple[Value, Value, Value, Value]:
         return g.op(
             'TRT::EfficientNMS_TRT',
             boxes,
@@ -118,7 +117,7 @@ def forward(
         score_activation: int = 0,
         class_agnostic: int = 1,
         plugin_version: str = '1',
-    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         batch_size, num_boxes, num_classes = scores.shape
         num_dets = torch.randint(0, max_output_boxes, (batch_size, 1), dtype=torch.int32)
         det_boxes = torch.randn(batch_size, max_output_boxes, 5, dtype=torch.float32)
@@ -140,7 +139,7 @@ def symbolic(
         score_activation: int = 0,
         class_agnostic: int = 1,
         plugin_version: str = '1',
-    ) -> Tuple[Value, Value, Value, Value]:
+    ) -> tuple[Value, Value, Value, Value]:
         return g.op(
             'TRT::EfficientRotatedNMS_TRT',
             boxes,
@@ -173,7 +172,7 @@ def forward(
         score_activation: int = 0,
         class_agnostic: int = 1,
         plugin_version: str = '1',
-    ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
         batch_size, num_boxes, num_classes = scores.shape
         num_dets = torch.randint(0, max_output_boxes, (batch_size, 1), dtype=torch.int32)
         det_boxes = torch.randn(batch_size, max_output_boxes, 4, dtype=torch.float32)
@@ -196,7 +195,7 @@ def symbolic(
         score_activation: int = 0,
         class_agnostic: int = 1,
         plugin_version: str = '1',
-    ) -> Tuple[Value, Value, Value, Value, Value]:
+    ) -> tuple[Value, Value, Value, Value, Value]:
         return g.op(
             'TRT::EfficientIdxNMS_TRT',
             boxes,

tensorrt_yolo/infer/__init__.py

@@ -1,41 +1,4 @@
-from .inference import (
-    DeployCGCls,
-    DeployCGDet,
-    DeployCGOBB,
-    DeployCGPose,
-    DeployCGSeg,
-    DeployCls,
-    DeployDet,
-    DeployOBB,
-    DeployPose,
-    DeploySeg,
-)
-from .result import Box, ClsResult, DetResult, KeyPoint, OBBResult, PoseResult, RotatedBox, SegResult
-from .timer import CpuTimer, GpuTimer
-from .utils import generate_labels_with_colors, image_batches, visualize
-
-__all__ = [
-    "DeployCGCls",
-    "DeployCGDet",
-    "DeployCGOBB",
-    "DeployCGPose",
-    "DeployCGSeg",
-    "DeployCls",
-    "DeployDet",
-    "DeployOBB",
-    "DeployPose",
-    "DeploySeg",
-    "Box",
-    "ClsResult",
-    "DetResult",
-    "KeyPoint",
-    "OBBResult",
-    "PoseResult",
-    "RotatedBox",
-    "SegResult",
-    "CpuTimer",
-    "GpuTimer",
-    "generate_labels_with_colors",
-    "image_batches",
-    "visualize",
-]
+from .inference import *  # noqa: F403
+from .result import *  # noqa: F403
+from .timer import *  # noqa: F403
+from .utils import *  # noqa: F403