diff --git a/atomai/models/sam.py b/atomai/models/sam.py index f7ff9f4..20dcfee 100644 --- a/atomai/models/sam.py +++ b/atomai/models/sam.py @@ -1,6 +1,5 @@ import numpy as np import cv2 -from segment_anything import SamAutomaticMaskGenerator, sam_model_registry import pandas as pd import matplotlib.pyplot as plt import torch @@ -26,7 +25,7 @@ class ParticleAnalyzer: >>> analyzer = ParticleAnalyzer(model_type="vit_h") >>> >>> # 2. Load image and run the analysis - >>> image = np.load(path_to_your_image) + >>> image = np.load(IMAGE_PATH) >>> result = analyzer.analyze(image) >>> >>> # 3. Print summary and visualize results @@ -96,6 +95,15 @@ def _download_model_if_needed(self, checkpoint_path, model_type): def _load_model(self, checkpoint_path, model_type): """Loads the SAM model from a checkpoint and moves it to the device.""" + try: + from segment_anything import sam_model_registry + except ImportError: + raise ImportError( + "The 'segment-anything' package is required to use this feature.\n" + "Please install it directly from the official repository:\n\n" + "pip install git+https://github.com/facebookresearch/segment-anything.git" + ) + try: sam = sam_model_registry[model_type](checkpoint=checkpoint_path) sam.to(device=self.device) @@ -179,6 +187,15 @@ def _preprocess_image(self, image_array, use_clahe): def _run_sam(self, image_rgb, preset_name): """Initializes and runs the SAM mask generator based on a preset.""" + try: + from segment_anything import SamAutomaticMaskGenerator + except ImportError: + raise ImportError( + "The 'segment-anything' package is required to use this feature.\n" + "Please install it directly from the official repository:\n\n" + "pip install git+https://github.com/facebookresearch/segment-anything.git" + ) + sam_param_presets = { "default": {}, "sensitive": { diff --git a/setup.py b/setup.py index 2407c50..c6f6eae 100644 --- a/setup.py +++ b/setup.py @@ -38,8 +38,7 @@ 'torchvision>=0.13.0', 'progressbar2>=3.38.0', 'gpytorch>=1.9.1', - 'pandas>=1.1.5', - 'segment-anything @ git+https://github.com/facebookresearch/segment-anything.git' + 'pandas>=1.1.5' ], classifiers=['Programming Language :: Python', 'Development Status :: 3 - Alpha',