def __call__(self, z, flow):
B, CH, T, Y, X = flow.shape
### back img generation
h = F.reshape(F.leaky_relu(self.bn_b0(self.l_b0(z))),
(B, self.gf_dim, 4, 4))
h = self.dc_b1(h)
h = self.dc_b2(h)
h = self.dc_b3(h)
h_back = F.tanh(self.dc_b4(h)) ### (B, CH, Y, X)
h_back = F.expand_dims(h_back, 2)
h_back = F.tile(h_back, (1, 1, T, 1, 1)) ### tile to (B, CH, T, Y, X)
### fore img generation
h_c = F.reshape(F.leaky_relu(self.bn_f0(self.l_f0(z))),
(B, self.gf_dim // 2, 2, 4, 4))
h_c = self.dc_f1(h_c)
h_c = self.dc_f2(h_c)
### colorize flow w U-net
## encode flow
h = flow
h_cm1 = self.c_m1(h)
h_cm2 = self.c_m2(h_cm1)
h = F.concat((h_cm2, h_c))
h_cm3 = self.c_m3(h)
h_cm4 = self.c_m4(h_cm3)
h_cm5 = self.c_m5(h_cm4)
## decode
h = self.dc_m1(h_cm5)
h = self.dc_m2(F.concat((h, h_cm4)))
h_dc3 = self.dc_m3(F.concat((h, h_cm3)))
h = self.dc_m4(F.concat((h_dc3, h_cm1)))
h_fore = F.tanh(self.dc_m5(h))
### make mask
h_mask = F.sigmoid(self.dc_mask(h_dc3))
h_mask = l1_penalty(h_mask, self.lamda)
h_mask = F.tile(h_mask, (1, 3, 1, 1, 1))
### calc video
x = h_mask * h_fore + (1 - h_mask) * h_back
if chainer.config.train:
return x
else:
return x, h_fore, h_back, h_mask
def __call__(self, z):
h = F.reshape(F.relu(self.bn0(self.l0(z))),
(z.data.shape[0], 512, (self.video_len // 16), 4, 4))
h = F.relu(self.bn1(self.dc1(h)))
h = F.relu(self.bn2(self.dc2(h)))
h = F.relu(self.bn3(self.dc3(h)))
h_fore = F.tanh(self.dc_fore(h))
h_mask = F.sigmoid(self.dc_mask(h))
h_mask = l1_penalty(h_mask)
h_mask_rep = F.tile(h_mask, (1, 2, 1, 1, 1))
x = h_mask_rep * h_fore
if chainer.config.train:
return x
else:
return x, h_fore, h_mask