def get_test_output_images(self, data):
r"""Get the visualization output of test function.
Args:
data (dict): Training data at the current iteration.
"""
vis_images = [
self.visualize_label(data['label'][:, -1]),
tensor2im(data['images'][:, -1]),
tensor2im(self.net_G_output['fake_images']),
]
return vis_images
def visualize_label(self, label):
r"""Visualize the input label when saving to image.
Args:
label (tensor): Input label tensor.
"""
cfgdata = self.cfg.data
if hasattr(cfgdata, 'for_pose_dataset'):
label = tensor2pose(self.cfg, label)
elif hasattr(cfgdata, 'input_labels') and \
'seg_maps' in cfgdata.input_labels:
for input_type in cfgdata.input_types:
if 'seg_maps' in input_type:
num_labels = input_type['seg_maps'].num_channels
label = tensor2label(label, num_labels)
elif getattr(cfgdata, 'label_channels', 1) > 3:
label = tensor2im(label.sum(1, keepdim=True))
else:
label = tensor2im(label)
return label
def get_images(data, net_G_output, return_first_frame=True,
for_model_average=False):
r"""Get the ourput images to save.
Args:
data (dict): Training data for current iteration.
net_G_output (dict): Generator output.
return_first_frame (bool): Return output for first frame in the
sequence.
for_model_average (bool): For model average output.
Return:
vis_images (list of numpy arrays): Visualization images.
"""
frame_idx = 0 if return_first_frame else -1
warped_idx = 0 if return_first_frame else 1
vis_images = []
if not for_model_average:
vis_images += [
tensor2im(data['few_shot_images'][:, frame_idx]),
self.visualize_label(data['label'][:, frame_idx]),
tensor2im(data['images'][:, frame_idx])
]
vis_images += [
tensor2im(net_G_output['fake_images']),
tensor2im(net_G_output['fake_raw_images'])]
if not for_model_average:
vis_images += [
tensor2im(net_G_output['warped_images'][warped_idx]),
tensor2flow(net_G_output['fake_flow_maps'][warped_idx]),
tensor2im(net_G_output['fake_occlusion_masks'][warped_idx],
normalize=False)
]
return vis_images
def get_test_output_images(self, data):
r"""Get the visualization output of test function.
Args:
data (dict): Training data at the current iteration.
"""
# Visualize labels.
label_lengths = self.val_data_loader.dataset.get_label_lengths()
labels = split_labels(data['label'], label_lengths)
vis_labels = []
for key, value in labels.items():
if key == 'seg_maps':
vis_labels.append(self.visualize_label(value[:, -1]))
else:
vis_labels.append(tensor2im(value[:, -1]))
# Get gt image.
im = tensor2im(data['images'][:, -1])
# Get guidance image and masks.
if self.net_G_output['guidance_images_and_masks'] is not None:
guidance_image = tensor2im(
self.net_G_output['guidance_images_and_masks'][:, :3])
guidance_mask = tensor2im(
self.net_G_output['guidance_images_and_masks'][:, 3:4],
normalize=False)
else:
guidance_image = [np.zeros_like(item) for item in im]
guidance_mask = [np.zeros_like(item) for item in im]
# Create output.
vis_images = [
*vis_labels,
im,
guidance_image,
guidance_mask,
tensor2im(self.net_G_output['fake_images']),
]
return vis_images
def test_single(self, data, output_dir=None, save_fake_only=False):
r"""The inference function. If output_dir exists, also save the
output image.
Args:
data (dict): Training data at the current iteration.
output_dir (str): Save image directory.
save_fake_only (bool): Only save the fake output image.
"""
if self.is_inference and self.cfg.trainer.model_average:
test_in_model_average_mode = True
else:
test_in_model_average_mode = getattr(self,
'test_in_model_average_mode',
False)
data_t = self.get_data_t(data, self.net_G_output, self.data_prev, 0)
if self.sequence_length > 1:
self.data_prev = data_t
# Generator forward.
# Reset renderer if first time step.
if self.t == 0:
self.net_G_module.reset_renderer(
is_flipped_input=data['is_flipped'])
with torch.no_grad():
if test_in_model_average_mode:
net_G = self.net_G.module.averaged_model
else:
net_G = self.net_G
self.net_G_output = net_G(data_t)
if output_dir is not None:
if save_fake_only:
image_grid = tensor2im(self.net_G_output['fake_images'])[0]
else:
vis_images = self.get_test_output_images(data)
image_grid = np.hstack(
[np.vstack(im) for im in vis_images if im is not None])
if 'img_name' in data:
save_name = data['img_name'].split('.')[0] + '.jpg'
else:
save_name = '%04d.jpg' % self.t
output_filename = os.path.join(output_dir, save_name)
os.makedirs(output_dir, exist_ok=True)
imageio.imwrite(output_filename, image_grid)
self.t += 1
return self.net_G_output
def renderer_update_point_cloud(self, image, point_info):
r"""Update the renderer's color dictionary."""
if point_info is None or len(point_info) == 0:
return
# print('Updating the renderer.')
_, _, h, w = image.size()
# Renderer expects (h, w, c) [0-255] RGB image.
if isinstance(image, torch.Tensor):
image = tensor2im(image.detach())[0]
# Flip this image to correspond to SfM camera pose.
if self.is_flipped_input:
image = np.fliplr(image).copy()
self.renderer.update_point_cloud(image, point_info)
self.renderer_num_forwards += 1
def test(self, test_data_loader, root_output_dir, inference_args):
r"""Run inference on all sequences.
Args:
test_data_loader (object): Test data loader.
root_output_dir (str): Location to dump outputs.
inference_args (optional): Optional args.
"""
# Go over all sequences.
loader = test_data_loader
num_inference_sequences = loader.dataset.num_inference_sequences()
for sequence_idx in range(num_inference_sequences):
loader.dataset.set_inference_sequence_idx(sequence_idx)
print('Seq id: %d, Seq length: %d' %
(sequence_idx + 1, len(loader)))
# Reset model at start of new inference sequence.
self.reset()
self.sequence_length = len(loader)
# Go over all frames of this sequence.
video = []
for idx, data in enumerate(tqdm(loader)):
key = data['key']['images'][0][0]
filename = key.split('/')[-1]
# Create output dir for this sequence.
if idx == 0:
output_dir, seq_name = \
self.create_sequence_output_dir(root_output_dir, key)
video_path = os.path.join(output_dir, '..', seq_name)
# Get output, and save all vis to all/.
data['img_name'] = filename
data = to_cuda(data)
output = self.test_single(data, output_dir=output_dir + '/all')
# Dump just the fake image here.
fake = tensor2im(output['fake_images'])[0]
video.append(fake)
imageio.imsave(output_dir + '/fake/%s.jpg' % (filename), fake)
# Save as mp4 and gif.
imageio.mimsave(video_path + '.mp4', video, fps=15)
def test_single(self, data, output_dir=None, inference_args=None):
r"""The inference function. If output_dir exists, also save the
output image.
Args:
data (dict): Training data at the current iteration.
output_dir (str): Save image directory.
inference_args (obj): Inference args.
"""
if getattr(inference_args, 'finetune', False):
if not getattr(self, 'has_finetuned', False):
self.finetune(data, inference_args)
net_G = self.net_G
if self.test_in_model_average_mode:
net_G = net_G.module.averaged_model
net_G.eval()
data_t = self.get_data_t(data, self.net_G_output, self.data_prev, 0)
if self.is_inference or self.sequence_length > 1:
self.data_prev = data_t
# Generator forward.
with torch.no_grad():
self.net_G_output = net_G(data_t)
if output_dir is None:
return self.net_G_output
save_fake_only = getattr(inference_args, 'save_fake_only', False)
if save_fake_only:
image_grid = tensor2im(self.net_G_output['fake_images'])[0]
else:
vis_images = self.get_test_output_images(data)
image_grid = np.hstack(
[np.vstack(im) for im in vis_images if im is not None])
if 'img_name' in data:
save_name = data['img_name'].split('.')[0] + '.jpg'
else:
save_name = '%04d.jpg' % self.t
output_filename = os.path.join(output_dir, save_name)
os.makedirs(output_dir, exist_ok=True)
imageio.imwrite(output_filename, image_grid)
self.t += 1
return image_grid
def save_image(self, path, data):
r"""Save the output images to path.
Note when the generate_raw_output is FALSE. Then,
first_net_G_output['fake_raw_images'] is None and will not be displayed.
In model average mode, we will plot the flow visualization twice.
Args:
path (str): Save path.
data (dict): Training data for current iteration.
"""
self.net_G.eval()
if self.cfg.trainer.model_average:
self.net_G.module.averaged_model.eval()
self.net_G_output = None
with torch.no_grad():
first_net_G_output, net_G_output, all_info = self.gen_frames(data)
if self.cfg.trainer.model_average:
first_net_G_output_avg, net_G_output_avg = self.gen_frames(
data, use_model_average=True)
# Visualize labels.
label_lengths = self.train_data_loader.dataset.get_label_lengths()
labels = split_labels(data['label'], label_lengths)
vis_labels_start, vis_labels_end = [], []
for key, value in labels.items():
if 'seg_maps' in key:
vis_labels_start.append(self.visualize_label(value[:, -1]))
vis_labels_end.append(self.visualize_label(value[:, 0]))
else:
normalize = self.train_data_loader.dataset.normalize[key]
vis_labels_start.append(
tensor2im(value[:, -1], normalize=normalize))
vis_labels_end.append(
tensor2im(value[:, 0], normalize=normalize))
if is_master():
vis_images = [
*vis_labels_start,
tensor2im(data['images'][:, -1]),
tensor2im(net_G_output['fake_images']),
tensor2im(net_G_output['fake_raw_images'])
]
if self.cfg.trainer.model_average:
vis_images += [
tensor2im(net_G_output_avg['fake_images']),
tensor2im(net_G_output_avg['fake_raw_images'])
]
if self.sequence_length > 1:
if net_G_output['guidance_images_and_masks'] is not None:
guidance_image = tensor2im(
net_G_output['guidance_images_and_masks'][:, :3])
guidance_mask = tensor2im(
net_G_output['guidance_images_and_masks'][:, 3:4],
normalize=False)
else:
im = tensor2im(data['images'][:, -1])
guidance_image = [np.zeros_like(item) for item in im]
guidance_mask = [np.zeros_like(item) for item in im]
vis_images += [guidance_image, guidance_mask]
vis_images_first = [
*vis_labels_end,
tensor2im(data['images'][:, 0]),
tensor2im(first_net_G_output['fake_images']),
tensor2im(first_net_G_output['fake_raw_images']),
[np.zeros_like(item) for item in guidance_image],
[np.zeros_like(item) for item in guidance_mask]
]
if self.cfg.trainer.model_average:
vis_images_first += [
tensor2im(first_net_G_output_avg['fake_images']),
tensor2im(first_net_G_output_avg['fake_raw_images'])
]
if self.use_flow:
flow_gt, conf_gt = self.criteria['Flow'].flowNet(
data['images'][:, -1], data['images'][:, -2])
warped_image_gt = resample(data['images'][:, -1], flow_gt)
vis_images_first += [
tensor2flow(flow_gt),
tensor2im(conf_gt, normalize=False),
tensor2im(warped_image_gt),
]
vis_images += [
tensor2flow(net_G_output['fake_flow_maps']),
tensor2im(net_G_output['fake_occlusion_masks'],
normalize=False),
tensor2im(net_G_output['warped_images']),
]
if self.cfg.trainer.model_average:
vis_images_first += [
tensor2flow(flow_gt),
tensor2im(conf_gt, normalize=False),
tensor2im(warped_image_gt),
]
vis_images += [
tensor2flow(net_G_output_avg['fake_flow_maps']),
tensor2im(net_G_output_avg['fake_occlusion_masks'],
normalize=False),
tensor2im(net_G_output_avg['warped_images'])
]
vis_images = [[
np.vstack((im_first, im))
for im_first, im in zip(imgs_first, imgs)
] for imgs_first, imgs in zip(vis_images_first, vis_images)
if imgs is not None]
image_grid = np.hstack(
[np.vstack(im) for im in vis_images if im is not None])
print('Save output images to {}'.format(path))
os.makedirs(os.path.dirname(path), exist_ok=True)
imageio.imwrite(path, image_grid)
# Gather all inputs and outputs for dumping into video.
if self.sequence_length > 1:
input_images, output_images, output_guidance = [], [], []
for item in all_info['inputs']:
input_images.append(tensor2im(item['image'])[0])
for item in all_info['outputs']:
output_images.append(tensor2im(item['fake_images'])[0])
if item['guidance_images_and_masks'] is not None:
output_guidance.append(
tensor2im(
item['guidance_images_and_masks'][:, :3])[0])
else:
output_guidance.append(np.zeros_like(
output_images[-1]))
imageio.mimwrite(os.path.splitext(path)[0] + '.mp4',
output_images,
fps=2,
macro_block_size=None)
imageio.mimwrite(os.path.splitext(path)[0] + '_guidance.mp4',
output_guidance,
fps=2,
macro_block_size=None)
# for idx, item in enumerate(output_guidance):
# imageio.imwrite(os.path.splitext(
# path)[0] + '_guidance_%d.jpg' % (idx), item)
# for idx, item in enumerate(input_images):
# imageio.imwrite(os.path.splitext(
# path)[0] + '_input_%d.jpg' % (idx), item)
self.net_G.float()
