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

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/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)

neuro_sam/training/__init__.py

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

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/training/utils/__init__.py

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

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

neuro_sam/training/utils/stream_dendrites.py

@@ -0,0 +1,299 @@
+import numpy as np
+import flammkuchen as fl
+from tensorflow.keras.utils import Sequence
+from .prompt_generation_dendrites import PromptGeneration
+import cv2
+import albumentations as A
+import random
+
+class DataGeneratorStream(Sequence):
+    def __init__(self, fn, batch_size, samples_per_epoch=50000, size=(1, 128, 128), target_resolution=None, augment=True, 
+        shuffle=True, seed=42, normalize=[-1, 1], resizing_size=1024, min_content=0., positive_points=1, negative_points=1
+    ):
+        """Data Generator that streams data dynamically for training DeepD3.
+
+        Args:
+            fn (str): The path to the training data file
+            batch_size (int): Batch size for training deep neural networks
+            samples_per_epoch (int, optional): Samples used in each epoch. Defaults to 50000.
+            size (tuple, optional): Shape of a single sample. Defaults to (1, 128, 128).
+            target_resolution (float, optional): Target resolution in microns. Defaults to None.
+            augment (bool, optional): Enables augmenting the data. Defaults to True.
+            shuffle (bool, optional): Enabled shuffling the data. Defaults to True.
+            seed (int, optional): Creates pseudorandom numbers for shuffling. Defaults to 42.
+            normalize (list, optional): Values range when normalizing data. Defaults to [-1, 1].
+            min_content (float, optional): Minimum content in image (annotated dendrite or spine), not considered if 0. Defaults to 0.
+        """
+        
+        # Save settings            
+        self.batch_size = batch_size
+        self.augment = augment
+        self.fn = fn
+        self.shuffle = shuffle
+        self.aug = self._get_augmenter()
+        self.add_gaus_noise = self._add_gaus_noise()
+        self.seed = seed
+        self.normalize = normalize
+        self.samples_per_epoch = samples_per_epoch
+        self.size = size
+        self.target_resolution = target_resolution
+        self.min_content = min_content
+        self.resizing_size = resizing_size
+        self.ppn = positive_points
+        self.pnn = negative_points
+
+        self.d = fl.load(self.fn)
+        self.data = self.d['data']
+        self.meta = self.d['meta']
+
+        # Seed randomness
+        random.seed(self.seed)
+        np.random.seed(self.seed)
+
+        self.on_epoch_end()
+        self.pg = PromptGeneration(random_seed=42, neg_range=(3, 9))
+
+    def __len__(self):
+        """Denotes the number of batches per epoch"""
+        return self.samples_per_epoch // self.batch_size
+
+    def __getitem__(self, index):
+        """Generate one batch of data
+        
+        Parameters
+        ----------
+        index : int
+            batch index in image/label id list
+        
+        Returns
+        -------
+        tuple
+            Contains two numpy arrays,
+            each of shape (batch_size, height, width, 1). 
+        """
+        X = []
+        Y0 = []
+        Y1 = []
+        pps = []
+        pts = []
+        eps = 1e-5
+
+        if self.shuffle is False:
+            np.random.seed(index)
+
+        # Create all pairs in a given batch
+        for i in range(self.batch_size):
+            # Retrieve a single sample pair
+            image, dendrite, spines = self.getSample()
+            image = cv2.resize(image, (self.resizing_size, self.resizing_size), interpolation=cv2.INTER_LINEAR)
+            dendrite = cv2.resize(dendrite.astype(np.uint8), (self.resizing_size, self.resizing_size), interpolation=cv2.INTER_LINEAR)
+            spines = cv2.resize(spines.astype(np.uint8), (self.resizing_size, self.resizing_size), interpolation=cv2.INTER_LINEAR)
+
+            # Augmenting the data
+            if self.augment:
+                augmented = self.aug(image=image, 
+                    mask1=dendrite, 
+                    mask2=spines) #augment image
+                
+                image = augmented['image']
+                dendrite = augmented['mask1']
+                spines = augmented['mask2']
+                try:
+                    pp, pt = self.extract_points(dendrite)
+
+                    augmented = self.add_gaus_noise(image=image, 
+                        mask1=dendrite, 
+                        mask2=spines) #augment image
+                        
+                    image = augmented['image']
+                    dendrite = augmented['mask1']
+                    spines = augmented['mask2']
+                except:
+                    pass
+
+            else:
+                dendrite = dendrite
+                spines = spines
+                try:
+                    pp, pt = self.extract_points(dendrite)     
+                except:
+                    return False  
+
+            # Min/max scaling
+            image = (image.astype(np.float32) - image.min()) / (image.max() - image.min() + eps)
+            # Shifting and scaling
+            # image = image * (self.normalize[1]-self.normalize[0]) + self.normalize[0]
+
+            X.append(cv2.cvtColor(image, cv2.COLOR_GRAY2RGB))
+            Y0.append(dendrite.astype(np.float32) / (dendrite.max() + eps))
+            Y1.append(spines.astype(np.float32) / (spines.max() + eps)) # to ensure binary targets
+            pps.append(pp)
+            pts.append(pt)
+
+        # return np.asarray(X, dtype=np.float32)[..., None], (np.asarray(Y0, dtype=np.float32)[..., None],
+        #         np.asarray(Y1, dtype=np.float32)[..., None]), (np.asarray(pps), np.asarray(pts))
+        return np.asarray(X, dtype=np.float32), np.asarray(Y0, dtype=np.float32), np.asarray(pps), np.asarray(pts)
+
+    def extract_points(self, dendrite):
+        ppp, ppn = self.ppn,self.pnn
+        coord_positive, coord_negative = self.pg.get_prompt_points(dendrite, ppp, ppn)
+        pp = coord_positive + coord_negative
+        pt = [1] * len(coord_positive) + [0] * len(coord_negative)
+        pp = np.array(pp, dtype=np.float32)
+        pt = np.array(pt, dtype=np.int32)
+        return pp, pt
+
+
+    def _get_augmenter(self):
+        """Defines used augmentations"""
+        aug = A.Compose([
+            A.RandomBrightnessContrast(p=0.25),    
+            A.Rotate(limit=10, border_mode=cv2.BORDER_REFLECT, p=0.5),
+            A.RandomRotate90(p=0.5),
+            A.HorizontalFlip(p=0.5),
+            A.VerticalFlip(p=0.5),
+            A.Blur(p=0.2)],p=1,
+            #A.GaussNoise(p=0.5)], p=1,
+            additional_targets={
+                'mask1': 'mask',
+                'mask2': 'mask'
+            })
+        return aug
+
+    def _add_gaus_noise(self):
+        aug = A.Compose([
+            A.GaussNoise(p=0.5)], p=1,
+            additional_targets={
+                'mask1': 'mask',
+                'mask2': 'mask'
+            })
+        return aug
+    
+
+    def getSample(self, squeeze=True):
+        """Get a sample from the provided data
+
+        Args:
+            squeeze (bool, optional): if plane is 2D, skip 3D. Defaults to True.
+
+        Returns:
+            list(np.ndarray, np.ndarray, np.ndarray): stack image with respective labels
+        """
+        while True:
+            r = self._getSample(squeeze)
+            
+            # If sample was successfully generated
+            # and we don't care about the content
+            if r is not None and self.min_content == 0:
+                return r
+
+            # If sample was successfully generated
+            # and we do care about the content
+            elif r is not None:
+                # In either or both annotation should be at least `min_content` pixels
+                # that are being labelled.
+                if (r[1]).sum() > self.min_content and (r[2]).sum() > self.min_content:
+                    return r
+                else:
+                    continue
+
+            else:
+                continue
+
+    def _getSample(self, squeeze=True):
+        """Retrieves a sample
+
+        Args:
+            squeeze (bool, optional): Squeezes return shape. Defaults to True.
+
+        Returns:
+            tuple: Tuple of stack (X), dendrite (Y0) and spines (Y1)
+        """
+        # Adjust for 2 images
+        if len(self.size) == 2:
+            size = (1,) + self.size
+
+        else:
+            size = self.size
+        
+        # sample random stack
+        r_stack = np.random.choice(len(self.meta))
+        
+        target_h = size[1]
+        target_w = size[2]
+
+
+        if self.target_resolution is None:
+            # Keep everything as is
+            scaling = 1
+            h = target_h
+            w = target_w
+
+        else:
+            # Computing scaling factor
+            scaling = self.target_resolution / self.meta.iloc[r_stack].Resolution_XY
+            
+            # Compute the height and width and random offsets
+            h = round(scaling * target_h)
+            w = round(scaling * target_w)
+        
+        # Correct for stack dimensions
+        if self.meta.iloc[r_stack].Width-w == 0:
+            x = 0
+            
+        elif self.meta.iloc[r_stack].Width-w < 0:
+            return
+        
+        else:
+            x = np.random.choice(self.meta.iloc[r_stack].Width-w)
+
+        # Correct for stack dimensions            
+        if self.meta.iloc[r_stack].Height-h == 0:
+            y = 0
+            
+        elif self.meta.iloc[r_stack].Height-h < 0:
+            return
+        
+        else:
+            y = np.random.choice(self.meta.iloc[r_stack].Height-h)
+            
+        ## Select random plane + range
+        r_plane = np.random.choice(self.meta.iloc[r_stack].Depth-size[0]+1)
+        
+        z_begin = r_plane
+        z_end   = r_plane+size[0]
+        
+        
+        # Scale if neccessary to the correct dimensions
+        tmp_stack = self.data['stacks'][f'x{r_stack}'][z_begin:z_end, y:y+h, x:x+w]
+        tmp_dendrites = self.data['dendrites'][f'x{r_stack}'][z_begin:z_end, y:y+h, x:x+w]
+        tmp_spines = self.data['spines'][f'x{r_stack}'][z_begin:z_end, y:y+h, x:x+w]
+
+        # Data needs to be rescaled
+        if scaling != 1:
+            return_stack = []
+            return_dendrites = []
+            return_spines = []
+        
+            # Do this for each plane
+            # and ensure that OpenCV is happy
+            for i in range(tmp_stack.shape[0]):
+                return_stack.append(cv2.resize(tmp_stack[i], (target_h, target_w)))
+                return_dendrites.append(cv2.resize(tmp_dendrites[i].astype(np.uint8), (target_h, target_w)).astype(bool))
+                return_spines.append(cv2.resize(tmp_spines[i].astype(np.uint8), (target_h, target_w)).astype(bool))
+                
+            return_stack = np.asarray(return_stack)
+            return_dendrites = np.asarray(return_dendrites)
+            return_spines = np.asarray(return_spines)
+            
+        else:
+            return_stack = tmp_stack
+            return_dendrites = tmp_dendrites
+            return_spines = tmp_spines
+                
+        if squeeze:
+            # Return sample
+            return return_stack.squeeze(), return_dendrites.squeeze(), return_spines.squeeze()
+        
+        else:
+            return return_stack, return_dendrites, return_spines

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: neuro-sam
-Version: 0.1.7
+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
@@ -55,6 +55,10 @@ 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.7.dist-info → neuro_sam-0.1.8.dist-info}/RECORD

@@ -30,12 +30,11 @@ neuro_sam/napari_utils/anisotropic_scaling.py,sha256=VA6Sd9zEhIAhzjAGto2cOjE9moc
 neuro_sam/napari_utils/color_utils.py,sha256=Hf5R8f0rh7b9CY1VT72o3tLGfGnnjRREkX8iWsiiu7k,4243
 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=FfnC6V_FErczkaQP5y3rp1YBMWPVx6YMI4TxEHah_Vo,16862
+neuro_sam/napari_utils/path_tracing_module.py,sha256=bDhSawWNMfY-Vs-Zdt8XTb90pLkD9jUBRqpbC2_6li0,44070
+neuro_sam/napari_utils/punet_widget.py,sha256=WCAND8YLn4CA_20YW4quwcuJ_xEqPAc7GHrSBk4Anw0,16928
 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
@@ -43,7 +42,12 @@ neuro_sam/punet/punet_inference.py,sha256=v5ufB2Zz5WfgfFZ5-rDjBEobpr5gy-HKPPWZpC
 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.7.dist-info/licenses/LICENSE,sha256=akmTIN8IuZn3Y7UK_8qVQnyKDWSDcVUwB8RPGNXCojw,1068
+neuro_sam/training/__init__.py,sha256=QISf7Tk0xphF08BWHd-pLI1CMHaEzDsTdmUkgzvwAUM,32
+neuro_sam/training/train_dendrites.py,sha256=TMG4YrrQV0Q784omMziXsB71tNICEmGLYSHr0zI9i6Y,11346
+neuro_sam/training/utils/__init__.py,sha256=4hbcx57NtRu8nryvXQYqmXK4hyUgUYNDz97kCw3Efs8,31
+neuro_sam/training/utils/prompt_generation_dendrites.py,sha256=_ntzXNV1lXPrpInRKaZ5CPpq3akF2IuD1naOXbTC8TU,3201
+neuro_sam/training/utils/stream_dendrites.py,sha256=qS_ZWrhJdW1Sg3RBjoRUJFlCV0u5X1Ns_tjYgJUjWJw,11024
+neuro_sam-0.1.8.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
@@ -87,8 +91,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.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.8.dist-info/METADATA,sha256=1vHp2trqfoQL7eP7aia_Xqa8H9f2b6JjsPECvGpsqrk,9746
+neuro_sam-0.1.8.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+neuro_sam-0.1.8.dist-info/entry_points.txt,sha256=EQg0SmFbnbGGcchHq5ROmhO9pkgby72Y5G5w90WyLZI,220
+neuro_sam-0.1.8.dist-info/top_level.txt,sha256=yPbWxFcw79sErTk8zohihUHMK9LL31i3bXir2MrS4OQ,15
+neuro_sam-0.1.8.dist-info/RECORD,,

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

@@ -1,6 +1,7 @@
 [console_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
 
 [napari.manifest]
 neuro-sam = neuro_sam:napari.yaml