def __init__(self, w_in, ch):
super().__init__()
self.w_in = w_in
self.ch = ch
with self.init_scope():
self.s = EqualizedLinear(w_in, ch, initial_bias=chainer.initializers.One(), gain=1)
self.b = EqualizedLinear(w_in, ch, initial_bias=chainer.initializers.Zero(), gain=1)
def __init__(self, in_ch, hidden_ch):
super(OccupancyNet, self).__init__()
net = [
EqualizedLinear(in_ch, hidden_ch), F.leaky_relu,
EqualizedLinear(hidden_ch, hidden_ch), F.leaky_relu,
EqualizedLinear(hidden_ch, 1)
]
with self.init_scope():
self.net = chainer.Sequential(*net)
def __init__(self):
super(CameraParamGenerator, self).__init__()
net = [
EqualizedLinear(8, 64), F.leaky_relu,
EqualizedLinear(64, 64), F.leaky_relu,
EqualizedLinear(64, 9)
]
with self.init_scope():
self.net = chainer.Sequential(*net)
def __init__(self,
in_ch=128,
ch=512,
enable_blur=False,
rgbd=False,
use_encoder=False,
use_occupancy_net=False,
initial_depth=None):
super(DCGANGenerator, self).__init__()
self.in_ch = in_ch
self.ch = ch
self.max_stage = 17
self.rgbd = rgbd
self.use_occupancy_net = use_occupancy_net
out_ch = 4 if rgbd else 3
if initial_depth is None:
initial_depth = 1.0
with self.init_scope():
if self.rgbd:
self.linear = EqualizedLinear(in_ch + 9, ch * 4 * 4)
else:
self.linear = EqualizedLinear(in_ch, ch * 4 * 4)
self.blocks = chainer.ChainList(
DCGANBlock(ch, ch, enable_blur=enable_blur), # 8
DCGANBlock(ch, ch, enable_blur=enable_blur), # 16
DCGANBlock(ch, ch, enable_blur=enable_blur), # 32
DCGANBlock(ch // 2, ch, enable_blur=enable_blur), # 64
DCGANBlock(ch // 4, ch // 2, enable_blur=enable_blur), # 128
)
self.outs = chainer.ChainList(
EqualizedConv2d(ch, out_ch, 1, 1, 0, gain=1),
EqualizedConv2d(ch, out_ch, 1, 1, 0, gain=1),
EqualizedConv2d(ch, out_ch, 1, 1, 0, gain=1),
EqualizedConv2d(ch // 2, out_ch, 1, 1, 0, gain=1),
EqualizedConv2d(ch // 4, out_ch, 1, 1, 0, gain=1),
)
if use_encoder:
self.enc = Encoder(ch, in_ch, enable_blur=enable_blur)
if use_occupancy_net:
self.occupancy = OccupancyNet(in_ch=in_ch + 3, hidden_ch=32)
# initialize depth weight to 0
for out in self.outs:
out.c.W.array[-1] = 0
out.c.b.array[-1] = math.log(math.e**initial_depth - 1)
self.n_blocks = len(self.blocks)
self.image_size = 128
self.enable_blur = enable_blur
def __init__(self, ch, dim_z=256):
super(EncoderBlockBase, self).__init__()
with self.init_scope():
self.c0 = EqualizedConv2d(ch, ch, 3, 1, 1)
self.c1 = EqualizedConv2d(ch, ch, 4, 1, 0)
self.l2 = EqualizedLinear(ch, dim_z, gain=1)
self.bn0 = L.BatchNormalization(ch)
self.bn1 = L.BatchNormalization(ch)
def __init__(self, ch=512, out_ch=64, enable_blur=False):
super(BackgroundFeatureGenerator, self).__init__()
self.img_size = 64
self.background_depth = 4 # virtual background distance
x_pos, y_pos = np.meshgrid(
np.arange(self.img_size) - self.img_size // 2,
np.arange(self.img_size) - self.img_size // 2)
depth_map = self.background_depth * \
self.img_size * 2 / np.sqrt((self.img_size * 2) ** 2 + x_pos ** 2 + y_pos ** 2)
with self.init_scope():
self.blocks = chainer.ChainList(
SynthesisBlock(ch, ch, ch, upsample=False), # 4
SynthesisBlock(ch,
ch,
ch,
upsample=True,
enable_blur=enable_blur), # 8
SynthesisBlock(ch,
ch,
ch,
upsample=True,
enable_blur=enable_blur), # 16
SynthesisBlock(ch,
ch,
ch,
upsample=True,
enable_blur=enable_blur), # 32
SynthesisBlock(ch // 2,
ch,
ch,
upsample=True,
enable_blur=enable_blur), # 64
)
self.conv = EqualizedConv2d(ch // 2, out_ch, 1, 1, 0, gain=1)
self.l1 = EqualizedLinear(ch + 9, ch)
self.l2 = EqualizedLinear(ch, ch)
self.add_persistent("depth_map", depth_map)
self.n_blocks = len(self.blocks)
self.image_size = 64
self.enable_blur = enable_blur
def __init__(self, ch, out_ch, sn=False):
super(MLP, self).__init__()
w = chainer.initializers.Uniform(1)
if sn:
net = [
L.Linear(ch, ch, initialW=w).add_hook(SpectralNormalization()),
F.leaky_relu,
L.Linear(ch, ch, initialW=w).add_hook(SpectralNormalization()),
F.leaky_relu,
L.Linear(ch, out_ch,
initialW=w).add_hook(SpectralNormalization()),
]
else:
net = [
EqualizedLinear(ch, ch),
F.leaky_relu,
EqualizedLinear(ch, ch),
F.leaky_relu,
EqualizedLinear(ch, out_ch),
]
with self.init_scope():
self.net = chainer.Sequential(*net)
def __init__(self, ch, dim_z, enable_blur=False, sn=False, res=False):
super(BigBiGANDiscriminator, self).__init__()
self.sn = sn
w = chainer.initializers.Uniform(1)
with self.init_scope():
self.f = Discriminator(ch,
ch,
enable_blur=enable_blur,
sn=sn,
res=res)
self.h_z = MLP(dim_z, ch, sn=sn)
self.h_cp = MLP(9, ch, sn=sn)
self.j = MLP(ch * 3, 1, sn=sn)
if sn:
self.s_x = L.Linear(ch, 1, initialW=w).add_hook(
SpectralNormalization())
self.s_z = L.Linear(ch, 1, initialW=w).add_hook(
SpectralNormalization())
self.s_cp = L.Linear(ch, 1, initialW=w).add_hook(
SpectralNormalization())
else:
self.s_x = EqualizedLinear(ch, 1)
self.s_z = EqualizedLinear(ch, 1)
self.s_cp = EqualizedLinear(ch, 1)
def __init__(self, ch, out_dim=1, sn=False):
super(DiscriminatorBlockBase, self).__init__()
with self.init_scope():
if not sn:
self.c0 = EqualizedConv2d(ch, ch, 3, 1, 1)
self.c1 = EqualizedConv2d(ch, ch, 4, 1, 0)
self.l2 = EqualizedLinear(ch, out_dim, gain=1)
else:
w = chainer.initializers.Uniform(1)
self.c0 = L.Convolution2D(ch, ch, 3, 1, 1,
initialW=w).add_hook(
SpectralNormalization())
self.c1 = L.Convolution2D(ch, ch, 4, 1, 0,
initialW=w).add_hook(
SpectralNormalization())
self.l2 = L.Linear(ch, out_dim, initialW=w).add_hook(
SpectralNormalization())
def __init__(self, ch=512):
super().__init__()
self.ch = ch
with self.init_scope():
self.l = chainer.ChainList(
EqualizedLinear(ch, ch),
LinkLeakyRelu(),
EqualizedLinear(ch, ch),
LinkLeakyRelu(),
EqualizedLinear(ch, ch),
LinkLeakyRelu(),
EqualizedLinear(ch, ch),
LinkLeakyRelu(),
EqualizedLinear(ch, ch),
LinkLeakyRelu(),
EqualizedLinear(ch, ch),
LinkLeakyRelu(),
EqualizedLinear(ch, ch),
LinkLeakyRelu(),
EqualizedLinear(ch, ch),
LinkLeakyRelu(),
)
self.ln = len(self.l)
def __init__(self, ch):
super(DiscriminatorBlockBase, self).__init__()
with self.init_scope():
self.c0 = EqualizedConv2d(ch, ch, 3, 1, 1)
self.c1 = EqualizedConv2d(ch, ch, 4, 1, 0)
self.l2 = EqualizedLinear(ch, 1, gain=1)
def __init__(self,
ch=512,
enable_blur=False,
rgbd=False,
rotate_conv_input=False,
use_encoder=False,
use_occupancy_net=False,
initial_depth=1.0):
super(StyleGenerator, self).__init__()
self.max_stage = 17
self.rgbd = rgbd
self.rotate_conv_input = rotate_conv_input
self.use_occupancy_net = use_occupancy_net
out_ch = 4 if rgbd else 3
with self.init_scope():
self.blocks = chainer.ChainList(
SynthesisBlock(ch, ch, ch, upsample=False), # 4
SynthesisBlock(ch,
ch,
ch,
upsample=True,
enable_blur=enable_blur), # 8
SynthesisBlock(ch,
ch,
ch,
upsample=True,
enable_blur=enable_blur), # 16
SynthesisBlock(ch,
ch,
ch,
upsample=True,
enable_blur=enable_blur), # 32
SynthesisBlock(ch // 2,
ch,
ch,
upsample=True,
enable_blur=enable_blur), # 64
SynthesisBlock(ch // 4,
ch // 2,
ch,
upsample=True,
enable_blur=enable_blur), # 128
# SynthesisBlock(ch // 8, ch // 4, ch, upsample=True, enable_blur=enable_blur), # 256
# SynthesisBlock(ch // 16, ch // 8, ch, upsample=True, enable_blur=enable_blur), # 512
# SynthesisBlock(ch // 32, ch // 16, ch, upsample=True, enable_blur=enable_blur) # 1024
)
self.outs = chainer.ChainList(
EqualizedConv2d(ch, out_ch, 1, 1, 0, gain=1),
EqualizedConv2d(ch, out_ch, 1, 1, 0, gain=1),
EqualizedConv2d(ch, out_ch, 1, 1, 0, gain=1),
EqualizedConv2d(ch, out_ch, 1, 1, 0, gain=1),
EqualizedConv2d(ch // 2, out_ch, 1, 1, 0, gain=1),
EqualizedConv2d(ch // 4, out_ch, 1, 1, 0, gain=1),
# EqualizedConv2d(ch // 8, out_ch, 1, 1, 0, gain=1),
# EqualizedConv2d(ch // 16, out_ch, 1, 1, 0, gain=1),
# EqualizedConv2d(ch // 32, out_ch, 1, 1, 0, gain=1)
)
if self.rgbd:
if self.rotate_conv_input:
self.l1 = EqualizedLinear(9, ch)
else:
self.l1 = EqualizedLinear(ch + 9, ch)
self.l2 = EqualizedLinear(ch, ch)
if rotate_conv_input:
self.rotate = StyleBlock(ch, ch)
if use_encoder:
self.enc = Encoder(ch, ch * 2, enable_blur=enable_blur)
if use_occupancy_net:
self.occupancy = OccupancyNet(in_ch=ch * 2 + 3, hidden_ch=32)
# initialize depth weight to 0
for out in self.outs:
out.c.W.array[-1] = 0
out.c.b.array[-1] = math.log(math.e**initial_depth - 1)
self.n_blocks = len(self.blocks)
self.image_size = 128
self.enable_blur = enable_blur
