neuro-sam 0.1.6__tar.gz → 0.1.8__tar.gz

{neuro_sam-0.1.6/src/neuro_sam.egg-info → neuro_sam-0.1.8}/PKG-INFO

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

{neuro_sam-0.1.6 → neuro_sam-0.1.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "neuro-sam"
-version = "0.1.6"
+version = "0.1.8"
 description = "Neuro-SAM: Foundation Models for Dendrite and Dendritic Spine Segmentation"
 readme = "README.md"
 authors = [
@@ -37,7 +37,12 @@ dependencies = [
     "numba",
     "PyQt5",
     "opencv-python-headless",
-    "matplotlib"
+    "matplotlib",
+    "requests",
+    "flammkuchen",
+    "albumentations",
+    "wandb",
+    "tensorflow"
 ]
 requires-python = ">=3.10"
 
@@ -55,6 +60,8 @@ include = ["neuro_sam*", "sam2*"]
 
 [project.scripts]
 neuro-sam = "neuro_sam.plugin:main"
+neuro-sam-download = "neuro_sam.utils:download_all_models"
+neuro-sam-train = "neuro_sam.training.train_dendrites:main"
 
 [project.entry-points."napari.manifest"]
 neuro-sam = "neuro_sam:napari.yaml"

{neuro_sam-0.1.6 → neuro_sam-0.1.8}/src/neuro_sam/napari_utils/path_tracing_module.py

@@ -509,10 +509,10 @@ class PathTracingWidget(QWidget):
             self.scaler.original_spacing_xyz[0] / x_nm   # X
         ])
         
-        self.scaling_status.setText(
-            f"Pending: X={x_nm:.1f}, Y={y_nm:.1f}, Z={z_nm:.1f} nm\n"
-            f"Scale factors (Z,Y,X): {temp_scale_factors[0]:.3f}, {temp_scale_factors[1]:.3f}, {temp_scale_factors[2]:.3f}"
-        )
+        # self.scaling_status.setText(
+        #     f"Pending: X={x_nm:.1f}, Y={y_nm:.1f}, Z={z_nm:.1f} nm\n"
+        #     f"Scale factors (Z,Y,X): {temp_scale_factors[0]:.3f}, {temp_scale_factors[1]:.3f}, {temp_scale_factors[2]:.3f}"
+        # )
         
     def _apply_scaling(self):
         """Apply current scaling settings"""

{neuro_sam-0.1.6 → neuro_sam-0.1.8}/src/neuro_sam/napari_utils/punet_widget.py

@@ -21,6 +21,8 @@ from neuro_sam.punet.punet_inference import run_inference_volume
 
 
 
+from neuro_sam.utils import get_weights_path
+
 class PunetSpineSegmentationWidget(QWidget):
     """
     Widget for spine segmentation using Probabilistic U-Net.
@@ -36,7 +38,7 @@ class PunetSpineSegmentationWidget(QWidget):
         
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model = None
-        self.model_path = "punet/punet_best.pth" # Default relative path
+        self.model_path = get_weights_path("punet_best.pth") # Auto-download default weights
         
         # Connect custom progress signal
         self.progress_signal.connect(self._on_worker_progress)
@@ -191,13 +193,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-0.1.6 → neuro_sam-0.1.8}/src/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-0.1.6 → neuro_sam-0.1.8}/src/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-0.1.6 → neuro_sam-0.1.8}/src/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-0.1.8/src/neuro_sam/training/__init__.py

@@ -0,0 +1 @@
+# Init file for training module

neuro_sam-0.1.8/src/neuro_sam/training/train_dendrites.py

@@ -0,0 +1,226 @@
+import numpy as np
+import torch
+import os
+import wandb
+from torch.onnx.symbolic_opset11 import hstack
+import torch.nn.functional as F
+from .utils.stream_dendrites import DataGeneratorStream
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+import time
+import argparse
+from neuro_sam.utils import get_weights_path
+
+def main():
+    parser = argparse.ArgumentParser(description="Train Neuro-SAM Dendrite Segmenter")
+    parser.add_argument("--ppn", type=int, required=True, help="Positive Points Number")
+    parser.add_argument("--pnn", type=int, required=True, help="Negative Points Number")
+    parser.add_argument("--model_name", type=str, required=True, choices=['small', 'base_plus', 'large', 'tiny'], help="SAM2 Model Size")
+    parser.add_argument("--batch_size", type=int, required=True, help="Batch Size")
+    parser.add_argument("--logger", type=str, default="False", help="Use WandB Logger (True/False)")
+    parser.add_argument("--data_dir", type=str, default="./data", help="Directory containing .d3set data files")
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    TRAINING_DATA_PATH = os.path.join(args.data_dir, "DeepD3_Training.d3set")
+    VALIDATION_DATA_PATH = os.path.join(args.data_dir, "DeepD3_Validation.d3set")
+
+    if not os.path.exists(TRAINING_DATA_PATH):
+        raise FileNotFoundError(f"Training data not found at {TRAINING_DATA_PATH}")
+
+    positive_points = args.ppn
+    negative_points = args.pnn
+    batch_size = args.batch_size
+    logger = (args.logger.lower() == "true")
+
+    print("Initializing Data Generator...")
+    dg_training = DataGeneratorStream(TRAINING_DATA_PATH,
+                                    batch_size=batch_size, # Data processed at once, depends on your GPU
+                                    target_resolution=0.094, # fixed to 94 nm, can be None for mixed resolution training
+                                    min_content=100,
+                                    resizing_size = 128,
+                                    positive_points = positive_points,
+                                    negative_points = negative_points) # images need to have at least 50 segmented px
+
+    dg_validation = DataGeneratorStream(VALIDATION_DATA_PATH,
+                                        batch_size=batch_size, 
+                                        target_resolution=0.094,
+                                        min_content=100, 
+                                        augment=False,
+                                        shuffle=False,
+                                        resizing_size = 128,
+                                        positive_points = positive_points,
+                                    negative_points = negative_points)
+
+    # Load model
+    model_name_map = {
+        'large': ("sam2.1_hiera_large.pt", "sam2.1_hiera_l.yaml"),
+        'base_plus': ("sam2.1_hiera_base_plus.pt", "sam2.1_hiera_b+.yaml"),
+        'small': ("sam2.1_hiera_small.pt", "sam2.1_hiera_s.yaml"),
+        'tiny': ("sam2.1_hiera_tiny.pt", "sam2.1_hiera_t.yaml"),
+    }
+    
+    ckpt_name, model_cfg = model_name_map[args.model_name]
+    
+    # Try to find weights via util or local defaults
+    try:
+        sam2_checkpoint = get_weights_path(ckpt_name)
+    except:
+        # Fallback if not downloadable or found
+        sam2_checkpoint = ckpt_name
+        print(f"Warning: Could not resolve weight path for {ckpt_name}, assuming local file.")
+
+    print(f"Loading SAM2 model: {args.model_name} from {sam2_checkpoint}")
+    sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device) # load model
+
+    predictor = SAM2ImagePredictor(sam2_model)
+
+    # Set training parameters
+
+    predictor.model.sam_mask_decoder.train(True) # enable training of mask decoder
+    predictor.model.sam_prompt_encoder.train(True) # enable training of prompt encoder
+    predictor.model.image_encoder.train(True) # enable training of image encoder: For this to work you need to scan the code for "no_grad" and remove them all
+    optimizer=torch.optim.AdamW(params=predictor.model.parameters(),lr=1e-5,weight_decay=4e-5)
+    scaler = torch.cuda.amp.GradScaler() # mixed precision
+
+
+    time_str = "-".join(["{:0>2}".format(x) for x in time.localtime(time.time())][:-3])
+    ckpt_path = f'results/samv2_{args.model_name}_{time_str}'
+    if not os.path.exists(ckpt_path): os.makedirs(ckpt_path)
+
+    if logger:
+        wandb.init(
+            # set the wandb project where this run will be logged
+            project="N-SAMv2",
+
+            # track hyperparameters and run metadata
+            config={
+            "architecture": "SAMv2",
+            "dataset": "DeepD3",
+            "model": args.model_name,
+            "epochs": 100000,
+            "ckpt_path": ckpt_path,
+            "image_size": (1,128,128),
+            "min_content": 100,
+            "positive_points": positive_points,
+            "negative_points": negative_points,
+            "batch_size": batch_size,
+            "prompt_seed":42
+            }
+        )
+
+
+    # add val code here
+
+    def perform_validation(predictor):
+        print('Performing Validation')
+        mean_iou = []
+        mean_dice = []
+        mean_loss = []
+        with torch.no_grad():
+            for i in range(20):
+                try:
+                    n = np.random.randint(len(dg_validation))
+                    image,mask,input_point, input_label = dg_validation[n]
+                    
+                except:
+                    print('Error in validation batch generation')
+                    continue
+                if mask.shape[0]==0: continue # ignore empty batches
+                predictor.set_image_batch(image) # apply SAM image encoder to the image
+
+                mask_input, unnorm_coords, labels, unnorm_box = predictor._prep_prompts(input_point, input_label, box=None, mask_logits=None, normalize_coords=True)
+                sparse_embeddings, dense_embeddings = predictor.model.sam_prompt_encoder(points=(unnorm_coords, labels),boxes=None,masks=None,)
+
+                # mask decoder
+
+                high_res_features = [feat_level[-1].unsqueeze(0) for feat_level in predictor._features["high_res_feats"]]
+                low_res_masks, prd_scores, _, _ = predictor.model.sam_mask_decoder(image_embeddings=predictor._features["image_embed"],image_pe=predictor.model.sam_prompt_encoder.get_dense_pe(),sparse_prompt_embeddings=sparse_embeddings,dense_prompt_embeddings=dense_embeddings,multimask_output=True,repeat_image=False,high_res_features=high_res_features,)
+                prd_masks = predictor._transforms.postprocess_masks(low_res_masks, predictor._orig_hw[-1])# Upscale the masks to the original image resolution
+
+                # Segmentaion Loss caclulation
+
+                gt_mask = torch.tensor(mask.astype(np.float32)).to(device)
+                prd_mask = torch.sigmoid(prd_masks[:, 0])# Turn logit map to probability map
+                seg_loss = (-gt_mask * torch.log(prd_mask + 0.00001) - (1 - gt_mask) * torch.log((1 - prd_mask) + 0.00001)).mean() # cross entropy loss
+
+                # Score loss calculation (intersection over union) IOU
+
+                inter = (gt_mask * (prd_mask > 0.5)).sum(1).sum(1)
+                iou = inter / (gt_mask.sum(1).sum(1) + (prd_mask > 0.5).sum(1).sum(1) - inter)
+                total_sum = gt_mask.sum(1).sum(1) + (prd_mask > 0.5).sum(1).sum(1)
+                dicee = ((2 * inter) / total_sum).mean()
+                score_loss = torch.abs(prd_scores[:, 0] - iou).mean()
+                loss=seg_loss+score_loss*0.05  # mix losses
+
+                mean_iou.append(iou.cpu().detach().numpy())
+                mean_dice.append(dicee.cpu().detach().numpy())
+                mean_loss.append(loss.cpu().detach().numpy())
+
+        return np.array(mean_iou).mean(), np.array(mean_dice).mean(), np.array(mean_loss).mean()
+    
+
+    for itr in range(100000):
+        epoch_dice = epoch_iou = epoch_loss = []
+        with torch.cuda.amp.autocast(): # cast to mix precision
+            n = np.random.randint(len(dg_training))
+            try:
+                image,mask,input_point, input_label = dg_training[n]
+            except:
+                print('Error in training batch')
+                continue
+            if mask.shape[0]==0: continue # ignore empty batches
+            predictor.set_image_batch(image) # apply SAM image encoder to the image
+            
+            # prompt encoding
+
+            mask_input, unnorm_coords, labels, unnorm_box = predictor._prep_prompts(input_point, input_label, box=None, mask_logits=None, normalize_coords=True)
+            sparse_embeddings, dense_embeddings = predictor.model.sam_prompt_encoder(points=(unnorm_coords, labels),boxes=None,masks=None,)
+
+            # mask decoder
+
+            high_res_features = [feat_level[-1].unsqueeze(0) for feat_level in predictor._features["high_res_feats"]]
+            low_res_masks, prd_scores, _, _ = predictor.model.sam_mask_decoder(image_embeddings=predictor._features["image_embed"],image_pe=predictor.model.sam_prompt_encoder.get_dense_pe(),sparse_prompt_embeddings=sparse_embeddings,dense_prompt_embeddings=dense_embeddings,multimask_output=True,repeat_image=False,high_res_features=high_res_features,)
+            prd_masks = predictor._transforms.postprocess_masks(low_res_masks, predictor._orig_hw[-1])# Upscale the masks to the original image resolution
+
+            gt_mask = torch.tensor(mask.astype(np.float32)).to(device)
+            prd_mask = torch.sigmoid(prd_masks[:, 0])# Turn logit map to probability map
+            seg_loss = (-gt_mask * torch.log(prd_mask + 0.00001) - (1 - gt_mask) * torch.log((1 - prd_mask) + 0.00001)).mean() # cross entropy loss
+
+            # Score loss calculation (intersection over union) IOU
+
+            inter = (gt_mask * (prd_mask > 0.5)).sum(1).sum(1)
+            iou = inter / (gt_mask.sum(1).sum(1) + (prd_mask > 0.5).sum(1).sum(1) - inter)
+            total_sum = gt_mask.sum(1).sum(1) + (prd_mask > 0.5).sum(1).sum(1)
+            dicee = ((2 * inter) / total_sum).mean()
+            score_loss = torch.abs(prd_scores[:, 0] - iou).mean()
+            loss=seg_loss+score_loss*0.05  # mix losses
+
+            predictor.model.zero_grad() # empty gradient 
+            scaler.scale(loss).backward()  # Backpropogate
+            scaler.step(optimizer)
+            scaler.update() # Mix precision
+
+
+            # Display results
+            if itr==0: mean_iou=0
+            mean_iou = mean_iou * 0.99 + 0.01 * np.mean(iou.cpu().detach().numpy())
+            if logger:
+                wandb.log({'step': itr, 'loss':loss, 'iou':mean_iou, 'dice_score':dicee})
+            print(f'step: {itr}, loss: {loss}, iou: {mean_iou}, dice_score: {dicee}')
+
+            if itr%500 == 0: 
+
+                val_iou, val_dice, val_loss = perform_validation(predictor)
+                if logger:
+                    wandb.log({'step': itr, 'val_loss':val_loss, 'val_iou':val_iou, 'val_dice_score':val_dice})
+                print(f'step: {itr}, val_loss: {val_loss}, val_iou: {val_iou}, val_dice_score: {val_dice}')
+
+                torch.save(predictor.model.state_dict(), f"{ckpt_path}/model_{itr}.torch") # save model
+    if logger:
+        wandb.finish()
+
+if __name__ == "__main__":
+    main()

neuro_sam-0.1.8/src/neuro_sam/training/utils/__init__.py

@@ -0,0 +1 @@
+# Init file for training utils

neuro_sam-0.1.8/src/neuro_sam/training/utils/prompt_generation_dendrites.py

@@ -0,0 +1,78 @@
+import numpy as np
+from scipy import ndimage
+import cv2, random
+from collections import *
+from itertools import *
+from functools import *
+
+class PromptGeneration:
+    def __init__(self, random_seed=0, neg_range=(3, 20), min_content = 20):
+        if random_seed:  
+            random.seed(random_seed)
+            np.random.seed(random_seed)
+
+        self.neg_range = neg_range
+
+        self.min_area = min_content
+
+    def get_labelmap(self, label):
+        structure = ndimage.generate_binary_structure(2, 2)
+        labelmaps, connected_num = ndimage.label(label, structure=structure)
+
+        label = np.zeros_like(labelmaps)
+        for i in range(1, 1+connected_num):
+            if np.sum(labelmaps==i) >= self.min_area: label += np.where(labelmaps==i, 255, 0)
+
+        structure = ndimage.generate_binary_structure(2, 2)
+        labelmaps, connected_num = ndimage.label(label, structure=structure)
+
+        return labelmaps, connected_num
+    
+    def search_negative_region_numpy(self, labelmap):
+        inner_range, outer_range = self.neg_range
+        def search(neg_range):
+            kernel = np.ones((neg_range * 2 + 1, neg_range * 2 + 1), np.uint8)
+            negative_region = cv2.dilate(labelmap, kernel, iterations=1)
+            mx = labelmap.max() + 1
+            labelmap_r = (mx - labelmap) * np.minimum(1, labelmap)
+            r = cv2.dilate(labelmap_r, kernel, iterations=1)
+            negative_region_r = (r.astype(np.int32) - mx) * np.minimum(1, r)
+            diff = negative_region.astype(np.int32) + negative_region_r
+            overlap = np.minimum(1, np.abs(diff).astype(np.uint8))
+            return negative_region - overlap - labelmap
+        return search(outer_range) - search(inner_range) 
+
+    def get_prompt_points(self, label_mask, ppp, ppn):
+        label_mask_cp = np.copy(label_mask)
+        label_mask_cp[label_mask_cp >= 1] = 1
+        labelmaps, connected_num = self.get_labelmap(label_mask_cp)
+
+        coord_positive, coord_negative = [], []
+
+        connected_components = list(range(1, connected_num+1))
+        random.shuffle(connected_components)
+
+        for i in connected_components:
+            cc = np.copy(labelmaps)
+            cc[cc!=i] = 0
+            cc[cc==i] = 1
+            if ppp:
+                coord_positive.append(random.choice([[y, x] for x, y in np.argwhere(cc == 1)]))
+                ppp -= 1
+        
+        random.shuffle(connected_components)
+        for i in connected_components:
+            cc = np.copy(labelmaps)
+            cc[cc!=i] = 0
+            cc[cc==i] = 1
+            negative_region = self.search_negative_region_numpy(cc.astype(np.uint8))
+            negative_region = negative_region * (1 - label_mask_cp)
+            if ppn:
+                coord_negative.append(random.choice([[y, x] for x, y in np.argwhere(negative_region == 1)]))
+                ppn -= 1
+
+        negative_region = self.search_negative_region_numpy(label_mask_cp)
+
+        if ppp: coord_positive += random.sample([[y, x] for x, y in np.argwhere(label_mask_cp == 1)], ppp)
+        if ppn: coord_negative += random.sample([[y, x] for x, y in np.argwhere(negative_region == 1)], ppn)
+        return coord_positive, coord_negative