Sync logic with bat_tracker v1.1.7

app.py

@@ -13,7 +13,7 @@ import yaml
 from bat_tracker.pipeline import run_pipeline
 
 
-APP_VERSION = "v1.1.5"
+APP_VERSION = "v1.1.7"
 APP_TITLE = f"Bat Tracker {APP_VERSION}"
 APP_DESCRIPTION = (
     "Sube un video IR monocromo para ejecutar el pipeline, revisar la region valida "

bat_tracker/pipeline.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import csv
+import heapq
 import json
 import sys
 from collections import Counter, defaultdict, deque
@@ -146,6 +147,24 @@ def _classify_direction(start_inside: bool, end_inside: bool) -> str:
     return "outside"
 
 
+def _classify_direction_full(
+    s_in: bool,
+    e_in: bool,
+    tps: List[TrackPoint],
+    valid_mask: np.ndarray | None,
+    frame_shape: tuple[int, int] | None = None,
+) -> str:
+    direction = _classify_direction(s_in, e_in)
+    if direction == "inside" and valid_mask is not None and len(tps) > 2:
+        for tp in tps[1:-1]:
+            if not _point_in_mask(tp, valid_mask):
+                direction = "exit"
+                break
+    if direction == "outside" and frame_shape is not None:
+        direction = _infer_outside_direction_from_motion(tps[0], tps[-1], frame_shape)
+    return direction
+
+
 def _infer_outside_direction_from_motion(
     start: TrackPoint,
     end: TrackPoint,
@@ -189,9 +208,7 @@ def _write_events_csv(
         if valid_mask is not None:
             s_in = _point_in_mask(start, valid_mask)
             e_in = _point_in_mask(end, valid_mask)
-            direction = _classify_direction(s_in, e_in)
-            if direction == "outside":
-                direction = _infer_outside_direction_from_motion(start, end, valid_mask.shape[:2])
+            direction = _classify_direction_full(s_in, e_in, tps, valid_mask, valid_mask.shape[:2])
         else:
             s_in = None
             e_in = None
@@ -477,15 +494,13 @@ def _filter_track_points(
         if require_start_or_end_in_valid_region and gate_mask is not None:
             s_in = _point_in_mask(start, gate_mask)
             e_in = _point_in_mask(end, gate_mask)
-            direction = _classify_direction(s_in, e_in)
+            direction = _classify_direction_full(s_in, e_in, track_points, valid_mask)
             if not (s_in or e_in):
                 reject_reasons.append("valid_region_gate")
         elif valid_mask is not None:
             s_in = _point_in_mask(start, valid_mask)
             e_in = _point_in_mask(end, valid_mask)
-            direction = _classify_direction(s_in, e_in)
-            if direction == "outside":
-                direction = _infer_outside_direction_from_motion(start, end, valid_mask.shape[:2])
+            direction = _classify_direction_full(s_in, e_in, track_points, valid_mask, valid_mask.shape[:2])
 
         accepted = not reject_reasons
         if not accepted:
@@ -565,6 +580,7 @@ def _auto_merge_track_points(points: List[TrackPoint], tracking_cfg: Dict) -> tu
     max_overlap_mean_dist = float(tracking_cfg.get("merge_overlap_max_mean_distance", 60.0))
     min_overlap_cos = float(tracking_cfg.get("merge_overlap_min_direction_cosine", 0.8))
     local_overlap_min_cos = max(0.65, min_overlap_cos - 0.15)
+    proximity_override_dist = float(tracking_cfg.get("merge_overlap_proximity_override_distance", 0.0))
 
     parent: Dict[int, int] = {track_id: track_id for track_id in by_track}
 
@@ -586,93 +602,150 @@ def _auto_merge_track_points(points: List[TrackPoint], tracking_cfg: Dict) -> tu
 
     merges_applied: List[Dict] = []
     track_ids = sorted(by_track.keys())
-    for idx, track_a_id in enumerate(track_ids):
-        a_pts = by_track[track_a_id]
-        a_start = a_pts[0]
-        a_end = a_pts[-1]
-        a_start_vec, a_end_vec = _track_edge_vectors(a_pts)
-        a_frames = {p.frame: p for p in a_pts}
-
-        for track_b_id in track_ids[idx + 1 :]:
-            b_pts = by_track[track_b_id]
-            b_start = b_pts[0]
-            b_end = b_pts[-1]
-            b_start_vec, b_end_vec = _track_edge_vectors(b_pts)
-
-            reason = None
-            reason_data: Dict[str, float | int] = {}
-
-            if a_end.frame < b_start.frame:
-                gap = b_start.frame - a_end.frame
-                dist = hypot(b_start.x - a_end.x, b_start.y - a_end.y)
-                if gap <= max_gap and dist <= max_endpoint_dist:
-                    reason = "handoff"
-                    reason_data = {"gap_frames": gap, "endpoint_distance": dist}
-            elif b_end.frame < a_start.frame:
-                gap = a_start.frame - b_end.frame
-                dist = hypot(a_start.x - b_end.x, a_start.y - b_end.y)
-                if gap <= max_gap and dist <= max_endpoint_dist:
-                    reason = "handoff"
-                    reason_data = {"gap_frames": gap, "endpoint_distance": dist}
-            else:
-                b_frames = {p.frame: p for p in b_pts}
-                common_frames = sorted(set(a_frames.keys()).intersection(b_frames.keys()))
-                if len(common_frames) >= 2:
-                    distances = []
-                    for frame in common_frames:
-                        pa = a_frames[frame]
-                        pb = b_frames[frame]
-                        distances.append(hypot(pa.x - pb.x, pa.y - pb.y))
-
-                    mean_distance = sum(distances) / len(distances)
-                    start_cos = _vector_cosine(a_start_vec, b_start_vec)
-                    end_cos = _vector_cosine(a_end_vec, b_end_vec)
-                    connector_cos = _vector_cosine(a_end_vec, b_start_vec)
-                    global_cosines = [c for c in (start_cos, end_cos) if c is not None]
-                    mean_cos = (sum(global_cosines) / len(global_cosines)) if global_cosines else None
-
-                    overlap_reason = None
-                    if len(common_frames) >= min_overlap_common:
-                        if mean_distance <= max_overlap_mean_dist and (
-                            mean_cos is None or mean_cos >= min_overlap_cos or connector_cos is not None and connector_cos >= min_overlap_cos
-                        ):
-                            overlap_reason = "overlap"
+
+    track_data: Dict[int, Dict] = {}
+    for tid in track_ids:
+        pts = by_track[tid]
+        start_vec, end_vec = _track_edge_vectors(pts)
+        track_data[tid] = {
+            "start": pts[0],
+            "end": pts[-1],
+            "start_vec": start_vec,
+            "end_vec": end_vec,
+            "frames": {p.frame: p for p in pts},
+        }
+
+    n_tracks = len(track_ids)
+    start_frames = [track_data[tid]["start"].frame for tid in track_ids]
+    end_frames = [track_data[tid]["end"].frame for tid in track_ids]
+    start_x = [track_data[tid]["start"].x for tid in track_ids]
+    start_y = [track_data[tid]["start"].y for tid in track_ids]
+    end_x = [track_data[tid]["end"].x for tid in track_ids]
+    end_y = [track_data[tid]["end"].y for tid in track_ids]
+    max_endpoint_dist_sq = max_endpoint_dist * max_endpoint_dist
+
+    sorted_positions = sorted(range(n_tracks), key=lambda idx: start_frames[idx])
+    heap: List[tuple[int, int]] = []
+    candidate_pairs: List[tuple[int, int]] = []
+
+    for b_pos in sorted_positions:
+        b_start_frame = start_frames[b_pos]
+        b_id = track_ids[b_pos]
+
+        while heap and heap[0][0] < b_start_frame - max_gap:
+            heapq.heappop(heap)
+
+        for a_end_frame, a_pos in heap:
+            a_id = track_ids[a_pos]
+            if a_end_frame < b_start_frame:
+                dx = end_x[a_pos] - start_x[b_pos]
+                dy = end_y[a_pos] - start_y[b_pos]
+                if dx * dx + dy * dy > max_endpoint_dist_sq:
+                    continue
+            candidate_pairs.append((min(a_id, b_id), max(a_id, b_id)))
+
+        heapq.heappush(heap, (end_frames[b_pos], b_pos))
+
+    for track_a_id, track_b_id in candidate_pairs:
+        td_a = track_data[track_a_id]
+        td_b = track_data[track_b_id]
+        a_start = td_a["start"]
+        a_end = td_a["end"]
+        a_start_vec = td_a["start_vec"]
+        a_end_vec = td_a["end_vec"]
+        a_frames = td_a["frames"]
+        b_start = td_b["start"]
+        b_end = td_b["end"]
+        b_start_vec = td_b["start_vec"]
+        b_end_vec = td_b["end_vec"]
+
+        reason = None
+        reason_data: Dict[str, float | int] = {}
+
+        if a_end.frame < b_start.frame:
+            gap = b_start.frame - a_end.frame
+            dist = hypot(b_start.x - a_end.x, b_start.y - a_end.y)
+            if gap <= max_gap and dist <= max_endpoint_dist:
+                reason = "handoff"
+                reason_data = {"gap_frames": gap, "endpoint_distance": dist}
+        elif b_end.frame < a_start.frame:
+            gap = a_start.frame - b_end.frame
+            dist = hypot(a_start.x - b_end.x, a_start.y - b_end.y)
+            if gap <= max_gap and dist <= max_endpoint_dist:
+                reason = "handoff"
+                reason_data = {"gap_frames": gap, "endpoint_distance": dist}
+        else:
+            b_frames = td_b["frames"]
+            common_frames = sorted(set(a_frames.keys()).intersection(b_frames.keys()))
+            if len(common_frames) >= 1:
+                distances = []
+                for frame in common_frames:
+                    pa = a_frames[frame]
+                    pb = b_frames[frame]
+                    distances.append(hypot(pa.x - pb.x, pa.y - pb.y))
+
+                mean_distance = sum(distances) / len(distances)
+                start_cos = _vector_cosine(a_start_vec, b_start_vec)
+                end_cos = _vector_cosine(a_end_vec, b_end_vec)
+                connector_cos = _vector_cosine(a_end_vec, b_start_vec)
+                global_cosines = [c for c in (start_cos, end_cos) if c is not None]
+                mean_cos = (sum(global_cosines) / len(global_cosines)) if global_cosines else None
+
+                overlap_reason = None
+                if (
+                    proximity_override_dist > 0.0
+                    and mean_distance <= proximity_override_dist
+                    and len(common_frames) >= 1
+                ):
+                    overlap_reason = "overlap_proximity"
+                elif len(common_frames) >= min_overlap_common:
+                    if mean_distance <= max_overlap_mean_dist and (
+                        mean_cos is None or mean_cos >= min_overlap_cos or connector_cos is not None and connector_cos >= min_overlap_cos
+                    ):
+                        overlap_reason = "overlap"
                     elif (
                         mean_distance <= max_overlap_mean_dist
-                        and connector_cos is not None
-                        and connector_cos >= local_overlap_min_cos
+                        and start_cos is not None
+                        and start_cos >= local_overlap_min_cos
                     ):
-                        overlap_reason = "overlap_local"
-
-                    if overlap_reason is not None:
-                        direction_score = connector_cos
-                        if direction_score is None:
-                            direction_score = mean_cos if mean_cos is not None else 1.0
-                        reason = overlap_reason
-                        reason_data = {
-                            "common_frames": len(common_frames),
-                            "mean_distance": mean_distance,
-                            "mean_direction_cosine": direction_score,
-                        }
-
-            if reason is None:
-                continue
-
-            ra = find(track_a_id)
-            rb = find(track_b_id)
-            if ra == rb:
-                continue
-            union(track_a_id, track_b_id)
-            merged_to = min(find(track_a_id), find(track_b_id))
-            merges_applied.append(
-                {
-                    "track_a": track_a_id,
-                    "track_b": track_b_id,
-                    "merged_to": merged_to,
-                    "reason": reason,
-                    **reason_data,
-                }
-            )
+                        overlap_reason = "overlap_start_aligned"
+                elif (
+                    len(common_frames) >= 1
+                    and mean_distance <= max_overlap_mean_dist
+                    and connector_cos is not None
+                    and connector_cos >= local_overlap_min_cos
+                ):
+                    overlap_reason = "overlap_local"
+
+                if overlap_reason is not None:
+                    direction_score = connector_cos
+                    if direction_score is None:
+                        direction_score = mean_cos if mean_cos is not None else 1.0
+                    reason = overlap_reason
+                    reason_data = {
+                        "common_frames": len(common_frames),
+                        "mean_distance": mean_distance,
+                        "mean_direction_cosine": direction_score,
+                    }
+
+        if reason is None:
+            continue
+
+        ra = find(track_a_id)
+        rb = find(track_b_id)
+        if ra == rb:
+            continue
+        union(track_a_id, track_b_id)
+        merged_to = min(find(track_a_id), find(track_b_id))
+        merges_applied.append(
+            {
+                "track_a": track_a_id,
+                "track_b": track_b_id,
+                "merged_to": merged_to,
+                "reason": reason,
+                **reason_data,
+            }
+        )
 
     remap: Dict[int, int] = {track_id: find(track_id) for track_id in track_ids}
     if all(src == dst for src, dst in remap.items()):

bat_tracker/tracker.py

@@ -3,6 +3,9 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Dict, List, Tuple
 
+import numpy as np
+from scipy.optimize import linear_sum_assignment
+
 from .detection import Detection
 
 
@@ -49,28 +52,34 @@ class GreedyTracker:
 
         unmatched_track_ids = set(self._active.keys())
         unmatched_det_idxs = set(range(len(detections)))
-
-        candidate_pairs: List[Tuple[float, int, int]] = []
-        max_distance_sq = self.max_distance_sq
-        for track_id, track in self._active.items():
-            dt_pred = max(1, frame_idx - track.last_frame) / self.fps
-            pred_x = track.x + track.vx * dt_pred
-            pred_y = track.y + track.vy * dt_pred
-            for det_idx, det in enumerate(detections):
-                dx = pred_x - det.x
-                dy = pred_y - det.y
-                d_sq = dx * dx + dy * dy
-                if d_sq <= max_distance_sq:
-                    candidate_pairs.append((d_sq, track_id, det_idx))
-
-        candidate_pairs.sort(key=lambda t: t[0])
         assignments: List[Tuple[int, int]] = []
-
-        for _, track_id, det_idx in candidate_pairs:
-            if track_id in unmatched_track_ids and det_idx in unmatched_det_idxs:
-                assignments.append((track_id, det_idx))
-                unmatched_track_ids.remove(track_id)
-                unmatched_det_idxs.remove(det_idx)
+        if self._active and detections:
+            track_ids_list = list(self._active.keys())
+            n_tracks = len(track_ids_list)
+            n_dets = len(detections)
+            max_dist = self.max_distance
+            inf_cost = max_dist * 1e6
+
+            cost = np.full((n_tracks, n_dets), inf_cost, dtype=np.float64)
+            for i, track_id in enumerate(track_ids_list):
+                track = self._active[track_id]
+                dt_pred = max(1, frame_idx - track.last_frame) / self.fps
+                pred_x = track.x + track.vx * dt_pred
+                pred_y = track.y + track.vy * dt_pred
+                for j, det in enumerate(detections):
+                    dx = pred_x - det.x
+                    dy = pred_y - det.y
+                    dist = (dx * dx + dy * dy) ** 0.5
+                    if dist <= max_dist:
+                        cost[i, j] = dist
+
+            row_ind, col_ind = linear_sum_assignment(cost)
+            for i, j in zip(row_ind, col_ind):
+                if cost[i, j] < inf_cost:
+                    track_id = track_ids_list[i]
+                    assignments.append((track_id, j))
+                    unmatched_track_ids.discard(track_id)
+                    unmatched_det_idxs.discard(j)
 
         for track_id, det_idx in assignments:
             track = self._active[track_id]

requirements.txt

@@ -1,4 +1,5 @@
 numpy>=1.24
+scipy>=1.10
 opencv-python>=4.8
 matplotlib>=3.7
 PyYAML>=6.0