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