neuro-sam 0.1.0__py3-none-any.whl

neuro_sam/plugin.py

@@ -0,0 +1,260 @@
+import napari
+import numpy as np
+import imageio.v2 as io
+from neuro_sam.napari_utils.main_widget import NeuroSAMWidget  # Updated with anisotropic scaling
+
+
+def pad_image_for_patches(image, patch_size=128, pad_value=0):
+    """
+    Pad the image so that its height and width are multiples of patch_size.
+    Handles various image dimensions including stacks of colored images.
+    
+    Parameters:
+    -----------
+    image (np.ndarray): Input image array:
+        - 2D: (H x W)
+        - 3D: (C x H x W) for grayscale stacks or (H x W x C) for colored image
+        - 4D: (Z x H x W x C) for stacks of colored images
+    patch_size (int): The patch size to pad to, default is 128.
+    pad_value (int or tuple): The constant value(s) for padding.
+    
+    Returns:
+    --------
+    padded_image (np.ndarray): The padded image.
+    padding_amounts (tuple): The amount of padding applied (pad_h, pad_w).
+    original_dims (tuple): The original dimensions (h, w).
+    """
+    # Determine the image format and dimensions
+    if image.ndim == 2:
+        # 2D grayscale image (H x W)
+        h, w = image.shape
+        is_color = False
+        is_stack = False
+    elif image.ndim == 3:
+        # This could be either:
+        # - A stack of 2D grayscale images (Z x H x W)
+        # - A single color image (H x W x C)
+        # We'll check the third dimension to decide
+        if image.shape[2] <= 4:  # Assuming color channels ≤ 4 (RGB, RGBA)
+            # Single color image (H x W x C)
+            h, w, c = image.shape
+            is_color = True
+            is_stack = False
+        else:
+            # Stack of grayscale images (Z x H x W)
+            z, h, w = image.shape
+            is_color = False
+            is_stack = True
+    elif image.ndim == 4:
+        # Stack of color images (Z x H x W x C)
+        z, h, w, c = image.shape
+        is_color = True
+        is_stack = True
+    else:
+        raise ValueError(f"Unsupported image dimension: {image.ndim}")
+    
+    # Compute necessary padding for height and width
+    pad_h = (patch_size - h % patch_size) % patch_size
+    pad_w = (patch_size - w % patch_size) % patch_size
+    
+    # Pad the image based on its format
+    if not is_stack and not is_color:
+        # 2D grayscale image
+        padding = ((0, pad_h), (0, pad_w))
+        padded_image = np.pad(image, padding, mode='constant', constant_values=pad_value)
+    
+    elif is_stack and not is_color:
+        # Stack of grayscale images (Z x H x W)
+        padding = ((0, 0), (0, pad_h), (0, pad_w))
+        padded_image = np.pad(image, padding, mode='constant', constant_values=pad_value)
+    
+    elif not is_stack and is_color:
+        # Single color image (H x W x C)
+        padding = ((0, pad_h), (0, pad_w), (0, 0))
+        padded_image = np.pad(image, padding, mode='constant', constant_values=pad_value)
+    
+    elif is_stack and is_color:
+        # Stack of color images (Z x H x W x C)
+        padding = ((0, 0), (0, pad_h), (0, pad_w), (0, 0))
+        padded_image = np.pad(image, padding, mode='constant', constant_values=pad_value)
+    
+    return padded_image, (pad_h, pad_w), (h, w)
+
+
+def run_neuro_sam(image=None, image_path=None, spacing_xyz=(94.0, 94.0, 500.0)):
+    """
+    Launch the NeuroSAM plugin with anisotropic scaling support
+    
+    Parameters:
+    -----------
+    image : numpy.ndarray, optional
+        3D or higher-dimensional image data. If None, image_path must be provided.
+    image_path : str, optional
+        Path to image file to load. If None, image must be provided.
+    spacing_xyz : tuple, optional
+        Original voxel spacing in (x, y, z) nanometers. 
+        Default: (94.0, 94.0, 500.0) - typical for two photon microscopy
+    
+    Returns:
+    --------
+    viewer : napari.Viewer
+        The napari viewer instance
+    """
+    # Validate inputs
+    if image is None and image_path is None:
+        raise ValueError("Either image or image_path must be provided")
+    
+    # Load image if path provided
+    if image is None:
+        try:
+            image = np.asarray(io.imread(image_path))
+            print(f"Loaded image from {image_path}")
+            print(f"Image shape: {image.shape}")
+            print(f"Image dtype: {image.dtype}")
+        except Exception as e:
+            raise ValueError(f"Failed to load image from {image_path}: {str(e)}")
+    
+    # Normalize image to 0-1 range for better visualization
+    if image.max() > 1:
+        image = image.astype(np.float32)
+        image_min, image_max = image.min(), image.max()
+        image = (image - image_min) / (image_max - image_min)
+        print(f"Normalized image from range [{image_min:.2f}, {image_max:.2f}] to [0, 1]")
+
+    # Pad image for patch-based processing
+    image, padding_amounts, original_dims = pad_image_for_patches(image, patch_size=128, pad_value=0)
+    if padding_amounts[0] > 0 or padding_amounts[1] > 0:
+        print(f"Padded image by {padding_amounts} pixels to be divisible by 128")
+        print(f"New image shape: {image.shape}")
+    
+    # Create a viewer
+    viewer = napari.Viewer()
+    
+    # Display spacing information
+    print(f"Original voxel spacing: X={spacing_xyz[0]:.1f}, Y={spacing_xyz[1]:.1f}, Z={spacing_xyz[2]:.1f} nm")
+    
+    # Create and add our widget with anisotropic scaling capabilities
+    neuro_sam_widget = NeuroSAMWidget(
+        viewer=viewer, 
+        image=image, 
+        original_spacing_xyz=spacing_xyz
+    )
+    
+    viewer.window.add_dock_widget(
+        neuro_sam_widget, name="Neuro-SAM", area="right"
+    )
+    
+    # Set initial view
+    if image.ndim >= 3:
+        # Start with a mid-slice view for 3D+ images
+        mid_slice = image.shape[0] // 2
+        viewer.dims.set_point(0, mid_slice)
+    
+    # Display startup information
+    napari.utils.notifications.show_info(
+        f"NeuroSAM loaded! Image shape: {image.shape}. "
+        f"Configure voxel spacing in the 'Path Tracing' tab first."
+    )
+    
+    return viewer
+
+
+def run_neuro_sam_with_metadata(image_path, metadata=None):
+    """
+    Launch NeuroSAM with metadata-derived spacing information
+    
+    Parameters:
+    -----------
+    image_path : str
+        Path to image file
+    metadata : dict, optional
+        Metadata dictionary with spacing information.
+        Expected keys: 'spacing_x_nm', 'spacing_y_nm', 'spacing_z_nm'
+        
+    Returns:
+    --------
+    viewer : napari.Viewer
+        The napari viewer instance
+    """
+    # Default spacing values
+    default_spacing = (94.0, 94.0, 500.0)  # (x, y, z) in nm
+    
+    if metadata is not None:
+        try:
+            # Extract spacing from metadata
+            x_spacing = metadata.get('spacing_x_nm', default_spacing[0])
+            y_spacing = metadata.get('spacing_y_nm', default_spacing[1])
+            z_spacing = metadata.get('spacing_z_nm', default_spacing[2])
+            
+            spacing_xyz = (float(x_spacing), float(y_spacing), float(z_spacing))
+            print(f"Using metadata-derived spacing: X={spacing_xyz[0]:.1f}, Y={spacing_xyz[1]:.1f}, Z={spacing_xyz[2]:.1f} nm")
+        except (ValueError, TypeError) as e:
+            print(f"Error parsing metadata spacing, using defaults: {e}")
+            spacing_xyz = default_spacing
+    else:
+        spacing_xyz = default_spacing
+        print(f"No metadata provided, using default spacing: X={spacing_xyz[0]:.1f}, Y={spacing_xyz[1]:.1f}, Z={spacing_xyz[2]:.1f} nm")
+    
+    return run_neuro_sam(image_path=image_path, spacing_xyz=spacing_xyz)
+
+
+def load_ome_tiff_with_spacing(image_path):
+    """
+    Load OME-TIFF file and extract voxel spacing from metadata
+    
+    Parameters:
+    -----------
+    image_path : str
+        Path to OME-TIFF file
+        
+    Returns:
+    --------
+    tuple : (image, spacing_xyz)
+        Image array and spacing tuple
+    """
+    image = np.asarray(io.imread(image_path))
+    return image, (94.0, 94.0, 500.0)
+
+
+# For direct execution from command line
+if __name__ == "__main__":
+    import sys
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Launch NeuroSAM with anisotropic scaling support")
+    parser.add_argument("--image_path", nargs="?", help="Path to image file")
+    parser.add_argument("--x-spacing", type=float, default=94.0, help="X voxel spacing in nm (default: 94.0)")
+    parser.add_argument("--y-spacing", type=float, default=94.0, help="Y voxel spacing in nm (default: 94.0)")
+    parser.add_argument("--z-spacing", type=float, default=500.0, help="Z voxel spacing in nm (default: 500.0)")
+    parser.add_argument("--ome", action="store_true", help="Try to extract spacing from OME-TIFF metadata")
+    
+    args = parser.parse_args()
+    
+    if args.image_path:
+        if args.ome:
+            # Try to load OME-TIFF with metadata
+            try:
+                image, spacing_xyz = load_ome_tiff_with_spacing(args.image_path)
+                viewer = run_neuro_sam(image=image, spacing_xyz=spacing_xyz)
+            except Exception as e:
+                print(f"Error loading OME-TIFF: {e}")
+                print("Falling back to standard loading...")
+                spacing_xyz = (args.x_spacing, args.y_spacing, args.z_spacing)
+                viewer = run_neuro_sam(image_path=args.image_path, spacing_xyz=spacing_xyz)
+        else:
+            # Use command line spacing arguments
+            spacing_xyz = (args.x_spacing, args.y_spacing, args.z_spacing)
+            viewer = run_neuro_sam(image_path=args.image_path, spacing_xyz=spacing_xyz)
+    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}")
+            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:
+            sys.exit(1)
+    
+    print("\nStarted NeuroSAM with anisotropic scaling support!")
+    print("Configure voxel spacing in the 'Path Tracing' tab before starting analysis.")
+    napari.run()  # Start the Napari event loop

neuro_sam/punet/deepd3_model.py

@@ -0,0 +1,231 @@
+"""
+DeepD3 U-Net model with dual decoders for dendrites and spines.
+
+Architecture:
+- Single encoder with residual connections
+- Dual decoders (one for dendrites, one for spines)
+- Skip connections from encoder to both decoders
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class EncoderBlock(nn.Module):
+    """
+    Encoder block with residual connection.
+    
+    Structure:
+    - 1x1 conv for identity mapping
+    - Two 3x3 convs with normalization and activation
+    - Residual addition
+    - Max pooling
+    """
+    def __init__(self, in_channels, out_channels, activation, use_batchnorm=True):
+        super().__init__()
+        self.identity_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
+        
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False)
+        self.bn1 = nn.GroupNorm(8, out_channels) if use_batchnorm else nn.Identity()
+        
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.GroupNorm(8, out_channels) if use_batchnorm else nn.Identity()
+        
+        self.activation = activation
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        
+    def forward(self, x):
+        identity = self.identity_conv(x)
+        
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.activation(x)
+        
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.activation(x + identity)
+        
+        pooled = self.pool(x)
+        return x, pooled
+
+
+class DecoderBlock(nn.Module):
+    """
+    Decoder block with two 3x3 convolutions.
+    """
+    def __init__(self, in_channels, out_channels, activation, use_batchnorm=True):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False)
+        self.bn1 = nn.GroupNorm(8, out_channels) if use_batchnorm else nn.Identity()
+        
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.GroupNorm(8, out_channels) if use_batchnorm else nn.Identity()
+        
+        self.activation = activation
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.activation(x)
+        
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.activation(x)
+        
+        return x
+
+
+class Decoder(nn.Module):
+    """
+    Decoder module with skip connections from encoder.
+    
+    For each level:
+    1. Upsample latent features
+    2. Concatenate with encoder features
+    3. Apply decoder block
+    """
+    def __init__(self, num_layers, base_filters, activation, use_batchnorm=True):
+        super().__init__()
+        self.num_layers = num_layers
+        self.blocks = nn.ModuleList()
+        
+        for i in range(num_layers):
+            k = num_layers - 1 - i
+            if i == 0:
+                in_ch = base_filters * (2 ** num_layers) + base_filters * (2 ** (num_layers - 1))
+            else:
+                in_ch = base_filters * (2 ** (k + 1)) + base_filters * (2 ** k)
+            out_ch = base_filters * (2 ** k)
+            self.blocks.append(DecoderBlock(in_ch, out_ch, activation, use_batchnorm))
+        
+        self.last_layer_features = None
+        self.final_conv = nn.Conv2d(base_filters, 1, kernel_size=1)
+        self.final_activation = nn.Sigmoid()
+        
+    def forward(self, x, encoder_features):
+        """
+        Args:
+            x: Latent representation
+            encoder_features: List of encoder skip connections (deepest first)
+        """
+        for block in self.blocks:
+            enc_feat = encoder_features.pop()
+            
+            x = F.interpolate(
+                x, 
+                size=enc_feat.shape[2:], 
+                mode='bilinear', 
+                align_corners=True
+            )
+            
+            x = torch.cat([x, enc_feat], dim=1)
+            x = block(x)
+        
+        self.last_layer_features = x
+        x = self.final_conv(x)
+        x = self.final_activation(x)
+        
+        return x
+
+
+class DeepD3Model(nn.Module):
+    """
+    U-Net with one encoder and dual decoders for dendrites and spines.
+    
+    Args:
+        in_channels: Number of input channels (default: 1 for grayscale)
+        base_filters: Base number of filters (doubled at each level)
+        num_layers: Network depth (number of encoder/decoder blocks)
+        activation: Activation function ('swish' or 'relu')
+        use_batchnorm: Whether to use batch normalization
+        apply_last_layer: Whether to apply final 1x1 conv and sigmoid
+    """
+    def __init__(
+        self, 
+        in_channels=1, 
+        base_filters=32, 
+        num_layers=4, 
+        activation="swish", 
+        use_batchnorm=True,
+        apply_last_layer=True
+    ):
+        super().__init__()
+        
+        self.apply_last_layer = apply_last_layer
+        
+        if activation == "swish":
+            act = nn.SiLU(inplace=True)
+        else:
+            act = nn.ReLU(inplace=True)
+        
+        self.activation = act
+        self.num_layers = num_layers
+        self.base_filters = base_filters
+        self.use_batchnorm = use_batchnorm
+        
+        self.encoder_blocks = nn.ModuleList()
+        current_in = in_channels
+        for i in range(num_layers):
+            out_channels = base_filters * (2 ** i)
+            self.encoder_blocks.append(
+                EncoderBlock(current_in, out_channels, self.activation, use_batchnorm)
+            )
+            current_in = out_channels
+            
+        latent_in = base_filters * (2 ** (num_layers - 1))
+        latent_out = base_filters * (2 ** num_layers)
+        self.latent_conv = nn.Sequential(
+            nn.Conv2d(latent_in, latent_out, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(latent_out) if use_batchnorm else nn.Identity(),
+            self.activation
+        )
+        
+        self.decoder_dendrites = Decoder(num_layers, base_filters, self.activation, use_batchnorm)
+        self.decoder_spines = Decoder(num_layers, base_filters, self.activation, use_batchnorm)
+        
+    def forward(self, x):
+        """
+        Forward pass through encoder and dual decoders.
+        
+        Args:
+            x: Input tensor [B, C, H, W]
+            
+        Returns:
+            Tuple of (dendrite_output, spine_output)
+        """
+        encoder_features = []
+        
+        for block in self.encoder_blocks:
+            feat, x = block(x)
+            encoder_features.append(feat)
+        
+        enc_feats_d = encoder_features.copy()
+        enc_feats_s = encoder_features.copy()
+        
+        x_latent = self.latent_conv(x)
+        
+        dendrites_features = self.decoder_dendrites.forward(x_latent, enc_feats_d)
+        spines_features = self.decoder_spines.forward(x_latent, enc_feats_s)
+        
+        if self.apply_last_layer:
+            return dendrites_features, spines_features
+        else:
+            return (
+                self.decoder_dendrites.last_layer_features, 
+                self.decoder_spines.last_layer_features
+            )
+
+
+if __name__ == '__main__':
+    x = torch.randn(1, 1, 48, 48)
+    model = DeepD3Model(
+        in_channels=1, 
+        base_filters=32, 
+        num_layers=4, 
+        activation="swish",
+        apply_last_layer=False
+    )
+    dendrites, spines = model(x)
+    print("Dendrites output shape:", dendrites.shape)
+    print("Spines output shape:", spines.shape)