fix bugs

brisvag · Aug 4, 2023 · 3f0339c · 3f0339c
1 parent 59c4972
commit 3f0339c
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/stemia/image/classify_densities.py b/stemia/image/classify_densities.py
@@ -42,7 +42,7 @@ def cli(stacks, max_classes):
 
     images = {}
 
-    print('Running with {max_classes} classes.')
+    print(f'Running with {max_classes} classes.')
 
     df = pd.DataFrame(columns=['total_density', 'radius_of_gyration'])
     df.index.name = 'image'
@@ -79,11 +79,12 @@ def cli(stacks, max_classes):
         progress.update(proc_task, advance=1)
 
         for cl, df_cl in progress.track(df.groupby('class'), description='Splitting classes...'):
-            stacked = []
+            stacked = {}
             for img in df_cl.index:
                 *img_name, idx = img.split('_')
                 img_name = '_'.join(img_name)
                 idx = int(idx)
-                stacked.append(images[img_name][idx])
-            mrc = mrcfile.new(f'{img_name}_class_{cl:04}.mrc', np.stack(stacked), overwrite=True)
-            mrc.close()
+                stacked[img_name] = images[img_name][idx]
+            for img_name, data in stacked.items():
+                mrc = mrcfile.new(f'{img_name}_class_{cl:04}.mrc', np.stack(data), overwrite=True)
+                mrc.close()