Skip to content

Commit

Permalink
some typing improvements in base tracker
Browse files Browse the repository at this point in the history
  • Loading branch information
philipqueen committed Apr 15, 2024
1 parent c3b6f52 commit 0c07f39
Showing 1 changed file with 5 additions and 5 deletions.
10 changes: 5 additions & 5 deletions skellytracker/trackers/base_tracker/base_tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@ class BaseTracker(ABC):

def __init__(
self,
tracked_object_names: List[str] = None,
recorder: BaseRecorder = None,
tracked_object_names: List[str] = [],
recorder: Optional[BaseRecorder] = None,
**data: Any,
):
self.recorder = recorder
Expand Down Expand Up @@ -66,15 +66,15 @@ def process_video(
output_video_filepath: Optional[Union[str, Path]] = None,
save_data_bool: bool = False,
use_tqdm: bool = True,
) -> np.ndarray:
) -> Union[np.ndarray, None]:
"""
Run the tracker on a video.
:param input_video_filepath: Path to video file.
:param output_video_filepath: Path to save annotated video to, does not save video if None.
:param save_data_bool: Whether to save the data to a file.
:param use_tqdm: Whether to use tqdm to show a progress bar
:return: Array of tracked keypoint data
:return: Array of tracked keypoint data if tracker has an associated recorder
"""

cap = cv2.VideoCapture(str(input_video_filepath))
Expand Down Expand Up @@ -131,7 +131,7 @@ def process_video(
output_array = self.recorder.process_tracked_objects(image_size=image_size)
if save_data_bool:
self.recorder.save(
file_path=Path(input_video_filepath).with_suffix(".npy")
file_path=str(Path(input_video_filepath).with_suffix(".npy"))
)
else:
output_array = None
Expand Down

0 comments on commit 0c07f39

Please sign in to comment.