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

computational-cell-analytics · Jul 4, 2024 · 3d8d879 · 3d8d879
1 parent 09af0a7
commit 3d8d879
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