Merge pull request #14 from aewebb80/main

Update cost matrix to only consider tracks in the match frame
kocherlab · Apr 8, 2024 · 0396273 · 0396273
2 parents fd6db7b + 1730b3e
commit 0396273
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
diff --git a/naps/__init__.py b/naps/__init__.py
@@ -6,7 +6,7 @@
 
 # Basic Information
 __name__ = "naps-track"
-__version__ = "1.0.6"
+__version__ = "1.0.7"
 __summary__ = "NAPS (NAPS is ArUco Plus SLEAP)"
 __url__ = "https://github.com/kocherlab/naps"
 __code__ = "https://github.com/kocherlab/naps"

diff --git a/naps/cost_matrix.py b/naps/cost_matrix.py
@@ -74,6 +74,12 @@ def assignTrackTagPairs(self):
             cost_dataframe = cost_dataframe.loc[~(cost_dataframe == 0).all(axis=1)]
             cost_dataframe = cost_dataframe.loc[:, ~(cost_dataframe == 0).all(axis=0)]
 
+            # Create a list of the tracks in the current frame
+            match_tracks = set(self.unmatched_dict[match_frame]) & set(cost_dataframe.columns)
+
+            # Subset the cost dataframe to only include the tracks in the current frame
+            cost_dataframe = cost_dataframe[list(match_tracks)]
+
             # Store the assignments
             for track_index, tag_index in self.assignment_method(cost_dataframe.values):
                 matched_dict[match_frame][

diff --git a/tests/data/example_naps_track_output.h5.slp b/tests/data/example_naps_track_output.h5.slp