#!/usr/bin/env python 
# -*- coding: utf-8 -*-
"""
@project: 
@File    : resunet
@Author  : qiqq
@create_time    : 2023/7/20 18:45
"""
import torch
import torch.nn as nn

from taihuyuan_pv.mitunet.model.resnet2 import resnet50
from taihuyuan_pv.mitunet.model.decoder2 import *



class unetUp(nn.Module):
    def __init__(self, in_size, out_size):
        super(unetUp, self).__init__()

        self.up = nn.UpsamplingBilinear2d(scale_factor=2)

        self.cbr = nn.Sequential(
            nn.Conv2d(in_size, out_size, 3, 1, 1,  bias=False),
            nn.BatchNorm2d(out_size),
            nn.ReLU(inplace=True)
        )



    def forward(self, inputs1, inputs2):
        outputs = torch.cat([inputs1, self.up(inputs2)], 1)
        outputs = self.cbr(outputs)

        return outputs


class resUnet(nn.Module):
    def __init__(self, num_classes=2, pretrained=True, backbone='resnet50'):
        super(resUnet, self).__init__()

        self.nclas = num_classes
        self.finnal_channel = 512
        self.backbone = resnet50(pretrained=pretrained)
        self.decoder = unetDecoder(in_filters=[ 512, 1024, 3072],out_filters=[128, 256, 512])

    def forward(self, inputs):

        feaureslist = self.backbone(inputs)  #2  4  8 16 32
        feaureslist=feaureslist[1:]
        out = self.decoder(feaureslist)
        out = F.interpolate(out, size=inputs.size()[2:], mode='bilinear', align_corners=True)
        return out



class resUnetPAM(nn.Module):
    def __init__(self, num_classes=2, pretrained=True, backbone='resnet50'):
        super(resUnetPAM, self).__init__()

        self.nclas = num_classes
        self.backbone = resnet50(pretrained=pretrained)
        self.decoder = unetpamDecoder(in_filters=[ 512, 1024, 1536],out_filters=[128, 256, 512])

    def forward(self, inputs):

        feaureslist = self.backbone(inputs)  #2  4  8 16 32
        feaureslist=feaureslist[1:]
        out = self.decoder(feaureslist)
        out = F.interpolate(out, size=inputs.size()[2:], mode='bilinear', align_corners=True)
        return out





class resUnetpamcarb_d16(nn.Module):
    '''除了下采样变成16倍数 其他都没变'''
    def __init__(self, num_classes=2, pretrained=True,replace_stride_with_dilation=[False,False,True] ):
        super(resUnetpamcarb_d16, self).__init__()

        self.nclas = num_classes
        self.finnal_channel = 512
        self.backbone = resnet50(pretrained=pretrained,replace_stride_with_dilation=replace_stride_with_dilation)
        self.decoder = unetpamCARBDecoder()

    def forward(self, inputs):

        feaureslist = self.backbone(inputs)  #2  4  8 16 32
        feaureslist=feaureslist[1:]
        out = self.decoder(feaureslist)
        out = F.interpolate(out, size=inputs.size()[2:], mode='bilinear', align_corners=True)
        return out


class resUnetpamcam(nn.Module):
    def __init__(self, num_classes=2, pretrained=True, backbone='resnet50'):
        super(resUnetpamcam, self).__init__()

        self.nclas = num_classes
        self.finnal_channel = 512
        self.backbone = resnet50(pretrained=pretrained)
        self.decoder = unetpamcamDecoder()

    def forward(self, inputs):

        feaureslist = self.backbone(inputs)  #2  4  8 16 32
        feaureslist=feaureslist[1:]
        out = self.decoder(feaureslist)
        out = F.interpolate(out, size=inputs.size()[2:], mode='bilinear', align_corners=True)
        return out



# 1/4  1/8 1/16 1/32  记作 1 2 3 4
class resUnetpamcarb_4(nn.Module):
    def __init__(self, num_classes=2, pretrained=True, backbone='resnet50'):
        super(resUnetpamcarb_4, self).__init__()

        self.nclas = num_classes
        self.finnal_channel = 512
        self.backbone = resnet50(pretrained=pretrained)
        self.decoder = unetpamDecoderzuhe()

    def forward(self, inputs):

        feaureslist = self.backbone(inputs)  #2  4  8 16 32
        feaureslist=feaureslist[1:]
        out = self.decoder(feaureslist)
        out = F.interpolate(out, size=inputs.size()[2:], mode='bilinear', align_corners=True)
        return out





if __name__ == '__main__':
    indd=torch.rand(2,3,512,512)
    modl=resUnetpamcarb_d16()
    out=modl(indd)
    print(type(out))