import numpy as np
import torch
import torch.distributed as dist
import os
import copy
import cv2
import hashlib
import tqdm
import scipy
import uuid
import random
from PIL import Image
from utils.commons.hparams import hparams
from utils.commons.tensor_utils import tensors_to_scalars, convert_to_np, move_to_cuda
from utils.nn.model_utils import not_requires_grad, num_params
from utils.nn.schedulers import NoneSchedule
from utils.commons.ckpt_utils import load_ckpt, get_last_checkpoint, restore_weights, restore_opt_state
from utils.commons.base_task import BaseTask
from tasks.os_avatar.loss_utils.vgg19_loss import VGG19Loss
from modules.real3d.facev2v_warp.losses import PerceptualLoss
from tasks.os_avatar.dataset_utils.motion2video_dataset import Img2Plane_Dataset
import lpips
from utils.commons.dataset_utils import data_loader
from modules.real3d.img2plane_baseline import OSAvatar_Img2plane
from modules.eg3ds.models.triplane import TriPlaneGenerator
from modules.eg3ds.models.dual_discriminator import DualDiscriminator
from modules.eg3ds.torch_utils.ops import conv2d_gradfix
from modules.eg3ds.torch_utils.ops import upfirdn2d
from modules.eg3ds.models.dual_discriminator import filtered_resizing
class ScheduleForImg2Plane(NoneSchedule):
def __init__(self, optimizer, lr, lr_d, warmup_updates=0):
self.optimizer = optimizer
self.constant_lr = self.lr = lr
self.lr_d = lr_d
self.warmup_updates = warmup_updates
self.step(0)
def step(self, num_updates):
constant_lr = self.constant_lr
if self.warmup_updates > 0 and num_updates <= self.warmup_updates:
warmup = min(num_updates / self.warmup_updates, 1.0)
self.lr = max(constant_lr * warmup, 1e-7)
else:
self.lr = constant_lr
for optim_i in range(len(self.optimizer)-1):
self.optimizer[optim_i].param_groups[0]['lr'] = max(1e-5, self.lr * (hparams.get("lr_decay_rate", 0.95)) ** (num_updates // hparams.get("lr_decay_interval", 5_000))) # secc_img2plane
self.optimizer[optim_i].param_groups[1]['lr'] = max(1e-5, self.lr * (hparams.get("lr_decay_rate", 0.95)) ** (num_updates // hparams.get("lr_decay_interval", 5_000))) if num_updates >= min(2_000, hparams['start_adv_iters']) else 0 # decoder
# fix住来自预训练EG3D的超分, 给img2plane 30000步的warmup时间
self.optimizer[optim_i].param_groups[2]['lr'] = max(1e-5, self.lr * (hparams.get("lr_decay_rate", 0.95)) ** (num_updates // hparams.get("lr_decay_interval", 5_000))) if num_updates >= hparams['start_adv_iters'] else 0 # SR module
self.optimizer[-1].param_groups[0]['lr'] = self.lr_d * 1 # for disc
return self.lr
class OSAvatarImg2PlaneTask(BaseTask):
def __init__(self):
super().__init__()
if hparams['lpips_mode'] == 'vgg19':
self.criterion_lpips = VGG19Loss()
elif hparams['lpips_mode'] in ['vgg16', 'vgg']:
hparams['lpips_mode'] = 'vgg'
self.criterion_lpips = lpips.LPIPS(net=hparams['lpips_mode'],lpips=True)
elif hparams['lpips_mode'] == 'vgg19_v2':
self.criterion_lpips = PerceptualLoss()
else:
raise NotImplementedError
self.gen_tmp_output = {}
if hparams.get("use_kv_dataset", True):
self.dataset_cls = Img2Plane_Dataset
else:
raise NotImplementedError()
self.start_adv_iters = hparams["start_adv_iters"]
self.resample_filter = upfirdn2d.setup_filter([1,3,3,1])
@data_loader
def train_dataloader(self):
train_dataset = self.dataset_cls(prefix='train')
self.train_dl = train_dataset.get_dataloader()
return self.train_dl
@data_loader
def val_dataloader(self):
val_dataset = self.dataset_cls(prefix='val')
self.val_dl = val_dataset.get_dataloader()
return self.val_dl
@data_loader
def test_dataloader(self):
val_dataset = self.dataset_cls(prefix='val')
self.val_dl = val_dataset.get_dataloader()
return self.val_dl
def build_model(self):
self.eg3d_model = TriPlaneGenerator()
load_ckpt(self.eg3d_model, hparams['pretrained_eg3d_ckpt'], strict=True)
self.model = OSAvatar_Img2plane()
load_ckpt(self.model, hparams['pretrained_eg3d_ckpt'], strict=False, verbose=False)
self.disc = DualDiscriminator()
load_ckpt(self.disc, hparams['pretrained_eg3d_ckpt'], strict=False, model_name='disc')
if hparams.get('init_from_ckpt', '') != '':
# load_ckpt(self.model.img2plane_backbone, ckpt_dir, model_name='model.img2plane_backbone', strict=True)
load_ckpt(self.model, hparams['init_from_ckpt'], model_name='model', strict=False)
# load_ckpt(self.model, hparams['init_from_ckpt'], model_name='model', strict=True)
load_ckpt(self.disc, hparams['init_from_ckpt'], model_name='disc', strict=True)
# restore_weights(self, get_last_checkpoint(hparams.get('init_from_ckpt', ''))[0])
print(f"restored weights from {hparams.get('init_from_ckpt', '')}")
# define groups of params, used to assign optimizer.
self.img2plane_backbone_params = [p for k, p in self.model.img2plane_backbone.named_parameters() if p.requires_grad]
self.decoder_params = [p for p in self.model.decoder.parameters() if p.requires_grad]
self.upsample_params = [p for p in self.model.superresolution.parameters() if p.requires_grad]
self.disc_params = [p for k, p in self.disc.named_parameters() if p.requires_grad]
return self.model
def build_optimizer(self, model):
self.optimizer_gen = optimizer_gen = torch.optim.Adam(
self.img2plane_backbone_params,
lr=hparams['lr_g'],
betas=(hparams['optimizer_adam_beta1_g'], hparams['optimizer_adam_beta2_g'])
)
self.optimizer_gen.add_param_group({
'params': self.decoder_params,
'lr': hparams['lr_g'],
'betas': (hparams['optimizer_adam_beta1_g'], hparams['optimizer_adam_beta2_g'])
})
self.optimizer_gen.add_param_group({
'params': self.upsample_params,
'lr': hparams['lr_g'],
'betas': (hparams['optimizer_adam_beta1_g'], hparams['optimizer_adam_beta2_g'])
})
mb_ratio_d = hparams['reg_interval_d'] / (hparams['reg_interval_d'] + 1)
self.optimizer_disc = optimizer_disc = torch.optim.Adam(
self.disc_params,
lr=hparams['lr_d'] * mb_ratio_d,
betas=(hparams['optimizer_adam_beta1_d'] ** mb_ratio_d, hparams['optimizer_adam_beta2_d'] ** mb_ratio_d))
optimizers = [optimizer_gen] * 2 + [optimizer_disc] # optim0-1: ref/mv; optim2: disc
return optimizers
def build_scheduler(self, optimizer):
mb_ratio_d = hparams['reg_interval_d'] / (hparams['reg_interval_d'] + 1)
return ScheduleForImg2Plane(optimizer, hparams['lr_g'], hparams['lr_d'] * mb_ratio_d, hparams['warmup_updates'])
def on_train_start(self):
print("==============================")
num_params(self.model, model_name="Generator")
for n, m in self.model.named_children():
num_params(m, model_name="|-- "+n)
print("==============================")
num_params(self.disc, model_name="Discriminator")
for n, m in self.disc.named_children():
num_params(m, model_name="|-- "+n)
print("==============================")
def forward_G(self, img, camera, cond=None, ret=None, update_emas=False, cache_backbone=True, use_cached_backbone=False):
"""
ref_img: [B, 3, W, H]
camera: [b, 25], 16 dim c2w, and 9 dim intrinsic
"""
G = self.model
gen_output = G.forward(img=img, camera=camera, cond=cond, ret=ret, update_emas=update_emas, cache_backbone=cache_backbone, use_cached_backbone=use_cached_backbone)
return gen_output
def forward_D(self, img, camera, update_emas=False):
D = self.disc
logits = D.forward(img, camera, update_emas=update_emas)
return logits
def prepare_batch(self, batch):
# get synthesized img as GT data
out_batch = {}
ws_camera = batch['ffhq_ws_cameras']
camera = batch['ffhq_ref_cameras']
fake_camera = batch['ffhq_mv_cameras']
z = torch.randn([camera.shape[0], hparams['z_dim']],device=camera.device)
with torch.no_grad():
ws = self.eg3d_model.mapping(z, ws_camera, update_emas=False, truncation_psi=0.7)
ref_img_gen = self.eg3d_model.synthesis(ws, camera, cache_backbone=True, use_cached_backbone=False)
mv_img_gen = self.eg3d_model.synthesis(ws, fake_camera, cache_backbone=False, use_cached_backbone=True)
ref_img_gen['image_raw'] = filtered_resizing(ref_img_gen['image'], size=hparams['neural_rendering_resolution'], f=self.resample_filter, filter_mode='antialiased')
mv_img_gen['image_raw'] = filtered_resizing(mv_img_gen['image'], size=hparams['neural_rendering_resolution'], f=self.resample_filter, filter_mode='antialiased')
out_batch.update({
'ffhq_planes': ref_img_gen['plane'],
'ffhq_ref_imgs_raw': ref_img_gen['image_raw'],
'ffhq_ref_imgs': ref_img_gen['image'],
'ffhq_ref_imgs_depth': ref_img_gen['image_depth'],
'ffhq_mv_imgs_raw': mv_img_gen['image_raw'],
'ffhq_mv_imgs': mv_img_gen['image'],
'ffhq_mv_imgs_depth': mv_img_gen['image_depth'],
'ffhq_mv_imgs_feature': mv_img_gen['image_feature'],
'ffhq_ref_cameras': batch['ffhq_ref_cameras'],
'ffhq_mv_cameras': batch['ffhq_mv_cameras'],
})
return out_batch
def run_G_reference_image(self, batch):
losses = {}
ret = {}
if self.global_step+1 < self.start_adv_iters:
# at early stage, don't train on the ref camera too frequently, to prevent bill-board
if (self.global_step+1) % 5 != 0:
return losses
elif self.global_step < self.start_adv_iters + 5000:
if (self.global_step+1) % 2 != 0:
return losses
else:
# after 37,500 steps, we train on ref image every step.
pass
with torch.autograd.profiler.record_function('G_ref_forward'):
camera = batch['ffhq_ref_cameras']
img = batch['ffhq_ref_imgs']
img_raw = batch['ffhq_ref_imgs_raw']
img_depth = batch['ffhq_ref_imgs_depth']
gen_img = self.forward_G(img, camera, cond={'ref_cameras': batch['ffhq_ref_cameras']}, ret=ret)
if 'losses' in ret: losses.update(ret['losses'])
self.gen_tmp_output['recon_ref_imgs'] = gen_img['image'].detach()
self.gen_tmp_output['recon_ref_imgs_raw'] = gen_img['image_raw'].detach()
losses['G_ref_plane_l1_mean'] = (gen_img['plane'][:,:]).detach().abs().mean()
losses['G_ref_plane_l1_std'] = (gen_img['plane'][:,:]).detach().abs().std()
if hparams['use_mse']:
losses['G_ref_img_mae'] = (gen_img['image'] - img).pow(2).mean()
losses['G_ref_img_mae_raw'] = (gen_img['image_raw'] - img_raw).pow(2).mean()
else:
losses['G_ref_img_mae'] = (gen_img['image'] - img).abs().mean()
losses['G_ref_img_mae_raw'] = (gen_img['image_raw'] - img_raw).abs().mean()
pred_img_01 = ((gen_img['image'] + 1) / 2).clamp(0, 1)
pred_img_01_raw = ((gen_img['image_raw'] + 1) / 2).clamp(0, 1)
gt_img_01 = ((img + 1) / 2).clamp(0, 1)
gt_img_01_raw = ((img_raw + 1) / 2).clamp(0, 1)
losses['G_ref_img_lpips'] = self.criterion_lpips(pred_img_01, gt_img_01).mean()
losses['G_ref_img_lpips_raw'] = self.criterion_lpips(pred_img_01_raw, gt_img_01_raw).mean()
losses['G_ref_img_mae_depth'] = (gen_img['image_depth'] - img_depth).detach().abs().mean()
disc_inp_img = {
'image': gen_img['image'],
'image_raw': gen_img['image_raw'],
}
gen_logits = self.forward_D(disc_inp_img, camera)
losses['G_ref_adv'] = torch.nn.functional.softplus(-gen_logits).mean()
return losses
def run_G_multiview_image(self, batch):
losses = {}
ret = {}
with torch.autograd.profiler.record_function('G_mv_forward'):
camera = batch['ffhq_mv_cameras']
img = batch['ffhq_mv_imgs']
img_raw = batch['ffhq_mv_imgs_raw']
img_depth = batch['ffhq_mv_imgs_depth']
gen_img = self.forward_G(batch['ffhq_ref_imgs'], camera, cond={'ref_cameras': batch['ffhq_ref_cameras']}, ret=ret)
if 'losses' in ret: losses.update(ret['losses'])
self.gen_tmp_output['recon_mv_imgs'] = gen_img['image'].detach()
self.gen_tmp_output['recon_mv_imgs_raw'] = gen_img['image_raw'].detach()
losses['G_mv_plane_l1_mean'] = (gen_img['plane'][:,:]).detach().abs().mean()
losses['G_mv_plane_l1_std'] = (gen_img['plane'][:,:]).detach().abs().std()
if hparams['use_mse']:
losses['G_mv_img_mae'] = (gen_img['image'] - img).pow(2).mean()
losses['G_mv_img_mae_raw'] = (gen_img['image_raw'] - img_raw).pow(2).mean()
else:
losses['G_mv_img_mae'] = (gen_img['image'] - img).abs().mean()
losses['G_mv_img_mae_raw'] = (gen_img['image_raw'] - img_raw).abs().mean()
pred_img_01 = (gen_img['image'] + 1) / 2
pred_img_01_raw = (gen_img['image_raw'] + 1) / 2
gt_img_01 = (img + 1) / 2
gt_img_01_raw = (img_raw + 1) / 2
losses['G_mv_img_lpips'] = self.criterion_lpips(pred_img_01, gt_img_01).mean()
losses['G_mv_img_lpips_raw'] = self.criterion_lpips(pred_img_01_raw, gt_img_01_raw).mean()
losses['G_mv_img_mae_depth'] = (gen_img['image_depth'] - img_depth).detach().abs().mean()
disc_inp_img = {
'image': gen_img['image'],
'image_raw': gen_img['image_raw'],
}
gen_logits = self.forward_D(disc_inp_img, camera)
losses['G_mv_adv'] = torch.nn.functional.softplus(-gen_logits).mean()
return losses
def run_G_reg(self, batch):
losses = {}
imgs = batch['ffhq_mv_imgs']
if (self.global_step+1) % hparams['reg_interval_g'] == 0:
with torch.autograd.profiler.record_function('G_regularize_forward'):
initial_coordinates = torch.rand((imgs.shape[0], 1000, 3), device=imgs.device) - 0.5 # [-0.5,0.5]
perturbed_coordinates = initial_coordinates + torch.randn_like(initial_coordinates) * 5e-3
all_coordinates = torch.cat([initial_coordinates, perturbed_coordinates], dim=1)
source_sigma = self.model.sample(coordinates=all_coordinates, directions=torch.randn_like(all_coordinates), img=imgs, cond={'ref_cameras': batch['ffhq_mv_cameras']}, update_emas=False)['sigma']
source_sigma_initial = source_sigma[:, :source_sigma.shape[1]//2]
source_sigma_perturbed = source_sigma[:, source_sigma.shape[1]//2:]
density_reg_loss = torch.nn.functional.l1_loss(source_sigma_initial, source_sigma_perturbed)
# we want the pertubed position has similar density
losses['G_regularize_density_l1'] = density_reg_loss
return losses
def forward_D_main(self, batch):
"""
we update ema this substep.
"""
losses = {}
with torch.autograd.profiler.record_function('D_minimize_fake_forward'):
if 'recon_ref_imgs' in self.gen_tmp_output:
camera = batch['ffhq_ref_cameras']
disc_inp_img = {
'image': self.gen_tmp_output['recon_ref_imgs'],
'image_raw': self.gen_tmp_output['recon_ref_imgs_raw'],
}
gen_logits = self.forward_D(disc_inp_img, camera, update_emas=True)
losses['D_minimize_ref_fake'] = torch.nn.functional.softplus(gen_logits).mean()
if 'recon_mv_imgs' in self.gen_tmp_output:
camera = batch['ffhq_mv_cameras']
disc_inp_img = {
'image': self.gen_tmp_output['recon_mv_imgs'],
'image_raw': self.gen_tmp_output['recon_mv_imgs_raw'],
}
gen_logits = self.forward_D(disc_inp_img, camera, update_emas=True)
losses['D_minimize_mv_fake'] = torch.nn.functional.softplus(gen_logits).mean()
# Maximize confidence on true samples
with torch.autograd.profiler.record_function('D_maximize_true_forward'):
if hparams.get("ffhq_disc_inp_mode", "eg3d_gen") == 'eg3d_gen':
ref_cameras = batch['ffhq_ref_cameras']
ref_img_tmp_image = batch['ffhq_ref_imgs'].detach().requires_grad_(True)
ref_img_tmp_image_raw = batch['ffhq_ref_imgs_raw'].detach().requires_grad_(True)
elif hparams.get("ffhq_disc_inp_mode", "eg3d_gen") == 'ffhq':
ref_cameras = batch['ffhq_gt_cameras']
ref_img_tmp_image = batch['ffhq_gt_imgs'].detach().requires_grad_(True)
ref_img_tmp_image_raw = batch['ffhq_gt_imgs_raw'].detach().requires_grad_(True)
ref_img_tmp = {'image': ref_img_tmp_image, 'image_raw': ref_img_tmp_image_raw}
ref_logits = self.forward_D(ref_img_tmp, ref_cameras)
losses['D_maximize_ref'] = torch.nn.functional.softplus(-ref_logits).mean()
if hparams.get("ffhq_disc_inp_mode", "eg3d_gen") == 'eg3d_gen':
mv_cameras = batch['ffhq_mv_cameras']
mv_img_tmp_image = batch['ffhq_mv_imgs'].detach().requires_grad_(True)
mv_img_tmp_image_raw = batch['ffhq_mv_imgs_raw'].detach().requires_grad_(True)
mv_img_tmp = {'image': mv_img_tmp_image, 'image_raw': mv_img_tmp_image_raw}
mv_logits = self.forward_D(mv_img_tmp, mv_cameras)
losses['D_maximize_mv'] = torch.nn.functional.softplus(-mv_logits).mean()
if (self.global_step+1) % hparams['reg_interval_d'] == 0 and self.training is True:
with torch.autograd.profiler.record_function('D_gradient_penalty_on_real_imgs_forward'), conv2d_gradfix.no_weight_gradients():
ref_r1_grads = torch.autograd.grad(outputs=[ref_logits.sum()], inputs=[ref_img_tmp['image'], ref_img_tmp['image_raw']], create_graph=True, only_inputs=True)
ref_r1_grads_image = ref_r1_grads[0]
ref_r1_grads_image_raw = ref_r1_grads[1]
ref_r1_penalty_raw = ref_r1_grads_image_raw.square().sum([1,2,3]).mean()
ref_r1_penalty_image = ref_r1_grads_image.square().sum([1,2,3]).mean()
losses['D_ref_gradient_penalty'] = (ref_r1_penalty_image + ref_r1_penalty_raw) / 2
if hparams.get("ffhq_disc_inp_mode", "eg3d_gen") == 'eg3d_gen':
mv_r1_grads = torch.autograd.grad(outputs=[mv_logits.sum()], inputs=[mv_img_tmp['image'], mv_img_tmp['image_raw']], create_graph=True, only_inputs=True)
mv_r1_grads_image = mv_r1_grads[0]
mv_r1_grads_image_raw = mv_r1_grads[1]
mv_r1_penalty_raw = mv_r1_grads_image_raw.square().sum([1,2,3]).mean()
mv_r1_penalty_image = mv_r1_grads_image.square().sum([1,2,3]).mean()
losses['D_mv_gradient_penalty'] = (mv_r1_penalty_image + mv_r1_penalty_raw) / 2
self.gen_tmp_output = {}
return losses
def _training_step(self, sample, batch_idx, optimizer_idx):
if len(sample) == 0:
return None
if optimizer_idx == 0:
sample = self.prepare_batch(sample)
self.cache_sample = sample
else:
sample = self.cache_sample
# self.update_ema()
losses = {}
if optimizer_idx == 0:
losses.update(self.run_G_reference_image(sample))
# after the early stage, we decrease the lambda of error-based loss to prevent oversmoothness
loss_weights = {
'G_ref_img_mae': 1.0,
'G_ref_img_mae_raw': 1.0,
'G_ref_img_mae_depth': hparams.get('lambda_mse_depth', 0),
'G_ref_img_lpips': 0.1,
'G_ref_img_lpips_raw': 0.1,
'G_ref_adv': 0.0 if self.global_step < self.start_adv_iters else 0.1,
}
elif optimizer_idx == 1:
losses.update(self.run_G_multiview_image(sample))
losses.update(self.run_G_reg(sample))
loss_weights = {
'G_mv_img_mae': 1.0,
'G_mv_img_mae_raw': 1.0,
'G_mv_img_mae_depth': hparams.get('lambda_mse_depth', 0),
'G_mv_img_lpips': 0.1,
'G_mv_img_lpips_raw': 0.1,
'G_mv_adv': 0.0 if self.global_step < self.start_adv_iters else 0.025,
'G_regularize_density_l1': hparams['lambda_density_reg'],
}
elif optimizer_idx == 2:
losses.update(self.forward_D_main(sample))
loss_weights = {
'D_maximize_ref': 1.0,
'D_minimize_ref_fake': 1.0,
'D_ref_gradient_penalty': hparams['lambda_gradient_penalty'] * hparams['reg_interval_d'],
'D_maximize_mv': 1.0,
'D_minimize_mv_fake': 1.0,
'D_mv_gradient_penalty': hparams['lambda_gradient_penalty'] * hparams['reg_interval_d'],
}
total_loss = sum([loss_weights[k] * v for k, v in losses.items() if isinstance(v, torch.Tensor) and v.requires_grad])
if len(losses) == 0:
return None
return total_loss, losses
#####################
# Validation
#####################
def validation_start(self):
if self.global_step % hparams['valid_infer_interval'] == 0:
self.gen_dir = os.path.join(hparams['work_dir'], f'validation_results')
os.makedirs(self.gen_dir, exist_ok=True)
@torch.no_grad()
def validation_step(self, sample, batch_idx):
outputs = {}
losses = {}
if len(sample) == 0:
return None
sample = self.prepare_batch(sample)
rank = 0 if len(set(os.environ['CUDA_VISIBLE_DEVICES'].split(","))) == 1 else dist.get_rank()
losses.update(self.run_G_reference_image(sample))
losses.update(self.run_G_multiview_image(sample))
losses.update(self.forward_D_main(sample))
outputs['losses'] = losses
outputs['total_loss'] = sum(outputs['losses'].values())
outputs = tensors_to_scalars(outputs)
if self.global_step % hparams['valid_infer_interval'] == 0 \
and batch_idx < hparams['num_valid_plots'] and rank == 0:
imgs_ref = sample['ffhq_ref_imgs']
dummy_cond = {
'cond_src': torch.zeros([imgs_ref.shape[0], 3, 512, 512], dtype=sample['ffhq_mv_cameras'].dtype, device=sample['ffhq_mv_cameras'].device),
'cond_drv': torch.zeros([imgs_ref.shape[0], 3, 512, 512], dtype=sample['ffhq_mv_cameras'].dtype, device=sample['ffhq_mv_cameras'].device),
'ref_cameras': sample['ffhq_ref_cameras'],
}
gen_img = self.model.forward(imgs_ref, sample['ffhq_mv_cameras'], cond=dummy_cond, noise_mode='const')
gen_img_recon = self.model.forward(imgs_ref, sample['ffhq_ref_cameras'], cond=dummy_cond, noise_mode='const')
imgs_recon = gen_img_recon['image'].permute(0,2,3,1)
imgs_recon_raw = filtered_resizing(gen_img_recon['image_raw'], size=512, f=self.resample_filter, filter_mode='antialiased').permute(0, 2,3,1)
imgs_recon_depth = gen_img_recon['image_depth'].permute(0,2,3,1)
imgs_pred = gen_img['image'].permute(0,2,3,1)
imgs_pred_raw = filtered_resizing(gen_img['image_raw'], size=512, f=self.resample_filter, filter_mode='antialiased').permute(0, 2,3,1)
imgs_pred_depth = gen_img['image_depth'].permute(0,2,3,1)
imgs_ref = imgs_ref.permute(0,2,3,1)
imgs_mv = sample['ffhq_mv_imgs'].permute(0,2,3,1) # [B, H, W, 3]
for i in range(len(imgs_pred)):
idx_string = format(i+batch_idx * hparams['batch_size'], "05d")
base_fn = f"{idx_string}"
img_ref_mv_recon_pred = torch.cat([imgs_ref[i], imgs_mv[i], imgs_recon_raw[i], imgs_pred_raw[i], imgs_recon[i], imgs_pred[i]], dim=1)
self.save_rgb_to_fname(img_ref_mv_recon_pred, f"{self.gen_dir}/images_rgb_iter{self.global_step}/ref_mv_recon_pred_{base_fn}.png")
img_depth_recon_pred = torch.cat([imgs_recon_depth[i], imgs_pred_depth[i]], dim=1)
self.save_depth_to_fname(img_depth_recon_pred, f"{self.gen_dir}/images_depth_iter{self.global_step}/recon_pred_{base_fn}.png")
return outputs
def validation_end(self, outputs):
return super().validation_end(outputs)
@staticmethod
def save_rgb_to_fname(rgb, fname):
"""
rgb: [H, W, 3]
"""
os.makedirs(os.path.dirname(fname), exist_ok=True)
img = (rgb * 127.5 + 128).clamp(0, 255)
img = convert_to_np(img).astype(np.uint8)
Image.fromarray(img, 'RGB').save(fname)
@staticmethod
def save_depth_to_fname(depth, fname):
"""
depth: [H, W, 3]
"""
os.makedirs(os.path.dirname(fname), exist_ok=True)
low, high = depth.min(), depth.max()
img = (depth - low) * (255 / (high - low))
img = convert_to_np(img)
img = np.rint(img).clip(0, 255).astype(np.uint8)
Image.fromarray(img[:, :, 0], 'L').save(fname)