Error in pulling sam2.1 vs sam2 when SAM-2 is installed (#1782)

Signed-off-by: Sachidanand Alle <[email protected]>
Project-MONAI · Nov 13, 2024 · a5733eb · a5733eb
1 parent b885a59
commit a5733eb
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 31 deletions.
diff --git a/monailabel/config.py b/monailabel/config.py
@@ -9,15 +9,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
-from importlib.util import find_spec
+from importlib.metadata import distributions
 from typing import Any, Dict, List, Optional
 
 from pydantic import AnyHttpUrl
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
 def is_package_installed(name):
-    return False if find_spec(name) is None else True
+    return name in sorted(x.name for x in distributions())
 
 
 class Settings(BaseSettings):

diff --git a/sample-apps/radiology/lib/configs/segmentation.py b/sample-apps/radiology/lib/configs/segmentation.py
@@ -21,7 +21,7 @@
 from monailabel.interfaces.config import TaskConfig
 from monailabel.interfaces.tasks.infer_v2 import InferTask
 from monailabel.interfaces.tasks.train import TrainTask
-from monailabel.utils.others.generic import download_file, strtobool
+from monailabel.utils.others.generic import download_file, remove_file, strtobool
 
 _, has_cp = optional_import("cupy")
 _, has_cucim = optional_import("cucim")
@@ -34,33 +34,38 @@ def init(self, name: str, model_dir: str, conf: Dict[str, str], planner: Any, **
         super().init(name, model_dir, conf, planner, **kwargs)
 
         # Labels
-        self.labels = {
-            "spleen": 1,
-            "kidney_right": 2,
-            "kidney_left": 3,
-            "gallbladder": 4,
-            "liver": 5,
-            "stomach": 6,
-            "aorta": 7,
-            "inferior_vena_cava": 8,
-            "portal_vein_and_splenic_vein": 9,
-            "pancreas": 10,
-            "adrenal_gland_right": 11,
-            "adrenal_gland_left": 12,
-            "lung_upper_lobe_left": 13,
-            "lung_lower_lobe_left": 14,
-            "lung_upper_lobe_right": 15,
-            "lung_middle_lobe_right": 16,
-            "lung_lower_lobe_right": 17,
-            "esophagus": 42,
-            "trachea": 43,
-            "heart_myocardium": 44,
-            "heart_atrium_left": 45,
-            "heart_ventricle_left": 46,
-            "heart_atrium_right": 47,
-            "heart_ventricle_right": 48,
-            "pulmonary_artery": 49,
-        }
+        conf_labels = self.conf.get("labels")
+        self.labels = (
+            {label: idx for idx, label in enumerate(conf_labels.split(","), start=1)}
+            if conf_labels
+            else {
+                "spleen": 1,
+                "kidney_right": 2,
+                "kidney_left": 3,
+                "gallbladder": 4,
+                "liver": 5,
+                "stomach": 6,
+                "aorta": 7,
+                "inferior_vena_cava": 8,
+                "portal_vein_and_splenic_vein": 9,
+                "pancreas": 10,
+                "adrenal_gland_right": 11,
+                "adrenal_gland_left": 12,
+                "lung_upper_lobe_left": 13,
+                "lung_lower_lobe_left": 14,
+                "lung_upper_lobe_right": 15,
+                "lung_middle_lobe_right": 16,
+                "lung_lower_lobe_right": 17,
+                "esophagus": 42,
+                "trachea": 43,
+                "heart_myocardium": 44,
+                "heart_atrium_left": 45,
+                "heart_ventricle_left": 46,
+                "heart_atrium_right": 47,
+                "heart_ventricle_right": 48,
+                "pulmonary_artery": 49,
+            }
+        )
 
         # Model Files
         self.path = [
@@ -69,11 +74,15 @@ def init(self, name: str, model_dir: str, conf: Dict[str, str], planner: Any, **
         ]
 
         # Download PreTrained Model
-        if strtobool(self.conf.get("use_pretrained_model", "true")):
+        if not conf_labels and strtobool(self.conf.get("use_pretrained_model", "true")):
             url = f"{self.conf.get('pretrained_path', self.PRE_TRAINED_PATH)}"
             url = f"{url}/radiology_segmentation_segresnet_multilabel.pt"
             download_file(url, self.path[0])
 
+        # Remove pre-trained pt if user is using his/her custom labels.
+        if conf_labels:
+            remove_file(self.path[0])
+
         self.target_spacing = (1.5, 1.5, 1.5)  # target space for image
         # Setting ROI size - This is for the image padding
         self.roi_size = (96, 96, 96)