Fix type

nmhaddad · Feb 25, 2024 · 8587b3a · 8587b3a
1 parent 2f41051
commit 8587b3a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/fast_track/detectors/third_party/yolov8/yolov8_onnx.py b/fast_track/detectors/third_party/yolov8/yolov8_onnx.py
@@ -42,7 +42,7 @@ def __init__(self,
                  visualize: bool = False,
                  conf_thres: float = 0.25,
                  iou_thres: float = 0.45,
-                 classes: Optional[int] = None,
+                 classes: Optional[List[int]] = None,
                  agnostic: bool = False,
                  multi_label: bool = False,
                  labels: List[List[Union[int, float, torch.Tensor]]] = (),

diff --git a/fast_track/detectors/third_party/yolov9/yolov9_onnx.py b/fast_track/detectors/third_party/yolov9/yolov9_onnx.py
@@ -41,7 +41,7 @@ def __init__(self,
                  visualize: bool = False,
                  conf_thres: float = 0.25,
                  iou_thres: float = 0.45,
-                 classes: Optional[int] = None,
+                 classes: Optional[List[int]] = None,
                  agnostic: bool = False,
                  multi_label: bool = False,
                  labels: List[List[Union[int, float, torch.Tensor]]] = (),