fixed toml sam-2 dependency; clean-up after post-processing

juglab · Nov 7, 2024 · 9d4ca5e · 9d4ca5e
1 parent 8f1c8d9
commit 9d4ca5e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -60,7 +60,7 @@ dependencies = [
     "iopath>=0.1.10",
     "timm==1.0.9",
     "segment-anything-py",
-    "sam2 @ git+https://github.com/facebookresearch/sam2.git"
+    "sam-2 @ git+https://github.com/facebookresearch/sam2.git"
 ]
 [project.optional-dependencies]
 # development dependencies and tooling

diff --git a/src/featureforest/postprocess/postprocess_with_sam.py b/src/featureforest/postprocess/postprocess_with_sam.py
@@ -193,6 +193,10 @@ def postprocess_with_sam(
         # put the final label mask into final result mask
         final_mask[sam_label_mask] = label
 
+    # clean-up
+    del predictor
+    torch.cuda.empty_cache()
+
     return final_mask
 
 
@@ -232,6 +236,10 @@ def get_sam_auto_masks(input_image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]
     sam_masks = np.array([mask["segmentation"] for mask in sam_generated_masks])
     sam_areas = np.array([mask["area"] for mask in sam_generated_masks])
 
+    # clean-up
+    del mask_generator
+    torch.cuda.empty_cache()
+
     return sam_masks, sam_areas