init

README.md

@@ -1,8 +1,8 @@
 ---
 title: NeurIPS CellSeg
-emoji: 📊
-colorFrom: indigo
-colorTo: pink
+emoji: 🔥
+colorFrom: yellow
+colorTo: red
 sdk: gradio
 sdk_version: 3.4.1
 app_file: app.py

app.py

@@ -0,0 +1,144 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Author: Yao
+# Mail: zhangyao215@mails.ucas.ac.cn
+
+import gradio as gr
+
+import os
+join = os.path.join
+import time
+import numpy as np
+# from skimage.filters import threshold_otsu
+# from skimage.measure import label
+import torch
+import monai
+from monai.inferers import sliding_window_inference
+from unetr2d import UNETR2D
+import time
+from skimage import io, segmentation, morphology, measure, exposure
+
+
+def visualize_instance_seg_mask(mask):
+    image = np.zeros((mask.shape[0], mask.shape[1], 3))
+    labels = np.unique(mask)
+    label2color = {label: (random.randint(0, 1), random.randint(0, 255), random.randint(0, 255)) for label in labels}
+    for i in range(image.shape[0]):
+      for j in range(image.shape[1]):
+        image[i, j, :] = label2color[mask[i, j]]
+    image = image / 255
+    return image
+
+
+def load_model(model_name, custom_model_path):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    if model_name == 'unet':
+        model = monai.networks.nets.UNet(
+                   spatial_dims=2,
+                   in_channels=3,
+                   out_channels=3,
+                   channels=(16, 32, 64, 128, 256),
+                   strides=(2, 2, 2, 2),
+                   num_res_units=2,
+               )
+
+    elif model_name == 'unetr':
+        model = UNETR2D(
+                    in_channels=3,
+                    out_channels=3,
+                    img_size=(256, 256),
+                    feature_size=16,
+                    hidden_size=768,
+                    mlp_dim=3072,
+                    num_heads=12,
+                    pos_embed="perceptron",
+                    norm_name="instance",
+                    res_block=True,
+                    dropout_rate=0.0,
+                )
+    elif model_name == 'swinunetr':
+        model = monai.networks.nets.SwinUNETR(
+                        img_size=(256, 256), 
+                        in_channels=3, 
+                        out_channels=3,
+                        feature_size=24, # should be divisible by 12
+                        spatial_dims=2
+                    )
+        if os.path.isfile(custom_model_path):
+            checkpoint = torch.load(custom_model_path.resolve(), map_location=torch.device(device))
+        elif os.path.isfile(join(os.path.dirname(__file__), 'best_Dice_model.pth')):
+            checkpoint = torch.load(join(os.path.dirname(__file__), 'best_Dice_model.pth'), map_location=torch.device(device))
+        else:
+            torch.hub.download_url_to_file('https://zenodo.org/record/6792177/files/best_Dice_model.pth?download=1', join(os.path.dirname(__file__), 'work_dir/swinunetr/best_Dice_model.pth'))
+            checkpoint = torch.load(join(os.path.dirname(__file__), 'best_Dice_model.pth'), map_location=torch.device(device))
+
+    model.load_state_dict(checkpoint['model_state_dict'])
+
+    model = model.to(device)
+    model.eval()
+
+    return model
+
+
+def normalize_channel(img, lower=1, upper=99):
+    non_zero_vals = img[np.nonzero(img)]
+    percentiles = np.percentile(non_zero_vals, [lower, upper])
+    if percentiles[1] - percentiles[0] > 0.001:
+        img_norm = exposure.rescale_intensity(img, in_range=(percentiles[0], percentiles[1]), out_range='uint8')
+    else:
+        img_norm = img
+    return img_norm.astype(np.uint8)
+
+
+def preprocess(img_data):
+    if len(img_data.shape) == 2:
+        img_data = np.repeat(np.expand_dims(img_data, axis=-1), 3, axis=-1)
+    elif len(img_data.shape) == 3 and img_data.shape[-1] > 3:
+        img_data = img_data[:,:, :3]
+    else:
+        pass
+    pre_img_data = np.zeros(img_data.shape, dtype=np.uint8)
+    for i in range(3):
+        img_channel_i = img_data[:,:,i]
+        if len(img_channel_i[np.nonzero(img_channel_i)])>0:
+            pre_img_data[:,:,i] = normalize_channel(img_channel_i, lower=1, upper=99)
+    return pre_img_data
+
+
+def get_seg(pre_img_data, model_name, custom_model_path, threshold):
+    model = load_model(model_name, custom_model_path)
+    #%%
+    roi_size = (256, 256)
+    sw_batch_size = 4
+    with torch.no_grad():
+        t0 = time.time()
+        test_npy01 = pre_img_data/np.max(pre_img_data)
+        # test_tensor = torch.from_numpy(np.expand_dims(test_npy01, 0)).permute(0,3,1,2).type(torch.FloatTensor).to(device)
+        test_tensor = torch.from_numpy(np.expand_dims(test_npy01, 0)).permute(0,3,1,2).type(torch.FloatTensor)
+        test_pred_out = sliding_window_inference(test_tensor, roi_size, sw_batch_size, model)
+        test_pred_out = torch.nn.functional.softmax(test_pred_out, dim=1) # (B, C, H, W)
+        test_pred_npy = test_pred_out[0,1].cpu().numpy()
+        # convert probability map to binary mask and apply morphological postprocessing
+        test_pred_mask = measure.label(morphology.remove_small_objects(morphology.remove_small_holes(test_pred_npy>threshold),16))
+        # tif.imwrite(join(output_path, img_name.split('.')[0]+'_label.tiff'), test_pred_mask, compression='zlib')
+        t1 = time.time()
+        # print(f'Prediction finished: {img_layer.name}; img size = {pre_img_data.shape}; costing: {t1-t0:.2f}s')
+    return test_pred_mask
+
+
+def predict(img):
+    seg_labels = get_seg(preprocess(img), 'swinunetr', './best_Dice_model.pth', 0.5)
+    return seg_labels
+
+
+demo = gr.Interface(
+    predict, 
+    inputs=[gr.Image()], 
+    outputs="image",
+    title="NeurIPS CellSeg Demo",
+    examples=[["cell_00225.png"]]
+)
+
+demo.launch()
+

best_Dice_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:764db0c53184da5bf743db84d8837b1f9b2c7f3b0236b43c75ff747e47c75e5a
+size 75949863

requirements.txt

@@ -0,0 +1,6 @@
+numpy
+scikit-image
+numpy
+torch
+monai
+einops

unetr2d.py

@@ -0,0 +1,202 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Sun Mar 20 14:23:19 2022
+Author: MONAI
+"""
+
+from typing import Tuple, Union
+import torch
+import torch.nn as nn
+
+from monai.networks.blocks.dynunet_block import UnetOutBlock
+from monai.networks.blocks import UnetrBasicBlock, UnetrPrUpBlock, UnetrUpBlock
+from monai.networks.nets import ViT
+
+class UNETR2D(nn.Module):
+    """
+    UNETR based on: "Hatamizadeh et al.,
+    UNETR: Transformers for 3D Medical Image Segmentation <https://arxiv.org/abs/2103.10504>"
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        img_size: Tuple[int, int],
+        feature_size: int = 16,
+        hidden_size: int = 768,
+        mlp_dim: int = 3072,
+        num_heads: int = 12,
+        pos_embed: str = "perceptron",
+        norm_name: Union[Tuple, str] = "instance",
+        conv_block: bool = False,
+        res_block: bool = True,
+        dropout_rate: float = 0.0,
+        debug: bool = False
+    ) -> None:
+
+        super().__init__()
+
+        if not (0 <= dropout_rate <= 1):
+            raise AssertionError("dropout_rate should be between 0 and 1.")
+
+        if hidden_size % num_heads != 0:
+            raise AssertionError("hidden size should be divisible by num_heads.")
+
+        if pos_embed not in ["conv", "perceptron"]:
+            raise KeyError(f"Position embedding layer of type {pos_embed} is not supported.")
+
+        self.num_layers = 12
+        self.patch_size = (16, 16)
+        self.feat_size = tuple(img_d // p_d for img_d, p_d in zip(img_size, self.patch_size))
+        self.hidden_size = hidden_size
+        self.classification = False
+        self.debug = debug
+        self.vit = ViT(
+            in_channels=in_channels,
+            img_size=img_size,
+            patch_size=self.patch_size,
+            hidden_size=hidden_size,
+            mlp_dim=mlp_dim,
+            num_layers=self.num_layers,
+            num_heads=num_heads,
+            pos_embed=pos_embed,
+            classification=self.classification,
+            dropout_rate=dropout_rate,
+            spatial_dims=2
+        )
+        self.encoder1 = UnetrBasicBlock(
+            spatial_dims=2,
+            in_channels=in_channels,
+            out_channels=feature_size,
+            kernel_size=3,
+            stride=1,
+            norm_name=norm_name,
+            res_block=res_block,
+        )
+        self.encoder2 = UnetrPrUpBlock(
+            spatial_dims=2,
+            in_channels=hidden_size,
+            out_channels=feature_size * 2,
+            num_layer=2,
+            kernel_size=3,
+            stride=1,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            conv_block=conv_block,
+            res_block=res_block,
+        )
+        self.encoder3 = UnetrPrUpBlock(
+            spatial_dims=2,
+            in_channels=hidden_size,
+            out_channels=feature_size * 4,
+            num_layer=1,
+            kernel_size=3,
+            stride=1,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            conv_block=conv_block,
+            res_block=res_block,
+        )
+        self.encoder4 = UnetrPrUpBlock(
+            spatial_dims=2,
+            in_channels=hidden_size,
+            out_channels=feature_size * 8,
+            num_layer=0,
+            kernel_size=3,
+            stride=1,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            conv_block=conv_block,
+            res_block=res_block,
+        )
+        self.decoder5 = UnetrUpBlock(
+            spatial_dims=2,
+            in_channels=hidden_size,
+            out_channels=feature_size * 8,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            res_block=res_block,
+        )
+        self.decoder4 = UnetrUpBlock(
+            spatial_dims=2,
+            in_channels=feature_size * 8,
+            out_channels=feature_size * 4,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            res_block=res_block,
+        )
+        self.decoder3 = UnetrUpBlock(
+            spatial_dims=2,
+            in_channels=feature_size * 4,
+            out_channels=feature_size * 2,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            res_block=res_block,
+        )
+        self.decoder2 = UnetrUpBlock(
+            spatial_dims=2,
+            in_channels=feature_size * 2,
+            out_channels=feature_size,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            res_block=res_block,
+        )
+        self.out = UnetOutBlock(spatial_dims=2, in_channels=feature_size, out_channels=out_channels)  # type: ignore
+
+    def proj_feat(self, x, hidden_size, feat_size): # x: (B, 256, 768)
+        x = x.view(x.size(0), feat_size[0], feat_size[1], hidden_size) # (B, 16, 16, 768)
+        x = x.permute(0, 3, 1, 2).contiguous() # (B, 768, 16, 16)
+        return x
+
+    def forward(self, x_in):
+        x, hidden_states_out = self.vit(x_in) # x: (B, 256,768), hidden_states_out: list, 12 elements, (B,256,768)
+        enc1 = self.encoder1(x_in) # (1, 16, 256, 256)
+        x2 = hidden_states_out[3] # (B, 256, 768)
+        # self.proj_feat(x2, self.hidden_size, self.feat_size): (B, 768, 16,16) -> enc2: (B,32,128,128)
+        enc2 = self.encoder2(self.proj_feat(x2, self.hidden_size, self.feat_size)) # hidden_size=768, self.feat_size=16
+        x3 = hidden_states_out[6] # (B, 256, 768)
+        enc3 = self.encoder3(self.proj_feat(x3, self.hidden_size, self.feat_size)) #(B, 768, 16,16) -> (B, 64, 64, 64)
+        x4 = hidden_states_out[9] # (B, 256, 768)
+        enc4 = self.encoder4(self.proj_feat(x4, self.hidden_size, self.feat_size)) # (B, 768, 16, 16) -> (B, 128, 32, 32)
+        dec4 = self.proj_feat(x, self.hidden_size, self.feat_size) # (B, 768, 16, 16)
+        dec3 = self.decoder5(dec4, enc4) # up -> cat -> ResConv; (B, 128, 32, 32)
+        dec2 = self.decoder4(dec3, enc3) # (B, 64, 64, 64)
+        dec1 = self.decoder3(dec2, enc2) # (B, 32, 128, 128)
+        out = self.decoder2(dec1, enc1) # (B, 16, 256, 256)
+        logits = self.out(out)
+        
+        if self.debug:
+            return x, x2, x3,x4, hidden_states_out, enc1, enc2, enc3, enc4, dec4, dec3, dec2, dec1, logits 
+        else:
+            return logits
+    
+
+# model = UNETR2D(
+#     in_channels=3, # 3 channels, R,G,B
+#     out_channels=3,
+#     img_size=(256, 256),
+#     feature_size=16,
+#     hidden_size=768,
+#     mlp_dim=3072,
+#     num_heads=12,
+#     pos_embed="perceptron",
+#     norm_name="instance",
+#     res_block=True,
+#     dropout_rate=0.0,
+#     debug=True
+# ).cuda()
+
+# from torchinfo import summary
+
+# batch_size = 1
+# summary(model, input_size=(batch_size, 3, 256, 256))
+
+# x = torch.rand((1,3,256,256)).cuda()
+# x, x2, x3,x4, hidden_states_out, enc1, enc2, enc3, enc4, dec4, dec3, dec2, dec1, logits  = model(x)
+# print(logits.shape) # torch.Size([1, 3, 256, 256])