pipline.py

#### This is an implmentation of deeplabv3 plus for retina detection
from skimage.measure import label, regionprops
import torch
import torchvision
from torch.nn import functional as F
import torch.nn as nn
import numpy as np
import cv2
import torch
from collections import namedtuple

# check you have the right version of timm
# assert timm.__version__ == "0.3.2"
from timm.models.swin_transformer import swin_base_patch4_window12_384_in22k, SwinTransformer

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"
pad_value = 10

def forward_features(self, x):
    x = self.patch_embed(x)
    if self.absolute_pos_embed is not None:
        x = x + self.absolute_pos_embed
    x = self.pos_drop(x)

    hide=[]
    for layer in self.layers:
        x = layer(x)
        #print(x.shape)
        hide.append(x)

    #x = self.layers(x)
    x = self.norm(x)  # B L C
    return hide

def forward(self, x):
        x = self.forward_features(x)
        #x = self.forward_head(x)
        return x

SwinTransformer.forward_features = forward_features
SwinTransformer.forward = forward




def extract_regions_Last(img_test, ytruth, pad1=pad_value, pad2=pad_value, pad3=pad_value, pad4=pad_value):

    y_truth_copy = ytruth.copy()
    y_truth_copy[y_truth_copy == 2] = 1
    label_img = label(y_truth_copy)

    regions = regionprops(label_img)
    max_Area = -1
    cropped_results = dict()
    for props in regions:
        if props.area > max_Area:
            max_Area = props.area
            minr, minc, maxr, maxc = props.bbox
            bx = (minc, maxc, maxc, minc, minc)
            by = (minr, minr, maxr, maxr, minr)
            # print(minr,maxr)
            # print(bx)
            # ax.plot(bx, by, '-b', linewidth=2.5)
            # cropped_image= pred_class[minr-pad:maxr+pad, minc-pad:maxc+pad]
            # cropped_pred_mask = pred_class[minr - pad:maxr + pad, minc - pad:maxc + pad]
            if minr - pad1 < 0:
                pad1 = 5
                if minr - pad1 < 0:
                    pad1 = 0

            if minc - pad2 < 0:
                pad2 = 5
                if minc - pad2 < 0:
                    pad2 = 0
            if maxr + pad3 > label_img.shape[0]:
                pad3 = 5
                if maxr + pad3 > label_img.shape[0]:
                    pad3 = 0

            if maxc + pad4 > label_img.shape[1]:
                pad4 = 5
                if maxc + pad4 > label_img.shape[1]:
                    pad4 = 0

            cropped_image = img_test[minr - pad1:maxr + pad3, minc - pad2:maxc + pad4, :]
            cropped_truth = ytruth[minr - pad1:maxr + pad3, minc - pad2:maxc + pad4]
            txcordi = []
            txcordi.append(minr - pad1)
            txcordi.append(maxr + pad3)
            txcordi.append(minc - pad2)
            txcordi.append(maxc + pad4)
            cropped_results['image'] = cropped_image
            cropped_results['truth'] = cropped_truth
            cropped_results['cord'] = txcordi

    return cropped_results


class BasicBlock(nn.Module):
    def __init__(self, channel_num):
        super(BasicBlock, self).__init__()
        # TODO: 3x3 convolution -> relu
        # the input and output channel number is channel_num
        self.conv_block1 = nn.Sequential(
            nn.Conv2d(channel_num, 48, 1, padding=0),
            nn.GroupNorm(num_groups=8, num_channels=48),
            nn.GELU(),
        )
        self.conv_block2 = nn.Sequential(
            nn.Conv2d(48, channel_num, 3, padding=1),
            nn.GroupNorm(num_groups=8, num_channels=channel_num),
            nn.GELU(),
        )
        self.relu = nn.GELU()

    def forward(self, x):
        # TODO: forward
        residual = x
        x = self.conv_block1(x)
        x = self.conv_block2(x)
        x = x + residual
        return x


class ASPP(nn.Module):
    def __init__(self, image_dim=384, head=1):
        super(ASPP, self).__init__()
        self.image_dim = image_dim
        self.Residual2 = BasicBlock(channel_num=head)
        self.pixel_shuffle = nn.PixelShuffle(2)
        self.head = head

    def forward(self, x):
        x21 = F.interpolate(x, size=(self.image_dim, self.image_dim), mode='bilinear',
                            align_corners=True)
        return x21



class Transformer_Regression(nn.Module):
    def __init__(self, image_dim=224, dim_patch=24, num_classes=3, scale=1, feat_dim=192):
        super(Transformer_Regression, self).__init__()
        self.backbone = swin_base_patch4_window12_384_in22k(pretrained=True)
        self.aux = 1
        self.dim_patch = dim_patch
        self.image_dim = image_dim
        self.num_classes = num_classes
        self.ASPP1 = ASPP(image_dim, head=128)
        self.ASPP2 = ASPP(image_dim, head=128)
        # self.ASPP3=ASPP(image_dim,scale,feat_dim)
        self.feat_dim = feat_dim
        # self.scale=1
        self.Classifier_main = nn.Sequential(
            # nn.Dropout(0.1),
            nn.Conv2d(128, self.num_classes, 3, bias=True, padding=1),
        )
        self.Classifier_aux1 = nn.Sequential(
            # nn.Dropout(0.1),
            nn.Conv2d(128, self.num_classes, 3, bias=True, padding=1),
        )

        self.conv1 = nn.Sequential(nn.Conv2d(448, 128, kernel_size=(1, 1), padding=1), nn.GELU())
        self.pixelshufler1 = nn.PixelShuffle(2)
        self.pixelshufler2 = nn.PixelShuffle(4)

    def forward(self, x):
        hide1 = self.backbone(x)
        x1 = []
        x1.append((hide1[0][:, 0:].reshape(-1, 48, 48, 256)))
        x1.append((hide1[1][:, 0:].reshape(-1, 24, 24, 512)))
        x1.append((hide1[2][:, 0:].reshape(-1, 12, 12, 1024)))
        for jk in range(len(x1)):
            x1[jk] = x1[jk].permute(0, 3, 1, 2)
        x1[1] = self.pixelshufler1(x1[1])
        x1[2] = self.pixelshufler2(x1[2])

        x1[0] = torch.cat((x1[0], x1[1], x1[2]), 1)

        x1[0] = self.conv1(x1[0])
        Score = dict()
        x_main1 = self.ASPP1(x1[0])
        x_main = self.Classifier_main(x_main1)
        x_aux_1 = self.ASPP2(x1[0])
        x_aux_1 = self.Classifier_aux1(x_aux_1)  ####### x_aux_1

        Score['seg'] = x_main
        Score['seg_aux_1'] = x_aux_1
        # Score['seg_aux_2'] = x_aux_2

        return Score


Ratios = namedtuple("Ratios", 'cdr hcdr vcdr')
eps = np.finfo(np.float32).eps


def compute_ratios(mask_image):
    '''
    Given an input image containing the cup and disc masks the function returns
    a tuple with the area, horizontal, and vertical cup-to-disc ratios
    Input:
        mask_image: an image with values (0,1,2) or (255,128,0)
                    for bg, disc, cup respectively
    Output:
        Ratios(cdr,hcdr,vcdr): a named tuple containing the computed ratios
    '''

    # if mask_image.max() == 2:
    # make sure correct values are provided in the image
    #    if np.setdiff1d(np.unique(mask_image),np.array([0,1,2])).shape[0]>0:
    #        raise ValueError(('Mask values can only be (0,1,2) '
    #                          'or (255,128,0) for bg, disc, cup'))
    #    disc = np.uint8(mask_image > 0)
    #    cup = np.uint8(mask_image > 1)
    # elif mask_image.max() == 255:
    #    # make sure correct values are provided in the image
    #    if np.setdiff1d(np.unique(mask_image),np.array([0,128,255])).shape[0]>0:
    #        raise ValueError(('Mask values can only be (0,1,2) '
    #                          'or (255,128,0) for bg, disc, cup'))
    #    disc = np.uint8(mask_image < 255)
    #    cup = np.uint8(mask_image == 0)
    # else:
    #    raise ValueError(("Mask values can only be (0,1,2) or (255,128,0) "
    #                      "for bg, disc, cup"))

    # get the area
    disc = 0
    cup = 0
    disc = disc + np.uint8(mask_image > 0)
    cup = cup + np.uint8(mask_image > 1)

    disc_area = np.sum(disc)
    cup_area = np.sum(cup)
    # get the vertical and horizontal mesure of the cup
    cup_vert = np.sum(cup, axis=0).max().astype(np.int32)
    cup_horz = np.sum(cup, axis=1).max().astype(np.int32)
    # get the vertical and horizontal mesure of the disc
    disc_vert = np.sum(disc, axis=0).max().astype(np.int32)
    disc_horz = np.sum(disc, axis=1).max().astype(np.int32)
    # calculate the cup to disc ratio
    cdr = (cup_area + eps) / (disc_area + eps)  # add eps to avoid div by 0
    # calculate the horizontal and vertical cup to disc ration
    hcdr = (cup_horz + eps) / (disc_horz + eps)
    vcdr = (cup_vert + eps) / (disc_vert + eps)

    return Ratios(cdr, hcdr, vcdr)