Merge pull request #46 from chanzuckerberg/tomo-fix

fix: Make sure AMG Parameters are properly read for Tomography Workflow

Author: jtschwar

docs/tutorials/inference.md

@@ -74,11 +74,11 @@ saber segment tomograms \
     --config copick_config.json \
     --model-config results/model_config.yaml \
     --model-weights results/best_model.pth \
-    --segmentation-name organelles \
+    --seg-name organelles \
     --target-class 2 
 ```
 
-When no `--run-ids` are provided, SABER will segment the entire project and save the results under provided `--segmentation-name` and `--segmentation-session-id` flags (by default, the `user-id` will always be saber. )
+When no `--run-ids` are provided, SABER will segment the entire project and save the results under provided `--seg-name` and `--seg-session-id` flags (by default, the `user-id` will always be saber. )
 
 **Parameters**:
 
@@ -88,8 +88,8 @@ When no `--run-ids` are provided, SABER will segment the entire project and save
 | `--voxel-size` | Resolution of tomograms to process | `10` |
 | `--tomo-alg` | Reconstruction algorithm to query | `denoised` |
 | `--slab-thickness` | Thickness of slab for initial segmentation | `10` |
-| `--segmentation-name` | Name of segmentation session | `organelles` |
-| `--segmentation-session-id` | Session ID for segmentation mask | `1` |
+| `--seg-name` | Name of segmentated object | `organelles` |
+| `--seg-session-id` | Session ID for segmentation mask | `1` |
 | `--model-config` | Path to classifier model config | `results/model_config.yaml` |
 | `--model-weights` | Path to trained classifier weights | `results/best_model.pth` |
 | `--target-class` | Target class (-1 for semantic, N>0 for instance) | `2` |

saber/classifier/preprocess/tomogram_training_prep.py

@@ -46,7 +46,7 @@ def extract_sam2_candidates(
 
     # Get the Global Zarr Writer
     zwriter = zarr_writer.get_zarr_writer(output)
-    zwriter.set_dict_attr('amg', segmenter.amg_params)
+    zwriter.set_dict_attr('amg', segmenter.cfg)
 
     # Get Tomogram
     vol = readers.tomogram(run, voxel_size, tomogram_algorithm)
@@ -55,6 +55,7 @@ def extract_sam2_candidates(
         return
 
     # Hard coded conversion from Angstroms to nanometers
+    # Copick tomograms are typically stored in Angstroms
     voxel_size /= 10
 
     # Process Multiple Slabs or Single Slab at the Center of the Volume

saber/entry_points/inference_core.py

@@ -102,7 +102,7 @@ def segment_micrograph_core(
 
     # Get the Global Zarr Writer
     zwriter = zarr_writer.get_zarr_writer(output) 
-    zwriter.set_dict_attr('amg', segmenter.amg_params)
+    zwriter.set_dict_attr('amg', segmenter.cfg)
 
     # Ensure we're on the correct GPU
     torch.cuda.set_device(gpu_id)

saber/entry_points/run_micrograph_segment.py

@@ -51,7 +51,7 @@ def micrographs(
     if target_resolution is not None and scale_factor is not None:
         raise ValueError("Please provide either target_resolution OR scale_factor input, not both.")
 
-    # Get All Files in the Directory
+    # Get All Files in the Directoryclear
     files = glob.glob(input)
     if len(files) == 0:
         raise ValueError(f"No files found in {input}")

saber/segmenters/base.py

@@ -60,7 +60,7 @@ def __init__(self,
 
             # Use Default AMG Config
             self.cfg = amg.get_default()
-            self.cfg['cfg'] = sam2_cfg
+            self.cfg['sam2_cfg'] = sam2_cfg
 
         # Build SAM2 Automatic Mask Generator
         self.mask_generator = amg.build_amg(
@@ -238,9 +238,9 @@ def __init__(self,
         super().__init__(sam2_cfg, deviceID, classifier, target_class, min_mask_area, min_rel_box_size)
 
         # Build Tomogram Predictor (VOS Optimized)
-        (cfg, checkpoint) = pretrained_weights.get_sam2_checkpoint(sam2_cfg)
+        (cfg, checkpoint) = pretrained_weights.get_sam2_checkpoint(self.cfg['sam2_cfg'])
         self.video_predictor = tomogram_predictor.TomogramSAM2Adapter(cfg, checkpoint, self.device)  
-
+        
         # Initialize Inference State
         self.inference_state = None