141 lines
4.2 KiB
Python
141 lines
4.2 KiB
Python
#!/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)) |