Upload 2 files

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)

Predictor.py

@@ -0,0 +1,234 @@
+import torch
+import numpy as np
+import os, sys
+from monai.inferers import sliding_window_inference
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), "../../")))
+
+from core.BasePredictor import BasePredictor
+from core.MEDIAR.utils import compute_masks
+
+__all__ = ["Predictor"]
+
+
+class Predictor(BasePredictor):
+    def __init__(
+        self,
+        model,
+        device,
+        input_path,
+        output_path,
+        make_submission=False,
+        exp_name=None,
+        algo_params=None,
+    ):
+        super(Predictor, self).__init__(
+            model,
+            device,
+            input_path,
+            output_path,
+            make_submission,
+            exp_name,
+            algo_params,
+        )
+        self.hflip_tta = HorizontalFlip()
+        self.vflip_tta = VerticalFlip()
+
+    @torch.no_grad()
+    def _inference(self, img_data):
+        """Conduct model prediction"""
+
+        img_data = img_data.to(self.device)
+        img_base = img_data
+        outputs_base = self._window_inference(img_base)
+        outputs_base = outputs_base.cpu().squeeze()
+        img_base.cpu()
+
+        if not self.use_tta:
+            pred_mask = outputs_base
+            return pred_mask
+
+        else:
+            # HorizontalFlip TTA
+            img_hflip = self.hflip_tta.apply_aug_image(img_data, apply=True)
+            outputs_hflip = self._window_inference(img_hflip)
+            outputs_hflip = self.hflip_tta.apply_deaug_mask(outputs_hflip, apply=True)
+            outputs_hflip = outputs_hflip.cpu().squeeze()
+            img_hflip = img_hflip.cpu()
+
+            # VertricalFlip TTA
+            img_vflip = self.vflip_tta.apply_aug_image(img_data, apply=True)
+            outputs_vflip = self._window_inference(img_vflip)
+            outputs_vflip = self.vflip_tta.apply_deaug_mask(outputs_vflip, apply=True)
+            outputs_vflip = outputs_vflip.cpu().squeeze()
+            img_vflip = img_vflip.cpu()
+
+            # Merge Results
+            pred_mask = torch.zeros_like(outputs_base)
+            pred_mask[0] = (outputs_base[0] + outputs_hflip[0] - outputs_vflip[0]) / 3
+            pred_mask[1] = (outputs_base[1] - outputs_hflip[1] + outputs_vflip[1]) / 3
+            pred_mask[2] = (outputs_base[2] + outputs_hflip[2] + outputs_vflip[2]) / 3
+
+        return pred_mask
+
+    def _window_inference(self, img_data, aux=False):
+        """Inference on RoI-sized window"""
+        outputs = sliding_window_inference(
+            img_data,
+            roi_size=512,
+            sw_batch_size=4,
+            predictor=self.model if not aux else self.model_aux,
+            padding_mode="constant",
+            mode="gaussian",
+            overlap=0.6,
+        )
+
+        return outputs
+
+    def _post_process(self, pred_mask):
+        """Generate cell instance masks."""
+        dP, cellprob = pred_mask[:2], self._sigmoid(pred_mask[-1])
+        H, W = pred_mask.shape[-2], pred_mask.shape[-1]
+
+        if np.prod(H * W) < (5000 * 5000):
+            pred_mask = compute_masks(
+                dP,
+                cellprob,
+                use_gpu=True,
+                flow_threshold=0.4,
+                device=self.device,
+                cellprob_threshold=0.5,
+            )[0]
+
+        else:
+            print("\n[Whole Slide] Grid Prediction starting...")
+            roi_size = 2000
+
+            # Get patch grid by roi_size
+            if H % roi_size != 0:
+                n_H = H // roi_size + 1
+                new_H = roi_size * n_H
+            else:
+                n_H = H // roi_size
+                new_H = H
+
+            if W % roi_size != 0:
+                n_W = W // roi_size + 1
+                new_W = roi_size * n_W
+            else:
+                n_W = W // roi_size
+                new_W = W
+
+            # Allocate values on the grid
+            pred_pad = np.zeros((new_H, new_W), dtype=np.uint32)
+            dP_pad = np.zeros((2, new_H, new_W), dtype=np.float32)
+            cellprob_pad = np.zeros((new_H, new_W), dtype=np.float32)
+
+            dP_pad[:, :H, :W], cellprob_pad[:H, :W] = dP, cellprob
+
+            for i in range(n_H):
+                for j in range(n_W):
+                    print("Pred on Grid (%d, %d) processing..." % (i, j))
+                    dP_roi = dP_pad[
+                        :,
+                        roi_size * i : roi_size * (i + 1),
+                        roi_size * j : roi_size * (j + 1),
+                    ]
+                    cellprob_roi = cellprob_pad[
+                        roi_size * i : roi_size * (i + 1),
+                        roi_size * j : roi_size * (j + 1),
+                    ]
+
+                    pred_mask = compute_masks(
+                        dP_roi,
+                        cellprob_roi,
+                        use_gpu=True,
+                        flow_threshold=0.4,
+                        device=self.device,
+                        cellprob_threshold=0.5,
+                    )[0]
+
+                    pred_pad[
+                        roi_size * i : roi_size * (i + 1),
+                        roi_size * j : roi_size * (j + 1),
+                    ] = pred_mask
+
+            pred_mask = pred_pad[:H, :W]
+
+        return pred_mask
+
+    def _sigmoid(self, z):
+        return 1 / (1 + np.exp(-z))
+
+
+"""
+Adapted from the following references:
+[1] https://github.com/qubvel/ttach/blob/master/ttach/transforms.py
+
+"""
+
+
+def hflip(x):
+    """flip batch of images horizontally"""
+    return x.flip(3)
+
+
+def vflip(x):
+    """flip batch of images vertically"""
+    return x.flip(2)
+
+
+class DualTransform:
+    identity_param = None
+
+    def __init__(
+        self, name: str, params,
+    ):
+        self.params = params
+        self.pname = name
+
+    def apply_aug_image(self, image, *args, **params):
+        raise NotImplementedError
+
+    def apply_deaug_mask(self, mask, *args, **params):
+        raise NotImplementedError
+
+
+class HorizontalFlip(DualTransform):
+    """Flip images horizontally (left -> right)"""
+
+    identity_param = False
+
+    def __init__(self):
+        super().__init__("apply", [False, True])
+
+    def apply_aug_image(self, image, apply=False, **kwargs):
+        if apply:
+            image = hflip(image)
+        return image
+
+    def apply_deaug_mask(self, mask, apply=False, **kwargs):
+        if apply:
+            mask = hflip(mask)
+        return mask
+
+
+class VerticalFlip(DualTransform):
+    """Flip images vertically (up -> down)"""
+
+    identity_param = False
+
+    def __init__(self):
+        super().__init__("apply", [False, True])
+
+    def apply_aug_image(self, image, apply=False, **kwargs):
+        if apply:
+            image = vflip(image)
+
+        return image
+
+    def apply_deaug_mask(self, mask, apply=False, **kwargs):
+        if apply:
+            mask = vflip(mask)
+
+        return mask