from multiprocessing.spawn import prepare from turtle import forward import torch import torch.nn as nn import torch.nn.functional as F import torchvision.models as models import nvdiffrast.torch as dr import numpy as np import matplotlib.pyplot as plt import os import os.path as osp from video3d.render.regularizer import get_edge_length, normal_consistency from . import networks from .renderer import * from .utils import misc, meters, flow_viz, arap, custom_loss from .dataloaders import get_sequence_loader, get_image_loader from .cub_dataloaders import get_cub_loader from .utils.skinning_v4 import estimate_bones, skinning import lpips from einops import rearrange from .geometry.dmtet import DMTetGeometry from .geometry.dlmesh import DLMesh from .render import renderutils as ru from .render import material from .render import mlptexture from .render import util from .render import mesh from .render import light from .render import render EPS = 1e-7 def get_optimizer(model, lr=0.0001, betas=(0.9, 0.999), weight_decay=0): return torch.optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr=lr, betas=betas, weight_decay=weight_decay) def set_requires_grad(model, requires_grad): if model is not None: for param in model.parameters(): param.requires_grad = requires_grad def forward_to_matrix(vec_forward, up=[0,1,0]): up = torch.FloatTensor(up).to(vec_forward.device) # vec_forward = nn.functional.normalize(vec_forward, p=2, dim=-1) # x right, y up, z forward vec_right = up.expand_as(vec_forward).cross(vec_forward, dim=-1) vec_right = nn.functional.normalize(vec_right, p=2, dim=-1) vec_up = vec_forward.cross(vec_right, dim=-1) vec_up = nn.functional.normalize(vec_up, p=2, dim=-1) rot_mat = torch.stack([vec_right, vec_up, vec_forward], -2) return rot_mat def sample_pose_hypothesis_from_quad_prediction(poses_raw, total_iter, batch_size, num_frames, pose_xflip_recon=False, input_image_xflip_flag=None, rot_temp_scalar=1., num_hypos=4, naive_probs_iter=2000, best_pose_start_iter=6000, random_sample=True): rots_pred = poses_raw[..., :num_hypos*4].view(-1, num_hypos, 4) rots_logits = rots_pred[..., 0] # Nx4 temp = 1 / np.clip(total_iter / 1000 / rot_temp_scalar, 1., 100.) rots_probs = torch.nn.functional.softmax(-rots_logits / temp, dim=1) # N x K # naive_probs = torch.FloatTensor([10] + [1] * (num_hypos - 1)).to(rots_logits.device) naive_probs = torch.ones(num_hypos).to(rots_logits.device) naive_probs = naive_probs / naive_probs.sum() naive_probs_weight = np.clip(1 - (total_iter - naive_probs_iter) / 2000, 0, 1) rots_probs = naive_probs.view(1, num_hypos) * naive_probs_weight + rots_probs * (1 - naive_probs_weight) rots_pred = rots_pred[..., 1:4] trans_pred = poses_raw[..., -3:] best_rot_idx = torch.argmax(rots_probs, dim=1) # N if random_sample: # rand_rot_idx = torch.randint(0, 4, (batch_size * num_frames,), device=poses_raw.device) # N rand_rot_idx = torch.randperm(batch_size * num_frames, device=poses_raw.device) % num_hypos # N # rand_rot_idx = torch.randperm(batch_size, device=poses_raw.device)[:,None].repeat(1, num_frames).view(-1) % 4 # N best_flag = (torch.randperm(batch_size * num_frames, device=poses_raw.device) / (batch_size * num_frames) < np.clip((total_iter - best_pose_start_iter)/2000, 0, 0.8)).long() rand_flag = 1 - best_flag # best_flag = torch.zeros_like(best_rot_idx) rot_idx = best_rot_idx * best_flag + rand_rot_idx * (1 - best_flag) else: rand_flag = torch.zeros_like(best_rot_idx) rot_idx = best_rot_idx rot_pred = torch.gather(rots_pred, 1, rot_idx[:, None, None].expand(-1, 1, 3))[:, 0] # Nx3 pose_raw = torch.cat([rot_pred, trans_pred], -1) rot_prob = torch.gather(rots_probs, 1, rot_idx[:, None].expand(-1, 1))[:, 0] # N rot_logit = torch.gather(rots_logits, 1, rot_idx[:, None].expand(-1, 1))[:, 0] # N if pose_xflip_recon: raise NotImplementedError rot_mat = forward_to_matrix(pose_raw[:, :3], up=[0, 1, 0]) pose = torch.cat([rot_mat.view(batch_size * num_frames, -1), pose_raw[:, 3:]], -1) return pose_raw, pose, rot_idx, rot_prob, rot_logit, rots_probs, rand_flag class PriorPredictor(nn.Module): def __init__(self, cfgs): super().__init__() dmtet_grid = cfgs.get('dmtet_grid', 64) grid_scale = cfgs.get('grid_scale', 5) prior_sdf_mode = cfgs.get('prior_sdf_mode', 'mlp') num_layers_shape = cfgs.get('num_layers_shape', 5) hidden_size = cfgs.get('hidden_size', 64) embedder_freq_shape = cfgs.get('embedder_freq_shape', 8) embed_concat_pts = cfgs.get('embed_concat_pts', True) init_sdf = cfgs.get('init_sdf', None) jitter_grid = cfgs.get('jitter_grid', 0.) perturb_sdf_iter = cfgs.get('perturb_sdf_iter', 10000) sym_prior_shape = cfgs.get('sym_prior_shape', False) self.netShape = DMTetGeometry(dmtet_grid, grid_scale, prior_sdf_mode, num_layers=num_layers_shape, hidden_size=hidden_size, embedder_freq=embedder_freq_shape, embed_concat_pts=embed_concat_pts, init_sdf=init_sdf, jitter_grid=jitter_grid, perturb_sdf_iter=perturb_sdf_iter, sym_prior_shape=sym_prior_shape) mlp_hidden_size = cfgs.get('hidden_size', 64) tet_bbox = self.netShape.getAABB() self.render_dino_mode = cfgs.get('render_dino_mode', None) num_layers_dino = cfgs.get("num_layers_dino", 5) dino_feature_recon_dim = cfgs.get('dino_feature_recon_dim', 64) sym_dino = cfgs.get("sym_dino", False) dino_min = torch.zeros(dino_feature_recon_dim) + cfgs.get('dino_min', 0.) dino_max = torch.zeros(dino_feature_recon_dim) + cfgs.get('dino_max', 1.) min_max = torch.stack((dino_min, dino_max), dim=0) if self.render_dino_mode is None: pass elif self.render_dino_mode == 'feature_mlpnv': self.netDINO = mlptexture.MLPTexture3D(tet_bbox, channels=dino_feature_recon_dim, internal_dims=mlp_hidden_size, hidden=num_layers_dino-1, feat_dim=0, min_max=min_max, bsdf=None, perturb_normal=False, symmetrize=sym_dino) elif self.render_dino_mode == 'feature_mlp': embedder_scaler = 2 * np.pi / grid_scale * 0.9 # originally (-0.5*s, 0.5*s) rescale to (-pi, pi) * 0.9 embed_concat_pts = cfgs.get('embed_concat_pts', True) self.netDINO = networks.MLPTextureSimple( 3, # x, y, z coordinates dino_feature_recon_dim, num_layers_dino, nf=mlp_hidden_size, dropout=0, activation="sigmoid", min_max=min_max, n_harmonic_functions=cfgs.get('embedder_freq_dino', 8), omega0=embedder_scaler, extra_dim=0, embed_concat_pts=embed_concat_pts, perturb_normal=False, symmetrize=sym_dino ) elif self.render_dino_mode == 'cluster': num_layers_dino = cfgs.get("num_layers_dino", 5) dino_cluster_dim = cfgs.get('dino_cluster_dim', 64) self.netDINO = mlptexture.MLPTexture3D(tet_bbox, channels=dino_cluster_dim, internal_dims=mlp_hidden_size, hidden=num_layers_dino-1, feat_dim=0, min_max=None, bsdf=None, perturb_normal=False, symmetrize=sym_dino) else: raise NotImplementedError def forward(self, perturb_sdf=False, total_iter=None, is_training=True): prior_shape = self.netShape.getMesh(perturb_sdf=perturb_sdf, total_iter=total_iter, jitter_grid=is_training) return prior_shape, self.netDINO class InstancePredictor(nn.Module): def __init__(self, cfgs, tet_bbox=None): super().__init__() self.cfgs = cfgs self.grid_scale = cfgs.get('grid_scale', 5) self.enable_encoder = cfgs.get('enable_encoder', False) if self.enable_encoder: encoder_latent_dim = cfgs.get('latent_dim', 256) encoder_pretrained = cfgs.get('encoder_pretrained', False) encoder_frozen = cfgs.get('encoder_frozen', False) encoder_arch = cfgs.get('encoder_arch', 'simple') in_image_size = cfgs.get('in_image_size', 256) self.dino_feature_input = cfgs.get('dino_feature_input', False) dino_feature_dim = cfgs.get('dino_feature_dim', 64) if encoder_arch == 'simple': if self.dino_feature_input: self.netEncoder = networks.EncoderWithDINO(cin_rgb=3, cin_dino=dino_feature_dim, cout=encoder_latent_dim, in_size=in_image_size, zdim=None, nf=64, activation=None) else: self.netEncoder = networks.Encoder(cin=3, cout=encoder_latent_dim, in_size=in_image_size, zdim=None, nf=64, activation=None) elif encoder_arch == 'vgg': self.netEncoder = networks.VGGEncoder(cout=encoder_latent_dim, pretrained=encoder_pretrained) elif encoder_arch == 'resnet': self.netEncoder = networks.ResnetEncoder(cout=encoder_latent_dim, pretrained=encoder_pretrained) elif encoder_arch == 'vit': which_vit = cfgs.get('which_vit', 'dino_vits8') vit_final_layer_type = cfgs.get('vit_final_layer_type', 'conv') self.netEncoder = networks.ViTEncoder(cout=encoder_latent_dim, which_vit=which_vit, pretrained=encoder_pretrained, frozen=encoder_frozen, in_size=in_image_size, final_layer_type=vit_final_layer_type) else: raise NotImplementedError else: encoder_latent_dim = 0 mlp_hidden_size = cfgs.get('hidden_size', 64) bsdf = cfgs.get("bsdf", 'diffuse') num_layers_tex = cfgs.get("num_layers_tex", 5) feat_dim = cfgs.get("latent_dim", 64) if self.enable_encoder else 0 perturb_normal = cfgs.get("perturb_normal", False) sym_texture = cfgs.get("sym_texture", False) kd_min = torch.FloatTensor(cfgs.get('kd_min', [0., 0., 0., 0.])) kd_max = torch.FloatTensor(cfgs.get('kd_max', [1., 1., 1., 1.])) ks_min = torch.FloatTensor(cfgs.get('ks_min', [0., 0., 0.])) ks_max = torch.FloatTensor(cfgs.get('ks_max', [0., 0., 0.])) nrm_min = torch.FloatTensor(cfgs.get('nrm_min', [-1., -1., 0.])) nrm_max = torch.FloatTensor(cfgs.get('nrm_max', [1., 1., 1.])) mlp_min = torch.cat((kd_min[0:3], ks_min, nrm_min), dim=0) mlp_max = torch.cat((kd_max[0:3], ks_max, nrm_max), dim=0) min_max = torch.stack((mlp_min, mlp_max), dim=0) out_chn = 9 # TODO: if the tet verts are deforming, we need to recompute tet_bbox texture_mode = cfgs.get("texture_mode", 'mlp') if texture_mode == 'mlpnv': self.netTexture = mlptexture.MLPTexture3D(tet_bbox, channels=out_chn, internal_dims=mlp_hidden_size, hidden=num_layers_tex-1, feat_dim=feat_dim, min_max=min_max, bsdf=bsdf, perturb_normal=perturb_normal, symmetrize=sym_texture) elif texture_mode == 'mlp': embedder_scaler = 2 * np.pi / self.grid_scale * 0.9 # originally (-0.5*s, 0.5*s) rescale to (-pi, pi) * 0.9 embed_concat_pts = cfgs.get('embed_concat_pts', True) self.netTexture = networks.MLPTextureSimple( 3, # x, y, z coordinates out_chn, num_layers_tex, nf=mlp_hidden_size, dropout=0, activation="sigmoid", min_max=min_max, n_harmonic_functions=cfgs.get('embedder_freq_tex', 10), omega0=embedder_scaler, extra_dim=feat_dim, embed_concat_pts=embed_concat_pts, perturb_normal=perturb_normal, symmetrize=sym_texture ) self.rot_rep = cfgs.get('rot_rep', 'euler_angle') self.enable_pose = cfgs.get('enable_pose', False) if self.enable_pose: cam_pos_z_offset = cfgs.get('cam_pos_z_offset', 10.) fov = cfgs.get('crop_fov_approx', 25) half_range = np.tan(fov /2 /180 * np.pi) * cam_pos_z_offset # 2.22 self.max_trans_xy_range = half_range * cfgs.get('max_trans_xy_range_ratio', 1.) self.max_trans_z_range = half_range * cfgs.get('max_trans_z_range_ratio', 1.) self.lookat_init = cfgs.get('lookat_init', None) self.lookat_zeroy = cfgs.get('lookat_zeroy', False) self.rot_temp_scalar = cfgs.get('rot_temp_scalar', 1.) self.naive_probs_iter = cfgs.get('naive_probs_iter', 2000) self.best_pose_start_iter = cfgs.get('best_pose_start_iter', 6000) if self.rot_rep == 'euler_angle': pose_cout = 6 elif self.rot_rep == 'quaternion': pose_cout = 7 elif self.rot_rep == 'lookat': pose_cout = 6 elif self.rot_rep == 'quadlookat': self.num_pose_hypos = 4 pose_cout = (3 + 1) * self.num_pose_hypos + 3 # 4 forward vectors for 4 quadrants, 4 quadrant classification logits, 3 for translation self.orthant_signs = torch.FloatTensor([[1,1,1], [-1,1,1], [-1,1,-1], [1,1,-1]]) elif self.rot_rep == 'octlookat': self.num_pose_hypos = 8 pose_cout = (3 + 1) * self.num_pose_hypos + 3 # 4 forward vectors for 8 octants, 8 octant classification logits, 3 for translation self.orthant_signs = torch.stack(torch.meshgrid([torch.arange(1, -2, -2)] *3), -1).view(-1, 3) # 8x3 else: raise NotImplementedError self.pose_arch = cfgs.get('pose_arch', 'mlp') if self.pose_arch == 'mlp': num_layers_pose = cfgs.get('num_layers_pose', 5) self.netPose = networks.MLP( encoder_latent_dim, pose_cout, num_layers_pose, nf=mlp_hidden_size, dropout=0, activation=None ) elif self.pose_arch == 'encoder': if self.dino_feature_input: dino_feature_dim = cfgs.get('dino_feature_dim', 64) self.netPose = networks.EncoderWithDINO(cin_rgb=3, cin_dino=dino_feature_dim, cout=pose_cout, in_size=in_image_size, zdim=None, nf=64, activation=None) else: self.netPose = networks.Encoder(cin=3, cout=pose_cout, in_size=in_image_size, zdim=None, nf=64, activation=None) elif self.pose_arch in ['encoder_dino_patch_out', 'encoder_dino_patch_key']: if which_vit == 'dino_vits8': dino_feat_dim = 384 elif which_vit == 'dinov2_vits14': dino_feat_dim = 384 elif which_vit == 'dino_vitb8': dino_feat_dim = 768 self.netPose = networks.Encoder32(cin=dino_feat_dim, cout=pose_cout, nf=256, activation=None) elif self.pose_arch == 'vit': encoder_pretrained = cfgs.get('encoder_pretrained', False) encoder_frozen = cfgs.get('encoder_frozen', False) which_vit = cfgs.get('which_vit', 'dino_vits8') vit_final_layer_type = cfgs.get('vit_final_layer_type', 'conv') self.netPose = networks.ViTEncoder(cout=encoder_latent_dim, which_vit=which_vit, pretrained=encoder_pretrained, frozen=encoder_frozen, in_size=in_image_size, final_layer_type=vit_final_layer_type) else: raise NotImplementedError self.enable_deform = cfgs.get('enable_deform', False) if self.enable_deform: embedder_scaler = 2 * np.pi / self.grid_scale * 0.9 # originally (-0.5*s, 0.5*s) rescale to (-pi, pi) * 0.9 embed_concat_pts = cfgs.get('embed_concat_pts', True) num_layers_deform = cfgs.get('num_layers_deform', 5) self.deform_epochs = np.arange(*cfgs.get('deform_epochs', [0, 0])) sym_deform = cfgs.get("sym_deform", False) self.netDeform = networks.MLPWithPositionalEncoding( 3, # x, y, z coordinates 3, # dx, dy, dz deformation num_layers_deform, nf=mlp_hidden_size, dropout=0, activation=None, n_harmonic_functions=cfgs.get('embedder_freq_deform', 10), omega0=embedder_scaler, extra_dim=encoder_latent_dim, embed_concat_pts=embed_concat_pts, symmetrize=sym_deform ) self.enable_articulation = cfgs.get('enable_articulation', False) if self.enable_articulation: self.num_body_bones = cfgs.get('num_body_bones', 4) self.articulation_multiplier = cfgs.get('articulation_multiplier', 1) self.static_root_bones = cfgs.get('static_root_bones', False) self.skinning_temperature = cfgs.get('skinning_temperature', 1) self.articulation_epochs = np.arange(*cfgs.get('articulation_epochs', [0, 0])) self.num_legs = cfgs.get('num_legs', 0) self.num_leg_bones = cfgs.get('num_leg_bones', 0) self.body_bones_type = cfgs.get('body_bones_type', 'z_minmax') self.perturb_articulation_epochs = np.arange(*cfgs.get('perturb_articulation_epochs', [0, 0])) self.num_bones = self.num_body_bones + self.num_legs * self.num_leg_bones self.constrain_legs = cfgs.get('constrain_legs', False) self.attach_legs_to_body_epochs = np.arange(*cfgs.get('attach_legs_to_body_epochs', [0, 0])) self.max_arti_angle = cfgs.get('max_arti_angle', 60) num_layers_arti = cfgs.get('num_layers_arti', 5) which_vit = cfgs.get('which_vit', 'dino_vits8') if which_vit == 'dino_vits8': dino_feat_dim = 384 elif which_vit == 'dino_vitb8': dino_feat_dim = 768 self.articulation_arch = cfgs.get('articulation_arch', 'mlp') self.articulation_feature_mode = cfgs.get('articulation_feature_mode', 'sample') embedder_freq_arti = cfgs.get('embedder_freq_arti', 8) if self.articulation_feature_mode == 'global': feat_dim = encoder_latent_dim elif self.articulation_feature_mode == 'sample': feat_dim = dino_feat_dim elif self.articulation_feature_mode == 'sample+global': feat_dim = encoder_latent_dim + dino_feat_dim if self.articulation_feature_mode == 'attention': arti_feat_attn_zdim = cfgs.get('arti_feat_attn_zdim', 128) pos_dim = 1 + 2 + 3*2 self.netFeatureAttn = networks.FeatureAttention(which_vit, pos_dim, embedder_freq_arti, arti_feat_attn_zdim, img_size=in_image_size) embedder_scaler = np.pi * 0.9 # originally (-1, 1) rescale to (-pi, pi) * 0.9 self.netArticulation = networks.ArticulationNetwork(self.articulation_arch, feat_dim, 1+2+3*2, num_layers_arti, mlp_hidden_size, n_harmonic_functions=embedder_freq_arti, omega0=embedder_scaler) self.kinematic_tree_epoch = -1 self.enable_lighting = cfgs.get('enable_lighting', False) if self.enable_lighting: num_layers_light = cfgs.get('num_layers_light', 5) amb_diff_min = torch.FloatTensor(cfgs.get('amb_diff_min', [0., 0.])) amb_diff_max = torch.FloatTensor(cfgs.get('amb_diff_max', [1., 1.])) intensity_min_max = torch.stack((amb_diff_min, amb_diff_max), dim=0) self.netLight = light.DirectionalLight(encoder_latent_dim, num_layers_light, mlp_hidden_size, intensity_min_max=intensity_min_max) self.cam_pos_z_offset = cfgs.get('cam_pos_z_offset', 10.) self.crop_fov_approx = cfgs.get("crop_fov_approx", 25) def forward_encoder(self, images, dino_features=None): images_in = images.view(-1, *images.shape[2:]) * 2 - 1 # rescale to (-1, 1) patch_out = patch_key = None if self.dino_feature_input and self.cfgs.get('encoder_arch', 'simple') != 'vit': dino_features_in = dino_features.view(-1, *dino_features.shape[2:]) * 2 - 1 # rescale to (-1, 1) feat_out = self.netEncoder(images_in, dino_features_in) # Shape: (B, latent_dim) elif self.cfgs.get('encoder_arch', 'simple') == 'vit': feat_out, feat_key, patch_out, patch_key = self.netEncoder(images_in, return_patches=True) else: feat_out = self.netEncoder(images_in) # Shape: (B, latent_dim) return feat_out, feat_key, patch_out, patch_key def forward_pose(self, images, feat, patch_out, patch_key, dino_features): if self.pose_arch == 'mlp': pose = self.netPose(feat) elif self.pose_arch == 'encoder': images_in = images.view(-1, *images.shape[2:]) * 2 - 1 # rescale to (-1, 1) if self.dino_feature_input: dino_features_in = dino_features.view(-1, *dino_features.shape[2:]) * 2 - 1 # rescale to (-1, 1) pose = self.netPose(images_in, dino_features_in) # Shape: (B, latent_dim) else: pose = self.netPose(images_in) # Shape: (B, latent_dim) elif self.pose_arch == 'vit': images_in = images.view(-1, *images.shape[2:]) * 2 - 1 # rescale to (-1, 1) pose = self.netPose(images_in) elif self.pose_arch == 'encoder_dino_patch_out': pose = self.netPose(patch_out) # Shape: (B, latent_dim) elif self.pose_arch == 'encoder_dino_patch_key': pose = self.netPose(patch_key) # Shape: (B, latent_dim) else: raise NotImplementedError trans_pred = pose[...,-3:].tanh() * torch.FloatTensor([self.max_trans_xy_range, self.max_trans_xy_range, self.max_trans_z_range]).to(pose.device) if self.rot_rep == 'euler_angle': multiplier = 1. if self.gradually_expand_yaw: # multiplier += (min(iteration, 20000) // 500) * 0.25 multiplier *= 1.2 ** (min(iteration, 20000) // 500) # 1.125^40 = 111.200 rot_pred = torch.cat([pose[...,:1], pose[...,1:2]*multiplier, pose[...,2:3]], -1).tanh() rot_pred = rot_pred * torch.FloatTensor([self.max_rot_x_range, self.max_rot_y_range, self.max_rot_z_range]).to(pose.device) /180 * np.pi elif self.rot_rep == 'quaternion': quat_init = torch.FloatTensor([0.01,0,0,0]).to(pose.device) rot_pred = pose[...,:4] + quat_init rot_pred = nn.functional.normalize(rot_pred, p=2, dim=-1) # rot_pred = torch.cat([rot_pred[...,:1].abs(), rot_pred[...,1:]], -1) # make real part non-negative rot_pred = rot_pred * rot_pred[...,:1].sign() # make real part non-negative elif self.rot_rep == 'lookat': vec_forward_raw = pose[...,:3] if self.lookat_init is not None: vec_forward_raw = vec_forward_raw + torch.FloatTensor(self.lookat_init).to(pose.device) if self.lookat_zeroy: vec_forward_raw = vec_forward_raw * torch.FloatTensor([1,0,1]).to(pose.device) vec_forward_raw = nn.functional.normalize(vec_forward_raw, p=2, dim=-1) # x right, y up, z forward rot_pred = vec_forward_raw elif self.rot_rep in ['quadlookat', 'octlookat']: rots_pred = pose[..., :self.num_pose_hypos*4].view(-1, self.num_pose_hypos, 4) # (B, T, K, 4) rots_logits = rots_pred[..., :1] vec_forward_raw = rots_pred[..., 1:4] xs, ys, zs = vec_forward_raw.unbind(-1) margin = 0. xs = nn.functional.softplus(xs, beta=np.log(2)/(0.5+margin)) - margin # initialize to 0.5 if self.rot_rep == 'octlookat': ys = nn.functional.softplus(ys, beta=np.log(2)/(0.5+margin)) - margin # initialize to 0.5 if self.lookat_zeroy: ys = ys * 0 zs = nn.functional.softplus(zs, beta=2*np.log(2)) # initialize to 0.5 vec_forward_raw = torch.stack([xs, ys, zs], -1) vec_forward_raw = vec_forward_raw * self.orthant_signs.to(pose.device) vec_forward_raw = nn.functional.normalize(vec_forward_raw, p=2, dim=-1) # x right, y up, z forward rot_pred = torch.cat([rots_logits, vec_forward_raw], -1).view(-1, self.num_pose_hypos*4) else: raise NotImplementedError pose = torch.cat([rot_pred, trans_pred], -1) return pose def forward_deformation(self, shape, feat=None): original_verts = shape.v_pos num_verts = original_verts.shape[1] if feat is not None: deform_feat = feat[:, None, :].repeat(1, num_verts, 1) # Shape: (B, num_verts, latent_dim) original_verts = original_verts.repeat(len(feat),1,1) deformation = self.netDeform(original_verts, deform_feat) * 0.1 # Shape: (B, num_verts, 3) shape = shape.deform(deformation) return shape, deformation def forward_articulation(self, shape, feat, patch_feat, mvp, w2c, batch_size, num_frames, epoch): """ Forward propagation of articulation. For each bone, the network takes: 1) the 3D location of the bone; 2) the feature of the patch which the bone is projected to; and 3) an encoding of the bone's index to predict the bone's rotation (represented by an Euler angle). Args: shape: a Mesh object, whose v_pos has batch size BxF or 1. feat: the feature of the patches. Shape: (BxF, feat_dim, num_patches_per_axis, num_patches_per_axis) mvp: the model-view-projection matrix. Shape: (BxF, 4, 4) Returns: shape: a Mesh object, whose v_pos has batch size BxF (collapsed). articulation_angles: the predicted bone rotations. Shape: (B, F, num_bones, 3) aux: a dictionary containing auxiliary information. """ verts = shape.v_pos if len(verts) == 1: verts = verts[None] else: verts = verts.view(batch_size, num_frames, *verts.shape[1:]) if self.kinematic_tree_epoch != epoch: # if (epoch == self.articulation_epochs[0]) and (self.kinematic_tree_epoch != epoch): # if (epoch in [self.articulation_epochs[0], self.articulation_epochs[0]+2, self.articulation_epochs[0]+4]) and (self.kinematic_tree_epoch != epoch): attach_legs_to_body = epoch in self.attach_legs_to_body_epochs bones, self.kinematic_tree, self.bone_aux = estimate_bones(verts.detach(), self.num_body_bones, n_legs=self.num_legs, n_leg_bones=self.num_leg_bones, body_bones_type=self.body_bones_type, compute_kinematic_chain=True, attach_legs_to_body=attach_legs_to_body) self.kinematic_tree_epoch = epoch else: bones = estimate_bones(verts.detach(), self.num_body_bones, n_legs=self.num_legs, n_leg_bones=self.num_leg_bones, body_bones_type=self.body_bones_type, compute_kinematic_chain=False, aux=self.bone_aux) bones_pos = bones # Shape: (B, F, K, 2, 3) if batch_size > bones_pos.shape[0] or num_frames > bones_pos.shape[1]: assert bones_pos.shape[0] == 1 and bones_pos.shape[1] == 1, "If there is a mismatch, then there must be only one canonical mesh." bones_pos = bones_pos.repeat(batch_size, num_frames, 1, 1, 1) num_bones = bones_pos.shape[2] bones_pos = bones_pos.view(batch_size*num_frames, num_bones, 2, 3) # NxKx2x3 bones_mid_pos = bones_pos.mean(2) # NxKx3 bones_idx = torch.arange(num_bones).to(bones_pos.device) bones_mid_pos_world4 = torch.cat([bones_mid_pos, torch.ones_like(bones_mid_pos[..., :1])], -1) # NxKx4 bones_mid_pos_clip4 = bones_mid_pos_world4 @ mvp.transpose(-1, -2) bones_mid_pos_uv = bones_mid_pos_clip4[..., :2] / bones_mid_pos_clip4[..., 3:4] bones_mid_pos_uv = bones_mid_pos_uv.detach() bones_pos_world4 = torch.cat([bones_pos, torch.ones_like(bones_pos[..., :1])], -1) # NxKx2x4 bones_pos_cam4 = bones_pos_world4 @ w2c[:,None].transpose(-1, -2) bones_pos_cam3 = bones_pos_cam4[..., :3] / bones_pos_cam4[..., 3:4] bones_pos_cam3 = bones_pos_cam3 + torch.FloatTensor([0, 0, self.cam_pos_z_offset]).to(bones_pos_cam3.device).view(1, 1, 1, 3) bones_pos_in = bones_pos_cam3.view(batch_size*num_frames, num_bones, 2*3) / self.grid_scale * 2 # (-1, 1), NxKx(2*3) bones_idx_in = ((bones_idx[None, :, None] + 0.5) / num_bones * 2 - 1).repeat(batch_size * num_frames, 1, 1) # (-1, 1) bones_pos_in = torch.cat([bones_mid_pos_uv, bones_pos_in, bones_idx_in], -1).detach() if self.articulation_feature_mode == 'global': bones_patch_features = feat[:, None].repeat(1, num_bones, 1) # (BxF, K, feat_dim) elif self.articulation_feature_mode == 'sample': bones_patch_features = F.grid_sample(patch_feat, bones_mid_pos_uv.view(batch_size * num_frames, 1, -1, 2), mode='bilinear').squeeze(dim=-2).permute(0, 2, 1) # (BxF, K, feat_dim) elif self.articulation_feature_mode == 'sample+global': bones_patch_features = F.grid_sample(patch_feat, bones_mid_pos_uv.view(batch_size * num_frames, 1, -1, 2), mode='bilinear').squeeze(dim=-2).permute(0, 2, 1) # (BxF, K, feat_dim) bones_patch_features = torch.cat([feat[:, None].repeat(1, num_bones, 1), bones_patch_features], -1) elif self.articulation_feature_mode == 'attention': bones_patch_features = self.netFeatureAttn(bones_pos_in, patch_feat) else: raise NotImplementedError articulation_angles = self.netArticulation(bones_patch_features, bones_pos_in).view(batch_size, num_frames, num_bones, 3) * self.articulation_multiplier if self.static_root_bones: root_bones = [self.num_body_bones // 2 - 1, self.num_body_bones - 1] tmp_mask = torch.ones_like(articulation_angles) tmp_mask[:, :, root_bones] = 0 articulation_angles = articulation_angles * tmp_mask articulation_angles = articulation_angles.tanh() if self.constrain_legs: leg_bones_posx = [self.num_body_bones + i for i in range(self.num_leg_bones * self.num_legs // 2)] leg_bones_negx = [self.num_body_bones + self.num_leg_bones * self.num_legs // 2 + i for i in range(self.num_leg_bones * self.num_legs // 2)] tmp_mask = torch.zeros_like(articulation_angles) tmp_mask[:, :, leg_bones_posx + leg_bones_negx, 2] = 1 articulation_angles = tmp_mask * (articulation_angles * 0.3) + (1 - tmp_mask) * articulation_angles # no twist tmp_mask = torch.zeros_like(articulation_angles) tmp_mask[:, :, leg_bones_posx + leg_bones_negx, 1] = 1 articulation_angles = tmp_mask * (articulation_angles * 0.3) + (1 - tmp_mask) * articulation_angles # (-0.4, 0.4), limit side bending if epoch in self.perturb_articulation_epochs: articulation_angles = articulation_angles + torch.randn_like(articulation_angles) * 0.1 articulation_angles = articulation_angles * self.max_arti_angle / 180 * np.pi verts_articulated, aux = skinning(verts, bones, self.kinematic_tree, articulation_angles, output_posed_bones=True, temperature=self.skinning_temperature) verts_articulated = verts_articulated.view(batch_size*num_frames, *verts_articulated.shape[2:]) v_tex = shape.v_tex if len(v_tex) != len(verts_articulated): v_tex = v_tex.repeat(len(verts_articulated), 1, 1) shape = mesh.make_mesh( verts_articulated, shape.t_pos_idx, v_tex, shape.t_tex_idx, shape.material) return shape, articulation_angles, aux def get_camera_extrinsics_from_pose(self, pose, znear=0.1, zfar=1000.): N = len(pose) cam_pos_offset = torch.FloatTensor([0, 0, -self.cam_pos_z_offset]).to(pose.device) pose_R = pose[:, :9].view(N, 3, 3).transpose(2, 1) pose_T = pose[:, -3:] + cam_pos_offset[None, None, :] pose_T = pose_T.view(N, 3, 1) pose_RT = torch.cat([pose_R, pose_T], axis=2) # Nx3x4 w2c = torch.cat([pose_RT, torch.FloatTensor([0, 0, 0, 1]).repeat(N, 1, 1).to(pose.device)], axis=1) # Nx4x4 # We assume the images are perfect square. proj = util.perspective(self.crop_fov_approx / 180 * np.pi, 1, znear, zfar)[None].to(pose.device) mvp = torch.matmul(proj, w2c) campos = -torch.matmul(pose_R.transpose(2, 1), pose_T).view(N, 3) return mvp, w2c, campos def forward(self, images=None, prior_shape=None, epoch=None, dino_features=None, dino_clusters=None, total_iter=None, is_training=True): batch_size, num_frames = images.shape[:2] if self.enable_encoder: feat_out, feat_key, patch_out, patch_key = self.forward_encoder(images, dino_features) else: feat_out = feat_key = patch_out = patch_key = None shape = prior_shape texture = self.netTexture multi_hypothesis_aux = {} if self.enable_pose: poses_raw = self.forward_pose(images, feat_out, patch_out, patch_key, dino_features) pose_raw, pose, rot_idx, rot_prob, rot_logit, rots_probs, rand_pose_flag = sample_pose_hypothesis_from_quad_prediction(poses_raw, total_iter, batch_size, num_frames, rot_temp_scalar=self.rot_temp_scalar, num_hypos=self.num_pose_hypos, naive_probs_iter=self.naive_probs_iter, best_pose_start_iter=self.best_pose_start_iter, random_sample=is_training) multi_hypothesis_aux['rot_idx'] = rot_idx multi_hypothesis_aux['rot_prob'] = rot_prob multi_hypothesis_aux['rot_logit'] = rot_logit multi_hypothesis_aux['rots_probs'] = rots_probs multi_hypothesis_aux['rand_pose_flag'] = rand_pose_flag else: raise NotImplementedError mvp, w2c, campos = self.get_camera_extrinsics_from_pose(pose) deformation = None if self.enable_deform and epoch in self.deform_epochs: shape, deformation = self.forward_deformation(shape, feat_key) arti_params, articulation_aux = None, {} if self.enable_articulation and epoch in self.articulation_epochs: shape, arti_params, articulation_aux = self.forward_articulation(shape, feat_key, patch_key, mvp, w2c, batch_size, num_frames, epoch) if self.enable_lighting: light = self.netLight else: light = None aux = articulation_aux aux.update(multi_hypothesis_aux) return shape, pose_raw, pose, mvp, w2c, campos, texture, feat_out, deformation, arti_params, light, aux class Unsup3D: def __init__(self, cfgs): self.cfgs = cfgs self.device = cfgs.get('device', 'cpu') self.in_image_size = cfgs.get('in_image_size', 128) self.out_image_size = cfgs.get('out_image_size', 128) self.num_epochs = cfgs.get('num_epochs', 10) self.lr = cfgs.get('lr', 1e-4) self.use_scheduler = cfgs.get('use_scheduler', False) if self.use_scheduler: scheduler_milestone = cfgs.get('scheduler_milestone', [1,2,3,4,5]) scheduler_gamma = cfgs.get('scheduler_gamma', 0.5) self.make_scheduler = lambda optim: torch.optim.lr_scheduler.MultiStepLR(optim, milestones=scheduler_milestone, gamma=scheduler_gamma) self.cam_pos_z_offset = cfgs.get('cam_pos_z_offset', 10.) self.full_size_h = cfgs.get('full_size_h', 1080) self.full_size_w = cfgs.get('full_size_w', 1920) # self.fov_w = cfgs.get('fov_w', 60) # self.fov_h = np.arctan(np.tan(self.fov_w /2 /180*np.pi) / self.full_size_w * self.full_size_h) *2 /np.pi*180 # 36 self.crop_fov_approx = cfgs.get("crop_fov_approx", 25) self.mesh_regularization_mode = cfgs.get('mesh_regularization_mode', 'seq') self.enable_prior = cfgs.get('enable_prior', False) if self.enable_prior: self.netPrior = PriorPredictor(self.cfgs) self.prior_lr = cfgs.get('prior_lr', self.lr) self.prior_weight_decay = cfgs.get('prior_weight_decay', 0.) self.prior_only_epochs = cfgs.get('prior_only_epochs', 0) self.netInstance = InstancePredictor(self.cfgs, tet_bbox=self.netPrior.netShape.getAABB()) self.perturb_sdf = cfgs.get('perturb_sdf', False) self.blur_mask = cfgs.get('blur_mask', False) self.blur_mask_iter = cfgs.get('blur_mask_iter', 1) self.seqshape_epochs = np.arange(*cfgs.get('seqshape_epochs', [0, self.num_epochs])) self.avg_texture_epochs = np.arange(*cfgs.get('avg_texture_epochs', [0, 0])) self.swap_texture_epochs = np.arange(*cfgs.get('swap_texture_epochs', [0, 0])) self.swap_priorshape_epochs = np.arange(*cfgs.get('swap_priorshape_epochs', [0, 0])) self.avg_seqshape_epochs = np.arange(*cfgs.get('avg_seqshape_epochs', [0, 0])) self.swap_seqshape_epochs = np.arange(*cfgs.get('swap_seqshape_epochs', [0, 0])) self.pose_epochs = np.arange(*cfgs.get('pose_epochs', [0, 0])) self.pose_iters = cfgs.get('pose_iters', 0) self.deform_type = cfgs.get('deform_type', None) self.mesh_reg_decay_epoch = cfgs.get('mesh_reg_decay_epoch', 0) self.sdf_reg_decay_start_iter = cfgs.get('sdf_reg_decay_start_iter', 0) self.mesh_reg_decay_rate = cfgs.get('mesh_reg_decay_rate', 1) self.texture_epochs = np.arange(*cfgs.get('texture_epochs', [0, self.num_epochs])) self.zflip_epochs = np.arange(*cfgs.get('zflip_epochs', [0, self.num_epochs])) self.lookat_zflip_loss_epochs = np.arange(*cfgs.get('lookat_zflip_loss_epochs', [0, self.num_epochs])) self.lookat_zflip_no_other_losses = cfgs.get('lookat_zflip_no_other_losses', False) self.flow_loss_epochs = np.arange(*cfgs.get('flow_loss_epochs', [0, self.num_epochs])) self.sdf_inflate_reg_loss_epochs = np.arange(*cfgs.get('sdf_inflate_reg_loss_epochs', [0, self.num_epochs])) self.arti_reg_loss_epochs = np.arange(*cfgs.get('arti_reg_loss_epochs', [0, self.num_epochs])) self.background_mode = cfgs.get('background_mode', 'background') self.shape_prior_type = cfgs.get('shape_prior_type', 'deform') self.backward_prior = cfgs.get('backward_prior', True) self.resume_prior_optim = cfgs.get('resume_prior_optim', True) self.dmtet_grid_smaller_epoch = cfgs.get('dmtet_grid_smaller_epoch', 0) self.dmtet_grid_smaller = cfgs.get('dmtet_grid_smaller', 128) self.dmtet_grid = cfgs.get('dmtet_grid', 256) self.pose_xflip_recon_epochs = np.arange(*cfgs.get('pose_xflip_recon_epochs', [0, 0])) self.rot_rand_quad_epochs = np.arange(*cfgs.get('rot_rand_quad_epochs', [0, 0])) self.rot_all_quad_epochs = np.arange(*cfgs.get('rot_all_quad_epochs', [0, 0])) ## perceptual loss if cfgs.get('perceptual_loss_weight', 0.) > 0: self.perceptual_loss_use_lin = cfgs.get('perceptual_loss_use_lin', True) self.perceptual_loss = lpips.LPIPS(net='vgg', lpips=self.perceptual_loss_use_lin) self.glctx = dr.RasterizeGLContext() self.render_flow = self.cfgs.get('flow_loss_weight', 0.) > 0. self.extra_renders = cfgs.get('extra_renders', []) self.renderer_spp = cfgs.get('renderer_spp', 1) self.dino_feature_recon_dim = cfgs.get('dino_feature_recon_dim', 64) self.total_loss = 0. self.all_scores = torch.Tensor() self.checkpoint_dir = cfgs.get('checkpoint_dir', 'results') @staticmethod def get_data_loaders(cfgs, dataset, in_image_size=256, out_image_size=256, batch_size=64, num_workers=4, run_train=False, run_test=False, train_data_dir=None, val_data_dir=None, test_data_dir=None): train_loader = val_loader = test_loader = None color_jitter_train = cfgs.get('color_jitter_train', None) color_jitter_val = cfgs.get('color_jitter_val', None) random_flip_train = cfgs.get('random_flip_train', False) ## video dataset if dataset == 'video': data_loader_mode = cfgs.get('data_loader_mode', 'n_frame') skip_beginning = cfgs.get('skip_beginning', 4) skip_end = cfgs.get('skip_end', 4) num_sample_frames = cfgs.get('num_sample_frames', 2) min_seq_len = cfgs.get('min_seq_len', 10) max_seq_len = cfgs.get('max_seq_len', 10) debug_seq = cfgs.get('debug_seq', False) random_sample_train_frames = cfgs.get('random_sample_train_frames', False) shuffle_train_seqs = cfgs.get('shuffle_train_seqs', False) random_sample_val_frames = cfgs.get('random_sample_val_frames', False) load_background = cfgs.get('background_mode', 'none') == 'background' rgb_suffix = cfgs.get('rgb_suffix', '.png') load_dino_feature = cfgs.get('load_dino_feature', False) load_dino_cluster = cfgs.get('load_dino_cluster', False) dino_feature_dim = cfgs.get('dino_feature_dim', 64) get_loader = lambda **kwargs: get_sequence_loader( mode=data_loader_mode, batch_size=batch_size, num_workers=num_workers, in_image_size=in_image_size, out_image_size=out_image_size, debug_seq=debug_seq, skip_beginning=skip_beginning, skip_end=skip_end, num_sample_frames=num_sample_frames, min_seq_len=min_seq_len, max_seq_len=max_seq_len, load_background=load_background, rgb_suffix=rgb_suffix, load_dino_feature=load_dino_feature, load_dino_cluster=load_dino_cluster, dino_feature_dim=dino_feature_dim, **kwargs) if run_train: assert osp.isdir(train_data_dir), f"Training data directory does not exist: {train_data_dir}" print(f"Loading training data from {train_data_dir}") train_loader = get_loader(data_dir=train_data_dir, is_validation=False, random_sample=random_sample_train_frames, shuffle=shuffle_train_seqs, dense_sample=True, color_jitter=color_jitter_train, random_flip=random_flip_train) if val_data_dir is not None: assert osp.isdir(val_data_dir), f"Validation data directory does not exist: {val_data_dir}" print(f"Loading validation data from {val_data_dir}") val_loader = get_loader(data_dir=val_data_dir, is_validation=True, random_sample=random_sample_val_frames, shuffle=False, dense_sample=False, color_jitter=color_jitter_val, random_flip=False) if run_test: assert osp.isdir(test_data_dir), f"Testing data directory does not exist: {test_data_dir}" print(f"Loading testing data from {test_data_dir}") test_loader = get_loader(data_dir=test_data_dir, is_validation=True, dense_sample=False, color_jitter=None, random_flip=False) ## CUB dataset elif dataset == 'cub': get_loader = lambda **kwargs: get_cub_loader( batch_size=batch_size, num_workers=num_workers, image_size=in_image_size, **kwargs) if run_train: assert osp.isdir(train_data_dir), f"Training data directory does not exist: {train_data_dir}" print(f"Loading training data from {train_data_dir}") train_loader = get_loader(data_dir=train_data_dir, split='train', is_validation=False) val_loader = get_loader(data_dir=val_data_dir, split='val', is_validation=True) if run_test: assert osp.isdir(test_data_dir), f"Testing data directory does not exist: {test_data_dir}" print(f"Loading testing data from {test_data_dir}") test_loader = get_loader(data_dir=test_data_dir, split='test', is_validation=True) ## other datasets else: get_loader = lambda **kwargs: get_image_loader( batch_size=batch_size, num_workers=num_workers, image_size=in_image_size, **kwargs) if run_train: assert osp.isdir(train_data_dir), f"Training data directory does not exist: {train_data_dir}" print(f"Loading training data from {train_data_dir}") train_loader = get_loader(data_dir=train_data_dir, is_validation=False, color_jitter=color_jitter_train) if val_data_dir is not None: assert osp.isdir(val_data_dir), f"Validation data directory does not exist: {val_data_dir}" print(f"Loading validation data from {val_data_dir}") val_loader = get_loader(data_dir=val_data_dir, is_validation=True, color_jitter=color_jitter_val) if run_test: assert osp.isdir(test_data_dir), f"Testing data directory does not exist: {test_data_dir}" print(f"Loading testing data from {test_data_dir}") test_loader = get_loader(data_dir=test_data_dir, is_validation=True, color_jitter=None) return train_loader, val_loader, test_loader def load_model_state(self, cp): self.netInstance.load_state_dict(cp["netInstance"]) if self.enable_prior: self.netPrior.load_state_dict(cp["netPrior"]) def load_optimizer_state(self, cp): self.optimizerInstance.load_state_dict(cp["optimizerInstance"]) if self.use_scheduler: if 'schedulerInstance' in cp: self.schedulerInstance.load_state_dict(cp["schedulerInstance"]) if self.enable_prior and self.resume_prior_optim: self.optimizerPrior.load_state_dict(cp["optimizerPrior"]) if self.use_scheduler: if 'schedulerPrior' in cp: self.schedulerPrior.load_state_dict(cp["schedulerPrior"]) def get_model_state(self): state = {"netInstance": self.netInstance.state_dict()} if self.enable_prior: state["netPrior"] = self.netPrior.state_dict() return state def get_optimizer_state(self): state = {"optimizerInstance": self.optimizerInstance.state_dict()} if self.use_scheduler: state["schedulerInstance"] = self.schedulerInstance.state_dict() if self.enable_prior: state["optimizerPrior"] = self.optimizerPrior.state_dict() if self.use_scheduler: state["schedulerPrior"] = self.schedulerPrior.state_dict() return state def to(self, device): self.device = device self.netInstance.to(device) if self.enable_prior: self.netPrior.to(device) if hasattr(self, 'perceptual_loss'): self.perceptual_loss.to(device) def set_train(self): self.netInstance.train() if self.enable_prior: self.netPrior.train() def set_eval(self): self.netInstance.eval() if self.enable_prior: self.netPrior.eval() def reset_optimizers(self): print("Resetting optimizers...") self.optimizerInstance = get_optimizer(self.netInstance, self.lr) if self.use_scheduler: self.schedulerInstance = self.make_scheduler(self.optimizerInstance) if self.enable_prior: self.optimizerPrior = get_optimizer(self.netPrior, lr=self.prior_lr, weight_decay=self.prior_weight_decay) if self.use_scheduler: self.schedulerPrior = self.make_scheduler(self.optimizerPrior) def backward(self): self.optimizerInstance.zero_grad() if self.backward_prior: self.optimizerPrior.zero_grad() self.total_loss.backward() self.optimizerInstance.step() if self.backward_prior: self.optimizerPrior.step() self.total_loss = 0. def scheduler_step(self): if self.use_scheduler: self.schedulerInstance.step() if self.enable_prior: self.schedulerPrior.step() def zflip_pose(self, pose): if self.rot_rep == 'lookat': vec_forward = pose[:,:,6:9] vec_forward = vec_forward * torch.FloatTensor([1,1,-1]).view(1,1,3).to(vec_forward.device) up = torch.FloatTensor([0,1,0]).to(pose.device).view(1,1,3) vec_right = up.expand_as(vec_forward).cross(vec_forward, dim=-1) vec_right = nn.functional.normalize(vec_right, p=2, dim=-1) vec_up = vec_forward.cross(vec_right, dim=-1) vec_up = nn.functional.normalize(vec_up, p=2, dim=-1) rot_mat = torch.stack([vec_right, vec_up, vec_forward], 2) rot_pred = rot_mat.reshape(*pose.shape[:-1], -1) pose_zflip = torch.cat([rot_pred, pose[:,:,9:]], -1) else: raise NotImplementedError return pose_zflip def render(self, shape, texture, mvp, w2c, campos, resolution, background='none', im_features=None, light=None, prior_shape=None, render_flow=True, dino_pred=None, render_mode='diffuse', two_sided_shading=True, num_frames=None, spp=1): h, w = resolution N = len(mvp) if background in ['none', 'black']: bg_image = torch.zeros((N, h, w, 3), device=mvp.device) elif background == 'white': bg_image = torch.ones((N, h, w, 3), device=mvp.device) elif background == 'checkerboard': bg_image = torch.FloatTensor(util.checkerboard((h, w), 8), device=self.device).repeat(N, 1, 1, 1) # NxHxWxC else: raise NotImplementedError frame_rendered = render.render_mesh( self.glctx, shape, mtx_in=mvp, w2c=w2c, view_pos=campos, material=texture, lgt=light, resolution=resolution, spp=spp, msaa=True, background=bg_image, bsdf=render_mode, feat=im_features, prior_mesh=prior_shape, two_sided_shading=two_sided_shading, render_flow=render_flow, dino_pred=dino_pred, num_frames=num_frames) shaded = frame_rendered['shaded'].permute(0, 3, 1, 2) image_pred = shaded[:, :3, :, :] mask_pred = shaded[:, 3, :, :] albedo = frame_rendered['kd'].permute(0, 3, 1, 2)[:, :3, :, :] if 'shading' in frame_rendered: shading = frame_rendered['shading'].permute(0, 3, 1, 2)[:, :1, :, :] else: shading = None if render_flow: flow_pred = frame_rendered['flow'] flow_pred = flow_pred.permute(0, 3, 1, 2)[:, :2, :, :] else: flow_pred = None if dino_pred is not None: dino_feat_im_pred = frame_rendered['dino_feat_im_pred'] dino_feat_im_pred = dino_feat_im_pred.permute(0, 3, 1, 2)[:, :-1] else: dino_feat_im_pred = None return image_pred, mask_pred, flow_pred, dino_feat_im_pred, albedo, shading def compute_reconstruction_losses(self, image_pred, image_gt, mask_pred, mask_gt, mask_dt, mask_valid, flow_pred, flow_gt, dino_feat_im_gt, dino_feat_im_pred, background_mode='none', reduce=False): losses = {} batch_size, num_frames, _, h, w = image_pred.shape # BxFxCxHxW # image_loss = (image_pred - image_gt) ** 2 image_loss = (image_pred - image_gt).abs() ## silhouette loss mask_pred_valid = mask_pred * mask_valid # mask_pred_valid = mask_pred # losses["silhouette_loss"] = ((mask_pred - mask_gt) ** 2).mean() # mask_loss_mask = (image_loss.mean(2).detach() > 0.05).float() mask_loss = (mask_pred_valid - mask_gt) ** 2 # mask_loss = nn.functional.mse_loss(mask_pred, mask_gt) # num_mask_pixels = mask_loss_mask.reshape(batch_size*num_frames, -1).sum(1).clamp(min=1) # losses["silhouette_loss"] = (mask_loss.reshape(batch_size*num_frames, -1).sum(1) / num_mask_pixels).mean() losses['silhouette_loss'] = mask_loss.view(batch_size, num_frames, -1).mean(2) losses['silhouette_dt_loss'] = (mask_pred * mask_dt[:,:,1]).view(batch_size, num_frames, -1).mean(2) losses['silhouette_inv_dt_loss'] = ((1-mask_pred) * mask_dt[:,:,0]).view(batch_size, num_frames, -1).mean(2) mask_pred_binary = (mask_pred_valid > 0.).float().detach() mask_both_binary = (mask_pred_binary * mask_gt).view(batch_size*num_frames, 1, *mask_pred.shape[2:]) mask_both_binary = (nn.functional.avg_pool2d(mask_both_binary, 3, stride=1, padding=1).view(batch_size, num_frames, *mask_pred.shape[2:]) > 0.99).float().detach() # erode by 1 pixel ## reconstruction loss # image_loss_mask = (mask_pred*mask_gt).unsqueeze(2).expand_as(image_gt) # image_loss = image_loss * image_loss_mask # num_mask_pixels = image_loss_mask.reshape(batch_size*num_frames, -1).sum(1).clamp(min=1) # losses["rgb_loss"] = (image_loss.reshape(batch_size*num_frames, -1).sum(1) / num_mask_pixels).mean() if background_mode in ['background', 'input']: pass else: image_loss = image_loss * mask_both_binary.unsqueeze(2) losses['rgb_loss'] = image_loss.reshape(batch_size, num_frames, -1).mean(2) if self.cfgs.get('perceptual_loss_weight', 0.) > 0: if background_mode in ['background', 'input']: perc_image_pred = image_pred perc_image_gt = image_gt else: perc_image_pred = image_pred * mask_pred_binary.unsqueeze(2) + 0.5 * (1-mask_pred_binary.unsqueeze(2)) perc_image_gt = image_gt * mask_pred_binary.unsqueeze(2) + 0.5 * (1-mask_pred_binary.unsqueeze(2)) losses['perceptual_loss'] = self.perceptual_loss(perc_image_pred.view(-1, *image_pred.shape[2:]) *2-1, perc_image_gt.view(-1, *image_gt.shape[2:]) *2-1).view(batch_size, num_frames) ## flow loss - between first and second frame if flow_pred is not None: flow_loss = (flow_pred - flow_gt).abs() flow_loss_mask = mask_both_binary[:,:-1].unsqueeze(2).expand_as(flow_gt).detach() ## ignore frames where GT flow is too large (likely inaccurate) large_flow = (flow_gt.abs() > 0.5).float() * flow_loss_mask large_flow = (large_flow.view(batch_size, num_frames-1, -1).sum(2) > 0).float() self.large_flow = large_flow flow_loss = flow_loss * flow_loss_mask * (1 - large_flow[:,:,None,None,None]) num_mask_pixels = flow_loss_mask.reshape(batch_size, num_frames-1, -1).sum(2).clamp(min=1) losses['flow_loss'] = (flow_loss.reshape(batch_size, num_frames-1, -1).sum(2) / num_mask_pixels) # losses["flow_loss"] = flow_loss.mean() if dino_feat_im_pred is not None: dino_feat_loss = (dino_feat_im_pred - dino_feat_im_gt) ** 2 dino_feat_loss = dino_feat_loss * mask_both_binary.unsqueeze(2) losses['dino_feat_im_loss'] = dino_feat_loss.reshape(batch_size, num_frames, -1).mean(2) if reduce: for k, v in losses.item(): losses[k] = v.mean() return losses def compute_pose_xflip_reg_loss(self, input_image, dino_feat_im, pose_raw, input_image_xflip_flag=None): image_xflip = input_image.flip(4) if dino_feat_im is not None: dino_feat_im_xflip = dino_feat_im.flip(4) else: dino_feat_im_xflip = None feat_xflip, _ = self.netInstance.forward_encoder(image_xflip, dino_feat_im_xflip) batch_size, num_frames = input_image.shape[:2] pose_xflip_raw = self.netInstance.forward_pose(image_xflip, feat_xflip, dino_feat_im_xflip) if input_image_xflip_flag is not None: pose_xflip_raw_xflip = pose_xflip_raw * torch.FloatTensor([-1,1,1,-1,1,1]).to(pose_raw.device) # forward x, trans x pose_xflip_raw = pose_xflip_raw * (1 - input_image_xflip_flag.view(batch_size * num_frames, 1)) + pose_xflip_raw_xflip * input_image_xflip_flag.view(batch_size * num_frames, 1) rot_rep = self.netInstance.rot_rep if rot_rep == 'euler_angle' or rot_rep == 'soft_calss': pose_xflip_xflip = pose_xflip * torch.FloatTensor([1,-1,-1,-1,1,1]).to(pose_xflip.device) # rot y+z, trans x pose_xflip_reg_loss = ((pose_xflip_xflip - pose) ** 2.).mean() elif rot_rep == 'quaternion': rot_euler = pytorch3d.transforms.matrix_to_euler_angles(pytorch3d.transforms.quaternion_to_matrix(pose[...,:4]), convention='XYZ') pose_euler = torch.cat([rot_euler, pose[...,4:]], -1) rot_xflip_euler = pytorch3d.transforms.matrix_to_euler_angles(pytorch3d.transforms.quaternion_to_matrix(pose_xflip[...,:4]), convention='XYZ') pose_xflip_euler = torch.cat([rot_xflip_euler, pose_xflip[...,4:]], -1) pose_xflip_euler_xflip = pose_xflip_euler * torch.FloatTensor([1,-1,-1,-1,1,1]).to(pose_xflip.device) # rot y+z, trans x pose_xflip_reg_loss = ((pose_xflip_euler_xflip - pose_euler) ** 2.).mean() elif rot_rep == 'lookat': pose_xflip_raw_xflip = pose_xflip_raw * torch.FloatTensor([-1,1,1,-1,1,1]).to(pose_raw.device) # forward x, trans x pose_xflip_reg_loss = ((pose_xflip_raw_xflip - pose_raw)[...,0] ** 2.) # compute x only # if epoch >= self.nolookat_zflip_loss_epochs and self.lookat_zflip_no_other_losses: # pose_xflip_reg_loss = pose_xflip_reg_loss.mean(1) * is_pose_1_better pose_xflip_reg_loss = pose_xflip_reg_loss.mean() return pose_xflip_reg_loss, pose_xflip_raw def compute_edge_length_reg_loss(self, mesh, prior_mesh): prior_edge_lengths = get_edge_length(prior_mesh.v_pos, prior_mesh.t_pos_idx) max_length = prior_edge_lengths.max().detach() *1.1 edge_lengths = get_edge_length(mesh.v_pos, mesh.t_pos_idx) mesh_edge_length_loss = ((edge_lengths - max_length).clamp(min=0)**2).mean() return mesh_edge_length_loss, edge_lengths def compute_regularizers(self, mesh, prior_mesh, input_image, dino_feat_im, pose_raw, input_image_xflip_flag=None, arti_params=None, deformation=None): losses = {} aux = {} if self.enable_prior: losses.update(self.netPrior.netShape.get_sdf_reg_loss()) if self.cfgs.get('pose_xflip_reg_loss_weight', 0.) > 0: losses["pose_xflip_reg_loss"], aux['pose_xflip_raw'] = self.compute_pose_xflip_reg_loss(input_image, dino_feat_im, pose_raw, input_image_xflip_flag) b, f = input_image.shape[:2] if b >= 2: vec_forward = pose_raw[..., :3] losses['pose_entropy_loss'] = (vec_forward[:b//2] * vec_forward[b//2:(b//2)*2]).sum(-1).mean() else: losses['pose_entropy_loss'] = 0. losses['mesh_normal_consistency_loss'] = normal_consistency(mesh.v_pos, mesh.t_pos_idx) losses['mesh_edge_length_loss'], aux['edge_lengths'] = self.compute_edge_length_reg_loss(mesh, prior_mesh) if arti_params is not None: losses['arti_reg_loss'] = (arti_params ** 2).mean() if deformation is not None: losses['deformation_reg_loss'] = (deformation ** 2).mean() # losses['deformation_reg_loss'] = deformation.abs().mean() return losses, aux def forward(self, batch, epoch, iter, is_train=True, viz_logger=None, total_iter=None, save_results=False, save_dir=None, which_data='', logger_prefix='', is_training=True): batch = [x.to(self.device) if x is not None else None for x in batch] input_image, mask_gt, mask_dt, mask_valid, flow_gt, bbox, bg_image, dino_feat_im, dino_cluster_im, seq_idx, frame_idx = batch batch_size, num_frames, _, h0, w0 = input_image.shape # BxFxCxHxW h = w = self.out_image_size def collapseF(x): return None if x is None else x.view(batch_size * num_frames, *x.shape[2:]) def expandF(x): return None if x is None else x.view(batch_size, num_frames, *x.shape[1:]) if flow_gt.dim() == 2: # dummy tensor for not loading flow flow_gt = None if dino_feat_im.dim() == 2: # dummy tensor for not loading dino features dino_feat_im = None dino_feat_im_gt = None else: dino_feat_im_gt = expandF(torch.nn.functional.interpolate(collapseF(dino_feat_im), size=[h, w], mode="bilinear"))[:, :, :self.dino_feature_recon_dim] if dino_cluster_im.dim() == 2: # dummy tensor for not loading dino clusters dino_cluster_im = None dino_cluster_im_gt = None else: dino_cluster_im_gt = expandF(torch.nn.functional.interpolate(collapseF(dino_cluster_im), size=[h, w], mode="nearest")) seq_idx = seq_idx.squeeze(1) # seq_idx = seq_idx * 0 # single sequnce model frame_id, crop_x0, crop_y0, crop_w, crop_h, full_w, full_h, sharpness = bbox.unbind(2) # BxFx7 bbox = torch.stack([crop_x0, crop_y0, crop_w, crop_h], 2) mask_gt = (mask_gt[:, :, 0, :, :] > 0.9).float() # BxFxHxW mask_dt = mask_dt / self.in_image_size if which_data != 'video': flow_gt = None aux_viz = {} ## GT image_gt = input_image if self.out_image_size != self.in_image_size: image_gt = expandF(torch.nn.functional.interpolate(collapseF(image_gt), size=[h, w], mode='bilinear')) if flow_gt is not None: flow_gt = torch.nn.functional.interpolate(flow_gt.view(batch_size*(num_frames-1), 2, h0, w0), size=[h, w], mode="bilinear").view(batch_size, num_frames-1, 2, h, w) self.train_pose_only = False if epoch in self.pose_epochs: if (total_iter // self.pose_iters) % 2 == 0: self.train_pose_only = True ## flip input and pose if epoch in self.pose_xflip_recon_epochs: input_image_xflip = input_image.flip(-1) input_image_xflip_flag = torch.randint(0, 2, (batch_size, num_frames), device=input_image.device) input_image = input_image * (1 - input_image_xflip_flag[:,:,None,None,None]) + input_image_xflip * input_image_xflip_flag[:,:,None,None,None] else: input_image_xflip_flag = None ## 1st pose hypothesis with original predictions # ============================================================================================== # Predict prior mesh. # ============================================================================================== if self.enable_prior: if epoch < self.dmtet_grid_smaller_epoch: if self.netPrior.netShape.grid_res != self.dmtet_grid_smaller: self.netPrior.netShape.load_tets(self.dmtet_grid_smaller) else: if self.netPrior.netShape.grid_res != self.dmtet_grid: self.netPrior.netShape.load_tets(self.dmtet_grid) perturb_sdf = self.perturb_sdf if is_train else False prior_shape, dino_pred = self.netPrior(perturb_sdf=perturb_sdf, total_iter=total_iter, is_training=is_training) else: prior_shape = None raise NotImplementedError shape, pose_raw, pose, mvp, w2c, campos, texture, im_features, deformation, arti_params, light, forward_aux = self.netInstance(input_image, prior_shape, epoch, dino_feat_im, dino_cluster_im, total_iter, is_training=is_training) # frame dim collapsed N=(B*F) rot_logit = forward_aux['rot_logit'] rot_idx = forward_aux['rot_idx'] rot_prob = forward_aux['rot_prob'] aux_viz.update(forward_aux) if self.train_pose_only: safe_detach = lambda x: x.detach() if x is not None else None prior_shape = safe_detach(prior_shape) shape = safe_detach(shape) im_features = safe_detach(im_features) arti_params = safe_detach(arti_params) deformation = safe_detach(deformation) set_requires_grad(texture, False) set_requires_grad(light, False) set_requires_grad(dino_pred, False) else: set_requires_grad(texture, True) set_requires_grad(light, True) set_requires_grad(dino_pred, True) render_flow = self.render_flow and num_frames > 1 image_pred, mask_pred, flow_pred, dino_feat_im_pred, albedo, shading = self.render(shape, texture, mvp, w2c, campos, (h, w), background=self.background_mode, im_features=im_features, light=light, prior_shape=prior_shape, render_flow=render_flow, dino_pred=dino_pred, num_frames=num_frames, spp=self.renderer_spp) image_pred, mask_pred, flow_pred, dino_feat_im_pred = map(expandF, (image_pred, mask_pred, flow_pred, dino_feat_im_pred)) if flow_pred is not None: flow_pred = flow_pred[:, :-1] # Bx(F-1)x2xHxW if self.blur_mask: sigma = max(0.5, 3 * (1 - total_iter / self.blur_mask_iter)) if sigma > 0.5: mask_gt = util.blur_image(mask_gt, kernel_size=9, sigma=sigma, mode='gaussian') # mask_pred = util.blur_image(mask_pred, kernel_size=7, mode='average') losses = self.compute_reconstruction_losses(image_pred, image_gt, mask_pred, mask_gt, mask_dt, mask_valid, flow_pred, flow_gt, dino_feat_im_gt, dino_feat_im_pred, background_mode=self.background_mode, reduce=False) ## TODO: assume flow loss is not used logit_loss_target = torch.zeros_like(expandF(rot_logit)) final_losses = {} for name, loss in losses.items(): loss_weight_logit = self.cfgs.get(f"{name}_weight", 0.) # if (name in ['flow_loss'] and epoch not in self.flow_loss_epochs) or (name in ['rgb_loss', 'perceptual_loss'] and epoch not in self.texture_epochs): # if name in ['flow_loss', 'rgb_loss', 'perceptual_loss']: # loss_weight_logit = 0. if name in ['sdf_bce_reg_loss', 'sdf_gradient_reg_loss', 'sdf_inflate_reg_loss']: if total_iter >= self.sdf_reg_decay_start_iter: decay_rate = max(0, 1 - (total_iter-self.sdf_reg_decay_start_iter) / 10000) loss_weight_logit = max(loss_weight_logit * decay_rate, self.cfgs.get(f"{name}_min_weight", 0.)) if name in ['dino_feat_im_loss']: loss_weight_logit = loss_weight_logit * self.cfgs.get("logit_loss_dino_feat_im_loss_multiplier", 1.) if loss_weight_logit > 0: logit_loss_target += loss * loss_weight_logit if self.netInstance.rot_rep in ['quadlookat', 'octlookat']: loss = loss * rot_prob.detach().view(batch_size, num_frames)[:, :loss.shape[1]] *self.netInstance.num_pose_hypos if name == 'flow_loss' and num_frames > 1: ri = rot_idx.view(batch_size, num_frames) same_rot_idx = (ri[:, 1:] == ri[:, :-1]).float() loss = loss * same_rot_idx final_losses[name] = loss.mean() final_losses['logit_loss'] = ((expandF(rot_logit) - logit_loss_target.detach())**2.).mean() ## regularizers regularizers, aux = self.compute_regularizers(shape, prior_shape, input_image, dino_feat_im, pose_raw, input_image_xflip_flag, arti_params, deformation) final_losses.update(regularizers) aux_viz.update(aux) total_loss = 0 for name, loss in final_losses.items(): loss_weight = self.cfgs.get(f"{name}_weight", 0.) if loss_weight <= 0: continue if self.train_pose_only: if name not in ['silhouette_loss', 'silhouette_dt_loss', 'silhouette_inv_dt_loss', 'flow_loss', 'pose_xflip_reg_loss', 'lookat_zflip_loss', 'dino_feat_im_loss']: continue if epoch not in self.flow_loss_epochs: if name in ['flow_loss']: continue if epoch not in self.texture_epochs: if name in ['rgb_loss', 'perceptual_loss']: continue if epoch not in self.lookat_zflip_loss_epochs: if name in ['lookat_zflip_loss']: continue if name in ['mesh_laplacian_smoothing_loss', 'mesh_normal_consistency_loss']: if total_iter < self.cfgs.get('mesh_reg_start_iter', 0): continue if epoch >= self.mesh_reg_decay_epoch: decay_rate = self.mesh_reg_decay_rate ** (epoch - self.mesh_reg_decay_epoch) loss_weight = max(loss_weight * decay_rate, self.cfgs.get(f"{name}_min_weight", 0.)) if epoch not in self.sdf_inflate_reg_loss_epochs: if name in ['sdf_inflate_reg_loss']: continue if epoch not in self.arti_reg_loss_epochs: if name in ['arti_reg_loss']: continue if name in ['sdf_bce_reg_loss', 'sdf_gradient_reg_loss', 'sdf_inflate_reg_loss']: if total_iter >= self.sdf_reg_decay_start_iter: decay_rate = max(0, 1 - (total_iter-self.sdf_reg_decay_start_iter) / 10000) loss_weight = max(loss_weight * decay_rate, self.cfgs.get(f"{name}_min_weight", 0.)) total_loss += loss * loss_weight self.total_loss += total_loss # reset to 0 in backward step if torch.isnan(self.total_loss): print("NaN in loss...") import ipdb; ipdb.set_trace() final_losses['logit_loss_target'] = logit_loss_target.mean() metrics = {'loss': total_loss, **final_losses} ## log visuals if viz_logger is not None: b0 = max(min(batch_size, 16//num_frames), 1) viz_logger.add_image(logger_prefix+'image/image_gt', misc.image_grid(image_gt.detach().cpu()[:b0,:].reshape(-1,*input_image.shape[2:]).clamp(0,1)), total_iter) viz_logger.add_image(logger_prefix+'image/image_pred', misc.image_grid(image_pred.detach().cpu()[:b0,:].reshape(-1,*image_pred.shape[2:]).clamp(0,1)), total_iter) # viz_logger.add_image(logger_prefix+'image/flow_loss_mask', misc.image_grid(flow_loss_mask[:b0,:,:1].reshape(-1,1,*flow_loss_mask.shape[3:]).repeat(1,3,1,1).clamp(0,1)), total_iter) viz_logger.add_image(logger_prefix+'image/mask_gt', misc.image_grid(mask_gt.detach().cpu()[:b0,:].reshape(-1,*mask_gt.shape[2:]).unsqueeze(1).repeat(1,3,1,1).clamp(0,1)), total_iter) viz_logger.add_image(logger_prefix+'image/mask_pred', misc.image_grid(mask_pred.detach().cpu()[:b0,:].reshape(-1,*mask_pred.shape[2:]).unsqueeze(1).repeat(1,3,1,1).clamp(0,1)), total_iter) if self.render_flow and flow_gt is not None: flow_gt = flow_gt.detach().cpu() flow_gt_viz = torch.cat([flow_gt[:b0], torch.zeros_like(flow_gt[:b0,:,:1])], 2) + 0.5 # -0.5~1.5 flow_gt_viz = torch.nn.functional.pad(flow_gt_viz, pad=[0, 0, 0, 0, 0, 0, 0, 1]) ## draw marker on large flow frames large_flow_marker_mask = torch.zeros_like(flow_gt_viz) large_flow_marker_mask[:,:,:,:8,:8] = 1. large_flow = torch.cat([self.large_flow, self.large_flow[:,:1] *0.], 1).detach().cpu()[:b0] large_flow_marker_mask = large_flow_marker_mask * large_flow[:,:,None,None,None] red = torch.FloatTensor([1,0,0])[None,None,:,None,None] flow_gt_viz = large_flow_marker_mask * red + (1-large_flow_marker_mask) * flow_gt_viz viz_logger.add_image(logger_prefix+'image/flow_gt', misc.image_grid(flow_gt_viz.reshape(-1,*flow_gt_viz.shape[2:])), total_iter) if self.render_flow and flow_pred is not None: flow_pred = flow_pred.detach().cpu() flow_pred_viz = torch.cat([flow_pred[:b0], torch.zeros_like(flow_pred[:b0,:,:1])], 2) + 0.5 # -0.5~1.5 flow_pred_viz = torch.nn.functional.pad(flow_pred_viz, pad=[0, 0, 0, 0, 0, 0, 0, 1]) viz_logger.add_image(logger_prefix+'image/flow_pred', misc.image_grid(flow_pred_viz.reshape(-1,*flow_pred_viz.shape[2:])), total_iter) if light is not None: param_names = ['dir_x', 'dir_y', 'dir_z', 'int_ambient', 'int_diffuse'] for name, param in zip(param_names, light.light_params.unbind(-1)): viz_logger.add_histogram(logger_prefix+'light/'+name, param, total_iter) viz_logger.add_image( logger_prefix + f'image/albedo', misc.image_grid(expandF(albedo)[:b0, ...].view(-1, *albedo.shape[1:])), total_iter) viz_logger.add_image( logger_prefix + f'image/shading', misc.image_grid(expandF(shading)[:b0, ...].view(-1, *shading.shape[1:]).repeat(1, 3, 1, 1) /2.), total_iter) viz_logger.add_histogram(logger_prefix+'sdf', self.netPrior.netShape.get_sdf(perturb_sdf=False), total_iter) viz_logger.add_histogram(logger_prefix+'coordinates', shape.v_pos, total_iter) if arti_params is not None: viz_logger.add_histogram(logger_prefix+'arti_params', arti_params, total_iter) viz_logger.add_histogram(logger_prefix+'edge_lengths', aux_viz['edge_lengths'], total_iter) if deformation is not None: viz_logger.add_histogram(logger_prefix+'deformation', deformation, total_iter) rot_rep = self.netInstance.rot_rep if rot_rep == 'euler_angle' or rot_rep == 'soft_calss': for i, name in enumerate(['rot_x', 'rot_y', 'rot_z', 'trans_x', 'trans_y', 'trans_z']): viz_logger.add_histogram(logger_prefix+'pose/'+name, pose[...,i], total_iter) elif rot_rep == 'quaternion': for i, name in enumerate(['qt_0', 'qt_1', 'qt_2', 'qt_3', 'trans_x', 'trans_y', 'trans_z']): viz_logger.add_histogram(logger_prefix+'pose/'+name, pose[...,i], total_iter) rot_euler = pytorch3d.transforms.matrix_to_euler_angles(pytorch3d.transforms.quaternion_to_matrix(pose.detach().cpu()[...,:4]), convention='XYZ') for i, name in enumerate(['rot_x', 'rot_y', 'rot_z']): viz_logger.add_histogram(logger_prefix+'pose/'+name, rot_euler[...,i], total_iter) elif rot_rep in ['lookat', 'quadlookat', 'octlookat']: for i, name in enumerate(['fwd_x', 'fwd_y', 'fwd_z']): viz_logger.add_histogram(logger_prefix+'pose/'+name, pose_raw[...,i], total_iter) for i, name in enumerate(['trans_x', 'trans_y', 'trans_z']): viz_logger.add_histogram(logger_prefix+'pose/'+name, pose_raw[...,-3+i], total_iter) if rot_rep in ['quadlookat', 'octlookat']: for i, rp in enumerate(forward_aux['rots_probs'].unbind(-1)): viz_logger.add_histogram(logger_prefix+'pose/rot_prob_%d'%i, rp, total_iter) if 'pose_xflip_raw' in aux_viz: pose_xflip_raw = aux_viz['pose_xflip_raw'] if rot_rep == 'euler_angle' or rot_rep == 'soft_calss': for i, name in enumerate(['rot_x', 'rot_y', 'rot_z', 'trans_x', 'trans_y', 'trans_z']): viz_logger.add_histogram(logger_prefix+'pose_xflip/'+name, pose_xflip[...,i], total_iter) elif rot_rep == 'quaternion': for i, name in enumerate(['qt_0', 'qt_1', 'qt_2', 'qt_3', 'trans_x', 'trans_y', 'trans_z']): viz_logger.add_histogram(logger_prefix+'pose_xflip/'+name, pose_xflip[...,i], total_iter) rot_euler = pytorch3d.transforms.matrix_to_euler_angles(pytorch3d.transforms.quaternion_to_matrix(pose_xflip.detach().cpu()[...,:4]), convention='XYZ') for i, name in enumerate(['rot_x', 'rot_y', 'rot_z']): viz_logger.add_histogram(logger_prefix+'pose_xflip/'+name, rot_euler[...,i], total_iter) elif rot_rep in ['lookat', 'quadlookat', 'octlookat']: for i, name in enumerate(['fwd_x', 'fwd_y', 'fwd_z']): viz_logger.add_histogram(logger_prefix+'pose_xflip/'+name, pose_xflip_raw[...,i], total_iter) for i, name in enumerate(['trans_x', 'trans_y', 'trans_z']): viz_logger.add_histogram(logger_prefix+'pose_xflip/'+name, pose_xflip_raw[...,-3+i], total_iter) if dino_feat_im_gt is not None: dino_feat_im_gt_first3 = dino_feat_im_gt[:,:,:3] viz_logger.add_image(logger_prefix+'image/dino_feat_im_gt', misc.image_grid(dino_feat_im_gt_first3.detach().cpu()[:b0,:].reshape(-1,*dino_feat_im_gt_first3.shape[2:]).clamp(0,1)), total_iter) if dino_cluster_im_gt is not None: viz_logger.add_image(logger_prefix+'image/dino_cluster_im_gt', misc.image_grid(dino_cluster_im_gt.detach().cpu()[:b0,:].reshape(-1,*dino_cluster_im_gt.shape[2:]).clamp(0,1)), total_iter) if dino_feat_im_pred is not None: dino_feat_im_pred_first3 = dino_feat_im_pred[:,:,:3] viz_logger.add_image(logger_prefix+'image/dino_feat_im_pred', misc.image_grid(dino_feat_im_pred_first3.detach().cpu()[:b0,:].reshape(-1,*dino_feat_im_pred_first3.shape[2:]).clamp(0,1)), total_iter) for which_shape, modes in self.extra_renders.items(): # This is wrong # if which_shape == "prior": # shape_to_render = prior_shape.extend(im_features.shape[0]) # needed_im_features = None if which_shape == "instance": shape_to_render = shape needed_im_features = im_features else: raise NotImplementedError for mode in modes: rendered, _, _, _, _, _ = self.render(shape_to_render, texture, mvp, w2c, campos, (h, w), background=self.background_mode, im_features=needed_im_features, prior_shape=prior_shape, render_mode=mode, render_flow=False, dino_pred=None) if 'kd' in mode: rendered = util.rgb_to_srgb(rendered) rendered = rendered.detach().cpu() if 'posed_bones' in aux_viz: rendered_bone_image = self.render_bones(mvp, aux_viz['posed_bones'], (h, w)) rendered_bone_image_mask = (rendered_bone_image < 1).any(1, keepdim=True).float() # viz_logger.add_image(logger_prefix+'image/articulation_bones', misc.image_grid(self.render_bones(mvp, aux_viz['posed_bones'])), total_iter) rendered = rendered_bone_image_mask*0.8 * rendered_bone_image + (1-rendered_bone_image_mask*0.8) * rendered if rot_rep in ['quadlookat', 'octlookat']: rand_pose_flag = forward_aux['rand_pose_flag'].detach().cpu() rand_pose_marker_mask = torch.zeros_like(rendered) rand_pose_marker_mask[:,:,:16,:16] = 1. rand_pose_marker_mask = rand_pose_marker_mask * rand_pose_flag[:,None,None,None] red = torch.FloatTensor([1,0,0])[None,:,None,None] rendered = rand_pose_marker_mask * red + (1-rand_pose_marker_mask) * rendered viz_logger.add_image( logger_prefix + f'image/{which_shape}_{mode}', misc.image_grid(expandF(rendered)[:b0, ...].view(-1, *rendered.shape[1:])), total_iter) viz_logger.add_video( logger_prefix + f'animation/{which_shape}_{mode}', self.render_rotation_frames(shape_to_render, texture, light, (h, w), background=self.background_mode, im_features=needed_im_features, prior_shape=prior_shape, num_frames=15, render_mode=mode, b=1).detach().cpu().unsqueeze(0), total_iter, fps=2) viz_logger.add_video( logger_prefix+'animation/prior_image_rotation', self.render_rotation_frames(prior_shape, texture, light, (h, w), background=self.background_mode, im_features=im_features, num_frames=15, b=1).detach().cpu().unsqueeze(0).clamp(0,1), total_iter, fps=2) viz_logger.add_video( logger_prefix+'animation/prior_normal_rotation', self.render_rotation_frames(prior_shape, texture, light, (h, w), background=self.background_mode, im_features=im_features, num_frames=15, render_mode='geo_normal', b=1).detach().cpu().unsqueeze(0), total_iter, fps=2) if save_results: b0 = self.cfgs.get('num_saved_from_each_batch', batch_size*num_frames) fnames = [f'{total_iter:07d}_{fid:10d}' for fid in collapseF(frame_id.int())][:b0] misc.save_images(save_dir, collapseF(image_gt)[:b0].clamp(0,1).detach().cpu().numpy(), suffix='image_gt', fnames=fnames) misc.save_images(save_dir, collapseF(image_pred)[:b0].clamp(0,1).detach().cpu().numpy(), suffix='image_pred', fnames=fnames) misc.save_images(save_dir, collapseF(mask_gt)[:b0].unsqueeze(1).repeat(1,3,1,1).clamp(0,1).detach().cpu().numpy(), suffix='mask_gt', fnames=fnames) misc.save_images(save_dir, collapseF(mask_pred)[:b0].unsqueeze(1).repeat(1,3,1,1).clamp(0,1).detach().cpu().numpy(), suffix='mask_pred', fnames=fnames) # tmp_shape = shape.first_n(b0).clone() # tmp_shape.material = texture # feat = im_features[:b0] if im_features is not None else None # misc.save_obj(save_dir, tmp_shape, save_material=False, feat=feat, suffix="mesh", fnames=fnames) # Save the first mesh. # if self.render_flow and flow_gt is not None: # flow_gt_viz = torch.cat([flow_gt, torch.zeros_like(flow_gt[:,:,:1])], 2) + 0.5 # -0.5~1.5 # flow_gt_viz = flow_gt_viz.view(-1, *flow_gt_viz.shape[2:]) # misc.save_images(save_dir, flow_gt_viz[:b0].clamp(0,1).detach().cpu().numpy(), suffix='flow_gt', fnames=fnames) # if flow_pred is not None: # flow_pred_viz = torch.cat([flow_pred, torch.zeros_like(flow_pred[:,:,:1])], 2) + 0.5 # -0.5~1.5 # flow_pred_viz = flow_pred_viz.view(-1, *flow_pred_viz.shape[2:]) # misc.save_images(save_dir, flow_pred_viz[:b0].clamp(0,1).detach().cpu().numpy(), suffix='flow_pred', fnames=fnames) misc.save_txt(save_dir, pose[:b0].detach().cpu().numpy(), suffix='pose', fnames=fnames) return metrics def save_scores(self, path): header = 'mask_mse, \ mask_iou, \ image_mse, \ flow_mse' mean = self.all_scores.mean(0) std = self.all_scores.std(0) header = header + '\nMean: ' + ',\t'.join(['%.8f'%x for x in mean]) header = header + '\nStd: ' + ',\t'.join(['%.8f'%x for x in std]) misc.save_scores(path, self.all_scores, header=header) print(header) def render_rotation_frames(self, mesh, texture, light, resolution, background='none', im_features=None, prior_shape=None, num_frames=36, render_mode='diffuse', b=None): frames = [] if b is None: b = len(mesh) else: mesh = mesh.first_n(b) feat = im_features[:b] if im_features is not None else None delta_angle = np.pi / num_frames * 2 delta_rot_matrix = torch.FloatTensor([ [np.cos(delta_angle), 0, np.sin(delta_angle), 0], [0, 1, 0, 0], [-np.sin(delta_angle), 0, np.cos(delta_angle), 0], [0, 0, 0, 1], ]).to(self.device).repeat(b, 1, 1) w2c = torch.FloatTensor(np.diag([1., 1., 1., 1])) w2c[:3, 3] = torch.FloatTensor([0, 0, -self.cam_pos_z_offset *1.1]) w2c = w2c.repeat(b, 1, 1).to(self.device) proj = util.perspective(self.crop_fov_approx / 180 * np.pi, 1, n=0.1, f=1000.0).repeat(b, 1, 1).to(self.device) mvp = torch.bmm(proj, w2c) campos = -w2c[:, :3, 3] def rotate_pose(mvp, campos): mvp = torch.matmul(mvp, delta_rot_matrix) campos = torch.matmul(delta_rot_matrix[:,:3,:3].transpose(2,1), campos[:,:,None])[:,:,0] return mvp, campos for _ in range(num_frames): image_pred, _, _, _, _, _ = self.render(mesh, texture, mvp, w2c, campos, resolution, background=background, im_features=feat, light=light, prior_shape=prior_shape, render_flow=False, dino_pred=None, render_mode=render_mode, two_sided_shading=False) frames += [misc.image_grid(image_pred)] mvp, campos = rotate_pose(mvp, campos) return torch.stack(frames, dim=0) # Shape: (T, C, H, W) def render_bones(self, mvp, bones_pred, size=(256, 256)): bone_world4 = torch.concat([bones_pred, torch.ones_like(bones_pred[..., :1]).to(bones_pred.device)], dim=-1) b, f, num_bones = bone_world4.shape[:3] bones_clip4 = (bone_world4.view(b, f, num_bones*2, 1, 4) @ mvp.transpose(-1, -2).reshape(b, f, 1, 4, 4)).view(b, f, num_bones, 2, 4) bones_uv = bones_clip4[..., :2] / bones_clip4[..., 3:4] # b, f, num_bones, 2, 2 dpi = 32 fx, fy = size[1] // dpi, size[0] // dpi rendered = [] for b_idx in range(b): for f_idx in range(f): frame_bones_uv = bones_uv[b_idx, f_idx].cpu().numpy() fig = plt.figure(figsize=(fx, fy), dpi=dpi, frameon=False) ax = plt.Axes(fig, [0., 0., 1., 1.]) ax.set_axis_off() for bone in frame_bones_uv: ax.plot(bone[:, 0], bone[:, 1], marker='o', linewidth=8, markersize=20) ax.set_xlim(-1, 1) ax.set_ylim(-1, 1) ax.invert_yaxis() # Convert to image fig.add_axes(ax) fig.canvas.draw_idle() image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) w, h = fig.canvas.get_width_height() image.resize(h, w, 3) rendered += [image / 255.] return torch.from_numpy(np.stack(rendered, 0).transpose(0, 3, 1, 2)) def render_deformation_frames(self, mesh, texture, batch_size, num_frames, resolution, background='none', im_features=None, render_mode='diffuse', b=None): # frames = [] # if b is None: # b = batch_size # im_features = im_features[] # mesh = mesh.first_n(num_frames * b) # for i in range(b): # tmp_mesh = mesh.get_m_to_n(i*num_frames:(i+1)*num_frames) pass