indexing change

Kapoorlabs-CAPED · Aug 14, 2024 · 406962b · 406962b
1 parent d705a08
commit 406962b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 21 deletions.
diff --git a/src/napatrackmater/Trackmate.py b/src/napatrackmater/Trackmate.py
@@ -2046,7 +2046,7 @@ def _update_spot_fate(self, TrackIds, fate_label):
                         {self.fate_key: fate_label}
                     )
 
-    def _get_trackmate_ids_by_location(self, dataframe: pd.DataFrame, tracklet_length: int = None):
+    def _get_trackmate_ids_by_location(self, dataframe: pd.DataFrame, tracklet_length= None):
         trackmate_track_ids = []
         t = int(self.tend)
         dataframe.columns = map(str.lower, dataframe.columns)
@@ -2072,7 +2072,7 @@ def _get_trackmate_ids_by_location(self, dataframe: pd.DataFrame, tracklet_lengt
                 if self.trackid_key in spot_properties_dict.keys():
                     trackmate_track_id = spot_properties_dict[self.trackid_key]
                     if tracklet_length is None:
-                       trackmate_track_ids.append(trackmate_track_id)
+                        trackmate_track_ids.append(trackmate_track_id)
                     else:
                            track_duration = spot_properties_dict[self.track_duration_key]
                            if track_duration > tracklet_length:

diff --git a/src/napatrackmater/Trackvector.py b/src/napatrackmater/Trackvector.py
@@ -4050,33 +4050,19 @@ def inception_model_prediction(
 
     def sample_subarrays(data, num_samples, tracklet_length, total_duration):
 
-        interval = max(1, (total_duration - tracklet_length) // num_samples)
+
+        max_start_index = total_duration - tracklet_length
+
+        start_indices = random.sample(range(max_start_index + 1), num_samples)
         subarrays = []
-        for i in range(num_samples):
-            start_min = i * interval
-            start_max = min((i + 1) * interval, total_duration - tracklet_length)
-
-            if start_max <= start_min:
-               start_max = start_min + 1
-
-            start_index = np.random.randint(start_min, start_max)
+        for start_index in start_indices:
             end_index = start_index + tracklet_length
             if end_index <= total_duration:
                sub_data = data[start_index:end_index, :]
                if sub_data.shape[0] == tracklet_length:
                   subarrays.append(sub_data)
 
 
-        if len(subarrays) < num_samples:
-            additional_subarrays = []
-            for i in range(num_samples - len(subarrays)):
-                start_index = np.random.randint(0, total_duration - tracklet_length)
-                end_index = start_index + tracklet_length
-                sub_data = data[start_index:end_index, :]
-                if sub_data.shape[0] == tracklet_length:
-                    additional_subarrays.append(sub_data)
-            subarrays.extend(additional_subarrays[: num_samples - len(subarrays)])
-
         return subarrays
 
     sub_arrays_shape = sample_subarrays(