UPDATE: Docstrings

Author: Ashp116

trackers/core/ksp/solver.py

@@ -73,19 +73,40 @@ def __init__(
         self.reset()
 
     def reset(self) -> None:
+        """
+        Reset the solver state, clearing all buffered detections and the graph.
+        """
         self.detection_per_frame = []
         self.graph = nx.DiGraph()
 
     def append_frame(self, detections: sv.Detections) -> None:
+        """
+        Add detections for the current frame to the solver's buffer.
+
+        Args:
+            detections (sv.Detections): Detections for the current frame.
+        """
         self.detection_per_frame.append(detections)
 
     def _get_center(self, bbox: np.ndarray) -> np.ndarray:
+        """
+        Compute the center (x, y) of a bounding box.
+
+        Args:
+            bbox (np.ndarray): Bounding box as [x1, y1, x2, y2].
+
+        Returns:
+            np.ndarray: Center coordinates as (x, y).
+        """
         x1, y1, x2, y2 = bbox
         return np.array([(x1 + x2) / 2, (y1 + y2) / 2])
 
     def set_entry_exit_regions(self, regions: List[Tuple[int, int, int, int]]) -> None:
         """
         Set rectangular entry/exit zones (x1, y1, x2, y2).
+
+        Args:
+            regions (List[Tuple[int, int, int, int]]): List of rectangular regions.
         """
         self.entry_exit_regions = regions
 
@@ -101,9 +122,9 @@ def set_border_entry_exit(
 
         Args:
             use_border (bool): Enable/disable border-based entry/exit.
-            borders (set): Set of borders to use. {"left", "right", "top", "bottom"}
-            margin (int): Border thickness in pixels.
-            frame_size (Tuple[int, int]): Size of the image (width, height).
+            borders (Optional[Set[str]]): Set of borders to use.
+            margin (Optional[int]): Border thickness in pixels.
+            frame_size (Optional[Tuple[int, int]]): Size of the image (width, height).
         """
         self.use_border_regions = use_border
         self.active_borders = (
@@ -113,6 +134,15 @@ def set_border_entry_exit(
         self.frame_size = frame_size
 
     def _in_door(self, node: TrackNode) -> bool:
+        """
+        Check if a node is inside any entry/exit region (rectangular or border).
+
+        Args:
+            node (TrackNode): The node to check.
+
+        Returns:
+            bool: True if in any entry/exit region, else False.
+        """
         x, y = node.position
 
         # Check custom rectangular regions
@@ -137,6 +167,16 @@ def _in_door(self, node: TrackNode) -> bool:
         return False
 
     def _edge_cost(self, nodeU: TrackNode, nodeV: TrackNode) -> float:
+        """
+        Compute the cost of linking two detections (nodes) in the graph.
+
+        Args:
+            nodeU (TrackNode): Source node.
+            nodeV (TrackNode): Destination node.
+
+        Returns:
+            float: Edge cost based on IoU, distance, size, and confidence weights.
+        """
         bboxU, bboxV = nodeU.bbox, nodeV.bbox
         conf_u, conf_v = nodeU.confidence, nodeV.confidence
 
@@ -159,6 +199,10 @@ def _edge_cost(self, nodeU: TrackNode, nodeV: TrackNode) -> float:
         )
 
     def _build_graph(self):
+        """
+        Build the directed graph of detections for KSP computation.
+        Nodes represent detections; edges represent possible associations.
+        """
         G = nx.DiGraph()
         G.add_node(self.source)
         G.add_node(self.sink)
@@ -198,6 +242,22 @@ def _build_graph(self):
         self.graph = G
 
     def solve(self, k: Optional[int] = None) -> List[List[TrackNode]]:
+        """
+        Solve the K-Shortest Paths problem on the constructed detection graph.
+
+        This method extracts up to k node-disjoint paths from the source to the sink in
+        the detection graph, assigning each path as a unique object track. Edge reuse is
+        penalized to encourage distinct tracks. The cost of each edge is determined by
+        the configured weights for IoU, distance, size, and confidence.
+
+        Args:
+            k (Optional[int]): The number of tracks (paths) to extract. If None, uses
+                the maximum number of detections in any frame as the default.
+
+        Returns:
+            List[List[TrackNode]]: A list of tracks, each track is a list of TrackNode
+                objects representing the detections assigned to that track.
+        """
         self._build_graph()
 
         G_base = self.graph.copy()

trackers/core/ksp/tracker.py

@@ -20,60 +20,57 @@ class KSPTracker(BaseOfflineTracker):
     def __init__(
         self,
         path_overlap_penalty: Optional[int] = 40,
-        iou_weight: Optional[int] = 0.9,
-        dist_weight: Optional[int] = 0.1,
-        size_weight: Optional[int] = 0.1,
-        conf_weight: Optional[int] = 0.1,
+        iou_weight: Optional[float] = 0.9,
+        dist_weight: Optional[float] = 0.1,
+        size_weight: Optional[float] = 0.1,
+        conf_weight: Optional[float] = 0.1,
         entry_exit_regions: Optional[List[Tuple[int, int, int, int]]] = None,
         use_border: Optional[bool] = True,
         borders: Optional[Set[str]] = None,
         border_margin: Optional[int] = 40,
         frame_size: Optional[Tuple[int, int]] = (1920, 1080),
     ) -> None:
         """
-        Initialize the KSPTracker and its solver.
+        Initialize the KSPTracker and its underlying solver with region and cost configuration.
 
         Args:
-            path_overlap_penalty (Optional[int]): Penalty for reusing the same edge
-                (detection pairing) in multiple tracks. Increasing this value encourages
-                the tracker to produce more distinct, non-overlapping tracks by
-                discouraging shared detections between tracks.
-
-            iou_weight (Optional[int]): Weight for the Intersection-over-Union (IoU)
-                penalty in the edge cost. Higher values make the tracker favor linking
-                detections with greater spatial overlap, which helps maintain track
-                continuity for objects that move smoothly.
-
-            dist_weight (Optional[int]): Weight for the Euclidean distance between
-                detection centers in the edge cost. Increasing this value penalizes
-                large jumps between detections in consecutive frames, promoting
-                smoother, more physically plausible tracks.
-
-            size_weight (Optional[int]): Weight for the size difference penalty in the
-                edge cost. Higher values penalize linking detections with significantly
-                different bounding box areas, which helps prevent identity switches when
-                object size changes abruptly.
-
-            conf_weight (Optional[int]): Weight for the confidence penalty in the edge
-                cost. Higher values penalize edges between detections with lower
-                confidence scores, making the tracker prefer more reliable detections
-                and reducing the impact of false positives.
-
-            entry_exit_regions (Optional[List[Tuple[int, int, int, int]]]): List of rectangular entry/exit regions.
-            use_border (Optional[bool]): Enable/disable border-based entry/exit.
-            borders (Optional[Set[str]]): Set of borders to use. {"left", "right", "top", "bottom"}
-            border_margin (Optional[int]): Border thickness in pixels.
-            frame_size (Optional[Tuple[int, int]]): Size of the image (width, height).
+            path_overlap_penalty (Optional[int]): Penalty for reusing the same edge (detection pairing) in multiple tracks.
+                Higher values encourage the tracker to produce more distinct, non-overlapping tracks by discouraging shared
+                detections between tracks. Default is 40.
+            iou_weight (Optional[float]): Weight for the Intersection-over-Union (IoU) penalty in the edge cost.
+                Higher values make the tracker favor linking detections with greater spatial overlap, which helps maintain
+                track continuity for objects that move smoothly. Default is 0.9.
+            dist_weight (Optional[float]): Weight for the Euclidean distance between detection centers in the edge cost.
+                Increasing this value penalizes large jumps between detections in consecutive frames, promoting smoother,
+                more physically plausible tracks. Default is 0.1.
+            size_weight (Optional[float]): Weight for the size difference penalty in the edge cost.
+                Higher values penalize linking detections with significantly different bounding box areas, which helps
+                prevent identity switches when object size changes abruptly. Default is 0.1.
+            conf_weight (Optional[float]): Weight for the confidence penalty in the edge cost.
+                Higher values penalize edges between detections with lower confidence scores, making the tracker prefer
+                more reliable detections and reducing the impact of false positives. Default is 0.1.
+            entry_exit_regions (Optional[List[Tuple[int, int, int, int]]]): List of rectangular entry/exit regions,
+                each defined as (x1, y1, x2, y2) in pixel coordinates. These regions are used to determine when objects
+                enter or exit the scene. Default is an empty list.
+            use_border (Optional[bool]): Whether to enable border-based entry/exit logic. If True, objects entering or
+                exiting through the image borders (as defined by `borders` and `border_margin`) are considered for
+                entry/exit events. Default is True.
+            borders (Optional[Set[str]]): Set of border sides to use for entry/exit logic. Valid values are any subset
+                of {"left", "right", "top", "bottom"}. Default is all four borders.
+            border_margin (Optional[int]): Thickness of the border region (in pixels) used for entry/exit detection.
+                Default is 40.
+            frame_size (Optional[Tuple[int, int]]): Size of the image frames as (width, height). Used to determine
+                border region extents. Default is (1920, 1080).
         """
-        self.entry_exit_regions = (
+        self.entry_exit_regions: List[Tuple[int, int, int, int]] = (
             entry_exit_regions if entry_exit_regions is not None else []
         )
-        self.use_border = use_border
-        self.borders = (
+        self.use_border: bool = use_border if use_border is not None else True
+        self.borders: Set[str] = (
             borders if borders is not None else {"left", "right", "top", "bottom"}
         )
-        self.border_margin = border_margin
-        self.frame_size = frame_size
+        self.border_margin: int = border_margin if border_margin is not None else 40
+        self.frame_size: Tuple[int, int] = frame_size if frame_size is not None else (1920, 1080)
         self._solver = KSPSolver(
             path_overlap_penalty=path_overlap_penalty,
             iou_weight=iou_weight,
@@ -113,7 +110,10 @@ def _update(self, detections: sv.Detections) -> sv.Detections:
 
     def set_entry_exit_regions(self, regions: List[Tuple[int, int, int, int]]) -> None:
         """
-        Set rectangular entry/exit zones (x1, y1, x2, y2).
+        Set rectangular entry/exit zones (x1, y1, x2, y2) and update both the tracker and solver.
+
+        Args:
+            regions (List[Tuple[int, int, int, int]]): List of rectangular regions for entry/exit logic.
         """
         self.entry_exit_regions = regions
         self._solver.set_entry_exit_regions(regions)
@@ -125,12 +125,19 @@ def set_border_entry_exit(
         margin: Optional[int] = 40,
         frame_size: Optional[Tuple[int, int]] = (1920, 1080),
     ) -> None:
-        self.use_border = use_border
-        self.borders = (
-            borders if borders is not None else {"left", "right", "top", "bottom"}
-        )
-        self.border_margin = margin
-        self.frame_size = frame_size
+        """
+        Configure border-based entry/exit zones and update both the tracker and solver.
+
+        Args:
+            use_border (Optional[bool]): Enable/disable border-based entry/exit.
+            borders (Optional[Set[str]]): Set of borders to use. {"left", "right", "top", "bottom"}
+            margin (Optional[int]): Border thickness in pixels.
+            frame_size (Optional[Tuple[int, int]]): Size of the image (width, height).
+        """
+        self.use_border = use_border if use_border is not None else True
+        self.borders = borders if borders is not None else {"left", "right", "top", "bottom"}
+        self.border_margin = margin if margin is not None else 40
+        self.frame_size = frame_size if frame_size is not None else (1920, 1080)
         self._solver.set_border_entry_exit(
             use_border=self.use_border,
             borders=self.borders,