def test_assign_corner_interpolate():
inp = tensor(np.arange(1, 5, dtype=np.float32).reshape(1, 1, 2, 2))
out = F.interpolate(inp, [4, 4], align_corners=True)
out2 = F.interpolate(inp, scale_factor=2.0, align_corners=True)
assertTensorClose(out.numpy(), out2.numpy())
def test_many_batch_interpolate():
inp = tensor(np.arange(1, 9, dtype=np.float32).reshape(2, 1, 2, 2))
out = F.interpolate(inp, [4, 4])
out2 = F.interpolate(inp, scale_factor=2.0)
assertTensorClose(out.numpy(), out2.numpy())
def test_generator_batch(image, pre_S, pre_D, pre_S_hat, pre_D_hat, pre_SD, *,
netG):
netG.eval()
_, _, H, W = image.shape
image_S = F.interpolate(image, scale_factor=[0.25, 0.25])
image_S = F.interpolate(image_S, size=[H, W])
image_D = image - image_S
outputs = netG(image, image_S, image_D, pre_S, pre_D, pre_S_hat, pre_D_hat,
pre_SD)
return list(outputs)[:4] + [image_S, image_D]
def test_linear_interpolate():
inp = tensor(np.arange(1, 3, dtype=np.float32).reshape(1, 1, 2))
out = F.interpolate(inp, scale_factor=2.0, mode="LINEAR")
out2 = F.interpolate(inp, 4, mode="LINEAR")
assertTensorClose(out.numpy(),
np.array([[[1.0, 1.25, 1.75, 2.0]]], dtype=np.float32))
assertTensorClose(out2.numpy(),
np.array([[[1.0, 1.25, 1.75, 2.0]]], dtype=np.float32))
def test_generator_batch(LR, pre_S_hat, pre_D_hat, pre_SD, *, netG):
netG.eval()
B, T, _, H, W = LR.shape
# image
image_S = LR.reshape((B * T, -1, H, W))
image_S = F.interpolate(image_S, scale_factor=[0.25, 0.25])
image_S = F.interpolate(image_S, size=[H, W])
image_S = image_S.reshape((B, T, -1, H, W))
image_D = LR - image_S
outputs = netG(LR, image_S, image_D, pre_S_hat, pre_D_hat, pre_SD)
return list(outputs)[:4]
def train_generator_batch(image, label, *, opt, netG, netloss):
netG.train()
B, T, _, H, W = image.shape
# image
image_S = image.reshape((B * T, -1, H, W))
image_S = F.interpolate(image_S, scale_factor=[0.25, 0.25])
image_S = F.interpolate(image_S, size=[H, W])
image_S = image_S.reshape((B, T, -1, H, W))
image_D = image - image_S
# label
label_S = label.reshape((B * T, -1, 4 * H, 4 * W))
label_S = F.interpolate(label_S, scale_factor=[0.25, 0.25])
label_S = F.interpolate(label_S, size=[4 * H, 4 * W])
label_S = label_S.reshape((B, T, -1, 4 * H, 4 * W))
label_D = label - label_S
HR_G = []
HR_D = []
HR_S = []
pre_S_hat = mge.tensor(
np.zeros((B, hidden_channels, H, W), dtype=np.float32))
pre_D_hat = F.zeros_like(pre_S_hat)
pre_SD = F.zeros_like(pre_S_hat)
imgHR, pre_SD, pre_S_hat, pre_D_hat, img_S, img_D = netG(
image[:, 0, ...], image_S[:, 0, ...], image_D[:, 0, ...],
image_S[:, 1, ...], image_D[:, 1, ...], pre_S_hat, pre_D_hat, pre_SD)
HR_G.append(F.add_axis(imgHR, axis=1))
HR_D.append(F.add_axis(img_D, axis=1))
HR_S.append(F.add_axis(img_S, axis=1))
for t in range(1, T):
imgHR, pre_SD, pre_S_hat, pre_D_hat, img_S, img_D = netG(
image[:, t, ...], image_S[:, t, ...], image_D[:, t, ...],
image_S[:, t - 1, ...], image_D[:, t - 1,
...], pre_S_hat, pre_D_hat, pre_SD)
HR_G.append(F.add_axis(imgHR, axis=1))
HR_D.append(F.add_axis(img_S, axis=1))
HR_S.append(F.add_axis(img_D, axis=1))
HR_G = F.concat(HR_G, axis=1)
HR_D = F.concat(HR_D, axis=1)
HR_S = F.concat(HR_S, axis=1)
# assert HR_G.shape == HR_D.shape and HR_D.shape == HR_S.shape # [B,T,C,H,W]
loss = netloss(HR_G, HR_D, HR_S, label, label_D, label_S)
opt.backward(loss)
if dist.is_distributed():
# do all reduce mean
pass
return loss
def forward(self, x):
layers = self.backbone.extract_features(x)
up0 = self.aspp(layers["res5"])
up0 = self.dropout(up0)
up0 = F.interpolate(up0, scale_factor=self.sub_output_stride)
up1 = self.upstage1(layers["res2"])
up1 = F.concat([up0, up1], 1)
up2 = self.upstage2(up1)
out = self.convout(up2)
out = F.interpolate(out, scale_factor=4)
return out
def forward(self, x):
bottom_up_features = self.bottom_up.extract_features(x)
x = [bottom_up_features[f] for f in self.in_features[::-1]]
results = []
prev_features = self.lateral_convs[0](x[0])
results.append(self.output_convs[0](prev_features))
for features, lateral_conv, output_conv in zip(x[1:],
self.lateral_convs[1:],
self.output_convs[1:]):
top_down_features = F.interpolate(prev_features,
scale_factor=2,
mode="BILINEAR")
lateral_features = lateral_conv(features)
prev_features = lateral_features + top_down_features
results.insert(0, output_conv(prev_features))
if self.top_block is not None:
top_block_in_feature = bottom_up_features.get(
self.top_block.in_feature, None)
if top_block_in_feature is None:
top_block_in_feature = results[self._out_features.index(
self.top_block.in_feature)]
results.extend(self.top_block(top_block_in_feature, results[-1]))
return dict(zip(self._out_features, results))
def _upsample_conv(self, x, conv):
r"""
Helper function for performing convolution after upsampling.
"""
return conv(
F.interpolate(x,
scale_factor=2,
mode='bilinear',
align_corners=False))
def forward(self, x):
conv1 = self.conv1(x)
conv31 = self.conv2(x)
conv32 = self.conv3(x)
conv33 = self.conv4(x)
gp = F.mean(x, 2, True)
gp = F.mean(gp, 3, True)
gp = self.convgp(gp)
gp = F.interpolate(gp, (x.shapeof(2), x.shapeof(3)))
out = F.concat([conv1, conv31, conv32, conv33, gp], axis=1)
out = self.convout(out)
return out
def forward(self, x):
c1 = self.conv1(x)
p1 = self.pool(c1)
c2 = self.conv2(p1)
p2 = self.pool(c2)
c3 = self.conv3(p2)
p3 = self.pool(c3)
c4 = self.conv4(p3)
p4 = self.pool(c4)
c5 = self.conv5(p4)
# d5 = self.dropout(c5) # 崩
up_6 = F.interpolate(c5, [32, 32], align_corners=False)
merge6 = F.concat([up_6, c4], axis=1)
c6 = self.conv6(merge6)
up_7 = F.interpolate(c6, [64, 64], align_corners=False)
merge7 = F.concat([up_7, c3], axis=1)
c7 = self.conv7(merge7)
up_8 = F.interpolate(c7, [128, 128], align_corners=False)
merge8 = F.concat([up_8, c2], axis=1)
c8 = self.conv8(merge8)
up_9 = F.interpolate(c8, [256, 256], align_corners=False)
merge9 = F.concat([up_9, c1], axis=1)
c9 = self.conv9(merge9)
c10 = self.conv10(c9)
return c10
def forward(self, x):
bottom_up_features = self.bottom_up(x)
bottom_up_features = bottom_up_features[::-1]
results = []
prev_features = self.lateral_convs[0](bottom_up_features[0])
results.append(self.output_convs[0](prev_features))
for features, lateral_conv, output_conv in zip(bottom_up_features[1:],
self.lateral_convs[1:],
self.output_convs[1:]):
top_down_features = F.interpolate(prev_features,
scale_factor=2,
mode="BILINEAR")
lateral_features = lateral_conv(features)
prev_features = lateral_features + top_down_features
results.append(output_conv(prev_features))
# p6
last_p6 = F.max_pool2d(results[0], kernel_size=1, stride=2, padding=0)
results.insert(0, last_p6)
return results
def forward(self, x):
f = self.head(x)
outputs = []
for i in range(self.nr_stg):
multi_scale_features = self.stages["Stage_{}_body".format(i)](f)
multi_scale_heatmaps = []
for j in range(4):
out = self.stages["Stage_{}_tail".format(i)]["tail_{}".format(
j)](multi_scale_features[j])
out = F.interpolate(out, scale_factor=2**(3 - j))
multi_scale_heatmaps.append(out)
if i < self.nr_stg - 1:
f = self.stages["Stage_{}_next".format(i)](
multi_scale_features[-1])
outputs.append(multi_scale_heatmaps)
return outputs
def test_step(self, batchdata, **kwargs):
"""test step.
need to know whether the first frame for video, and every step restore some hidden state.
Args:
batchdata: list for train_batch, numpy.ndarray, length up to Collect class.
Returns:
list: outputs
"""
epoch = kwargs.get('epoch', 0)
image = batchdata[0] # [B,C,H,W]
image = ensemble_forward(image, Type=epoch) # for ensemble
H, W = image.shape[-2], image.shape[-1]
scale = getattr(self.generator, 'upscale_factor', 4)
padding_multi = self.eval_cfg.get('padding_multi', 1)
# padding for H and W
image = img_multi_padding(image,
padding_multi=padding_multi,
pad_value=-0.5) # [B,C,H,W]
assert image.shape[0] == 1 # only support batchsize 1
assert len(batchdata[1].shape) == 1 # first frame flag
if batchdata[1][0] > 0.5: # first frame
print("first frame")
self.now_test_num = 1
B, _, now_H, now_W = image.shape
print("use now_H : {} and now_W: {}".format(now_H, now_W))
self.pre_S_hat = np.zeros((B, hidden_channels, now_H, now_W),
dtype=np.float32)
self.pre_D_hat = np.zeros_like(self.pre_S_hat)
self.pre_SD = np.zeros_like(self.pre_S_hat)
self.pre_S = F.interpolate(mge.tensor(image),
scale_factor=[0.25, 0.25])
self.pre_S = F.interpolate(self.pre_S, size=[now_H, now_W]).numpy()
self.pre_D = image - self.pre_S
outputs = test_generator_batch(image,
self.pre_S,
self.pre_D,
self.pre_S_hat,
self.pre_D_hat,
self.pre_SD,
netG=self.generator)
outputs = list(outputs)
outputs[0] = img_de_multi_padding(outputs[0],
origin_H=H * scale,
origin_W=W * scale)
for i in range(0, 6):
outputs[i] = outputs[i].numpy()
# update hidden state
G, self.pre_SD, self.pre_S_hat, self.pre_D_hat, self.pre_S, self.pre_D = outputs
# back ensemble for G
G = ensemble_back(G, Type=epoch)
save_image_flag = kwargs.get('save_image')
if save_image_flag:
save_path = kwargs.get('save_path', None)
start_id = kwargs.get('sample_id', None)
if save_path is None or start_id is None:
raise RuntimeError(
"if save image in test_step, please set 'save_path' and 'sample_id' parameters"
)
for idx in range(G.shape[0]):
if epoch == 0:
imwrite(tensor2img(G[idx], min_max=(-0.5, 0.5)),
file_path=os.path.join(
save_path,
"idx_{}.png".format(start_id + idx)))
else:
imwrite(tensor2img(G[idx], min_max=(-0.5, 0.5)),
file_path=os.path.join(
save_path, "idx_{}_epoch_{}.png".format(
start_id + idx, epoch)))
print("now test num: {}".format(self.now_test_num))
self.now_test_num += 1
return outputs
def test_inappropriate_scale_linear_interpolate():
inp = tensor(np.arange(1, 3, dtype=np.float32).reshape(1, 1, 2))
with pytest.raises(ValueError):
F.interpolate(inp, scale_factor=[2.0, 3.0], mode="LINEAR")
def test_error_shape_linear_interpolate():
inp = tensor(np.arange(1, 5, dtype=np.float32).reshape(1, 1, 2, 2))
with pytest.raises(ValueError):
F.interpolate(inp, scale_factor=2.0, mode="LINEAR")