UnitPrint/scan-DocRes
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

init + inference.py патч

Browse Source
This commit is contained in:
TraceLumen
2026-09-12 03:29:38 +03:00
commit 6f7f6031d1
50 changed files with 7410 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

742
data/MBD/model/gienet.py Normal file
View File

@ -0,0 +1,742 @@
from math import log
import torch
import torch.nn as nn
from torch.nn import init
import functools
from model.cbam import CBAM
# Defines the Unet generator.
# |num_downs|: number of downsamplings in UNet. For example,
# if |num_downs| == 7, image of size 128x128 will become of size 1x1
# at the bottleneck
class SingleConv(nn.Module):
"""(convolution => [BN] => ReLU) * 2"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.double_conv = nn.Sequential(
nn.ReflectionPad2d(1),
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0,stride=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
# nn.ReflectionPad2d(1),
# nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0,stride=1),
# nn.BatchNorm2d(out_channels),
# nn.ReLU(inplace=True)
)
def forward(self, x):
return self.double_conv(x)
class Down_single(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.maxpool_conv = nn.Sequential(
nn.MaxPool2d(2),
SingleConv(in_channels, out_channels)
)
def forward(self, x):
return self.maxpool_conv(x)
class Up_single(nn.Module):
"""Upscaling then double conv"""
def __init__(self, in_channels, out_channels, bilinear=True):
super().__init__()
self.up = nn.Upsample(scale_factor=2, mode='nearest')
self.conv = SingleConv(in_channels, out_channels)
self.deconv = nn.ConvTranspose2d(in_channels, out_channels,kernel_size=4, stride=2,padding=1, bias=True)
def forward(self, x1, x2):
x1 = self.deconv(x1)
# input is BCHW
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
class DoubleConv(nn.Module):
"""(convolution => [BN] => ReLU) * 2"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.double_conv = nn.Sequential(
nn.ReflectionPad2d(1),
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0,stride=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.ReflectionPad2d(1),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0,stride=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.double_conv(x)
class Down(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.maxpool_conv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_channels, out_channels)
)
def forward(self, x):
return self.maxpool_conv(x)
class Up(nn.Module):
"""Upscaling then double conv"""
def __init__(self, in_channels, out_channels, bilinear=True):
super().__init__()
self.up = nn.Upsample(scale_factor=2, mode='nearest')
self.conv = DoubleConv(in_channels, out_channels)
self.deconv = nn.ConvTranspose2d(in_channels, out_channels,kernel_size=4, stride=2,padding=1, bias=True)
def forward(self, x1, x2):
x1 = self.deconv(x1)
# input is BCHW
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
class OutConv(nn.Module):
def __init__(self, in_channels, out_channels):
super(OutConv, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
self.tanh = nn.Tanh()
self.hardtanh = nn.Hardtanh()
self.sigmoid = nn.Sigmoid()
def forward(self, x1):
x = self.conv(x1)
# x = self.sigmoid(x)
# x = self.hardtanh(x)
# x = (x+1)/2
return x
class GiemaskGenerator(nn.Module):
"""Create a Unet-based generator"""
def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
"""Construct a Unet generator
Parameters:
input_nc (int) -- the number of channels in input images
output_nc (int) -- the number of channels in output images
num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
image of size 128x128 will become of size 1x1 # at the bottleneck
ngf (int) -- the number of filters in the last conv layer
norm_layer -- normalization layer
We construct the U-Net from the innermost layer to the outermost layer.
It is a recursive process.
"""
super(GiemaskGenerator, self).__init__()
self.init_channel =32
self.inc = DoubleConv(3,self.init_channel)
self.down1 = Down(self.init_channel, self.init_channel*2)
self.down2 = Down(self.init_channel*2, self.init_channel*4)
self.down3 = Down(self.init_channel*4, self.init_channel*8)
self.down4 = Down(self.init_channel*8, self.init_channel*16)
self.down5 = Down(self.init_channel*16, self.init_channel*32)
self.up1 = Up(self.init_channel*32, self.init_channel*16)
self.up2 = Up(self.init_channel*16, self.init_channel*8)
self.up3 = Up(self.init_channel*8, self.init_channel*4)
self.up4 = Up(self.init_channel*4,self.init_channel*2)
self.up5 = Up(self.init_channel*2, self.init_channel)
self.outc = OutConv(self.init_channel, 1)
self.up1_1 = Up_single(self.init_channel*32, self.init_channel*16)
self.up2_1 = Up_single(self.init_channel*16, self.init_channel*8)
self.up3_1 = Up_single(self.init_channel*8, self.init_channel*4)
self.up4_1 = Up_single(self.init_channel*4,self.init_channel*2)
self.up5_1 = Up_single(self.init_channel*2, self.init_channel)
self.outc_1 = OutConv(self.init_channel, 1)
# self.dropout = nn.Dropout(p=0.5)
def forward(self, input):
x1 = self.inc(input)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x6 = self.down5(x5)
x_1 = self.up1_1(x6, x5)
x_1 = self.up2_1(x_1, x4)
x_1 = self.up3_1(x_1, x3)
x_1 = self.up4_1(x_1, x2)
x_1 = self.up5_1(x_1, x1)
mask = self.outc_1(x_1)
x = self.up1(x6, x5)
# x = self.dropout(x)
x = self.up2(x, x4)
# x = self.dropout(x)
x = self.up3(x, x3)
# x = self.dropout(x)
x = self.up4(x, x2)
# x = self.dropout(x)
x = self.up5(x, x1)
# x = self.dropout(x)
depth = self.outc(x)
return depth,mask
"""Create a Unet-based generator"""
class Giemask2Generator(nn.Module):
"""Create a Unet-based generator"""
def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
"""Construct a Unet generator
Parameters:
input_nc (int) -- the number of channels in input images
output_nc (int) -- the number of channels in output images
num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
image of size 128x128 will become of size 1x1 # at the bottleneck
ngf (int) -- the number of filters in the last conv layer
norm_layer -- normalization layer
We construct the U-Net from the innermost layer to the outermost layer.
It is a recursive process.
"""
super(Giemask2Generator, self).__init__()
self.init_channel =32
self.inc = DoubleConv(3,self.init_channel)
self.down1 = Down(self.init_channel, self.init_channel*2)
self.down2 = Down(self.init_channel*2, self.init_channel*4)
self.down3 = Down(self.init_channel*4, self.init_channel*8)
self.down4 = Down(self.init_channel*8, self.init_channel*16)
self.down5 = Down(self.init_channel*16, self.init_channel*32)
self.up1 = Up(self.init_channel*32, self.init_channel*16)
self.up2 = Up(self.init_channel*16, self.init_channel*8)
self.up3 = Up(self.init_channel*8, self.init_channel*4)
self.up4 = Up(self.init_channel*4,self.init_channel*2)
self.up5 = Up(self.init_channel*2, self.init_channel)
self.outc = OutConv(self.init_channel, 1)
self.up1_1 = Up_single(self.init_channel*32, self.init_channel*16)
self.up2_1 = Up_single(self.init_channel*16, self.init_channel*8)
self.up3_1 = Up_single(self.init_channel*8, self.init_channel*4)
self.up4_1 = Up_single(self.init_channel*4,self.init_channel*2)
self.up5_1 = Up_single(self.init_channel*2, self.init_channel)
self.outc_1 = OutConv(self.init_channel, 1)
self.outc_2 = OutConv(self.init_channel, 1)
# self.dropout = nn.Dropout(p=0.5)
def forward(self, input):
x1 = self.inc(input)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x6 = self.down5(x5)
x_1 = self.up1_1(x6, x5)
x_1 = self.up2_1(x_1, x4)
x_1 = self.up3_1(x_1, x3)
x_1 = self.up4_1(x_1, x2)
x_1 = self.up5_1(x_1, x1)
mask = self.outc_1(x_1)
edge = self.outc_2(x_1)
x = self.up1(x6, x5)
# x = self.dropout(x)
x = self.up2(x, x4)
# x = self.dropout(x)
x = self.up3(x, x3)
# x = self.dropout(x)
x = self.up4(x, x2)
# x = self.dropout(x)
x = self.up5(x, x1)
# x = self.dropout(x)
depth = self.outc(x)
return depth,mask,edge
"""Create a Unet-based generator"""
class GieGenerator(nn.Module):
def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
"""Construct a Unet generator
Parameters:
input_nc (int) -- the number of channels in input images
output_nc (int) -- the number of channels in output images
num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
image of size 128x128 will become of size 1x1 # at the bottleneck
ngf (int) -- the number of filters in the last conv layer
norm_layer -- normalization layer
We construct the U-Net from the innermost layer to the outermost layer.
It is a recursive process.
"""
super(GieGenerator, self).__init__()
self.init_channel =32
self.inc = DoubleConv(input_nc,self.init_channel)
self.down1 = Down(self.init_channel, self.init_channel*2)
self.down2 = Down(self.init_channel*2, self.init_channel*4)
self.down3 = Down(self.init_channel*4, self.init_channel*8)
self.down4 = Down(self.init_channel*8, self.init_channel*16)
self.down5 = Down(self.init_channel*16, self.init_channel*32)
self.up1 = Up(self.init_channel*32, self.init_channel*16)
self.up2 = Up(self.init_channel*16, self.init_channel*8)
self.up3 = Up(self.init_channel*8, self.init_channel*4)
self.up4 = Up(self.init_channel*4,self.init_channel*2)
self.up5 = Up(self.init_channel*2, self.init_channel)
self.outc = OutConv(self.init_channel, 2)
# self.dropout = nn.Dropout(p=0.5)
def forward(self, input):
x1 = self.inc(input)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x6 = self.down5(x5)
x = self.up1(x6, x5)
# x = self.dropout(x)
x = self.up2(x, x4)
# x = self.dropout(x)
x = self.up3(x, x3)
# x = self.dropout(x)
x = self.up4(x, x2)
# x = self.dropout(x)
x = self.up5(x, x1)
# x = self.dropout(x)
logits1 = self.outc(x)
return logits1
class GiecbamGenerator(nn.Module):
def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
"""Construct a Unet generator
Parameters:
input_nc (int) -- the number of channels in input images
output_nc (int) -- the number of channels in output images
num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
image of size 128x128 will become of size 1x1 # at the bottleneck
ngf (int) -- the number of filters in the last conv layer
norm_layer -- normalization layer
We construct the U-Net from the innermost layer to the outermost layer.
It is a recursive process.
"""
super(GiecbamGenerator, self).__init__()
self.init_channel =32
self.inc = DoubleConv(input_nc,self.init_channel)
self.down1 = Down(self.init_channel, self.init_channel*2)
self.down2 = Down(self.init_channel*2, self.init_channel*4)
self.down3 = Down(self.init_channel*4, self.init_channel*8)
self.down4 = Down(self.init_channel*8, self.init_channel*16)
self.down5 = Down(self.init_channel*16, self.init_channel*32)
self.cbam = CBAM(gate_channels=self.init_channel*32)
self.up1 = Up(self.init_channel*32, self.init_channel*16)
self.up2 = Up(self.init_channel*16, self.init_channel*8)
self.up3 = Up(self.init_channel*8, self.init_channel*4)
self.up4 = Up(self.init_channel*4,self.init_channel*2)
self.up5 = Up(self.init_channel*2, self.init_channel)
self.outc = OutConv(self.init_channel, 2)
self.dropout = nn.Dropout(p=0.1)
def forward(self, input):
x1 = self.inc(input)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x6 = self.down5(x5)
x6 = self.cbam(x6)
x = self.up1(x6, x5)
x = self.up2(x, x4)
x = self.up3(x, x3)
x = self.up4(x, x2)
x = self.up5(x, x1)
x = self.dropout(x)
logits1 = self.outc(x)
return logits1
class Gie2headGenerator(nn.Module):
def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
"""Construct a Unet generator
Parameters:
input_nc (int) -- the number of channels in input images
output_nc (int) -- the number of channels in output images
num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
image of size 128x128 will become of size 1x1 # at the bottleneck
ngf (int) -- the number of filters in the last conv layer
norm_layer -- normalization layer
We construct the U-Net from the innermost layer to the outermost layer.
It is a recursive process.
"""
super(Gie2headGenerator, self).__init__()
self.init_channel =32
self.inc = DoubleConv(input_nc,self.init_channel)
self.down1 = Down(self.init_channel, self.init_channel*2)
self.down2 = Down(self.init_channel*2, self.init_channel*4)
self.down3 = Down(self.init_channel*4, self.init_channel*8)
self.down4 = Down(self.init_channel*8, self.init_channel*16)
self.down5 = Down(self.init_channel*16, self.init_channel*32)
self.up1_1 = Up(self.init_channel*32, self.init_channel*16)
self.up2_1 = Up(self.init_channel*16, self.init_channel*8)
self.up3_1 = Up(self.init_channel*8, self.init_channel*4)
self.up4_1 = Up(self.init_channel*4,self.init_channel*2)
self.up5_1 = Up(self.init_channel*2, self.init_channel)
self.outc_1 = OutConv(self.init_channel, 1)
self.up1_2 = Up(self.init_channel*32, self.init_channel*16)
self.up2_2 = Up(self.init_channel*16, self.init_channel*8)
self.up3_2 = Up(self.init_channel*8, self.init_channel*4)
self.up4_2 = Up(self.init_channel*4,self.init_channel*2)
self.up5_2 = Up(self.init_channel*2, self.init_channel)
self.outc_2 = OutConv(self.init_channel, 1)
def forward(self, input):
x1 = self.inc(input)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x6 = self.down5(x5)
x_1 = self.up1_1(x6, x5)
x_1 = self.up2_1(x_1, x4)
x_1 = self.up3_1(x_1, x3)
x_1 = self.up4_1(x_1, x2)
x_1 = self.up5_1(x_1, x1)
logits_1 = self.outc_1(x_1)
x_2 = self.up1_2(x6, x5)
x_2 = self.up2_2(x_2, x4)
x_2 = self.up3_2(x_2, x3)
x_2 = self.up4_2(x_2, x2)
x_2 = self.up5_2(x_2, x1)
logits_2 = self.outc_2(x_2)
logits = torch.cat((logits_1,logits_2),1)
return logits
class BmpGenerator(nn.Module):
def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
"""Construct a Unet generator
Parameters:
input_nc (int) -- the number of channels in input images
output_nc (int) -- the number of channels in output images
num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
image of size 128x128 will become of size 1x1 # at the bottleneck
ngf (int) -- the number of filters in the last conv layer
norm_layer -- normalization layer
We construct the U-Net from the innermost layer to the outermost layer.
It is a recursive process.
"""
super(BmpGenerator, self).__init__()
self.init_channel =32
self.output_nc = output_nc
self.inc = DoubleConv(input_nc,self.init_channel)
self.down1 = Down(self.init_channel, self.init_channel*2)
self.down2 = Down(self.init_channel*2, self.init_channel*4)
self.down3 = Down(self.init_channel*4, self.init_channel*8)
self.down4 = Down(self.init_channel*8, self.init_channel*16)
self.down5 = Down(self.init_channel*16, self.init_channel*32)
self.up1 = Up(self.init_channel*32, self.init_channel*16)
self.up2 = Up(self.init_channel*16, self.init_channel*8)
self.up3 = Up(self.init_channel*8, self.init_channel*4)
self.up4 = Up(self.init_channel*4,self.init_channel*2)
self.up5 = Up(self.init_channel*2, self.init_channel)
self.outc = OutConv(self.init_channel, self.output_nc)
# self.dropout = nn.Dropout(p=0.5)
def forward(self, input):
x1 = self.inc(input)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x6 = self.down5(x5)
x = self.up1(x6, x5)
# x = self.dropout(x)
x = self.up2(x, x4)
# x = self.dropout(x)
x = self.up3(x, x3)
# x = self.dropout(x)
x = self.up4(x, x2)
# x = self.dropout(x)
x = self.up5(x, x1)
# x = self.dropout(x)
logits1 = self.outc(x)
return logits1
class Bmp2Generator(nn.Module):
"""Create a Unet-based generator"""
def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
"""Construct a Unet generator
Parameters:
input_nc (int) -- the number of channels in input images
output_nc (int) -- the number of channels in output images
num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
image of size 128x128 will become of size 1x1 # at the bottleneck
ngf (int) -- the number of filters in the last conv layer
norm_layer -- normalization layer
We construct the U-Net from the innermost layer to the outermost layer.
It is a recursive process.
"""
super(Bmp2Generator, self).__init__()
#gienet
self.init_channel =32
self.inc = DoubleConv(3,self.init_channel)
self.down1 = Down(self.init_channel, self.init_channel*2)
self.down2 = Down(self.init_channel*2, self.init_channel*4)
self.down3 = Down(self.init_channel*4, self.init_channel*8)
self.down4 = Down(self.init_channel*8, self.init_channel*16)
self.down5 = Down(self.init_channel*16, self.init_channel*32)
self.up1 = Up(self.init_channel*32, self.init_channel*16)
self.up2 = Up(self.init_channel*16, self.init_channel*8)
self.up3 = Up(self.init_channel*8, self.init_channel*4)
self.up4 = Up(self.init_channel*4,self.init_channel*2)
self.up5 = Up(self.init_channel*2, self.init_channel)
self.outc = OutConv(self.init_channel, 1)
self.up1_1 = Up_single(self.init_channel*32, self.init_channel*16)
self.up2_1 = Up_single(self.init_channel*16, self.init_channel*8)
self.up3_1 = Up_single(self.init_channel*8, self.init_channel*4)
self.up4_1 = Up_single(self.init_channel*4,self.init_channel*2)
self.up5_1 = Up_single(self.init_channel*2, self.init_channel)
self.outc_1 = OutConv(self.init_channel, 1)
self.outc_2 = OutConv(self.init_channel, 1)
#bpm net
self.inc_b = DoubleConv(4,self.init_channel)
self.down1_b = Down(self.init_channel, self.init_channel*2)
self.down2_b = Down(self.init_channel*2, self.init_channel*4)
self.down3_b = Down(self.init_channel*4, self.init_channel*8)
self.down4_b = Down(self.init_channel*8, self.init_channel*16)
self.down5_b = Down(self.init_channel*16, self.init_channel*32)
self.up1_b = Up(self.init_channel*32, self.init_channel*16)
self.up2_b = Up(self.init_channel*16, self.init_channel*8)
self.up3_b = Up(self.init_channel*8, self.init_channel*4)
self.up4_b = Up(self.init_channel*4,self.init_channel*2)
self.up5_b = Up(self.init_channel*2, self.init_channel)
self.outc_b = OutConv(self.init_channel, 2)
# self.dropout = nn.Dropout(p=0.5)
def forward(self, input):
#gienet
x1 = self.inc(input)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x6 = self.down5(x5)
x_1 = self.up1_1(x6, x5)
x_1 = self.up2_1(x_1, x4)
x_1 = self.up3_1(x_1, x3)
x_1 = self.up4_1(x_1, x2)
x_1 = self.up5_1(x_1, x1)
mask = self.outc_1(x_1)
edge = self.outc_2(x_1)
x = self.up1(x6, x5)
x = self.up2(x, x4)
x = self.up3(x, x3)
x = self.up4(x, x2)
x = self.up5(x, x1)
depth = self.outc(x)
#bmpnet
mask[mask>0.5]=1.
mask[mask<=0.5]=0.
image_cat_depth = torch.cat((input*mask,depth*mask),dim=1)
x1_b = self.inc_b(image_cat_depth)
x2_b = self.down1_b(x1_b)
x3_b = self.down2_b(x2_b)
x4_b = self.down3_b(x3_b)
x5_b = self.down4_b(x4_b)
x6_b = self.down5_b(x5_b)
x_b = self.up1_b(x6_b, x5_b)
x_b = self.up2_b(x_b, x4_b)
x_b = self.up3_b(x_b, x3_b)
x_b = self.up4_b(x_b, x2_b)
x_b = self.up5_b(x_b, x1_b)
bm = self.outc_b(x_b)
# return depth,mask,edge,bm
return bm
class UnetGenerator(nn.Module):
def __init__(self, input_nc, output_nc, num_downs, ngf=64,
norm_layer=nn.BatchNorm2d, use_dropout=False):
super(UnetGenerator, self).__init__()
# construct unet structure
unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
for i in range(num_downs - 5):
unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
self.model = unet_block
def forward(self, input):
return self.model(input)
#class GieGenerator(nn.Module):
# def __init__(self, input_nc, output_nc, num_downs, ngf=64,
# norm_layer=nn.BatchNorm2d, use_dropout=False):
# super(GieGenerator, self).__init__()
#
# # construct unet structure
# unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
# for i in range(num_downs - 5):
# unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
# unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
# unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
# unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
# unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
#
# self.model = unet_block
#
# def forward(self, input):
# return self.model(input)
# Defines the submodule with skip connection.
# X -------------------identity---------------------- X
# |-- downsampling -- |submodule| -- upsampling --|
class UnetSkipConnectionBlock(nn.Module):
def __init__(self, outer_nc, inner_nc, input_nc=None,
submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
super(UnetSkipConnectionBlock, self).__init__()
self.outermost = outermost
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
if input_nc is None:
input_nc = outer_nc
downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
stride=2, padding=1, bias=use_bias)
downrelu = nn.LeakyReLU(0.2, True)
downnorm = norm_layer(inner_nc)
uprelu = nn.ReLU(True)
upnorm = norm_layer(outer_nc)
if outermost:
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
kernel_size=4, stride=2,
padding=1)
down = [downconv]
up = [uprelu, upconv, nn.Tanh()]
model = down + [submodule] + up
elif innermost:
# resize = nn.Upsample(scale_factor=2)
# conv = nn.Conv2d(inner_nc,outer_nc,kernel_size=4,stride=2,padding=1,bias=use_bias)
upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
kernel_size=4, stride=2,
padding=1, bias=use_bias)
down = [downrelu, downconv]
up = [uprelu, upconv, upnorm]
#up = [uprelu, resize, conv, upnorm]
model = down + up
else:
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
kernel_size=4, stride=2,
padding=1, bias=use_bias)
down = [downrelu, downconv, downnorm]
up = [uprelu, upconv, upnorm]
if use_dropout:
model = down + [submodule] + up + [nn.Dropout(0.5)]
else:
model = down + [submodule] + up
self.model = nn.Sequential(*model)
def forward(self, x):
if self.outermost:
return self.model(x)
else:
return torch.cat([x, self.model(x)], 1)
##===================================================================================================
class DilatedDoubleConv(nn.Module):
"""(convolution => [BN] => ReLU) * 2"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=4,stride=1,dilation=4),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=4,stride=1,dilation=4),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.double_conv(x)
class DilatedDown(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.maxpool_conv = nn.Sequential(
nn.MaxPool2d(2),
DilatedDoubleConv(in_channels, out_channels)
)
def forward(self, x):
return self.maxpool_conv(x)
class DilatedUp(nn.Module):
"""Upscaling then double conv"""
def __init__(self, in_channels, out_channels, bilinear=True):
super().__init__()
self.up = nn.Upsample(scale_factor=2, mode='nearest')
self.conv = DilatedDoubleConv(in_channels, out_channels)
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=4,stride=1,dilation=4),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
)
# self.deconv = nn.ConvTranspose2d(in_channels, out_channels,kernel_size=4, stride=2,padding=1, bias=True)
def forward(self, x1, x2):
x1 = self.up(x1)
x1 = self.conv1(x1)
# x1 = self.deconv(x1)
# input is BCHW
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
class DilatedSingleUnet(nn.Module):
def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
super(DilatedSingleUnet, self).__init__()
self.init_channel = 32
self.inc = DilatedDoubleConv(input_nc,self.init_channel)
self.down1 = DilatedDown(self.init_channel, self.init_channel*2)
self.down2 = DilatedDown(self.init_channel*2, self.init_channel*4)
self.down3 = DilatedDown(self.init_channel*4, self.init_channel*8)
self.down4 = DilatedDown(self.init_channel*8, self.init_channel*16)
self.down5 = DilatedDown(self.init_channel*16, self.init_channel*32)
self.cbam = CBAM(gate_channels=self.init_channel*32)
self.up1 = DilatedUp(self.init_channel*32, self.init_channel*16)
self.up2 = DilatedUp(self.init_channel*16, self.init_channel*8)
self.up3 = DilatedUp(self.init_channel*8, self.init_channel*4)
self.up4 = DilatedUp(self.init_channel*4,self.init_channel*2)
self.up5 = DilatedUp(self.init_channel*2, self.init_channel)
self.outc = OutConv(self.init_channel, output_nc)
def forward(self, input):
x1 = self.inc(input)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x6 = self.down5(x5)
x6 = self.cbam(x6)
x = self.up1(x6, x5)
x = self.up2(x, x4)
x = self.up3(x, x3)
x = self.up4(x, x2)
x = self.up5(x, x1)
logits1 = self.outc(x)
return logits1