Awesomenous/CellSegmentation

1import numpy as np2import torch3from torch import nn4from torchvision.transforms.functional import center_crop5from torchvision.transforms import Resize6import gradio as gr7from PIL import Image8 9class CNNBlock(nn.Module):10    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=0):11        super(CNNBlock, self).__init__()12    13        self.seq_block = nn.Sequential(14            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False),15            nn.BatchNorm2d(out_channels),16            nn.ReLU(inplace=True)17        )18    19    def forward(self, x):20        return self.seq_block(x)21 22class CNNBlocks(nn.Module):23    def __init__(self, n_conv, in_channels, out_channels, padding):24        super(CNNBlocks, self).__init__()25        26        self.layers = nn.ModuleList()27        for i in range(n_conv):28            self.layers.append(CNNBlock(in_channels, out_channels, padding=padding))29            in_channels = out_channels30        31    def forward(self, x):32        for layers in self.layers:33            x = layers(x)34        return x35 36class Encoder(nn.Module):37    def __init__(self, in_channels, out_channels, padding, downhill=4):38        super(Encoder, self).__init__()39        40        self.enc_layers = nn.ModuleList()41        for i in range(downhill):42            self.enc_layers += [CNNBlocks(n_conv=2, in_channels=in_channels, out_channels=out_channels, padding=padding), 43                                nn.MaxPool2d(2,2)]44            in_channels = out_channels45            out_channels *= 246        self.enc_layers.append(CNNBlocks(n_conv=2, in_channels=in_channels, out_channels=out_channels, padding=padding))47    48    def forward(self, x):49        route_connection = []50        for layer in self.enc_layers:51            x = layer(x)52            if isinstance(layer, CNNBlocks):53                route_connection.append(x)54        return x, route_connection55 56class Decoder(nn.Module):57    def __init__(self, in_channels, out_channels, exit_channels, padding, uphill=4):58        super(Decoder, self).__init__()59        60        self.exit_channels = exit_channels61        self.layers = nn.ModuleList()62        for i in range(uphill):63 64            self.layers += [nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),65                            CNNBlocks(n_conv=2, in_channels=in_channels, out_channels=out_channels, padding=padding)]66            in_channels //= 267            out_channels //= 268        self.layers.append(nn.Conv2d(in_channels, exit_channels, kernel_size=1, padding=0))69 70    def forward(self, x, routes_connection):71        routes_connection.pop(-1)72        for layer in self.layers:73            if isinstance(layer, CNNBlocks):74                routes_connection[-1] = center_crop(routes_connection[-1], (x.shape[2], x.shape[3]))75                x = torch.cat([x, routes_connection.pop(-1)], dim=1)76                x = layer(x)77            else:78                x = layer(x)79        return x80 81class UNET(nn.Module):82    def __init__(self, in_channels, first_out_channels, exit_channels, downhill, padding=0):83        super(UNET, self).__init__()84        self.encoder = Encoder(in_channels, first_out_channels, padding=padding, downhill=downhill)85        self.decoder = Decoder(first_out_channels*(2**downhill), first_out_channels*(2**(downhill-1)),86                               exit_channels, padding=padding, uphill=downhill)87    def forward(self, x):88        enc_out, routes = self.encoder(x)89        out = self.decoder(enc_out, routes)90        return out91 92model = UNET(3, 64, 1, 4, 1)93model.load_state_dict(torch.load("epoch_27.pth", map_location=torch.device('cpu')))94 95def predict(input_image):96    transform1 = Resize(size=(256,256))97    numpy_in = np.array(input_image.convert("RGB"))98    img = transform1(torch.from_numpy(numpy_in).permute(2,0,1).unsqueeze(dim=0).type(torch.float32))99    model.eval()100    with torch.inference_mode():101        pred = model(img)102    transform2 = Resize(size=(numpy_in.shape[0], numpy_in.shape[1]))103    final_img = transform2(pred.squeeze(dim=0).sigmoid()).round().squeeze().numpy()104    return Image.fromarray((final_img * 255).astype(np.uint8))105 106interface = gr.Interface(107    fn=predict,108    inputs=gr.Image(type="pil", label="Upload Image"),109    outputs=gr.Image(type="pil", label="Segmentation Mask"),110    examples=["00ae65c1c6631ae6f2be1a449902976e6eb8483bf6b0740d00530220832c6d3e.png",111             "2f929b067a59f88530b6bfa6f6889bc3a38adf88d594895973d1c8b2549fd93d.png",112             "432f367a4c5b5674de2e2977744d10289a064e5704b21af6607b4975be47c580.png",113             "Picture2.png"]114)115 116interface.launch()