unet改进笔记

论文：U-Net: Convolutional Networks for Biomedical Image Segmentation

github：https://github.com/milesial/Pytorch-UNet

改进1：数据增强，新增augmentations.py

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
Image augmentation functions
"""

import math
import random

import cv2
import numpy as np
import imgaug.augmenters as iaa



class Albumentations:
    # YOLOv5 Albumentations class (optional, only used if package is installed)
    def __init__(self):
        self.transform = None
        try:
            #version 1.0.3
            import albumentations as A

            T = [
                #A.OneOf([
                    #A.IAAAdditiveGaussianNoise(),   # 将高斯噪声添加到输入图像
                    #A.GaussNoise(),    # 将高斯噪声应用于输入图像。
                #], p=0.2),   # 应用选定变换的概率
                A.OneOf([
                    A.MotionBlur(p=0.2),   # 使用随机大小的内核将运动模糊应用于输入图像。
                    A.MedianBlur(blur_limit=3, p=0.01),    # 中值滤波
                    A.Blur(blur_limit=3, p=0.01),   # 使用随机大小的内核模糊输入图像。
                ], p=0.2),
                # 随机应用仿射变换：平移，缩放和旋转输入
                A.RandomBrightnessContrast(p=0.2),   # 随机明亮对比度
                A.CLAHE(p=0.01),
                A.RandomGamma(p=0.0),
                A.ImageCompression(quality_lower=75, p=0.0)]  # transforms
            self.transform = A.Compose(T)

            print('albumentations: ' + ', '.join(f'{x}' for x in self.transform.transforms if x.p))
        except ImportError:  # package not installed, skip
            pass
        except Exception as e:
            print('albumentations: '+ f'{e}')

    def __call__(self, im, p=0.8):
        if self.transform and random.random() < p:
            new = self.transform(image=im)  # transformed
            im = new['image']


        if random.random() > p:
            im = augment_hsv(im, hgain=0.5, sgain=0.5, vgain=0.5)
        if random.random() > p:
            im = hist_equalize(im, clahe=True, bgr=True) 
            
            
            
            
        #数据增强

        range_num = random.randint(1, 2)
        for iii in range(range_num):
            rand_num = random.randint(0, 20)
            if rand_num == 0: 
                im = channel_shuffle(im)
            #elif rand_num == 1:
                #im = random_noise(im, limit=[0, 0.2], p=0.5)
            elif rand_num == 2:
                im = random_brightness(im, brightness=0.3)
            elif rand_num == 3:
                im = random_contrast(im,contrast=0.3)
            elif rand_num == 4:
                im = random_saturation(im, saturation=0.5)
            elif rand_num == 5:
                im = EqExtension(im)
            elif rand_num == 6:
                im = HueExtension(im)
            elif rand_num == 7:
                im = zaodian(im)
            elif rand_num == 8:
                im = superpixelsaug(im)
            elif rand_num == 9:
                im = fogaug(im)
            elif rand_num == 10:
                im = cloudsaug(im)
            #elif rand_num == 11:
                #im = fnaug(im)
            elif rand_num == 11:
                im = Coarseaug(im)
            elif rand_num == 12:
                im = random_hue(im)
            else :
                im = im
        ####aug
        #数据增强
            

        return im


def augment_hsv(im, hgain=0.5, sgain=0.5, vgain=0.5):
    # HSV color-space augmentation
    if hgain or sgain or vgain:
        r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains
        hue, sat, val = cv2.split(cv2.cvtColor(im, cv2.COLOR_BGR2HSV))
        dtype = im.dtype  # uint8

        x = np.arange(0, 256, dtype=r.dtype)
        lut_hue = ((x * r[0]) % 180).astype(dtype)
        lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
        lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

        im_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
        return cv2.cvtColor(im_hsv, cv2.COLOR_HSV2BGR)  # no return needed


def hist_equalize(im, clahe=True, bgr=True):
    # Equalize histogram on BGR image 'im' with im.shape(n,m,3) and range 0-255
    yuv = cv2.cvtColor(im, cv2.COLOR_BGR2YUV if bgr else cv2.COLOR_RGB2YUV)
    if clahe:
        c = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        yuv[:, :, 0] = c.apply(yuv[:, :, 0])
    else:
        yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0])  # equalize Y channel histogram
    return cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR if bgr else cv2.COLOR_YUV2RGB)  # convert YUV image to RGB




def channel_shuffle(img):
    if(img.shape[2] == 3):
        ch_arr = [0, 1, 2]
        np.random.shuffle(ch_arr)
        img = img[..., ch_arr]
    return img




def EqExtension(src):
    I_backup = src.copy()
    b, g, r = cv2.split(I_backup)
    b = cv2.equalizeHist(b)
    g = cv2.equalizeHist(g)
    r = cv2.equalizeHist(r)
    I_eq = cv2.merge([b, g, r])
    return I_eq.astype(np.uint8)


def HueExtension(src):
    img_hsv = cv2.cvtColor(src, cv2.COLOR_BGR2HSV)
    cre = random.randint(90, 100)
    cre = float(cre) / 100
    img_hsv[:,:,2] = img_hsv[:,:,2] * cre
 
    # print(img_hsv[:,:,0])
    dst = cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR)
    return dst.astype(np.uint8)




# 随机生成500个椒盐噪声
def zaodian(img):
    height,weight,channel = img.shape
    img_zao = img.copy()
    for i in range(500):
        x = np.random.randint(0,height)
        y = np.random.randint(0,weight)
        img_zao[x ,y ,:] = 255
    return img_zao.astype(np.uint8)


#def random_noise(img, limit=[0, 0.2], p=0.5):
    #if random.random() < p:
        #H, W = img.shape[:2]
        #noise = np.random.uniform(limit[0], limit[1], size=(H,W)) * 255

        #img = img + noise[:,:,np.newaxis]*np.array([1,1,1])
        #img = np.clip(img, 0, 255).astype(np.uint8)
        
    #return img

def random_brightness(image, brightness=0.3):
    alpha = 1 + np.random.uniform(-brightness, brightness)
    img = alpha * image
    img = np.clip(img, 0, 255).astype(np.uint8)
    return img

def random_contrast(img,contrast=0.3):
    coef = np.array([[[0.114, 0.587,  0.299]]])   # rgb to gray (YCbCr)
    alpha = 1.0 + np.random.uniform(-contrast, contrast)
    gray = img * coef
    gray = (3.0 * (1.0 - alpha) / gray.size) * np.sum(gray)
    img = alpha*img  + gray
    img = np.clip(img, 0, 255).astype(np.uint8)
    return img

def random_saturation(img, saturation=0.5):
    coef = np.array([[[0.299, 0.587, 0.114]]])
    alpha = np.random.uniform(-saturation, saturation)
    gray  = img * coef
    gray  = np.sum(gray,axis=2, keepdims=True)
    img = alpha*img  + (1.0 - alpha)*gray
    img = np.clip(img, 0, 255).astype(np.uint8)
    return img

def random_hue(image,hue=0.5):
    h = int(np.random.uniform(-hue, hue)*180)

    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    hsv[:, :, 0] = (hsv[:, :, 0].astype(int) + h) % 180
    image = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    return image



def superpixelsaug(img):
    images = np.expand_dims(img, axis=0)
    aug = iaa.Superpixels(p_replace=0.005,max_size=50)
    images_aug = aug(images=images)
    img_aug = np.squeeze(images_aug)
    return img_aug.astype(np.uint8)



def fogaug(img):
    images = np.expand_dims(img, axis=0)
    aug = iaa.Fog()
    images_aug = aug(images=images)
    img_aug = np.squeeze(images_aug)
    return img_aug.astype(np.uint8)


def cloudsaug(img):
    images = np.expand_dims(img, axis=0)
    aug = iaa.Clouds()
    images_aug = aug(images=images)
    img_aug = np.squeeze(images_aug)
    return img_aug.astype(np.uint8)
 
#def fnaug(img):
    #images = np.expand_dims(img, axis=0)
    #aug = iaa.FrequencyNoiseAlpha(first=iaa.EdgeDetect(0.5))
    #images_aug = aug(images=images)
    #img_aug = np.squeeze(images_aug)
    #return img_aug.astype(np.uint8)


def Coarseaug(img):
    images = np.expand_dims(img, axis=0)
    aug = iaa.CoarseDropout(0.02, size_percent=0.5)
    images_aug = aug(images=images)
    img_aug = np.squeeze(images_aug)
    return img_aug.astype(np.uint8)

调用方式，在utils/data_loading.py中__getitem__函数增加

#albumentations
            if self.albumentations is not None and self.mode == 'train':  
                image = self.albumentations(image)
            #albumentations

改进2：基于图片拼接的数据增强，需要修改图片和标签，在utils/data_loading.py中__getitem__函数增加

def __getitem__(self, idx):
        #训练数据使用bgr格式
        max_height = 0
        max_width = 0
        height_up_left, width_up_left, height_up_right, width_up_right, height_down_left, width_down_left, height_down_right, width_down_right = 0,0,0,0,0,0,0,0
        num1, num2, num3, num4 = idx, random.randint(0,self.img_num-1),random.randint(0,self.img_num-1),random.randint(0,self.img_num-1)
        
        
        img_up_left = cv2.imread(self.img_list[num1], 1)
        #随机4个图片拼接为一个大图片
        if self.train:
            img_up_right  = cv2.imread(self.img_list[num2], 1)
            img_down_left  = cv2.imread(self.img_list[num3], 1)
            img_down_right  = cv2.imread(self.img_list[num4], 1)
        
            height_up_left, width_up_left = img_up_left.shape[:2]
            height_up_right, width_up_right = img_up_right.shape[:2]
            height_down_left, width_down_left = img_down_left.shape[:2]
            height_down_right, width_down_right = img_down_right.shape[:2]
            
            max_height = max([height_up_left, height_up_right, height_down_left, height_down_right])
            max_width = max([width_up_left, width_up_right, width_down_left, width_down_right])
            
            #假设训练数据大小不一样，类比fasterrcnn的方式，采用最大长宽的进行拼接，其余补黑边
            for num, im in enumerate([img_up_left, img_up_right, img_down_left, img_down_right]):
                img_template = np.zeros((max_height, max_width, 3),np.uint8)
                if num ==0: 
                    img_template[:height_up_left, :width_up_left, :] = im.copy()
                    img_up_left = img_template.copy()
                elif num==1:
                    img_template[:height_up_right, :width_up_right, :] = im.copy()
                    img_up_right = img_template.copy()
                elif num==2:
                    img_template[:height_down_left, :width_down_left, :] = im.copy()
                    img_down_left = img_template.copy()
                elif num==3:
                    img_template[:height_down_right, :width_down_right, :] = im.copy()
                    img_down_right = img_template.copy()
                else:
                    pass
                 
            img = np.vstack([np.hstack([img_up_left,img_up_right]), np.hstack([img_down_left,img_down_right])])
        else:
            img = img_up_left
        

        
       
        
        

        mask_up_left = cv2.imread(self.mask[num1], 0)
        if self.train:
            mask_up_right = cv2.imread(self.mask[num2], 0)
            mask_down_left = cv2.imread(self.mask[num3], 0)
            mask_down_right = cv2.imread(self.mask[num4], 0)
        
        
            #假设训练数据大小不一样，类比fasterrcnn的方式，采用最大长宽的进行拼接，其余补黑边
            for num, ma in enumerate([mask_up_left, mask_up_right, mask_down_left, mask_down_right]):
                mask_template = np.zeros((max_height, max_width), np.uint8)
                if num ==0: 
                    mask_template[:height_up_left, :width_up_left] = ma.copy()
                    mask_up_left = mask_template.copy()
                elif num==1:
                    mask_template[:height_up_right, :width_up_right] = ma.copy()
                    mask_up_right = mask_template.copy()
                elif num==2:
                    mask_template[:height_down_left, :width_down_left] = ma.copy()
                    mask_down_left = mask_template.copy()
                elif num==3:
                    mask_template[:height_down_right, :width_down_right] = ma.copy()
                    mask_down_right = mask_template.copy()
                else:
                    pass
        
            mask = np.vstack([np.hstack([mask_up_left,mask_up_right]), np.hstack([mask_down_left,mask_down_right])])
        else:
            mask = mask_up_left
        mask = mask/ 255.0  # 前景==255，背景==0
        
        
        if self.transforms is not None:
            #opencv to pillow
            img= Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
            mask = Image.fromarray(mask)
            img, mask = self.transforms(img, mask)
            
         

        return img, mask

改进3：训练图片分布的随机扰动，在utils/data_loading.py中__getitem__函数增加

if self.train:
                #随机乘加操作
                img = img * random.randint(90,110)/100 + random.randint(-10,10)/100
                #随机乘加操作

改进4：增加预测推理代码，

import os
import time
import cv2

import torch
from torchvision import transforms
import numpy as np
from PIL import Image

from src import UNet




class ThyroidNoduleDetect():
    def __init__(self, weights_path = "save_weights/N-net3_weight_best.pth"):
        #超声波甲状腺结节检测
        self.use_cuda = True
        self.use_float16 = False
        # mean = (0.302, 0.302, 0.302)
        # std = (0.093, 0.093, 0.093)
        self.mean = [0.326, 0.327, 0.354]
        self.std = [0.149, 0.149, 0.157]

        # get devices
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        print("using {} device.".format(device))

        # create model
        classes = 1  # exclude background
        self.model = UNet(in_channels=3, num_classes=classes+1, base_c=32)

        # load weights
        self.model.load_state_dict(torch.load(weights_path, map_location='cpu')['model'])
        self.model.to(device)
        self.model.requires_grad_(False)
        self.model.eval()  # 进入验证模式

    def process(self, image):
    
        x = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
        x = (x/np.max(x) - self.mean)/self.std
        x = np.expand_dims(x, axis=0)

        if self.use_cuda:
            x = torch.from_numpy(x).cuda()
        else:
            x = torch.from_numpy(x)
    
        x = x.to(torch.float32 if not self.use_float16 else torch.float16).permute(0, 3, 1, 2)

        t1 = time.time()
        with torch.no_grad():
            output = self.model(x)
            
            prediction = output['out'].argmax(1).squeeze(0)
            prediction = prediction.to("cpu").numpy().astype(np.uint8)
            prediction_mask = prediction*255
            
         

        t2 = time.time()
        tact_time = (t2 - t1)
        print("ThyroidNoduleDetect ",f'{ tact_time} seconds, {1 / tact_time} FPS, @batch_size 1')
        return prediction_mask


    def draw_circles(self, image, prediction_mask):
        _,binary = cv2.threshold(prediction_mask,127,255,cv2.THRESH_BINARY) 
        contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) 
        cv2.drawContours(image,contours,-1,(0,0,255),3)
        
        for index,c in enumerate(contours):
            M = cv2.moments(c)
            print(index, M)
            if M["m00"]!=0.0:
                #计算重心
                cX = int(M["m10"] / M["m00"])
                cY = int(M["m01"] / M["m00"])
                cv2.circle(image, (cX, cY), 4, (0, 255, 0), -1)

        return image

def test_image():
    img_path = "TNUI/test/images/1079.png"
    roi_mask_path = "TNUI/test/masks/1079.png"
    image =cv2.imread(img_path,1)
    
    tnd = ThyroidNoduleDetect()
    prediction_mask = tnd.process(image)
    cv2.imwrite("pre_result/N-net2_test_1079.png", prediction_mask)


def eval_images():
    tnd = ThyroidNoduleDetect()
    image_path = "TNUI/test/images/"
    mask_path = "TNUI/test/masks/"

    for name in  os.listdir(image_path):
        image = cv2.imread(os.path.join(image_path,name),1)
        mask = cv2.imread(os.path.join(mask_path,name),1)
        
        prediction_mask = tnd.process(image)
        prediction_mask_bgr = cv2.cvtColor(prediction_mask,cv2.COLOR_GRAY2BGR)
        image = tnd.draw_circles(image, prediction_mask)
        
        cv2.imwrite(os.path.join("pre_result/",name), np.hstack([image, prediction_mask_bgr, mask]))




if __name__ == '__main__':
    test_image()
    #eval_images()

改进5：MobileV3Unet，新增mobilenet_unet.py

from collections import OrderedDict
from typing import Dict
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torchvision.models import mobilenet_v3_large
from .unet import Up, OutConv


class IntermediateLayerGetter(nn.ModuleDict):
    """
    Module wrapper that returns intermediate layers from a model

    It has a strong assumption that the modules have been registered
    into the model in the same order as they are used.
    This means that one should **not** reuse the same nn.Module
    twice in the forward if you want this to work.

    Additionally, it is only able to query submodules that are directly
    assigned to the model. So if `model` is passed, `model.feature1` can
    be returned, but not `model.feature1.layer2`.

    Args:
        model (nn.Module): model on which we will extract the features
        return_layers (Dict[name, new_name]): a dict containing the names
            of the modules for which the activations will be returned as
            the key of the dict, and the value of the dict is the name
            of the returned activation (which the user can specify).
    """
    _version = 2
    __annotations__ = {
        "return_layers": Dict[str, str],
    }

    def __init__(self, model: nn.Module, return_layers: Dict[str, str]) -> None:
        if not set(return_layers).issubset([name for name, _ in model.named_children()]):
            raise ValueError("return_layers are not present in model")
        orig_return_layers = return_layers
        return_layers = {str(k): str(v) for k, v in return_layers.items()}

        # 重新构建backbone，将没有使用到的模块全部删掉
        layers = OrderedDict()
        for name, module in model.named_children():
            layers[name] = module
            if name in return_layers:
                del return_layers[name]
            if not return_layers:
                break

        super(IntermediateLayerGetter, self).__init__(layers)
        self.return_layers = orig_return_layers

    def forward(self, x: Tensor) -> Dict[str, Tensor]:
        out = OrderedDict()
        for name, module in self.items():
            x = module(x)
            if name in self.return_layers:
                out_name = self.return_layers[name]
                out[out_name] = x
        return out


class MobileV3Unet(nn.Module):
    def __init__(self, num_classes, pretrain_backbone: bool = False):
        super(MobileV3Unet, self).__init__()
        backbone = mobilenet_v3_large(pretrained=pretrain_backbone)

        # if pretrain_backbone:
        #     # 载入mobilenetv3 large backbone预训练权重
        #     # https://download.pytorch.org/models/mobilenet_v3_large-8738ca79.pth
        #     backbone.load_state_dict(torch.load("mobilenet_v3_large.pth", map_location='cpu'))

        backbone = backbone.features

        stage_indices = [1, 3, 6, 12, 15]
        self.stage_out_channels = [backbone[i].out_channels for i in stage_indices]
        return_layers = dict([(str(j), f"stage{i}") for i, j in enumerate(stage_indices)])
        self.backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)

        c = self.stage_out_channels[4] + self.stage_out_channels[3]
        self.up1 = Up(c, self.stage_out_channels[3])
        c = self.stage_out_channels[3] + self.stage_out_channels[2]
        self.up2 = Up(c, self.stage_out_channels[2])
        c = self.stage_out_channels[2] + self.stage_out_channels[1]
        self.up3 = Up(c, self.stage_out_channels[1])
        c = self.stage_out_channels[1] + self.stage_out_channels[0]
        self.up4 = Up(c, self.stage_out_channels[0])
        self.conv = OutConv(self.stage_out_channels[0], num_classes=num_classes)

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        input_shape = x.shape[-2:]
        backbone_out = self.backbone(x)
        x = self.up1(backbone_out['stage4'], backbone_out['stage3'])
        x = self.up2(x, backbone_out['stage2'])
        x = self.up3(x, backbone_out['stage1'])
        x = self.up4(x, backbone_out['stage0'])
        x = self.conv(x)
        x = F.interpolate(x, size=input_shape, mode="bilinear", align_corners=False)

        return {"out": x}

改进6：VGG16UNet，新增vgg_unet.py

from collections import OrderedDict
from typing import Dict

import torch
import torch.nn as nn
from torch import Tensor
from torchvision.models import vgg16_bn
from .unet import Up, OutConv


class IntermediateLayerGetter(nn.ModuleDict):
    """
    Module wrapper that returns intermediate layers from a model

    It has a strong assumption that the modules have been registered
    into the model in the same order as they are used.
    This means that one should **not** reuse the same nn.Module
    twice in the forward if you want this to work.

    Additionally, it is only able to query submodules that are directly
    assigned to the model. So if `model` is passed, `model.feature1` can
    be returned, but not `model.feature1.layer2`.

    Args:
        model (nn.Module): model on which we will extract the features
        return_layers (Dict[name, new_name]): a dict containing the names
            of the modules for which the activations will be returned as
            the key of the dict, and the value of the dict is the name
            of the returned activation (which the user can specify).
    """
    _version = 2
    __annotations__ = {
        "return_layers": Dict[str, str],
    }

    def __init__(self, model: nn.Module, return_layers: Dict[str, str]) -> None:
        if not set(return_layers).issubset([name for name, _ in model.named_children()]):
            raise ValueError("return_layers are not present in model")
        orig_return_layers = return_layers
        return_layers = {str(k): str(v) for k, v in return_layers.items()}

        # 重新构建backbone，将没有使用到的模块全部删掉
        layers = OrderedDict()
        for name, module in model.named_children():
            layers[name] = module
            if name in return_layers:
                del return_layers[name]
            if not return_layers:
                break

        super(IntermediateLayerGetter, self).__init__(layers)
        self.return_layers = orig_return_layers

    def forward(self, x: Tensor) -> Dict[str, Tensor]:
        out = OrderedDict()
        for name, module in self.items():
            x = module(x)
            if name in self.return_layers:
                out_name = self.return_layers[name]
                out[out_name] = x
        return out


class VGG16UNet(nn.Module):
    def __init__(self, num_classes, pretrain_backbone: bool = False):
        super(VGG16UNet, self).__init__()
        backbone = vgg16_bn(pretrained=pretrain_backbone)

        # if pretrain_backbone:
        #     # 载入vgg16_bn预训练权重
        #     # https://download.pytorch.org/models/vgg16_bn-6c64b313.pth
        #     backbone.load_state_dict(torch.load("vgg16_bn.pth", map_location='cpu'))

        backbone = backbone.features

        stage_indices = [5, 12, 22, 32, 42]
        self.stage_out_channels = [64, 128, 256, 512, 512]
        return_layers = dict([(str(j), f"stage{i}") for i, j in enumerate(stage_indices)])
        self.backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)

        c = self.stage_out_channels[4] + self.stage_out_channels[3]
        self.up1 = Up(c, self.stage_out_channels[3])
        c = self.stage_out_channels[3] + self.stage_out_channels[2]
        self.up2 = Up(c, self.stage_out_channels[2])
        c = self.stage_out_channels[2] + self.stage_out_channels[1]
        self.up3 = Up(c, self.stage_out_channels[1])
        c = self.stage_out_channels[1] + self.stage_out_channels[0]
        self.up4 = Up(c, self.stage_out_channels[0])
        self.conv = OutConv(self.stage_out_channels[0], num_classes=num_classes)

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        backbone_out = self.backbone(x)
        x = self.up1(backbone_out['stage4'], backbone_out['stage3'])
        x = self.up2(x, backbone_out['stage2'])
        x = self.up3(x, backbone_out['stage1'])
        x = self.up4(x, backbone_out['stage0'])
        x = self.conv(x)

        return {"out": x}

改进7：HNet，新增HNet.py

from functools import partial
from typing import Dict
import torch
import torch.nn as nn
import torch.nn.functional as F

nonlinearity = partial(F.relu, inplace=True)


class SDCBlock(nn.Module):
    def __init__(self, channel):
        super(SDCBlock, self).__init__()
        self.dilate1 = nn.Conv2d(channel, channel, kernel_size=3, dilation=1, padding=1)
        self.dilate2 = nn.Conv2d(channel, channel, kernel_size=3, dilation=3, padding=3)
        self.dilate3 = nn.Conv2d(channel, channel, kernel_size=3, dilation=11, padding=11)
        self.conv1x1 = nn.Conv2d(channel, channel, kernel_size=1, dilation=1, padding=0)

    def forward(self, x):
        dilate1_out = nonlinearity(self.dilate1(x))
        dilate2_out = nonlinearity(self.conv1x1(self.dilate2(x)))
        dilate3_out = nonlinearity(self.conv1x1(self.dilate2(self.dilate1(x))))
        dilate4_out = nonlinearity(self.conv1x1(self.dilate2(self.dilate1(self.dilate3(x)))))
        out = x + dilate1_out + dilate2_out + dilate3_out + dilate4_out
        return out


class DoubleConv(nn.Sequential):
    def __init__(self, in_channels, out_channels, mid_channels=None):
        if mid_channels is None:
            mid_channels = out_channels
        super(DoubleConv, self).__init__(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )


class SingleConv(nn.Sequential):
    def __init__(self, in_channels, out_channels, mid_channels=None):
        if mid_channels is None:
            mid_channels = out_channels
        super(SingleConv, self).__init__(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
        )


class Up(nn.Module):
    def __init__(self, in_channels, out_channels, bilinear=True):
        super(Up, self).__init__()
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        x1 = self.up(x1)
        # [N, C, H, W]
        diff_y = x2.size()[2] - x1.size()[2]
        diff_x = x2.size()[3] - x1.size()[3]

        # padding_left, padding_right, padding_top, padding_bottom
        x1 = F.pad(x1, [diff_x // 2, diff_x - diff_x // 2,
                        diff_y // 2, diff_y - diff_y // 2])

        x = torch.cat([x2, x1], dim=1)
        x = self.conv(x)
        return x


class OutConv(nn.Sequential):
    def __init__(self, in_channels, num_classes):
        super(OutConv, self).__init__(
            nn.Conv2d(in_channels, num_classes, kernel_size=1)
        )

class Down(nn.Sequential):
    def __init__(self, in_channels, out_channels):
        super(Down, self).__init__(
            nn.MaxPool2d(2, stride=2),
            DoubleConv(in_channels, out_channels)
        )


class DownLast(nn.Sequential):
    def __init__(self, in_channels, out_channels):
        super(DownLast, self).__init__(
            nn.MaxPool2d(2, stride=2),
            SingleConv(in_channels, out_channels)
        )

class UNet(nn.Module):
    def __init__(self,
                 in_channels: int = 1,
                 num_classes: int = 2,
                 bilinear: bool = True,
                 base_c: int = 64):
        super(UNet, self).__init__()
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.bilinear = bilinear

        self.in_conv = DoubleConv(in_channels, base_c)
        self.down1 = Down(base_c, base_c * 2)
        self.down2 = Down(base_c * 2, base_c * 4)
        self.down3 = Down(base_c * 4, base_c * 8)
        factor = 2 if bilinear else 1
        self.down4 = DownLast(base_c * 8, base_c * 16 // factor)
        # self.down4 = Down(base_c * 8, base_c * 16 // factor)
        self.sdc = SDCBlock(base_c * 8)
        self.singleconv = SingleConv(base_c * 8, base_c * 16 // factor)
        self.up1 = Up(base_c * 16, base_c * 8 // factor, bilinear)
        self.up2 = Up(base_c * 8, base_c * 4 // factor, bilinear)
        self.up3 = Up(base_c * 4, base_c * 2 // factor, bilinear)
        self.up4 = Up(base_c * 2, base_c, bilinear)
        self.out_conv = OutConv(base_c, num_classes)

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        x1 = self.in_conv(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.singleconv(self.sdc(self.down4(x4)))
        # x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.out_conv(x)

        return {"out": logits}