diff --git a/src/laptrack/_tracking.py b/src/laptrack/_tracking.py index 0a16758e..a4201856 100644 --- a/src/laptrack/_tracking.py +++ b/src/laptrack/_tracking.py @@ -657,6 +657,7 @@ def predict_dataframe( validate_frame: bool = True, only_coordinate_cols: bool = True, connected_edges: Optional[List[Tuple[Int, Int]]] = None, + index_offset: Int = 0, ) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: """Shorthand for the tracking with the dataframe input / output. @@ -676,6 +677,8 @@ def predict_dataframe( The edges that is known to be connected. Must be a list of the tuples of the row numbers (not index, but `iloc`). If None, no edges are assumed to be connected. + index_offset : Int + The offset to add to the track and tree index. Returns ------- @@ -721,6 +724,16 @@ def predict_dataframe( track_df, split_df, merge_df = convert_tree_to_dataframe( tree, dataframe=df, frame_index=frame_index ) + + track_df["track_id"] = track_df["track_id"] + index_offset + track_df["tree_id"] = track_df["tree_id"] + index_offset + if not split_df.empty: + split_df["parent_track_id"] = split_df["parent_track_id"] + index_offset + split_df["child_track_id"] = split_df["child_track_id"] + index_offset + if not merge_df.empty: + merge_df["parent_track_id"] = merge_df["parent_track_id"] + index_offset + merge_df["child_track_id"] = merge_df["child_track_id"] + index_offset + return track_df, split_df, merge_df diff --git a/tests/test_tracking.py b/tests/test_tracking.py index 58794df6..d40120d2 100644 --- a/tests/test_tracking.py +++ b/tests/test_tracking.py @@ -137,6 +137,28 @@ def test_reproducing_trackmate(testdata, tracker_class, parallel_backend) -> Non assert all(split_df == split_df2) assert all(merge_df == merge_df2) + # check index offset + track_df3, split_df3, merge_df3 = lt.predict_dataframe( + df, ["x", "y"], index_offset=2 + ) + assert min(track_df3["track_id"]) == 2 + assert min(track_df3["tree_id"]) == 2 + + track_df4 = track_df2.copy() + split_df4 = split_df2.copy() + merge_df4 = merge_df2.copy() + track_df4["track_id"] = track_df4["track_id"] + 2 + track_df4["tree_id"] = track_df4["tree_id"] + 2 + if not split_df4.empty: + split_df4["parent_track_id"] = split_df4["parent_track_id"] + 2 + split_df4["child_track_id"] = split_df4["child_track_id"] + 2 + if not merge_df4.empty: + merge_df4["parent_track_id"] = merge_df4["parent_track_id"] + 2 + merge_df4["child_track_id"] = merge_df4["child_track_id"] + 2 + assert all(track_df3 == track_df2) + assert all(split_df3 == split_df2) + assert all(merge_df3 == merge_df2) + track_df, split_df, merge_df = lt.predict_dataframe( df, ["x", "y"], only_coordinate_cols=False )