diff --git a/skellytracker/trackers/mediapipe_tracker/mediapipe_model_info.py b/skellytracker/trackers/mediapipe_tracker/mediapipe_model_info.py index 5453694..10768a7 100644 --- a/skellytracker/trackers/mediapipe_tracker/mediapipe_model_info.py +++ b/skellytracker/trackers/mediapipe_tracker/mediapipe_model_info.py @@ -1,3 +1,4 @@ +from typing import Literal from mediapipe.python.solutions import holistic as mp_holistic from mediapipe.python.solutions.face_mesh import FACEMESH_NUM_LANDMARKS_WITH_IRISES @@ -35,3 +36,7 @@ class MediapipeTrackingParams(BaseTrackingParams): min_detection_confidence: float = 0.5 min_tracking_confidence: float = 0.5 static_image_mode: bool = False + bounding_box_buffer_percentage: float = 10 + buffer_size_method: Literal["buffer_by_box_size", "buffer_by_image_size"] = ( + "buffer_by_box_size" + ) diff --git a/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py b/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py index c5293b3..d6ff664 100644 --- a/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py +++ b/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py @@ -3,7 +3,7 @@ import copy import mediapipe as mp import torch -from typing import Dict, Tuple +from typing import Dict, Literal, Tuple from ultralytics import YOLO from skellytracker.trackers.base_tracker.base_tracker import BaseTracker @@ -28,8 +28,10 @@ def __init__( min_tracking_confidence: float = 0.5, static_image_mode: bool = False, smooth_landmarks: bool = True, - bounding_box_buffer_percentage: float = 0, - buffer_size_method: str = "buffer_by_image_size", + bounding_box_buffer_percentage: float = 10, + buffer_size_method: Literal[ + "buffer_by_box_size", "buffer_by_image_size" + ] = "buffer_by_box_size", ): super().__init__( tracked_object_names=MediapipeModelInfo.mediapipe_tracked_object_names,