diff --git a/src/neuron_proofreader/machine_learning/geometric_gnn_models.py b/src/neuron_proofreader/models/geometric_gnn_models.py similarity index 100% rename from src/neuron_proofreader/machine_learning/geometric_gnn_models.py rename to src/neuron_proofreader/models/geometric_gnn_models.py diff --git a/src/neuron_proofreader/machine_learning/gnn_models.py b/src/neuron_proofreader/models/gnn_models.py similarity index 100% rename from src/neuron_proofreader/machine_learning/gnn_models.py rename to src/neuron_proofreader/models/gnn_models.py diff --git a/src/neuron_proofreader/models/new_vision_models.py b/src/neuron_proofreader/models/new_vision_models.py new file mode 100644 index 00000000..9a413770 --- /dev/null +++ b/src/neuron_proofreader/models/new_vision_models.py @@ -0,0 +1,231 @@ +""" +Created on Thu July 2 13:00:00 2026 + +@author: Anna Grim +@email: anna.grim@alleninstitute.org + +Code for vision models that perform image classification tasks within +NeuronProofreader pipelines. + +""" + +import torch.nn as nn +import torch.nn.functional as F + +from neuron_proofreader.utils.ml_util import FeedForwardNet + + +class NewCNN3D(nn.Module): + """ + Convolutional neural network for 3D images. + """ + + def __init__( + self, + patch_shape, + in_channels=2, + first_out_channels=16, + output_dim=1, + channel_growth=2, + dropout=0.1, + max_channels=256, + num_blocks=5, + num_single_blocks=2, + use_double=True, + ): + """ + Instantiates a CNN3D object. + + Parameters + ---------- + patch_shape : Tuple[int] + Shape of input image patch. + output_dim : int, optional + Dimension of output. Default is 1. + dropout : float, optional + Fraction of values to randomly drop during training. Default is + 0.1. + num_blocks : int, optional + Number of convolutional blocks. Default is 5. + use_double : bool, optional + Indication of whether to use double convolution. Default is True. + """ + # Call parent class + nn.Module.__init__(self) + + # Instance attributes + self.dropout = dropout + self.patch_shape = patch_shape + + self.encode = Encoder3D( + in_channels, + first_out_channels, + num_blocks, + channel_growth=channel_growth, + use_double=use_double, + max_channels=max_channels, + num_single_blocks=num_single_blocks, + ) + self.output = FeedForwardNet(self.encode.out_channels, output_dim, 3) + + # Initialize weights + self.apply(self.init_weights) + + @staticmethod + def init_weights(m): + """ + Initializes the weights and biases of a given PyTorch layer. + + Parameters + ---------- + m : nn.Module + PyTorch layer or module. + """ + if isinstance(m, nn.Conv3d): + nn.init.xavier_normal_(m.weight) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.xavier_normal_(m.weight) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.GroupNorm): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + def forward(self, x): + """ + Passes the given input through this neural network. + + Parameters + ---------- + x : torch.Tensor + Input vector of features. + + Returns + ------- + x : torch.Tensor + Output of the neural network. + """ + x = self.encode(x) + x = F.adaptive_avg_pool3d(x, 1).flatten(1) + x = self.output(x) + return x + + +class Encoder3D(nn.Module): + """ + Sequence of ConvBlock3D blocks with growing (capped) channel width. + """ + def __init__( + self, + in_channels, + out_channels, + num_blocks, + channel_growth=2, + use_double=True, + max_channels=256, + num_single_blocks=2, + ): + """ + Instantiates an Encoder3D object. + + Parameters + ---------- + in_channels : int + Number of channels input to the first block. + out_channels : int + Number of channels output by the first block. + num_blocks : int + Number of conv blocks in the encoder. + channel_growth : float, optional + Multiplicative channel growth factor per layer. Default is 2. + use_double : bool, optional + Indication of whether blocks use double convolution. Default is True. + max_channels : int, optional + Cap on channel growth across layers. Default is 128. + """ + # Call parent class + super().__init__() + + # Create encoding blocks + blocks = list() + for i in range(num_blocks): + use_double = i > num_single_blocks + block = ConvBlock3D( + in_channels, out_channels, use_double=use_double + ) + blocks.append(block) + in_channels = block.out_channels + out_channels = int(min(out_channels * channel_growth, max_channels)) + + # Instance attributes + self.blocks = nn.ModuleList(blocks) + self.out_channels = self.blocks[-1].out_channels + + def forward(self, x): + for block in self.blocks: + x = block(x) + return x + + +class ConvBlock3D(nn.Module): + + def __init__( + self, in_channels, out_channels, kernel_size=3, use_double=True + ): + """ + Instantiates a ConvBlock3D object. + + Parameters + ---------- + in_channels : int + Number of input channels to this block. + out_channels : int + Number of output channels from this block. + kernel_size : int, optional + Size of kernel used on convolutional layers. Default is 3. + use_double : bool, optional + Indication of whether to apply a second conv+norm+act before + pooling. Default is True. + """ + # Call parent class + super().__init__() + + # Instance attributes + self.out_channels = out_channels + + # Create encoding layers + layers = self.create_unit(in_channels, out_channels, kernel_size) + if use_double: + layers.extend( + self.create_unit(out_channels, out_channels, kernel_size) + ) + + self.net = nn.Sequential(*layers, nn.MaxPool3d(kernel_size=2)) + + def forward(self, x): + return self.net(x) + + # --- Helpers --- + def create_unit(self, in_channels, out_channels, kernel_size): + padding = kernel_size // 2 + n_groups = self.get_num_groups(out_channels, 8) + unit = [ + nn.Conv3d( + in_channels, + out_channels, + kernel_size, + padding=padding, + ), + nn.GroupNorm(n_groups, out_channels), + nn.GELU() + ] + return unit + + @staticmethod + def get_num_groups(num_channels, max_groups=8): + for g in reversed(range(1, max_groups + 1)): + if num_channels % g == 0: + return g + return 1 diff --git a/src/neuron_proofreader/machine_learning/point_cloud_models.py b/src/neuron_proofreader/models/point_cloud_models.py similarity index 100% rename from src/neuron_proofreader/machine_learning/point_cloud_models.py rename to src/neuron_proofreader/models/point_cloud_models.py diff --git a/src/neuron_proofreader/machine_learning/vision_models.py b/src/neuron_proofreader/models/vision_models.py similarity index 100% rename from src/neuron_proofreader/machine_learning/vision_models.py rename to src/neuron_proofreader/models/vision_models.py