Merge pull request #157 from computational-cell-analytics/fix-amg-state-serialization

Serialize amg state to the cpu so that it can be loaded without a GPU

Author: constantinpape

micro_sam/instance_segmentation.py

@@ -288,6 +288,7 @@ def get_state(self) -> Dict[str, Any]:
         """
         if not self.is_initialized:
             raise RuntimeError("The state has not been computed yet. Call initialize first.")
+
         return {"crop_list": self.crop_list, "crop_boxes": self.crop_boxes, "original_size": self.original_size}
 
     def set_state(self, state: Dict[str, Any]) -> None:

micro_sam/precompute_state.py

@@ -9,6 +9,7 @@
 from typing import Optional, Tuple, Union
 
 import numpy as np
+import torch
 from segment_anything.predictor import SamPredictor
 from tqdm import tqdm
 
@@ -50,9 +51,21 @@ def cache_amg_state(
 
     if verbose:
         print("Precomputing the state for instance segmentation.")
+
     amg.initialize(raw, image_embeddings=image_embeddings, verbose=verbose)
+    amg_state = amg.get_state()
+
+    # put all state onto the cpu so that the state can be deserialized without a gpu
+    new_crop_list = []
+    for mask_data in amg_state["crop_list"]:
+        for k, v in mask_data.items():
+            if torch.is_tensor(v):
+                mask_data[k] = v.cpu()
+        new_crop_list.append(mask_data)
+    amg_state["crop_list"] = new_crop_list
+
     with open(save_path_amg, "wb") as f:
-        pickle.dump(amg.get_state(), f)
+        pickle.dump(amg_state, f)
 
     return amg