Fix deprecated sklear linear_assignment

see nwojke#220
jimtang123 · Jul 22, 2022 · f9de9ac · f9de9ac
1 parent 280b8bd
commit f9de9ac
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/deep_sort/linear_assignment.py b/deep_sort/linear_assignment.py
@@ -1,7 +1,7 @@
 # vim: expandtab:ts=4:sw=4
 from __future__ import absolute_import
 import numpy as np
-from sklearn.utils.linear_assignment_ import linear_assignment
+from scipy.optimize import linear_sum_assignment as linear_assignment
 from . import kalman_filter
 
 
@@ -55,7 +55,8 @@ def min_cost_matching(
     cost_matrix = distance_metric(
         tracks, detections, track_indices, detection_indices)
     cost_matrix[cost_matrix > max_distance] = max_distance + 1e-5
-    indices = linear_assignment(cost_matrix)
+    r_ind, c_ind = linear_assignment(cost_matrix)
+    indices = np.vstack((r_ind, c_ind)).T
 
     matches, unmatched_tracks, unmatched_detections = [], [], []
     for col, detection_idx in enumerate(detection_indices):