neuro-sam 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

neuro_sam/napari_utils/punet_widget.py

@@ -15,14 +15,8 @@ import napari
 from napari.qt.threading import thread_worker
 
 # Import the model class
-# Assuming running from root, so punet is a package
-try:
-    from neuro_sam.punet.punet_inference import run_inference_volume
-except ImportError:
-    # Fallback if running from a different context, try to append path
-    import sys
-    sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'punet'))
-    from neuro_sam.punet.punet_inference import run_inference_volume
+# Assumes installed package structure
+from neuro_sam.punet.punet_inference import run_inference_volume
 
 
 
@@ -197,13 +191,13 @@ class PunetSpineSegmentationWidget(QWidget):
     def _segmentation_worker(self, vol, params):
         import traceback
         try:
+            # Import the refactored inference function from the package
             # Import the refactored inference function from the package
             try:
                 from neuro_sam.punet.punet_inference import run_inference_volume
             except ImportError:
-                import sys
-                sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'punet'))
-                from neuro_sam.punet_inference import run_inference_volume
+                # Fallback should not be needed with proper package execution
+                raise ImportError("Could not import run_inference_volume from neuro_sam.punet.punet_inference")
             
             yield "Starting inference..."
             

neuro_sam/napari_utils/segmentation_model.py

@@ -8,22 +8,33 @@ from scipy.ndimage import label
 from matplotlib.path import Path
 
 
+from neuro_sam.utils import get_weights_path
+
 class DendriteSegmenter:
     """Class for segmenting dendrites from 3D image volumes using SAM2 with overlapping patches"""
     
-    def __init__(self, model_path="./Train-SAMv2/checkpoints/sam2.1_hiera_small.pt", config_path="sam2.1_hiera_s.yaml", weights_path="./Train-SAMv2/results/samv2_dendrite/dendrite_model.torch", device="cuda"):
+    def __init__(self, model_path=None, config_path="sam2.1_hiera_s.yaml", weights_path=None, device="cuda"):
         """
         Initialize the dendrite segmenter with overlapping patches.
         
         Args:
-            model_path: Path to SAM2 model checkpoint
+            model_path: Path to SAM2 model checkpoint (auto-downloaded if None)
             config_path: Path to model configuration
-            weights_path: Path to trained weights
+            weights_path: Path to trained weights (auto-downloaded if None)
             device: Device to run the model on (cpu or cuda)
         """
-        self.model_path = model_path
+        if model_path is None:
+            self.model_path = get_weights_path("sam2.1_hiera_small.pt")
+        else:
+            self.model_path = model_path
+            
         self.config_path = config_path
-        self.weights_path = weights_path
+        
+        if weights_path is None:
+            self.weights_path = get_weights_path("dendrite_model.torch")
+        else:
+            self.weights_path = weights_path
+            
         self.device = device
         self.predictor = None
 
@@ -35,8 +46,6 @@ class DendriteSegmenter:
             
             # Try importing first to catch import errors
             try:
-                import sys
-                sys.path.append('./Train-SAMv2')
                 from sam2.build_sam import build_sam2
                 from sam2.sam2_image_predictor import SAM2ImagePredictor
                 print("Successfully imported SAM2 modules")

neuro_sam/napari_utils/segmentation_module.py

@@ -349,11 +349,8 @@ class SegmentationWidget(QWidget):
             # Initialize segmenter if not already done
             if self.segmenter is None:
                 self.segmenter = DendriteSegmenter(
-                    model_path="./Train-SAMv2/checkpoints/sam2.1_hiera_small.pt",
-                    config_path="sam2.1_hiera_s.yaml",
-                    weights_path="./Train-SAMv2/results/samv2_dendrite/dendrite_model.torch",
                     device=device
-                )
+                ) # Paths are now handled automatically by default args
             
             # Load the model
             success = self.segmenter.load_model()

neuro_sam/plugin.py

@@ -247,11 +247,13 @@ def main():
     else:
         # Try to load a default benchmark image
         try:
-            default_path = './DeepD3_Benchmark.tif'
-            print(f"No image path provided, trying to load default: {default_path}")
+            from neuro_sam.utils import get_weights_path
+            default_path = get_weights_path('DeepD3_Benchmark.tif')
+            print(f"No image path provided, loading default: {default_path}")
             spacing_xyz = (args.x_spacing, args.y_spacing, args.z_spacing)
             viewer = run_neuro_sam(image_path=default_path, spacing_xyz=spacing_xyz)
-        except FileNotFoundError:
+        except Exception as e:
+            print(f"Failed to load default image: {e}")
             sys.exit(1)
     
     print("\nStarted NeuroSAM with anisotropic scaling support!")

neuro_sam/punet/punet_inference.py

@@ -7,7 +7,7 @@ import torch
 import tifffile as tiff
 from tqdm import tqdm
 
-from prob_unet_with_tversky import ProbabilisticUnetDualLatent
+from .prob_unet_with_tversky import ProbabilisticUnetDualLatent
 
 
 def pad_to_multiple(img: np.ndarray, multiple: int = 32):

neuro_sam/utils.py

@@ -0,0 +1,90 @@
+import os
+import sys
+import requests
+from pathlib import Path
+from tqdm import tqdm
+
+# Constants for Model URLs
+WEIGHTS_URLS = {
+    "DeepD3_Benchmark.tif": "https://github.com/nipunarora8/Neuro-SAM/releases/download/weights/DeepD3_Benchmark.tif",
+    "dendrite_model.torch": "https://github.com/nipunarora8/Neuro-SAM/releases/download/weights/dendrite_model.torch",
+    "sam2.1_hiera_small.pt": "https://github.com/nipunarora8/Neuro-SAM/releases/download/weights/sam2.1_hiera_small.pt",
+    "punet_best.pth": "https://github.com/nipunarora8/Neuro-SAM/releases/download/weights/punet_best.pth"
+}
+
+def get_weights_dir():
+    """Get the directory where weights are stored (~/.neuro_sam/checkpoints)."""
+    weights_dir = Path.home() / ".neuro_sam" / "checkpoints"
+    weights_dir.mkdir(parents=True, exist_ok=True)
+    return weights_dir
+
+def download_file(url, dest_path):
+    """Download a file from a URL to a destination path with a progress bar."""
+    print(f"Downloading {url} to {dest_path}...")
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        total_size = int(response.headers.get('content-length', 0))
+        
+        block_size = 1024 # 1 Kibibyte
+        t = tqdm(total=total_size, unit='iB', unit_scale=True)
+        
+        with open(dest_path, 'wb') as file:
+            for data in response.iter_content(block_size):
+                t.update(len(data))
+                file.write(data)
+        t.close()
+        
+        if total_size != 0 and t.n != total_size:
+            print("ERROR, something went wrong with the download")
+            return False
+        
+        print(f"Download complete: {dest_path}")
+        return True
+    except Exception as e:
+        print(f"Error downloading file: {e}")
+        if os.path.exists(dest_path):
+            os.remove(dest_path)
+        return False
+
+def get_weights_path(filename, url=None):
+    """
+    Get the path to a weights file. 
+    Checks local 'checkpoints' folder first, then ~/.neuro_sam/checkpoints.
+    If not found, downloads it.
+    """
+    # 1. Check local checkpoints folder (development mode)
+    local_path = Path("checkpoints") / filename
+    if local_path.exists():
+        return str(local_path.absolute())
+        
+    # 2. Check global cache directory
+    weights_dir = get_weights_dir()
+    cache_path = weights_dir / filename
+    
+    if cache_path.exists():
+        return str(cache_path.absolute())
+        
+    # 3. Download if not found
+    if url is None:
+        url = WEIGHTS_URLS.get(filename)
+        
+    if url:
+        print(f"Weights file {filename} not found locally. Downloading...")
+        success = download_file(url, cache_path)
+        if success:
+            return str(cache_path.absolute())
+        else:
+            raise RuntimeError(f"Failed to download {filename}")
+    else:
+        raise FileNotFoundError(f"Weights file {filename} not found and no URL provided.")
+
+def download_all_models():
+    """Download all known models to the cache directory."""
+    print("Downloading all Neuro-SAM models...")
+    for filename, url in WEIGHTS_URLS.items():
+        try:
+            get_weights_path(filename, url)
+        except Exception as e:
+            print(f"Failed to process {filename}: {e}")
+    print("All downloads processed.")

{neuro_sam-0.1.5.dist-info → neuro_sam-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: neuro-sam
-Version: 0.1.5
+Version: 0.1.7
 Summary: Neuro-SAM: Foundation Models for Dendrite and Dendritic Spine Segmentation
 Author-email: Nipun Arora <nipunarora8@yahoo.com>
 License: MIT License
@@ -54,6 +54,7 @@ Requires-Dist: numba
 Requires-Dist: PyQt5
 Requires-Dist: opencv-python-headless
 Requires-Dist: matplotlib
+Requires-Dist: requests
 Dynamic: license-file
 
 <div align="center">

{neuro_sam-0.1.5.dist-info → neuro_sam-0.1.7.dist-info}/RECORD

@@ -1,5 +1,6 @@
 neuro_sam/__init__.py,sha256=0pbMqpI_nQyhP0_pfTaIg97FVEcFkS5w8gQrsMiBcG4,34
-neuro_sam/plugin.py,sha256=35z7ybZQaFTxXL78z9oNOaHghdILJ8vaAtVShyQyISo,10116
+neuro_sam/plugin.py,sha256=qFEdbEFEBgs0RSQj4zdK3cmgnCxjTmz9_HLi7LI__Tk,10235
+neuro_sam/utils.py,sha256=epdMrxYBO__sVZk3Twk9rHHOtswAu-wmqUektdQNP-s,3307
 neuro_sam/brightest_path_lib/__init__.py,sha256=vU3VvX26D2c9B26Lid09uThzKMQJYPx622SkPhxNlDI,123
 neuro_sam/brightest_path_lib/connected_componen.py,sha256=x_kjDGZ_8U2ks9kZJOOyM8_ow84UugAsPUByI-NlXFk,12734
 neuro_sam/brightest_path_lib/algorithm/__init__.py,sha256=XFYxFyx97FG7lK3_j7uA8GgunpfkOcAo2NIJU0GUn40,170
@@ -30,19 +31,19 @@ neuro_sam/napari_utils/color_utils.py,sha256=Hf5R8f0rh7b9CY1VT72o3tLGfGnnjRREkX8
 neuro_sam/napari_utils/contrasting_color_system.py,sha256=a-lt_3zJLDL9YyIdWJhFDGMYzBb6yH85cV7BNCabbdI,6771
 neuro_sam/napari_utils/main_widget.py,sha256=yahfPLwmhBt_hImpRykIObzfMwbVZvVJTEKKzMZ11bw,48588
 neuro_sam/napari_utils/path_tracing_module.py,sha256=0mMAtrMmtgK_ujMzaWzIguYVDPr8nfzalaTAwgF3NaQ,44062
-neuro_sam/napari_utils/punet_widget.py,sha256=_kGI4bQGbRpjSyhfZrFAqRK5RdSkdwenSrFAFv53pvs,17064
-neuro_sam/napari_utils/segmentation_model.py,sha256=uAnNvMdZVrtlaST36TpCdyoYiaNtVsQLeaPcroD-aT0,33939
-neuro_sam/napari_utils/segmentation_module.py,sha256=TPn_x53uwf-WVC0vTleX2g9fCFXXWNAGtuB8Ht5vmAM,28909
+neuro_sam/napari_utils/punet_widget.py,sha256=FfnC6V_FErczkaQP5y3rp1YBMWPVx6YMI4TxEHah_Vo,16862
+neuro_sam/napari_utils/segmentation_model.py,sha256=mHXVjksqEcxHRH5KWp5-hXLEnRHgGhwPUxyUkV8eJGM,34141
+neuro_sam/napari_utils/segmentation_module.py,sha256=iObM5k8VkARtB_rcqAQGzKJ-PmaAKLeFJD14_Jy6xhs,28732
 neuro_sam/napari_utils/visualization_module.py,sha256=JtZlBoKlfIwVLa2Sqg7b2KTr07fNlAcwR0M7fHsn2oM,24723
 neuro_sam/punet/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 neuro_sam/punet/deepd3_model.py,sha256=nGVEqzCPz_E4cFA6QmknW2CffDcjxH7VsdYAyTdAtY0,7509
 neuro_sam/punet/prob_unet_deepd3.py,sha256=syXNleUVrfYtmVveN9G461oAhumxsijsavps8in4VRw,14698
 neuro_sam/punet/prob_unet_with_tversky.py,sha256=2dBbO_BEHbhYWBXW7rXQX6s2DnqoTgBKkgk6VkgN-Ds,12845
-neuro_sam/punet/punet_inference.py,sha256=U3ne-sP1y732PISFYlV50cPbXN6eEREQue1Jf6w1Io8,8032
+neuro_sam/punet/punet_inference.py,sha256=v5ufB2Zz5WfgfFZ5-rDjBEobpr5gy-HKPPWZpCALV7A,8033
 neuro_sam/punet/run_inference.py,sha256=c9ATKWJvhOzNEaww_sUCI5fFS1q0bQ4GYUwNUqxWcwA,5312
 neuro_sam/punet/unet_blocks.py,sha256=ZRNKay9P3OnJ0PmtKXw_iSgUyRE1DkkGefGXwSbYZGY,3171
 neuro_sam/punet/utils.py,sha256=ibwcpkqqZ3_3Afz2VYxzplz8_8FWQ5qYQqjJiKS8hIo,1786
-neuro_sam-0.1.5.dist-info/licenses/LICENSE,sha256=akmTIN8IuZn3Y7UK_8qVQnyKDWSDcVUwB8RPGNXCojw,1068
+neuro_sam-0.1.7.dist-info/licenses/LICENSE,sha256=akmTIN8IuZn3Y7UK_8qVQnyKDWSDcVUwB8RPGNXCojw,1068
 sam2/__init__.py,sha256=uHyh6VzVS4F2box0rPDpN5UmOVKeQNK0CIaTKG9JQZ4,395
 sam2/automatic_mask_generator.py,sha256=Zt8mbb4UQSMFrjOY8OwbshswOpMhaxAtdn5sTuXUw9c,18461
 sam2/benchmark.py,sha256=m3o1BriIQuwJAx-3zQ_B0_7YLhN84G28oQSV5sGA3ak,2811
@@ -86,8 +87,8 @@ sam2/utils/__init__.py,sha256=NL2AacVHZOe41zp4kF2-ZGcUCi9zFwh1Eo9spNjN0Ko,197
 sam2/utils/amg.py,sha256=t7MwkOKvcuBNu4FcjzKv9BpO0av5Zo9itZ8b3WQMpdg,12842
 sam2/utils/misc.py,sha256=AWAMAcFhzQedcQb7HU2oRc-RqjGrK87K-MsVG21tIKI,13090
 sam2/utils/transforms.py,sha256=ujpk9GAMYvIJIGpt87QOP88TPtrjL61liDG7DCptEUY,4885
-neuro_sam-0.1.5.dist-info/METADATA,sha256=04rkWKMPdFGRctGcQgHGe69JxMENwwkzNtMHOQiGn_Y,9618
-neuro_sam-0.1.5.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-neuro_sam-0.1.5.dist-info/entry_points.txt,sha256=K0SaLior6N1REi_K0O4962Wn4DhTSos_Sx6Enx0mIm4,105
-neuro_sam-0.1.5.dist-info/top_level.txt,sha256=yPbWxFcw79sErTk8zohihUHMK9LL31i3bXir2MrS4OQ,15
-neuro_sam-0.1.5.dist-info/RECORD,,
+neuro_sam-0.1.7.dist-info/METADATA,sha256=42Pp29GxNSpY1GeNvzeQ3XJyDDsEiMxsCLSSV2Ib-FA,9642
+neuro_sam-0.1.7.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+neuro_sam-0.1.7.dist-info/entry_points.txt,sha256=a1JXEgiM_QOPJdV8zvcIS60WAE62MeqgIVY2oSx81FY,162
+neuro_sam-0.1.7.dist-info/top_level.txt,sha256=yPbWxFcw79sErTk8zohihUHMK9LL31i3bXir2MrS4OQ,15
+neuro_sam-0.1.7.dist-info/RECORD,,

{neuro_sam-0.1.5.dist-info → neuro_sam-0.1.7.dist-info}/entry_points.txt

@@ -1,5 +1,6 @@
 [console_scripts]
 neuro-sam = neuro_sam.plugin:main
+neuro-sam-download = neuro_sam.utils:download_all_models
 
 [napari.manifest]
 neuro-sam = neuro_sam:napari.yaml