Fix bug in precompute for 3d data (#649)

computational-cell-analytics · Jul 3, 2024 · 6ad5e67 · 6ad5e67
1 parent 1e1cff9
commit 6ad5e67
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/micro_sam/precompute_state.py b/micro_sam/precompute_state.py
@@ -68,7 +68,7 @@ def cache_amg_state(
     if verbose:
         print("Precomputing the state for instance segmentation.")
 
-    amg.initialize(raw, image_embeddings=image_embeddings, verbose=verbose, i=i)
+    amg.initialize(raw if i is None else raw[i], image_embeddings=image_embeddings, verbose=verbose, i=i)
     amg_state = amg.get_state()
 
     # put all state onto the cpu so that the state can be deserialized without a gpu