spacr 0.2.68__py3-none-any.whl → 0.3.0__py3-none-any.whl

spacr/resources/MEDIAR/train_tools/models/MEDIARFormer.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn as nn
+
+from segmentation_models_pytorch import MAnet
+from segmentation_models_pytorch.base.modules import Activation
+
+__all__ = ["MEDIARFormer"]
+
+
+class MEDIARFormer(MAnet):
+    """MEDIAR-Former Model"""
+
+    def __init__(
+        self,
+        encoder_name="mit_b5",  # Default encoder
+        encoder_weights="imagenet",  # Pre-trained weights
+        decoder_channels=(1024, 512, 256, 128, 64),  # Decoder configuration
+        decoder_pab_channels=256,  # Decoder Pyramid Attention Block channels
+        in_channels=3,  # Number of input channels
+        classes=3,  # Number of output classes
+    ):
+        # Initialize the MAnet model with provided parameters
+        super().__init__(
+            encoder_name=encoder_name,
+            encoder_weights=encoder_weights,
+            decoder_channels=decoder_channels,
+            decoder_pab_channels=decoder_pab_channels,
+            in_channels=in_channels,
+            classes=classes,
+        )
+
+        # Remove the default segmentation head as it's not used in this architecture
+        self.segmentation_head = None
+
+        # Modify all activation functions in the encoder and decoder from ReLU to Mish
+        _convert_activations(self.encoder, nn.ReLU, nn.Mish(inplace=True))
+        _convert_activations(self.decoder, nn.ReLU, nn.Mish(inplace=True))
+
+        # Add custom segmentation heads for different segmentation tasks
+        self.cellprob_head = DeepSegmentationHead(
+            in_channels=decoder_channels[-1], out_channels=1
+        )
+        self.gradflow_head = DeepSegmentationHead(
+            in_channels=decoder_channels[-1], out_channels=2
+        )
+
+    def forward(self, x):
+        """Forward pass through the network"""
+        # Ensure the input shape is correct
+        self.check_input_shape(x)
+
+        # Encode the input and then decode it
+        features = self.encoder(x)
+        decoder_output = self.decoder(*features)
+
+        # Generate masks for cell probability and gradient flows
+        cellprob_mask = self.cellprob_head(decoder_output)
+        gradflow_mask = self.gradflow_head(decoder_output)
+
+        # Concatenate the masks for output
+        masks = torch.cat([gradflow_mask, cellprob_mask], dim=1)
+
+        return masks
+
+
+class DeepSegmentationHead(nn.Sequential):
+    """Custom segmentation head for generating specific masks"""
+
+    def __init__(
+        self, in_channels, out_channels, kernel_size=3, activation=None, upsampling=1
+    ):
+        # Define a sequence of layers for the segmentation head
+        layers = [
+            nn.Conv2d(
+                in_channels,
+                in_channels // 2,
+                kernel_size=kernel_size,
+                padding=kernel_size // 2,
+            ),
+            nn.Mish(inplace=True),
+            nn.BatchNorm2d(in_channels // 2),
+            nn.Conv2d(
+                in_channels // 2,
+                out_channels,
+                kernel_size=kernel_size,
+                padding=kernel_size // 2,
+            ),
+            nn.UpsamplingBilinear2d(scale_factor=upsampling)
+            if upsampling > 1
+            else nn.Identity(),
+            Activation(activation) if activation else nn.Identity(),
+        ]
+        super().__init__(*layers)
+
+
+def _convert_activations(module, from_activation, to_activation):
+    """Recursively convert activation functions in a module"""
+    for name, child in module.named_children():
+        if isinstance(child, from_activation):
+            setattr(module, name, to_activation)
+        else:
+            _convert_activations(child, from_activation, to_activation)

spacr/resources/MEDIAR/train_tools/models/__init__.py

@@ -0,0 +1 @@
+from .MEDIARFormer import *

spacr/resources/MEDIAR/train_tools/utils.py

@@ -0,0 +1,70 @@
+import torch
+import numpy as np
+import os, json, random
+from pprint import pprint
+
+__all__ = ["ConfLoader", "directory_setter", "random_seeder", "pprint_config"]
+
+
+class ConfLoader:
+    """
+    Load json config file using DictWithAttributeAccess object_hook.
+    ConfLoader(conf_name).opt attribute is the result of loading json config file.
+    """
+
+    class DictWithAttributeAccess(dict):
+        """
+        This inner class makes dict to be accessed same as class attribute.
+        For example, you can use opt.key instead of the opt['key'].
+        """
+
+        def __getattr__(self, key):
+            return self[key]
+
+        def __setattr__(self, key, value):
+            self[key] = value
+
+    def __init__(self, conf_name):
+        self.conf_name = conf_name
+        self.opt = self.__get_opt()
+
+    def __load_conf(self):
+        with open(self.conf_name, "r") as conf:
+            opt = json.load(
+                conf, object_hook=lambda dict: self.DictWithAttributeAccess(dict)
+            )
+        return opt
+
+    def __get_opt(self):
+        opt = self.__load_conf()
+        opt = self.DictWithAttributeAccess(opt)
+
+        return opt
+
+
+def directory_setter(path="./results", make_dir=False):
+    """
+    Make dictionary if not exists.
+    """
+    if not os.path.exists(path) and make_dir:
+        os.makedirs(path)  # make dir if not exist
+        print("directory %s is created" % path)
+
+    if not os.path.isdir(path):
+        raise NotADirectoryError(
+            "%s is not valid. set make_dir=True to make dir." % path
+        )
+
+
+def random_seeder(seed):
+    """Fix randomness."""
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+def pprint_config(opt):
+    print("\n" + "=" * 50 + " Configuration " + "=" * 50)
+    pprint(opt, compact=True)
+    print("=" * 115 + "\n")