def __init__(self, inch, nlayer, stride, act='lrelu', acto='tanh', norm=None, conv_mode='3D'):
super(RGBBlock, self).__init__()
seq = []
for i in range(nlayer):
outch = inch#//(2**(i+1))
seq += [('l%d'%i, ConvST3d(inch, outch, kernel_size=(3,3,3), stride=stride ,padding=(1,1,1), act=act, norm=norm, conv_mode=conv_mode))]
inch = outch
seq += [('l%d'%nlayer, ConvST3d(inch, 3, kernel_size=(1,1,1), act=acto, norm=None, conv_mode='3D'))]
self.seq = nn.Sequential(OrderedDict(seq))
def __init__(self,
ngf=64,
nbf=32,
nblock=5,
nup=2,
conv_mode='ST',
act='lrelu',
acto=None,
norm=None):
super(G4, self).__init__()
self.convin = ConvST3d(3,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=None,
norm=None,
conv_mode=conv_mode)
self.res_hall = nn.Sequential(
OrderedDict([('rrdb%d' % (i + 1),
RRDB(inch=ngf,
nbf=nbf,
norm=norm,
act=act,
conv_mode=conv_mode)) for i in range(nblock)]))
self.convbout = ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=None,
norm=norm,
conv_mode=conv_mode)
up = [('upconv%d' % (i + 1),
ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
scale_mode='trilinear',
act=act,
norm=None,
conv_mode=conv_mode)) for i in range(nup)]
self.up_hall = nn.Sequential(OrderedDict(up))
self.conv1 = ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=act,
norm=None,
conv_mode=conv_mode)
self.conv2 = ConvST3d(ngf,
3, (3, 3, 3),
padding=(1, 1, 1),
act=acto,
norm=None,
conv_mode=conv_mode)
def __init__(self,
inch=32,
outch=32,
scale_factor=[1, 2, 2],
scale_mode='trilinear',
norm=None,
actenc='lrelu',
actup='lrelu',
conv_mode='3D'):
super(Block, self).__init__()
self.enc = DenseBlock(inch,
outch,
res_scale=0.2,
kernel_size=(3, 3, 3),
norm=norm,
act=actenc,
conv_mode=conv_mode)
self.up = ConvST3d(outch,
outch, (3, 3, 3),
padding=(1, 1, 1),
scale_factor=scale_factor,
scale_mode=scale_mode,
act=actup,
norm=norm,
conv_mode=conv_mode)
self.dec = DenseBlock(outch,
outch,
res_scale=0.2,
kernel_size=(3, 3, 3),
norm=norm,
act=actup,
conv_mode=conv_mode)
def __init__(self,
nin,
nout,
res_scale=0.2,
act='lrelu',
norm='bn',
conv_mode='3D'):
super(EncoderBlock1, self).__init__()
#self.rrdb = RRDB(nin, nout, res_scale=0.2, kernel_size=(3,3,3), norm=None, act=act, conv_mode=conv_mode)
#self.conv1 = ConvST3d(nin, nout, (3,3,3), stride=stride, padding=(1,1,1), act=act, norm=norm, conv_mode=conv_mode)
kernel_size = (3, 3, 3)
self.res_scale = res_scale
self.RDB1 = ResidualDenseBlock(inch=nin,
outch=nin,
res_scale=res_scale,
kernel_size=kernel_size,
norm=norm,
act=act,
conv_mode=conv_mode)
self.conv1 = ConvST3d(nin,
nout, (3, 3, 3),
stride=(2, 2, 2),
padding=(1, 1, 1),
act=act,
norm=norm,
conv_mode=conv_mode)
self.RDB2 = ResidualDenseBlock(inch=nout,
outch=nout,
res_scale=res_scale,
kernel_size=kernel_size,
norm=norm,
act=act,
conv_mode=conv_mode)
def __init__(self,
nin,
nout,
scale_mode='trilinear',
scale_factor=(2, 2, 2),
res_scale=0.2,
act='lrelu',
norm='bn',
conv_mode='3D'):
super(DecoderBlock1, self).__init__()
self.RDB1 = ResidualDenseBlock(inch=nin,
outch=nin,
res_scale=res_scale,
kernel_size=(3, 3, 3),
norm=norm,
act=act,
conv_mode=conv_mode)
self.conv1 = ConvST3d(nin,
nout, (3, 3, 3),
padding=(1, 1, 1),
scale_mode=scale_mode,
scale_factor=scale_factor,
act=act,
norm=norm,
conv_mode=conv_mode)
def __init__(self,
inch,
outch,
res_scale=0.2,
kernel_size=(3, 3, 3),
norm=None,
act='lrelu',
conv_mode='ST'):
super(ResidualDenseBlock, self).__init__()
self.res_scale = res_scale
self.conv1 = ConvST3d(inch,
outch,
kernel_size,
padding=(1, 1, 1),
act=act,
norm=norm,
conv_mode=conv_mode)
self.conv2 = ConvST3d(inch + outch,
outch,
kernel_size,
padding=(1, 1, 1),
act=act,
norm=norm,
conv_mode=conv_mode)
self.conv3 = ConvST3d(inch + outch * 2,
outch,
kernel_size,
padding=(1, 1, 1),
act=act,
norm=norm,
conv_mode=conv_mode)
self.conv4 = ConvST3d(inch + outch * 3,
outch,
kernel_size,
padding=(1, 1, 1),
act=act,
norm=norm,
conv_mode=conv_mode)
self.conv5 = ConvST3d(inch + outch * 4,
inch,
kernel_size,
padding=(1, 1, 1),
act=None,
norm=norm,
conv_mode=conv_mode)
def __init__(self,
ndf=16,
nlayer=5,
actc='lrelu',
acto='sigmoid',
norm='spec',
conv_mode='3D'):
super(D4, self).__init__()
seq = [
('conv1',
ConvST3d(3,
ndf, (3, 3, 3),
padding=(1, 1, 1),
stride=(1, 1, 1),
act=actc,
norm=norm,
conv_mode=conv_mode)),
]
nin = ndf
nout = ndf
for i in range(nlayer - 2):
nout = nin * (2**i)
seq += [('conv%d' % (i + 2),
ConvST3d(nin,
nout, (3, 3, 3),
padding=(1, 1, 1),
stride=(1, 2, 2),
act=actc,
norm=norm,
conv_mode=conv_mode))]
nin = nout
seq += [('convout',
ConvST3d(nout,
1, (3, 3, 3),
padding=(1, 1, 1),
stride=(1, 2, 2),
act=acto,
norm=None,
conv_mode=conv_mode))]
self.model = nn.Sequential(OrderedDict(seq))
def __init__(self, ndf=16, nfc=256, nframe=16, nout=16, actc='lrelu', acto='sigmoid', norm='bn', conv_mode='ST'):
super(D3, self).__init__()
self.nout = nout
seq = [
('conv1', ConvST3d(3, ndf, (3,3,3), padding=(1,1,1), stride=(1,2,2), act=actc, norm=norm, conv_mode=conv_mode)),
('conv2', ConvST3d(ndf, ndf, (3,3,3), padding=(1,1,1), stride=(1,2,2), act=actc, norm=norm, conv_mode=conv_mode)),
('conv3', ConvST3d(ndf, ndf*2, (3,3,3), padding=(1,1,1), stride=(1,2,2), act=actc, norm=norm, conv_mode=conv_mode)),
('conv4', ConvST3d(ndf*2, ndf*4, (3,3,3), padding=(1,1,1), stride=(1,2,2), act=actc, norm=norm, conv_mode=conv_mode)),
('conv5', ConvST3d(ndf*4, ndf*8, (3,3,3), padding=(1,1,1), stride=(1,2,2), act=actc, norm=norm, conv_mode=conv_mode))
]
self.nfeat = 4*4 * nframe * ndf*8
lin = [
('fc1', FullyConnected(self.nfeat, nfc, act='relu')),
('fc2', FullyConnected(nfc, self.nout, act=acto))
]
self.model = nn.Sequential(OrderedDict(seq))
self.fc = nn.Sequential(OrderedDict(lin))
def __init__(self,
ndf=16,
actc='relu',
acto='sigmoid',
norm='bn',
conv_mode='ST'):
super(D2, self).__init__()
seq = [
('conv1',
ConvST3d(3,
ndf, (3, 3, 3),
padding=(1, 1, 1),
stride=(1, 2, 2),
act=actc,
norm=norm,
conv_mode=conv_mode)),
('conv2',
ConvST3d(ndf,
ndf * 2, (3, 3, 3),
padding=(1, 1, 1),
stride=(1, 2, 2),
act=actc,
norm=norm,
conv_mode=conv_mode)),
('conv3',
ConvST3d(ndf * 2,
ndf * 4, (3, 3, 3),
padding=(1, 1, 1),
stride=(1, 2, 2),
act=actc,
norm=norm,
conv_mode=conv_mode)),
('conv4',
ConvST3d(ndf * 4,
ndf * 8, (3, 3, 3),
padding=(1, 1, 1),
stride=(1, 2, 2),
act=actc,
norm=norm,
conv_mode=conv_mode)),
('conv5',
ConvST3d(ndf * 8,
ndf * 8, (3, 3, 3),
padding=(1, 1, 1),
stride=(1, 2, 2),
act=actc,
norm=norm,
conv_mode=conv_mode)),
('conv6',
ConvST3d(ndf * 8,
1, (1, 4, 4),
stride=(1, 4, 4),
act=acto,
conv_mode=conv_mode)),
]
self.model = nn.Sequential(OrderedDict(seq))
def __init__(self,
nbf=32,
act='lrelu',
acto='tanh',
norm='spec',
conv_mode='3D'):
super(G8, self).__init__()
self.convin = ConvST3d(3,
nbf, (3, 3, 3),
padding=(1, 1, 1),
act=act,
norm=norm,
conv_mode=conv_mode)
self.net = nn.Sequential(RRDB(nbf, nbf, act=act),
#RRDB(nbf, nbf, act=act)
)
self.convout = ConvST3d(nbf,
3, (3, 3, 3),
padding=(1, 1, 1),
act=acto,
conv_mode=conv_mode)
def __init__(self,
outch,
nlayer,
stride,
act='lrelu',
acto='tanh',
norm=None,
conv_mode='3D'):
super(RGB2FeatProjector1, self).__init__()
self.conv1 = ConvST3d(3,
outch,
kernel_size=(1, 1, 1),
act=acto,
norm=None,
conv_mode='3D')
def __init__(self,
inch,
outch,
nlayer,
stride=[1, 1, 1],
norm='spec',
act='lrelu',
conv_mode='3D'):
super(Block, self).__init__()
seq = []
for _ in range(nlayer):
seq += [
ConvST3d(inch,
outch, (3, 3, 3),
padding=(1, 1, 1),
stride=stride,
act=act,
norm=norm,
conv_mode=conv_mode)
]
self.model = nn.Sequential(*seq)
def __init__(self,
ngf=128,
actd='relu',
acto='tanh',
cnorm='bn',
unorm='bn'):
super(G2, self).__init__()
seq = [('conv1',
ConvST3d(3,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=actd,
norm=cnorm)),
('uconv1',
ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
scale_mode='trilinear',
act=actd,
norm=unorm)),
('conv2',
ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=actd,
norm=cnorm)),
('uconv2',
ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
scale_mode='trilinear',
act=actd,
norm=unorm)),
('conv3',
ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=actd,
norm=cnorm)),
('conv4',
ConvST3d(ngf, 3, (3, 3, 3), padding=(1, 1, 1), act=acto))]
self.model = nn.Sequential(OrderedDict(seq))
def __init__(self,
ngf=128,
actd='relu',
acto='tanh',
cnorm='bn',
unorm='bn'):
super(G3, self).__init__()
self.conv1 = ConvST3d(3,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=actd,
norm=cnorm)
self.uconv1 = ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
scale_mode='trilinear',
act=actd,
norm=unorm)
self.conv2 = ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=actd,
norm=cnorm)
self.uconv2 = ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
scale_mode='trilinear',
act=actd,
norm=unorm)
self.conv3 = ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=actd,
norm=cnorm)
self.conv4 = ConvST3d(ngf, 3, (3, 3, 3), padding=(1, 1, 1), act=acto)
def __init__(self,
ndf=16,
actc='lrelu',
acto='sigmoid',
norm='spec',
conv_mode='3D'):
super(PD1, self).__init__()
feat = [
ConvST3d(3,
ndf, (3, 3, 3),
padding=(1, 1, 1),
stride=(1, 1, 1),
act=actc,
norm=norm,
conv_mode=conv_mode),
ConvST3d(3,
ndf, (3, 3, 3),
padding=(1, 1, 1),
stride=(1, 1, 1),
act=actc,
norm=norm,
conv_mode=conv_mode),
ConvST3d(3,
ndf, (3, 3, 3),
padding=(1, 1, 1),
stride=(1, 1, 1),
act=actc,
norm=norm,
conv_mode=conv_mode)
]
blocks = [
Block(ndf,
ndf,
2,
stride=[1, 2, 2],
norm=norm,
act=actc,
conv_mode=conv_mode),
Block(ndf,
ndf,
1,
stride=[1, 2, 2],
norm=norm,
act=actc,
conv_mode=conv_mode),
Block(ndf,
ndf,
1,
stride=[1, 2, 2],
norm=norm,
act=actc,
conv_mode=conv_mode)
]
self.convout = ConvST3d(ndf,
1, (3, 3, 3),
padding=(1, 1, 1),
stride=(2, 1, 1),
act=acto,
norm=None,
conv_mode=conv_mode)
self.avgpool = nn.AvgPool3d(kernel_size=[1, 2, 2], stride=[1, 2, 2])
self.feat_proj = nn.ModuleList(feat)
self.layers = nn.ModuleList(blocks)
def __init__(self,
ngf=32,
actenc='lrelu',
actdec='lrelu',
acto='tanh',
norm='spec',
conv_mode='3D'):
super(PG1, self).__init__()
self.convin = ConvST3d(3,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=actenc,
norm=norm,
conv_mode=conv_mode)
blocks = [
Block(inch=ngf,
outch=ngf,
scale_factor=[1, 2, 2],
norm=norm,
actenc=actenc,
actup=actdec,
conv_mode=conv_mode),
Block(inch=ngf,
outch=ngf,
scale_factor=[1, 2, 2],
norm=norm,
actenc=actenc,
actup=actdec,
conv_mode=conv_mode),
Block(inch=ngf,
outch=ngf,
scale_factor=[1, 2, 2],
norm=norm,
actenc=actenc,
actup=actdec,
conv_mode=conv_mode)
]
proj = [
RGBBlock(ngf,
2,
stride=(2, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode),
RGBBlock(ngf,
1,
stride=(2, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode),
RGBBlock(ngf,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode)
]
self.layers = nn.ModuleList(blocks)
self.rgb_proj = nn.ModuleList(proj)
def __init__(self,
ngf=64,
nbf=32,
nblock=5,
nup=2,
attn_resblock_idx=None,
attn_upblock_idx=None,
attn_before_up=False,
conv_mode='3D',
act='lrelu',
acto='tanh',
norm='spec'):
super(G5, self).__init__()
self.attn_before_up = attn_before_up
# Initial convolution
self.convin = ConvST3d(3,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=None,
norm=norm,
conv_mode=conv_mode)
# RRDB blocks
rh = []
for i in range(nblock):
rh.append(('rrdb%d' % (i + 1),
RRDB(inch=ngf,
nbf=nbf,
norm=norm,
act=act,
conv_mode=conv_mode)))
if attn_resblock_idx is not None and attn_resblock_idx == i:
rh.append(('rb_attn', SelfAttention(ngf)))
self.res_hall = nn.Sequential(OrderedDict(rh))
# Convolution btw residual blocks and upsampling
self.convbout = ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=None,
norm=norm,
conv_mode=conv_mode)
if self.attn_before_up:
self.attn = SelfAttention(ngf)
# Upsampling layers
up = []
for i in range(nup):
up.append(('upconv%d' % (i + 1),
ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
scale_mode='trilinear',
act=act,
norm=norm,
conv_mode=conv_mode)))
if attn_upblock_idx is not None and attn_upblock_idx == i:
up.append(('up_attn', SelfAttention(ngf)))
self.up_hall = nn.Sequential(OrderedDict(up))
# Output convolutions
self.conv1 = ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=act,
norm=norm,
conv_mode=conv_mode)
self.conv2 = ConvST3d(ngf,
3, (3, 3, 3),
padding=(1, 1, 1),
act=acto,
norm=norm,
conv_mode=conv_mode)
self.apply(init_weights)
def __init__(self,
mode,
scale_factor,
ngf=32,
nbf=32,
nblock=4,
actenc='lrelu',
actdec='lrelu',
acto='tanh',
norm='spec',
conv_mode='3D'):
super(PG3, self).__init__()
assert mode in ['3D', 'ST']
assert scale_factor in [4, 8]
self.convin = ConvST3d(3,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=None,
norm=None,
conv_mode=conv_mode)
self.res_hall = nn.Sequential(
OrderedDict([('rrdb%d' % (i + 1),
RRDB(inch=ngf,
nbf=nbf,
norm=norm,
act=actenc,
conv_mode=conv_mode)) for i in range(nblock)]))
self.convbout = ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
act=None,
norm=norm,
conv_mode=conv_mode)
if mode == '3D':
self.enc_dec_ratio = 1
ups = [
nn.Sequential(
ConvST3d(nbf,
nbf,
kernel_size=(3, 3, 3),
padding=(1, 1, 1),
stride=(2, 1, 1),
act=actdec,
norm=norm,
conv_mode=conv_mode),
ConvST3d(nbf,
nbf,
kernel_size=(3, 3, 3),
padding=(1, 1, 1),
stride=(2, 1, 1),
act=actdec,
norm=norm,
conv_mode=conv_mode),
UpBlock(nbf,
nbf,
kernel=(3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
act=actdec,
norm=norm,
conv_mode=conv_mode)),
nn.Sequential(
ConvST3d(nbf,
nbf,
kernel_size=(3, 3, 3),
padding=(1, 1, 1),
stride=(2, 1, 1),
act=actdec,
norm=norm,
conv_mode=conv_mode),
UpBlock(nbf,
nbf,
kernel=(3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
act=actdec,
norm=norm,
conv_mode=conv_mode)),
UpBlock(nbf,
nbf,
kernel=(3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
act=actdec,
norm=norm,
conv_mode=conv_mode),
]
proj = [
RGBBlock(ngf,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode),
RGBBlock(ngf,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode),
RGBBlock(ngf,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode)
]
if scale_factor == 4:
ups, proj = ups[1:], proj[1:]
elif mode == 'ST':
self.enc_dec_ratio = 2
ups = [
nn.Sequential(
ConvST3d(nbf,
nbf,
kernel_size=(3, 3, 3),
padding=(1, 1, 1),
stride=(2, 1, 1),
act=actdec,
norm=norm,
conv_mode=conv_mode),
ConvST3d(nbf,
nbf,
kernel_size=(3, 3, 3),
padding=(1, 1, 1),
stride=(2, 1, 1),
act=actdec,
norm=norm,
conv_mode=conv_mode)),
nn.Sequential(
ConvST3d(nbf,
nbf,
kernel_size=(3, 3, 3),
padding=(1, 1, 1),
stride=(2, 1, 1),
act=actdec,
norm=norm,
conv_mode=conv_mode),
ConvST3d(nbf,
nbf,
kernel_size=(3, 3, 3),
padding=(1, 1, 1),
stride=(2, 1, 1),
act=actdec,
norm=norm,
conv_mode=conv_mode),
UpBlock(nbf,
nbf,
kernel=(3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
act=actdec,
norm=norm,
conv_mode=conv_mode)),
nn.Sequential(
ConvST3d(nbf,
nbf,
kernel_size=(3, 3, 3),
padding=(1, 1, 1),
stride=(2, 1, 1),
act=actdec,
norm=norm,
conv_mode=conv_mode)),
nn.Sequential(
ConvST3d(nbf,
nbf,
kernel_size=(3, 3, 3),
padding=(1, 1, 1),
stride=(2, 1, 1),
act=actdec,
norm=norm,
conv_mode=conv_mode),
UpBlock(nbf,
nbf,
kernel=(3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
act=actdec,
norm=norm,
conv_mode=conv_mode)),
ConvST3d(nbf,
nbf,
kernel_size=(3, 3, 3),
padding=(1, 1, 1),
stride=(1, 1, 1),
act=actdec,
norm=norm,
conv_mode=conv_mode),
UpBlock(nbf,
nbf,
kernel=(3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
act=actdec,
norm=norm,
conv_mode=conv_mode)
]
proj = [
RGBBlock(ngf,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode),
RGBBlock(ngf,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode),
RGBBlock(ngf,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode),
RGBBlock(ngf,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode),
RGBBlock(ngf,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode),
RGBBlock(ngf,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode)
]
if scale_factor == 4:
ups, proj = ups[2:], proj[2:]
self.up_layers = nn.ModuleList(ups)
self.rgb_proj = nn.ModuleList(proj)
def __init__(self, nin, nout, kernel, scale_factor, padding=(0,0,0), stride=(1,1,1), act='lrelu', norm='spec', conv_mode='3D'):
super(UpBlock, self).__init__()
self.c1 = ConvST3d(nin, nout, kernel, padding=padding, act=act, norm=norm, conv_mode=conv_mode)
self.c2 = ConvST3d(nin, nout, kernel, padding=padding, scale_factor=scale_factor, stride=stride, scale_mode='trilinear', act=act, norm=norm, conv_mode=conv_mode)
self.c3 = ConvST3d(nin, nout, kernel, padding=padding, act=act, norm=norm, conv_mode=conv_mode)
def __init__(self, mode, scale_factor, ngf=32, nbf=32, actenc='lrelu', actdec='lrelu', acto='tanh', norm='spec', conv_mode='3D'):
super(LOCPG2, self).__init__()
assert mode in ['3D', 'ST']
assert scale_factor in [4 , 8]
self.convin = ConvST3d(3, ngf, (3,3,3), padding=(1,1,1), act=None, norm=None, conv_mode=conv_mode)
blocks = [
RRDB(inch=ngf, nbf=nbf, norm=norm, act=actenc, conv_mode=conv_mode),
RRDB(inch=ngf, nbf=nbf, norm=norm, act=actenc, conv_mode=conv_mode),
RRDB(inch=ngf, nbf=nbf, norm=norm, act=actenc, conv_mode=conv_mode),
]
if scale_factor == 4:
blocks = blocks[1:]
if mode == '3D':
self.enc_dec_ratio = 1
ups = [
nn.Sequential(
ConvST3d(nbf, nbf, kernel_size=(3,3,3), padding=(1,1,1), stride=(2,1,1), act=actdec, norm=norm, conv_mode=conv_mode),
ConvST3d(nbf, nbf, kernel_size=(3,3,3), padding=(1,1,1), stride=(2,1,1), act=actdec, norm=norm, conv_mode=conv_mode),
UpBlock(nbf, nbf, kernel=(3,3,3), padding=(1,1,1), scale_factor=[1,2,2], act=actdec, norm=norm, conv_mode=conv_mode)
),
nn.Sequential(
ConvST3d(nbf, nbf, kernel_size=(3,3,3), padding=(1,1,1), stride=(2,1,1), act=actdec, norm=norm, conv_mode=conv_mode),
UpBlock(nbf, nbf, kernel=(3,3,3), padding=(1,1,1), scale_factor=[1,2,2], act=actdec, norm=norm, conv_mode=conv_mode)
),
UpBlock(nbf, nbf, kernel=(3,3,3), padding=(1,1,1), scale_factor=[1,2,2], act=actdec, norm=norm, conv_mode=conv_mode),
]
proj = [
RGBBlock(ngf, 1, stride=(1,1,1), act=actenc, acto=acto, norm=norm, conv_mode=conv_mode),
RGBBlock(ngf, 1, stride=(1,1,1), act=actenc, acto=acto, norm=norm, conv_mode=conv_mode),
RGBBlock(ngf, 1, stride=(1,1,1), act=actenc, acto=acto, norm=norm, conv_mode=conv_mode)
]
if scale_factor == 4:
ups, proj = ups[1:], proj[1:]
elif mode == 'ST':
self.enc_dec_ratio = 2
ups = [
nn.Sequential(
ConvST3d(nbf, nbf, kernel_size=(3,3,3), padding=(1,1,1), stride=(2,1,1), act=actdec, norm=norm, conv_mode=conv_mode),
ConvST3d(nbf, nbf, kernel_size=(3,3,3), padding=(1,1,1), stride=(2,1,1), act=actdec, norm=norm, conv_mode=conv_mode)
),
nn.Sequential(
ConvST3d(nbf, nbf, kernel_size=(3,3,3), padding=(1,1,1), stride=(2,1,1), act=actdec, norm=norm, conv_mode=conv_mode),
ConvST3d(nbf, nbf, kernel_size=(3,3,3), padding=(1,1,1), stride=(2,1,1), act=actdec, norm=norm, conv_mode=conv_mode),
UpBlock(nbf, nbf, kernel=(3,3,3), padding=(1,1,1), scale_factor=[1,2,2], act=actdec, norm=norm, conv_mode=conv_mode)
),
nn.Sequential(
ConvST3d(nbf, nbf, kernel_size=(3,3,3), padding=(1,1,1), stride=(2,1,1), act=actdec, norm=norm, conv_mode=conv_mode)
),
nn.Sequential(
ConvST3d(nbf, nbf, kernel_size=(3,3,3), padding=(1,1,1), stride=(2,1,1), act=actdec, norm=norm, conv_mode=conv_mode),
UpBlock(nbf, nbf, kernel=(3,3,3), padding=(1,1,1), scale_factor=[1,2,2], act=actdec, norm=norm, conv_mode=conv_mode)
),
ConvST3d(nbf, nbf, kernel_size=(3,3,3), padding=(1,1,1), stride=(1,1,1), act=actdec, norm=norm, conv_mode=conv_mode),
UpBlock(nbf, nbf, kernel=(3,3,3), padding=(1,1,1), scale_factor=[1,2,2], act=actdec, norm=norm, conv_mode=conv_mode)
]
proj = [
RGBBlock(ngf, 1, stride=(1,1,1), act=actenc, acto=acto, norm=norm, conv_mode=conv_mode),
RGBBlock(ngf, 1, stride=(1,1,1), act=actenc, acto=acto, norm=norm, conv_mode=conv_mode),
RGBBlock(ngf, 1, stride=(1,1,1), act=actenc, acto=acto, norm=norm, conv_mode=conv_mode),
RGBBlock(ngf, 1, stride=(1,1,1), act=actenc, acto=acto, norm=norm, conv_mode=conv_mode),
RGBBlock(ngf, 1, stride=(1,1,1), act=actenc, acto=acto, norm=norm, conv_mode=conv_mode),
RGBBlock(ngf, 1, stride=(1,1,1), act=actenc, acto=acto, norm=norm, conv_mode=conv_mode)
]
if scale_factor == 4:
ups, proj = ups[2:], proj[2:]
self.res_hall = nn.ModuleList(blocks)
self.up_layers = nn.ModuleList(ups)
self.rgb_proj = nn.ModuleList(proj)
if scale_factor == 8:
self.saliency_dec = nn.Sequential(
ConvST3d(nbf, nbf//2, (3,3,3), padding=(1,1,1), scale_mode='trilinear', scale_factor=(1,2,2), act=actdec, norm='bn', conv_mode=conv_mode),
ConvST3d(nbf//2, nbf//4, (3,3,3), padding=(1,1,1), scale_mode='trilinear', scale_factor=(1,2,2), act=actdec, norm='bn', conv_mode=conv_mode),
ConvST3d(nbf//4, 1, (3,3,3), padding=(1,1,1), scale_mode='trilinear', scale_factor=(1,2,2), act='sigmoid', norm=None, conv_mode=conv_mode)
)
elif scale_factor == 4:
self.saliency_dec = nn.Sequential(
ConvST3d(nbf, nbf//4, (3,3,3), padding=(1,1,1), scale_mode='trilinear', scale_factor=(1,2,2), act=actdec, norm='bn', conv_mode=conv_mode),
ConvST3d(nbf//4, 1, (3,3,3), padding=(1,1,1), scale_mode='trilinear', scale_factor=(1,2,2), act='sigmoid', norm=None, conv_mode=conv_mode)
)
def __init__(self, ngf=16, actc='relu', actd='relu', acto='tanh'):
super(G1, self).__init__()
seq = [('convin',
ConvST3d(3, ngf, (5, 5, 5), padding=(2, 2, 2), act=actc)),
('d1_uconv1',
ConvST3d(ngf,
ngf, (3, 3, 3),
padding=(1, 1, 1),
scale_factor=[2, 2, 2],
scale_mode='trilinear',
act=actd)),
('d1_conv1',
ConvST3d(ngf, ngf, (3, 3, 3), padding=(1, 1, 1), act=actd)),
('d1_uconv2',
ConvST3d(ngf,
ngf * 2, (3, 3, 3),
padding=(1, 1, 1),
scale_factor=[2, 2, 2],
scale_mode='trilinear',
act=actd)),
('d1_conv2',
ConvST3d(ngf * 2,
ngf * 2, (3, 3, 3),
padding=(1, 1, 1),
act=actd)),
('e1_conv1',
ConvST3d(ngf * 2,
ngf * 2, (3, 3, 3),
padding=(1, 1, 1),
stride=(2, 1, 1),
act=actc)),
('e1_conv2',
ConvST3d(ngf * 2,
ngf, (3, 3, 3),
padding=(1, 1, 1),
stride=(2, 2, 2),
act=actc)),
('e1_conv3',
ConvST3d(ngf,
int(ngf / 2), (3, 3, 3),
padding=(1, 1, 1),
stride=(2, 2, 2),
act=actc)),
('d2_uconv1',
ConvST3d(int(ngf / 2),
ngf, (3, 3, 3),
padding=(1, 1, 1),
scale_factor=[1, 2, 2],
scale_mode='trilinear',
act=actd)),
('d2_conv1',
ConvST3d(ngf, ngf, (3, 3, 3), padding=(1, 1, 1), act=actd)),
('d2_uconv2',
ConvST3d(ngf,
ngf * 2, (3, 3, 3),
padding=(1, 1, 1),
scale_factor=[2, 2, 2],
scale_mode='trilinear',
act=actd)),
('d2_conv2',
ConvST3d(ngf * 2,
ngf * 2, (3, 3, 3),
padding=(1, 1, 1),
act=actd)),
('convout',
ConvST3d(ngf * 2, 3, (3, 3, 3), padding=(1, 1, 1), act=acto))]
self.model = nn.Sequential(OrderedDict(seq))
def __init__(self,
nch=32,
actenc='lrelu',
actdec='lrelu',
acto='tanh',
norm='spec',
conv_mode='3D'):
super(LOCPG5, self).__init__()
self.encoder = nn.Sequential(*[
RGB2FeatProjector1(nch,
1,
stride=(1, 1, 1),
act=actenc,
acto=actenc,
norm=norm,
conv_mode=conv_mode),
EncoderBlock1(nch, 2 * nch),
EncoderBlock1(2 * nch, 4 * nch),
EncoderBlock1(4 * nch, 8 * nch),
])
self.decoder_blocks = nn.ModuleList([
DecoderBlock1(8 * nch, 4 * nch), # 0
DecoderBlock1(4 * nch, 2 * nch), # 1
DecoderBlock1(2 * nch, nch) # 2
])
self.feat2rgb = nn.ModuleList([
Feat2RGBProjector1(4 * nch,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode),
Feat2RGBProjector1(2 * nch,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode),
Feat2RGBProjector1(nch,
1,
stride=(1, 1, 1),
act=actenc,
acto=acto,
norm=norm,
conv_mode=conv_mode)
])
self.saliency_dec = nn.Sequential(
ConvST3d(8 * nch,
2 * nch, (3, 3, 3),
padding=(1, 1, 1),
scale_mode='trilinear',
scale_factor=(2, 2, 2),
act=actdec,
norm='bn',
conv_mode=conv_mode),
ConvST3d(2 * nch,
nch, (3, 3, 3),
padding=(1, 1, 1),
scale_mode='trilinear',
scale_factor=(2, 2, 2),
act=actdec,
norm='bn',
conv_mode=conv_mode),
ConvST3d(nch,
1, (3, 3, 3),
padding=(1, 1, 1),
scale_mode='trilinear',
scale_factor=(2, 2, 2),
act='sigmoid',
norm=None,
conv_mode=conv_mode))
