dewarpNet矫正扭曲
This commit is contained in:
24
dewarp/models/__init__.py
Normal file
24
dewarp/models/__init__.py
Normal file
@@ -0,0 +1,24 @@
|
||||
from dewarp.models.densenetccnl import dnetccnl
|
||||
from dewarp.models.unetnc import UnetGenerator
|
||||
|
||||
|
||||
def get_model(name, n_classes=1, in_channels=3):
|
||||
model = _get_model_instance(name)
|
||||
|
||||
if name == 'dnetccnl':
|
||||
model = model(img_size=128, in_channels=in_channels, out_channels=n_classes, filters=32)
|
||||
elif name == 'unetnc':
|
||||
model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
|
||||
else:
|
||||
model = model(n_classes=n_classes)
|
||||
return model
|
||||
|
||||
|
||||
def _get_model_instance(name):
|
||||
try:
|
||||
return {
|
||||
'dnetccnl': dnetccnl,
|
||||
'unetnc': UnetGenerator,
|
||||
}[name]
|
||||
except:
|
||||
print('Model {} not available'.format(name))
|
||||
243
dewarp/models/densenetccnl.py
Normal file
243
dewarp/models/densenetccnl.py
Normal file
@@ -0,0 +1,243 @@
|
||||
# Densenet decoder encoder with intermediate fully connected layers and dropout
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def add_coordConv_channels(t):
|
||||
n, c, h, w = t.size()
|
||||
xx_channel = np.ones((h, w))
|
||||
xx_range = np.array(range(h))
|
||||
xx_range = np.expand_dims(xx_range, -1)
|
||||
xx_coord = xx_channel * xx_range
|
||||
yy_coord = xx_coord.transpose()
|
||||
|
||||
xx_coord = xx_coord / (h - 1)
|
||||
yy_coord = yy_coord / (h - 1)
|
||||
xx_coord = xx_coord * 2 - 1
|
||||
yy_coord = yy_coord * 2 - 1
|
||||
xx_coord = torch.from_numpy(xx_coord).float()
|
||||
yy_coord = torch.from_numpy(yy_coord).float()
|
||||
|
||||
if t.is_cuda:
|
||||
xx_coord = xx_coord.cuda()
|
||||
yy_coord = yy_coord.cuda()
|
||||
|
||||
xx_coord = xx_coord.unsqueeze(0).unsqueeze(0).repeat(n, 1, 1, 1)
|
||||
yy_coord = yy_coord.unsqueeze(0).unsqueeze(0).repeat(n, 1, 1, 1)
|
||||
|
||||
t_cc = torch.cat((t, xx_coord, yy_coord), dim=1)
|
||||
|
||||
return t_cc
|
||||
|
||||
|
||||
class DenseBlockEncoder(nn.Module):
|
||||
def __init__(self, n_channels, n_convs, activation=nn.ReLU, args=[False]):
|
||||
super(DenseBlockEncoder, self).__init__()
|
||||
assert (n_convs > 0)
|
||||
|
||||
self.n_channels = n_channels
|
||||
self.n_convs = n_convs
|
||||
self.layers = nn.ModuleList()
|
||||
for i in range(n_convs):
|
||||
self.layers.append(nn.Sequential(
|
||||
nn.BatchNorm2d(n_channels),
|
||||
activation(*args),
|
||||
nn.Conv2d(n_channels, n_channels, 3, stride=1, padding=1, bias=False), ))
|
||||
|
||||
def forward(self, inputs):
|
||||
outputs = []
|
||||
|
||||
for i, layer in enumerate(self.layers):
|
||||
if i > 0:
|
||||
next_output = 0
|
||||
for no in outputs:
|
||||
next_output = next_output + no
|
||||
outputs.append(next_output)
|
||||
else:
|
||||
outputs.append(layer(inputs))
|
||||
return outputs[-1]
|
||||
|
||||
|
||||
# Dense block in encoder.
|
||||
class DenseBlockDecoder(nn.Module):
|
||||
def __init__(self, n_channels, n_convs, activation=nn.ReLU, args=[False]):
|
||||
super(DenseBlockDecoder, self).__init__()
|
||||
assert (n_convs > 0)
|
||||
|
||||
self.n_channels = n_channels
|
||||
self.n_convs = n_convs
|
||||
self.layers = nn.ModuleList()
|
||||
for i in range(n_convs):
|
||||
self.layers.append(nn.Sequential(
|
||||
nn.BatchNorm2d(n_channels),
|
||||
activation(*args),
|
||||
nn.ConvTranspose2d(n_channels, n_channels, 3, stride=1, padding=1, bias=False), ))
|
||||
|
||||
def forward(self, inputs):
|
||||
outputs = []
|
||||
|
||||
for i, layer in enumerate(self.layers):
|
||||
if i > 0:
|
||||
next_output = 0
|
||||
for no in outputs:
|
||||
next_output = next_output + no
|
||||
outputs.append(next_output)
|
||||
else:
|
||||
outputs.append(layer(inputs))
|
||||
return outputs[-1]
|
||||
|
||||
|
||||
class DenseTransitionBlockEncoder(nn.Module):
|
||||
def __init__(self, n_channels_in, n_channels_out, mp, activation=nn.ReLU, args=[False]):
|
||||
super(DenseTransitionBlockEncoder, self).__init__()
|
||||
self.n_channels_in = n_channels_in
|
||||
self.n_channels_out = n_channels_out
|
||||
self.mp = mp
|
||||
self.main = nn.Sequential(
|
||||
nn.BatchNorm2d(n_channels_in),
|
||||
activation(*args),
|
||||
nn.Conv2d(n_channels_in, n_channels_out, 1, stride=1, padding=0, bias=False),
|
||||
nn.MaxPool2d(mp),
|
||||
)
|
||||
|
||||
def forward(self, inputs):
|
||||
return self.main(inputs)
|
||||
|
||||
|
||||
class DenseTransitionBlockDecoder(nn.Module):
|
||||
def __init__(self, n_channels_in, n_channels_out, activation=nn.ReLU, args=[False]):
|
||||
super(DenseTransitionBlockDecoder, self).__init__()
|
||||
self.n_channels_in = n_channels_in
|
||||
self.n_channels_out = n_channels_out
|
||||
self.main = nn.Sequential(
|
||||
nn.BatchNorm2d(n_channels_in),
|
||||
activation(*args),
|
||||
nn.ConvTranspose2d(n_channels_in, n_channels_out, 4, stride=2, padding=1, bias=False),
|
||||
)
|
||||
|
||||
def forward(self, inputs):
|
||||
return self.main(inputs)
|
||||
|
||||
|
||||
## Dense encoders and decoders for image of size 128 128
|
||||
class waspDenseEncoder128(nn.Module):
|
||||
def __init__(self, nc=1, ndf=32, ndim=128, activation=nn.LeakyReLU, args=[0.2, False], f_activation=nn.Tanh,
|
||||
f_args=[]):
|
||||
super(waspDenseEncoder128, self).__init__()
|
||||
self.ndim = ndim
|
||||
|
||||
self.main = nn.Sequential(
|
||||
# input is (nc) x 128 x 128
|
||||
nn.BatchNorm2d(nc),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(nc, ndf, 4, stride=2, padding=1),
|
||||
|
||||
# state size. (ndf) x 64 x 64
|
||||
DenseBlockEncoder(ndf, 6),
|
||||
DenseTransitionBlockEncoder(ndf, ndf * 2, 2, activation=activation, args=args),
|
||||
|
||||
# state size. (ndf*2) x 32 x 32
|
||||
DenseBlockEncoder(ndf * 2, 12),
|
||||
DenseTransitionBlockEncoder(ndf * 2, ndf * 4, 2, activation=activation, args=args),
|
||||
|
||||
# state size. (ndf*4) x 16 x 16
|
||||
DenseBlockEncoder(ndf * 4, 16),
|
||||
DenseTransitionBlockEncoder(ndf * 4, ndf * 8, 2, activation=activation, args=args),
|
||||
|
||||
# state size. (ndf*4) x 8 x 8
|
||||
DenseBlockEncoder(ndf * 8, 16),
|
||||
DenseTransitionBlockEncoder(ndf * 8, ndf * 8, 2, activation=activation, args=args),
|
||||
|
||||
# state size. (ndf*8) x 4 x 4
|
||||
DenseBlockEncoder(ndf * 8, 16),
|
||||
DenseTransitionBlockEncoder(ndf * 8, ndim, 4, activation=activation, args=args),
|
||||
f_activation(*f_args),
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
input = add_coordConv_channels(input)
|
||||
output = self.main(input).view(-1, self.ndim)
|
||||
# print(output.size())
|
||||
return output
|
||||
|
||||
|
||||
class waspDenseDecoder128(nn.Module):
|
||||
def __init__(self, nz=128, nc=1, ngf=32, lb=0, ub=1, activation=nn.ReLU, args=[False], f_activation=nn.Hardtanh,
|
||||
f_args=[]):
|
||||
super(waspDenseDecoder128, self).__init__()
|
||||
self.main = nn.Sequential(
|
||||
# input is Z, going into convolution
|
||||
nn.BatchNorm2d(nz),
|
||||
activation(*args),
|
||||
nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False),
|
||||
|
||||
# state size. (ngf*8) x 4 x 4
|
||||
DenseBlockDecoder(ngf * 8, 16),
|
||||
DenseTransitionBlockDecoder(ngf * 8, ngf * 8),
|
||||
|
||||
# state size. (ngf*4) x 8 x 8
|
||||
DenseBlockDecoder(ngf * 8, 16),
|
||||
DenseTransitionBlockDecoder(ngf * 8, ngf * 4),
|
||||
|
||||
# state size. (ngf*2) x 16 x 16
|
||||
DenseBlockDecoder(ngf * 4, 12),
|
||||
DenseTransitionBlockDecoder(ngf * 4, ngf * 2),
|
||||
|
||||
# state size. (ngf) x 32 x 32
|
||||
DenseBlockDecoder(ngf * 2, 6),
|
||||
DenseTransitionBlockDecoder(ngf * 2, ngf),
|
||||
|
||||
# state size. (ngf) x 64 x 64
|
||||
DenseBlockDecoder(ngf, 6),
|
||||
DenseTransitionBlockDecoder(ngf, ngf),
|
||||
|
||||
# state size (ngf) x 128 x 128
|
||||
nn.BatchNorm2d(ngf),
|
||||
activation(*args),
|
||||
nn.ConvTranspose2d(ngf, nc, 3, stride=1, padding=1, bias=False),
|
||||
f_activation(*f_args),
|
||||
)
|
||||
# self.smooth=nn.Sequential(
|
||||
# nn.Conv2d(nc, nc, 1, stride=1, padding=0, bias=False),
|
||||
# f_activation(*f_args),
|
||||
# )
|
||||
|
||||
def forward(self, inputs):
|
||||
# return self.smooth(self.main(inputs))
|
||||
return self.main(inputs)
|
||||
|
||||
|
||||
class dnetccnl(nn.Module):
|
||||
# in_channels -> nc | encoder first layer
|
||||
# filters -> ndf | encoder first layer
|
||||
# img_size(h,w) -> ndim
|
||||
# out_channels -> optical flow (x,y)
|
||||
|
||||
def __init__(self, img_size=128, in_channels=1, out_channels=2, filters=32, fc_units=100):
|
||||
super(dnetccnl, self).__init__()
|
||||
self.nc = in_channels
|
||||
self.nf = filters
|
||||
self.ndim = img_size
|
||||
self.oc = out_channels
|
||||
self.fcu = fc_units
|
||||
|
||||
self.encoder = waspDenseEncoder128(nc=self.nc + 2, ndf=self.nf, ndim=self.ndim)
|
||||
self.decoder = waspDenseDecoder128(nz=self.ndim, nc=self.oc, ngf=self.nf)
|
||||
# self.fc_layers= nn.Sequential(nn.Linear(self.ndim, self.fcu),
|
||||
# nn.ReLU(True),
|
||||
# nn.Dropout(0.25),
|
||||
# nn.Linear(self.fcu,self.ndim),
|
||||
# nn.ReLU(True),
|
||||
# nn.Dropout(0.25),
|
||||
# )
|
||||
|
||||
def forward(self, inputs):
|
||||
encoded = self.encoder(inputs)
|
||||
encoded = encoded.unsqueeze(-1).unsqueeze(-1)
|
||||
decoded = self.decoder(encoded)
|
||||
# print torch.max(decoded)
|
||||
# print torch.min(decoded)
|
||||
|
||||
return decoded
|
||||
89
dewarp/models/unetnc.py
Normal file
89
dewarp/models/unetnc.py
Normal file
@@ -0,0 +1,89 @@
|
||||
import functools
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
# 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 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)
|
||||
|
||||
|
||||
# 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:
|
||||
upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
|
||||
kernel_size=4, stride=2,
|
||||
padding=1, bias=use_bias)
|
||||
down = [downrelu, downconv]
|
||||
up = [uprelu, upconv, 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)
|
||||
Reference in New Issue
Block a user