performance improvements for spot location refinement code; pinning some packages in requirements.txt

Author: ellenemerson

deepcell_spots/applications/polaris.py

@@ -436,19 +436,25 @@ def _predict(self,
             print('Refining spot locations.')
             dec_prob_im = np.zeros((spots_image.shape[:3]))
 
-            for i in range(len(df_results)):
-                gene = df_results.loc[i, 'predicted_name']
-                if gene in ['Background', 'Unknown']:
-                    continue
-                if 'Blank' in gene:
-                    continue
-                
-                x = df_results.loc[i, 'x']
-                y = df_results.loc[i, 'y']
-                b = df_results.loc[i, 'batch_id']
-                prob = max_proj_images[b, x, y]
-                
-                dec_prob_im[b, x, y] = prob
+            # Mask out unwanted genes
+            mask_valid = ~df_results['predicted_name'].isin(['Background', 'Unknown']) & \
+                         ~df_results['predicted_name'].str.contains('Blank', na=False)
+            valid_df = df_results[mask_valid]
+
+            b = valid_df['batch_id'].to_numpy()
+            x = valid_df['x'].to_numpy()
+            y = valid_df['y'].to_numpy()
+
+            dec_prob_im[b, x, y] = max_proj_images[b, x, y]
+
+
+            # Pre-create a dictionary for fast lookups
+            lookup_dict = {}
+            for idx, row in df_results.iterrows():
+                key = (row['batch_id'], row['x'], row['y'])
+                if key not in lookup_dict:
+                    lookup_dict[key] = []
+                lookup_dict[key].append(idx)
 
             mask = []
             for b in range(spots_image.shape[0]):
@@ -459,10 +465,10 @@ def _predict(self,
                     x = decoded_spots_locations[0][i, 0]
                     y = decoded_spots_locations[0][i, 1]
 
-                    mask.append(df_results.loc[(df_results.x==x) & 
-                                               (df_results.y==y) & 
-                                               (df_results.batch_id==b)].index[0])
-                
+                    key = (b, x, y)
+                    if key in lookup_dict:
+                        mask.append(lookup_dict[key][0])
+
             df_results = df_results.loc[mask]
 
         return df_results, segmentation_result

requirements.txt

@@ -8,11 +8,11 @@ jupyter>=1.0.0,<2
 networkx>=2.1
 opencv-python-headless<5
 deepcell>=0.12.7
-trackpy
+trackpy~=0.6.0
 tqdm
 plotly
-statsmodels
+statsmodels~=0.14.0
 --extra-index-url https://download.pytorch.org/whl/cpu # install the cpu only version of torch and torchvision
-torch
+torch~=2.3.0
 torchvision
-pyro-ppl
+pyro-ppl>=1.8.0