def __init__(self, in_channels=32, weight_num=None, init_type="xavier"):
super(MotionBlurArgumentLayer, self).__init__()
self.offset1 = M.conv2d_block(in_channels,
2,
1,
1,
0,
init_type=init_type,
activation=None,
norm_type=None)
self.offset2 = M.conv2d_block(in_channels,
2,
1,
1,
0,
init_type=init_type,
activation=None,
norm_type=None)
self.weight1 = M.conv2d_block(in_channels,
weight_num,
1,
1,
0,
init_type=init_type,
activation=None,
norm_type=None)
self.weight2 = M.conv2d_block(in_channels,
weight_num,
1,
1,
0,
init_type=init_type,
activation=None,
norm_type=None)
def __init__(self,
in_channels,
out_channels,
layers=3,
init_type="xavier",
activation=nn.ReLU(),
norm_type=None):
super(LayerConv, self).__init__()
convs = []
for i in range(layers):
if i == 0:
convs.append(
M.conv2d_block(in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=1,
init_type=init_type,
activation=activation,
norm_type=norm_type))
else:
convs.append(
M.conv2d_block(in_channels=out_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=1,
init_type=init_type,
activation=activation,
norm_type=norm_type))
self.main = nn.Sequential(*convs)
def __init__(self,
in_channels,
out_channels,
layers,
downsample_type="maxpool",
init_type="xavier",
activation=nn.ReLU(),
norm_type='bn'):
super(DownsampleConv, self).__init__()
if downsample_type == "maxpool":
self.downsample = nn.MaxPool2d(kernel_size=2)
elif downsample_type == "conv_stride2":
self.downsample = M.conv2d_block(in_channels=in_channels,
out_channels=in_channels,
kernel_size=4,
stride=2,
padding=1,
pad_type='reflect',
init_type=init_type,
activation=activation,
norm_type=norm_type)
else:
raise ValueError(
"invalid downsample type: {}".format(downsample_type))
convs = []
for i in range(layers):
if i == 0:
convs.append(
M.conv2d_block(in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=init_type,
activation=activation,
norm_type=norm_type))
else:
convs.append(
M.conv2d_block(in_channels=out_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=init_type,
activation=activation,
norm_type=norm_type))
self.main = nn.Sequential(*convs)
def __init__(self,
in_channels,
out_channels,
layers=2,
upsample_type="bilinear",
init_type="xavier",
activation=nn.ReLU(),
norm_type='bn'):
super(UpsampleConv, self).__init__()
convs = []
for i in range(layers):
if i == 0:
convs.append(
M.conv2d_block(in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=init_type,
activation=activation,
norm_type=norm_type))
else:
convs.append(
M.conv2d_block(in_channels=out_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=init_type,
activation=activation,
norm_type=norm_type))
self.main = nn.Sequential(*convs)
if upsample_type == "bilinear":
self.upsample = lambda x: F.interpolate(
x, scale_factor=2, mode='bilinear', align_corners=False)
elif upsample_type == "conv_transpose_stride2":
self.upsample = M.deconv2d_block(in_channels=out_channels,
out_channels=out_channels,
kernel_size=4,
stride=2,
padding=1,
pad_type='reflect',
init_type=init_type,
activation=activation,
norm_type=norm_type)
else:
raise ValueError("invalid upsample type: {}".format(upsample_type))
def __init__(self,
in_channels,
out_channels,
init_type="xavier",
use_batchnorm=True):
super(ReductionModuleA, self).__init__()
self.init_type = init_type
self.activation = nn.ReLU(inplace=True)
self.bn = use_batchnorm
self.branch0 = M.conv2d_block(in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=1,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn)
self.branch1 = nn.Sequential(
M.conv2d_block(in_channels=in_channels,
out_channels=out_channels // 2,
kernel_size=1,
stride=1,
padding=0,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn),
M.conv2d_block(in_channels=out_channels // 2,
out_channels=out_channels // 2,
kernel_size=3,
stride=1,
padding=1,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn),
M.conv2d_block(in_channels=out_channels // 2,
out_channels=out_channels,
kernel_size=1,
stride=1,
padding=0,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn))
self.branch2 = nn.Sequential(
M.conv2d_block(in_channels=in_channels,
out_channels=out_channels // 2,
kernel_size=1,
stride=1,
padding=0,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn),
M.conv2d_block(in_channels=out_channels // 2,
out_channels=out_channels,
kernel_size=7,
stride=1,
padding=3,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn))
def __init__(self,
in_channels=1,
layers=2,
init_type="xavier",
norm_type='WN'):
super(BurstDeblurMP, self).__init__()
self.init_type = init_type
self.bn = norm_type
self.activation = nn.ReLU(inplace=True)
self.layers = layers
self.input_conv = LayerConv(in_channels=in_channels,
out_channels=self.encoder_list[0],
layers=self.layers,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.fusion1 = M.conv2d_block(in_channels=2 * self.encoder_list[0],
out_channels=self.encoder_list[0],
kernel_size=1,
stride=1,
padding=0,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.downsample_conv1 = DownsampleConv(
in_channels=self.encoder_list[0],
out_channels=self.encoder_list[1],
layers=self.layers,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.fusion2 = M.conv2d_block(in_channels=2 * self.encoder_list[1],
out_channels=self.encoder_list[1],
kernel_size=1,
stride=1,
padding=0,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.downsample_conv2 = DownsampleConv(
in_channels=self.encoder_list[1],
out_channels=self.encoder_list[2],
layers=self.layers,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.fusion3 = M.conv2d_block(in_channels=2 * self.encoder_list[2],
out_channels=self.encoder_list[2],
kernel_size=1,
stride=1,
padding=0,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.mid = MidConv(in_channels=self.encoder_list[2],
out_channels=self.decoder_list[0],
layers=self.layers,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.fusion4 = M.conv2d_block(in_channels=2 * self.decoder_list[0],
out_channels=self.decoder_list[0],
kernel_size=1,
stride=1,
padding=0,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.upsample_conv1 = UpsampleConv(in_channels=self.decoder_list[0] +
self.encoder_list[2],
out_channels=self.decoder_list[1],
layers=self.layers,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.fusion5 = M.conv2d_block(in_channels=2 * self.decoder_list[1],
out_channels=self.decoder_list[1],
kernel_size=1,
stride=1,
padding=0,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.upsample_conv2 = UpsampleConv(in_channels=self.decoder_list[1] +
self.encoder_list[1],
out_channels=self.decoder_list[2],
layers=self.layers,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.fusion6 = M.conv2d_block(in_channels=2 * self.decoder_list[2],
out_channels=self.decoder_list[2],
kernel_size=1,
stride=1,
padding=0,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.output_conv1 = LayerConv(in_channels=self.decoder_list[2] +
self.encoder_list[0],
out_channels=self.decoder_list[3],
layers=self.layers,
init_type=self.init_type,
activation=self.activation,
norm_type=self.bn)
self.output_conv2 = LayerConv(in_channels=self.decoder_list[3],
out_channels=in_channels,
layers=1,
init_type=self.init_type,
activation=None,
norm_type=None)
def __init__(self,
in_channels,
out_channels,
layers,
downsample_type="maxpool",
upsample_type="nearest",
init_type="xavier",
activation=nn.ReLU(),
norm_type='BN'):
super(MidConv, self).__init__()
if downsample_type == "maxpool":
self.downsample = nn.MaxPool2d(kernel_size=2)
elif downsample_type == "conv_stride2":
self.downsample = M.conv2d_block(in_channels=in_channels,
out_channels=in_channels,
kernel_size=4,
stride=2,
padding=1,
init_type=init_type,
activation=activation,
norm_type=norm_type)
else:
raise ValueError(
"invalid downsample type: {}".format(downsample_type))
convs = []
for i in range(layers):
if i == 0:
convs.append(
M.conv2d_block(in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=1,
init_type=init_type,
activation=activation,
norm_type=norm_type))
else:
convs.append(
M.conv2d_block(in_channels=out_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=1,
init_type=init_type,
activation=activation,
norm_type=norm_type))
self.main = nn.Sequential(*convs)
if upsample_type == "nearest":
self.upsample = lambda x: nn.functional.interpolate(
x, scale_factor=2, mode='nearest')
elif upsample_type == "conv_transpose_stride2":
self.upsample = M.deconv2d_block(in_channels=out_channels,
out_channels=out_channels,
kernel_size=4,
stride=2,
padding=1,
init_type=init_type,
activation=activation,
norm_type=norm_type)
else:
raise ValueError("invalid upsample type: {}".format(upsample_type))
def __init__(self,
init_type="xavier",
use_batchnorm=False,
use_maxpool=False,
DEBUG=False):
super(GradientInferenceNetwork, self).__init__()
self.init_type = init_type
self.bn = use_batchnorm
self.use_maxpool = use_maxpool
#self.activation = nn.ReLU(inplace=True)
self.activation = Linear() ##change the activation
self.feature_maps = None
self.encoder = OrderedDict()
in_channels = 4
for i in range(len(self.encoder_list)):
self.encoder['conv{}'.format(i + 1)] = M.conv2d_block(
in_channels=in_channels,
out_channels=self.encoder_list[i],
kernel_size=3,
stride=1,
padding=1,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn)
in_channels = self.encoder_list[i]
if use_maxpool:
self.encoder['conv_next{}'.format(i + 1)] = M.conv2d_block(
in_channels=in_channels,
out_channels=self.encoder_list[i],
kernel_size=3,
stride=1,
padding=1,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn)
self.encoder['downsample{}'.format(i + 1)] = nn.MaxPool2d(
kernel_size=2)
else:
self.encoder['conv_downsample{}'.format(
i + 1)] = M.conv2d_block(in_channels=in_channels,
out_channels=self.encoder_list[i],
kernel_size=4,
stride=2,
padding=1,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn)
self.mid = nn.Sequential(
M.conv2d_block(in_channels=in_channels,
out_channels=1024,
kernel_size=7,
stride=1,
padding=3,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn),
M.conv2d_block(in_channels=1024,
out_channels=512,
kernel_size=1,
stride=1,
padding=0,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn))
in_channels = 512
self.decoder = OrderedDict()
for i in range(len(self.decoder_list)):
self.decoder['conv{}'.format(i + 1)] = M.conv2d_block(
in_channels=in_channels,
out_channels=self.decoder_list[i],
kernel_size=3,
stride=1,
padding=1,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn)
in_channels = self.decoder_list[i]
self.decoder['deconv{}'.format(i + 1)] = M.deconv2d_block(
in_channels=in_channels,
out_channels=self.decoder_list[i],
kernel_size=4,
stride=2,
padding=1,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn)
in_channels += self.encoder_list[3 - i] # concat
self.end = nn.Sequential(
M.conv2d_block(in_channels=in_channels,
out_channels=64,
kernel_size=4,
stride=2,
padding=1,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn),
M.conv2d_block(in_channels=64,
out_channels=1,
kernel_size=5,
stride=1,
padding=2,
init_type=self.init_type,
activation=None,
use_batchnorm=None))
if DEBUG:
print(self.encoder)
print(self.mid)
print(self.decoder)
print(self.end)
def __init__(self,
in_channels=6,
layers=3,
branch_num=[9, 17, 33],
init_type='xavier',
norm_type=None):
super(SynthesizeBlurBranch, self).__init__()
if torch.cuda.is_available():
self.device = 'cuda'
else:
self.device = 'cpu'
self.layers = layers
self.init_type = init_type
self.norm_type = norm_type
self.leak_rate = 0.2
self.activation = nn.LeakyReLU(self.leak_rate, inplace=True)
self.branch_num = branch_num
#
self.input_conv = LayerConv(in_channels=in_channels,
out_channels=self.encoder_list[0],
layers=self.layers,
activation=self.activation,
norm_type=self.norm_type)
self.downsample1 = DownsampleConv(in_channels=self.encoder_list[0],
out_channels=self.encoder_list[1],
layers=self.layers,
downsample_type="maxpool",
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.downsample2 = DownsampleConv(in_channels=self.encoder_list[1],
out_channels=self.encoder_list[2],
layers=self.layers,
downsample_type="maxpool",
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.downsample3 = DownsampleConv(in_channels=self.encoder_list[2],
out_channels=self.encoder_list[3],
layers=self.layers,
downsample_type="maxpool",
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.mid = MidConv(in_channels=self.encoder_list[3],
out_channels=self.decoder_list[0],
layers=self.layers,
downsample_type='maxpool',
upsample_type='bilinear',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.bottom1 = M.conv2d_block(in_channels=self.decoder_list[0],
out_channels=self.decoder_list[0],
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.upsample_conv1 = UpsampleConv(in_channels=self.decoder_list[0] +
self.encoder_list[3],
out_channels=self.decoder_list[0],
layers=2,
upsample_type='bilinear',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.bottom2 = M.conv2d_block(in_channels=self.decoder_list[0],
out_channels=self.decoder_list[1],
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.upsample_conv2 = UpsampleConv(in_channels=self.decoder_list[1] +
self.encoder_list[2],
out_channels=self.decoder_list[1],
layers=self.layers,
upsample_type='bilinear',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.bottom3 = M.conv2d_block(in_channels=self.decoder_list[1],
out_channels=self.decoder_list[2],
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.upsample_conv3 = UpsampleConv(in_channels=self.decoder_list[2] +
self.encoder_list[1],
out_channels=self.decoder_list[2],
layers=self.layers,
upsample_type='bilinear',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.bottom4 = M.conv2d_block(in_channels=self.decoder_list[2],
out_channels=self.decoder_list[3],
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.output_conv = LayerConv(in_channels=self.decoder_list[3] +
self.encoder_list[0],
out_channels=self.decoder_list[3],
layers=self.layers,
init_type='xavier',
activation=self.activation,
norm_type=self.norm_type)
self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
self.prob_layer = nn.Linear(32, len(branch_num))
self.output_layer = [
MotionBlurArgumentLayer(self.decoder_list[-1],
weight_num=x).to(self.device)
for x in branch_num
]
def __init__(self,
in_channels=6,
layers=3,
init_type='xavier',
norm_type=None):
super(SynthesizeBlur, self).__init__()
if torch.cuda.is_available():
self.device = 'cuda'
else:
self.device = 'cpu'
self.layers = layers
self.init_type = init_type
self.norm_type = norm_type
self.LeakRate = 0.2
self.activation = nn.LeakyReLU(self.LeakRate)
self.N = 17 # N evnely-spaced discrete samples
#
self.input_conv = LayerConv(in_channels=in_channels,
out_channels=self.encoder_list[0],
layers=self.layers,
activation=self.activation,
norm_type=self.norm_type)
self.downsample1 = DownsampleConv(in_channels=self.encoder_list[0],
out_channels=self.encoder_list[1],
layers=self.layers,
downsample_type="maxpool",
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.downsample2 = DownsampleConv(in_channels=self.encoder_list[1],
out_channels=self.encoder_list[2],
layers=self.layers,
downsample_type="maxpool",
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.downsample3 = DownsampleConv(in_channels=self.encoder_list[2],
out_channels=self.encoder_list[3],
layers=self.layers,
downsample_type="maxpool",
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.mid = MidConv(in_channels=self.encoder_list[3],
out_channels=self.decoder_list[0],
layers=self.layers,
downsample_type='maxpool',
upsample_type='bilinear',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.bottom1 = M.conv2d_block(in_channels=self.decoder_list[0],
out_channels=self.decoder_list[0],
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.upsample_conv1 = UpsampleConv(in_channels=self.decoder_list[0] +
self.encoder_list[3],
out_channels=self.decoder_list[0],
layers=2,
upsample_type='bilinear',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.bottom2 = M.conv2d_block(in_channels=self.decoder_list[0],
out_channels=self.decoder_list[1],
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.upsample_conv2 = UpsampleConv(in_channels=self.decoder_list[1] +
self.encoder_list[2],
out_channels=self.decoder_list[1],
layers=self.layers,
upsample_type='bilinear',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.bottom3 = M.conv2d_block(in_channels=self.decoder_list[1],
out_channels=self.decoder_list[2],
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.upsample_conv3 = UpsampleConv(in_channels=self.decoder_list[2] +
self.encoder_list[1],
out_channels=self.decoder_list[2],
layers=self.layers,
upsample_type='bilinear',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.bottom4 = M.conv2d_block(in_channels=self.decoder_list[2],
out_channels=self.decoder_list[3],
kernel_size=3,
stride=1,
padding=1,
pad_type='reflect',
init_type=self.init_type,
activation=self.activation,
norm_type=self.norm_type)
self.output_conv = LayerConv(in_channels=self.decoder_list[3] +
self.encoder_list[0],
out_channels=self.decoder_list[3],
layers=self.layers,
init_type='xavier',
activation=self.activation,
norm_type=self.norm_type)
'''
outputs: feeds into line prediction layer
'''
self.offset1 = nn.Conv2d(32, 2, kernel_size=1, stride=1, padding=0)
self.offset2 = nn.Conv2d(32, 2, kernel_size=1, stride=1, padding=0)
self.weight1 = nn.Conv2d(32,
self.N,
kernel_size=1,
stride=1,
padding=0)
self.weight2 = nn.Conv2d(32,
self.N,
kernel_size=1,
stride=1,
padding=0)
def __init__(self,
backbone_type=None,
use_batchnorm=True,
init_type="xavier"):
super(ImageInferenceNetwork, self).__init__()
if backbone_type not in self.backbone_list:
raise ValueError(
"not support backbone type: {}".format(backbone_type))
self.bn = use_batchnorm
self.maxpool = nn.MaxPool2d(kernel_size=2)
self.activation = nn.ReLU(inplace=True)
self.init_type = init_type
self.conv_count = 0
self.mp_count = 0
if backbone_type == 'vgg16' or backbone_type == 'vgg16_bn':
self.backbone = OrderedDict()
in_channels = 3
for v in self.vgg_cfg['D']:
if v == 'M':
self.backbone['mp{}'.format(
self._add_mp_count())] = self.maxpool
else:
self.backbone['conv{}'.format(
self._add_conv_count())] = (M.conv2d_block(
in_channels=in_channels,
out_channels=v,
kernel_size=3,
padding=1,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn))
in_channels = v
# print(self.backbone)
self.feature_extract = (M.conv2d_block(in_channels=in_channels,
out_channels=256,
kernel_size=3,
padding=1,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn))
self.reduction_a = ReductionModuleA(in_channels=256, out_channels=256)
self.inception1 = InceptionDeconvModule(in_channels=3 * 256,
out_channels=256)
self.inception2 = InceptionDeconvModule(in_channels=3 * 256 + 512 +
256,
out_channels=128)
self.reduction_b = ReductionModuleB(in_channels=3 * 128 + 256 + 128,
out_channels=128)
self.inception3 = InceptionDeconvModule(in_channels=3 * 128,
out_channels=64)
self.inception4 = InceptionDeconvModule(in_channels=3 * 64 + 128 + 64,
out_channels=32)
self.inception5 = InceptionDeconvModule(in_channels=3 * 32 + 64 + 32,
out_channels=16)
self.conv_transition = M.conv2d_block(in_channels=3 * 16,
out_channels=16,
kernel_size=3,
stride=1,
padding=1,
init_type=self.init_type,
activation=self.activation,
use_batchnorm=self.bn)
self.conv_estimate_B = M.conv2d_block(in_channels=16,
out_channels=3,
kernel_size=3,
stride=1,
padding=1,
init_type=self.init_type,
activation=None,
use_batchnorm=None)
