from lib.kits.basic import * import traceback from tqdm import tqdm from lib.body_models.common import make_SKEL from lib.body_models.skel_wrapper import SKELWrapper, SKELOutput from lib.body_models.abstract_skeletons import Skeleton_OpenPose25 from lib.utils.data import to_tensor, to_list from lib.utils.camera import perspective_projection from lib.utils.media import draw_kp2d_on_img, annotate_img, splice_img from lib.utils.vis import render_mesh_overlay_img from lib.modeling.losses import compute_poses_angle_prior_loss from .skelify.utils import get_kp_active_j_masks def compute_rel_change(prev_val: float, curr_val: float) -> float: ''' Compute the relative change between two values. Copied: from https://github.com/vchoutas/smplify-x ### Args: - prev_val: float - curr_val: float ### Returns: - float ''' return np.abs(prev_val - curr_val) / max([np.abs(prev_val), np.abs(curr_val), 1]) def gmof(x, sigma): ''' Geman-McClure error function, to be used as a robust loss function. ''' x_squared = x ** 2 sigma_squared = sigma ** 2 return (sigma_squared * x_squared) / (sigma_squared + x_squared) class SKELifyRefiner(): def __init__(self, cfg, name='SKELify', tb_logger=None, device='cuda:0'): self.cfg = cfg self.name = name self.eq_thre = cfg.early_quit_thresholds self.tb_logger = tb_logger self.device = device self.skel_model = instantiate(cfg.skel_model).to(device) # Dirty implementation for visualization. self.render_frames = [] def __call__( self, gt_kp2d : Union[torch.Tensor, np.ndarray], init_poses : Union[torch.Tensor, np.ndarray], init_betas : Union[torch.Tensor, np.ndarray], init_cam_t : Union[torch.Tensor, np.ndarray], img_patch : Optional[np.ndarray] = None, **kwargs ): ''' Use optimization to fit the SKEL parameters to the 2D keypoints. ### Args: - gt_kp2d : torch.Tensor or np.ndarray, (B, J, 3) - The last three dim means [x, y, conf]. - The 2D keypoints to fit, they are defined in [-0.5, 0.5], zero-centered space. - init_poses : torch.Tensor or np.ndarray, (B, 46) - init_betas : torch.Tensor or np.ndarray, (B, 10) - init_cam_t : torch.Tensor or np.ndarray, (B, 3) - img_patch : np.ndarray or None, (B, H, W, 3) - The image patch for visualization. H, W are defined in normalized bounding box space. - If None, the visualization will simply use a black image. ### Returns: - TODO: ''' # ⛩️ Prepare the input data. gt_kp2d = to_tensor(gt_kp2d, device=self.device).detach().float().clone() # (B, J, 3) init_poses = to_tensor(init_poses, device=self.device).detach().float().clone() # (B, 46) init_betas = to_tensor(init_betas, device=self.device).detach().float().clone() # (B, 10) init_cam_t = to_tensor(init_cam_t, device=self.device).detach().float().clone() # (B, 3) inputs = { 'poses_orient': init_poses[:, :3], # (B, 3) 'poses_body' : init_poses[:, 3:], # (B, 43) 'betas' : init_betas, # (B, 10) 'cam_t' : init_cam_t, # (B, 3) } focal_length = np.ones(2) * self.cfg.focal_length / self.cfg.img_patch_size focal_length = focal_length.reshape(1, 2).repeat(inputs['cam_t'].shape[0], 1) # ⛩️ Optimization phases, controlled by config file. prev_phase_steps = 0 # accumulate the steps are *supposed* to be done in the previous phases for phase_id, phase_name in enumerate(self.cfg.phases): phase_cfg = self.cfg.phases[phase_name] # Preparation. optim_params = [] for k in inputs.keys(): if k in phase_cfg.params_keys: inputs[k].requires_grad = True optim_params.append(inputs[k]) # (B, D) else: inputs[k].requires_grad = False optimizer = instantiate(phase_cfg.optimizer, optim_params, _recursive_=True) def closure(): optimizer.zero_grad() # Data preparation. cam_t = inputs['cam_t'] skel_params = { 'poses' : torch.cat([inputs['poses_orient'], inputs['poses_body']], dim=-1), # (B, 46) 'betas' : inputs['betas'], # (B, 10) 'skelmesh' : False, } # Optimize steps. skel_output = self.skel_model(**skel_params) pd_kp2d = perspective_projection( points = to_tensor(skel_output.joints, device=self.device), translation = to_tensor(cam_t, device=self.device), focal_length = to_tensor(focal_length, device=self.device), ) loss, losses = self._compute_losses( act_losses = phase_cfg.losses, act_parts = phase_cfg.get('parts', 'all'), gt_kp2d = gt_kp2d, pd_kp2d = pd_kp2d, pd_params = skel_params, **phase_cfg.get('weights', {}), ) # For visualize the optimization process. _conf = gt_kp2d[..., 2] # (B, J) metric = torch.sum((pd_kp2d - gt_kp2d[..., :2]) ** 2, dim=-1) * _conf # (B, J) metric = metric.sum(dim=-1) / (torch.sum(_conf, dim=-1) + 1e-6) # (B,) # Store logging data. if self.tb_logger is not None: log_data.update({ 'losses' : losses, 'pd_kp2d' : pd_kp2d[:self.cfg.logger.samples_per_record].detach().clone(), 'pd_verts' : skel_output.skin_verts[:self.cfg.logger.samples_per_record].detach().clone(), 'cam_t' : cam_t[:self.cfg.logger.samples_per_record].detach().clone(), 'metric' : metric[:self.cfg.logger.samples_per_record].detach().clone(), 'optim_betas' : inputs['betas'][:self.cfg.logger.samples_per_record].detach().clone(), }) loss.backward() return loss.item() # Optimization loop. prev_loss = None with tqdm(range(phase_cfg.max_loop)) as bar: bar.set_description(f'[{phase_name}] Loss: ???') for i in bar: log_data = {} curr_loss = optimizer.step(closure) # Logging. if self.tb_logger is not None: log_data.update({ 'img_patch' : img_patch[:self.cfg.logger.samples_per_record] if img_patch is not None else None, 'gt_kp2d' : gt_kp2d[:self.cfg.logger.samples_per_record].detach().clone(), }) self._tb_log(prev_phase_steps + i, log_data) # self._tb_log_for_report(prev_phase_steps + i, log_data) bar.set_description(f'[{phase_name}] Loss: {curr_loss:.4f}') if self._can_early_quit(optim_params, prev_loss, curr_loss): break prev_loss = curr_loss prev_phase_steps += phase_cfg.max_loop # ⛩️ Prepare the output data. outputs = { 'poses': torch.cat([inputs['poses_orient'], inputs['poses_body']], dim=-1).detach().clone(), # (B, 46) 'betas': inputs['betas'].detach().clone(), # (B, 10) 'cam_t': inputs['cam_t'].detach().clone(), # (B, 3) } return outputs def _compute_losses( self, act_losses : List[str], act_parts : List[str], gt_kp2d : torch.Tensor, pd_kp2d : torch.Tensor, pd_params : Dict, robust_sigma : float = 100, shape_prior_weight : float = 5, angle_prior_weight : float = 15.2, *args, **kwargs, ): ''' Compute the weighted losses according to the config file. Follow: https://github.com/nkolot/SPIN/blob/2476c436013055be5cb3905e4e4ecfa86966fac3/smplify/losses.py#L26-L58s ''' B = len(gt_kp2d) act_j_masks = get_kp_active_j_masks(act_parts, device=gt_kp2d.device) # (44,) # Reproject the 3D keypoints to image and compare the L2 error with the g.t. 2D keypoints. kp_conf = gt_kp2d[..., 2] # (B, J) gt_kp2d = gt_kp2d[..., :2] # (B, J, 2) reproj_err = gmof(pd_kp2d - gt_kp2d, robust_sigma) # (B, J, 2) reproj_loss = ((kp_conf ** 2) * reproj_err.sum(dim=-1) * act_j_masks[None]).sum(-1) # (B,) # Regularize the shape parameters. shape_prior_loss = (shape_prior_weight ** 2) * (pd_params['betas'] ** 2).sum(dim=-1) # (B,) # Use the SKEL angle prior knowledge (e.g., rotation limitation) to regularize the optimization process. # TODO: Is that necessary? angle_prior_loss = (angle_prior_weight ** 2) * compute_poses_angle_prior_loss(pd_params['poses']).mean() # (,) losses = { 'reprojection' : reproj_loss.mean(), # (,) 'shape_prior' : shape_prior_loss.mean(), # (,) 'angle_prior' : angle_prior_loss, # (,) } loss = torch.tensor(0., device=gt_kp2d.device) for k in act_losses: loss += losses[k] losses = {k: v.detach() for k, v in losses.items()} losses['sum'] = loss.detach() # (,) return loss, losses def _can_early_quit(self, opt_params, prev_loss, curr_loss): ''' Judge whether to early quit the optimization process. If yes, return True, otherwise False.''' if self.cfg.early_quit_thresholds is None: # Never early quit. return False # Relative change test. if prev_loss is not None: loss_rel_change = compute_rel_change(prev_loss, curr_loss) if loss_rel_change < self.cfg.early_quit_thresholds.rel: get_logger().info(f'Early quit due to relative change: {loss_rel_change:.4f} = rel({prev_loss}, {curr_loss})') return True # Absolute change test. if all([ torch.abs(param.grad.max()).item() < self.cfg.early_quit_thresholds.abs for param in opt_params if param.grad is not None ]): get_logger().info(f'Early quit due to absolute change.') return True return False @rank_zero_only def _tb_log(self, step_cnt:int, log_data:Dict, *args, **kwargs): ''' Write the logging information to the TensorBoard. ''' if step_cnt != 0 and (step_cnt + 1) % self.cfg.logger.interval != 0: return summary_writer = self.tb_logger.experiment # Save losses. for loss_name, loss_val in log_data['losses'].items(): summary_writer.add_scalar(f'skelify/{loss_name}', loss_val.detach().item(), step_cnt) # Visualization of the optimization process. TODO: Maybe we can make this more elegant. if log_data['img_patch'] is None: log_data['img_patch'] = [np.zeros((self.cfg.img_patch_size, self.cfg.img_patch_size, 3), dtype=np.uint8)] \ * len(log_data['gt_kp2d']) if len(self.render_frames) < 1: self.init_v = log_data['pd_verts'] self.init_metric = log_data['metric'] self.init_ct = log_data['cam_t'] # Overlay the skin mesh of the results on the original image. try: imgs_spliced = [] for i, img_patch in enumerate(log_data['img_patch']): metric = log_data['metric'][i].item() img_with_init = render_mesh_overlay_img( faces = self.skel_model.skin_f, verts = self.init_v[i], K4 = [self.cfg.focal_length, self.cfg.focal_length, 0, 0], img = img_patch, Rt = [torch.eye(3), self.init_ct[i]], mesh_color = 'pink', ) img_with_init = annotate_img(img_with_init, 'init') img_with_init = annotate_img(img_with_init, f'Quality: {self.init_metric[i].item()*1000:.3f}/1e3', pos='tl') img_with_mesh = render_mesh_overlay_img( faces = self.skel_model.skin_f, verts = log_data['pd_verts'][i], K4 = [self.cfg.focal_length, self.cfg.focal_length, 0, 0], img = img_patch, Rt = [torch.eye(3), log_data['cam_t'][i]], mesh_color = 'pink', ) img_with_mesh = annotate_img(img_with_mesh, 'pd_mesh') betas_max = log_data['optim_betas'][i].abs().max().item() img_with_mesh = annotate_img(img_with_mesh, f'Quality: {metric*1000:.3f}/1e3\nbetas_max: {betas_max:.3f}', pos='tl') img_patch_raw = annotate_img(img_patch, 'raw') log_data['gt_kp2d'][i][..., :2] = (log_data['gt_kp2d'][i][..., :2] + 0.5) * self.cfg.img_patch_size img_with_gt = annotate_img(img_patch, 'gt_kp2d') img_with_gt = draw_kp2d_on_img( img_with_gt, log_data['gt_kp2d'][i], Skeleton_OpenPose25.bones, Skeleton_OpenPose25.bone_colors, ) log_data['pd_kp2d'][i] = (log_data['pd_kp2d'][i] + 0.5) * self.cfg.img_patch_size img_with_pd = annotate_img(img_patch, 'pd_kp2d') img_with_pd = draw_kp2d_on_img( img_with_pd, log_data['pd_kp2d'][i], Skeleton_OpenPose25.bones, Skeleton_OpenPose25.bone_colors, ) img_spliced = splice_img( img_grids = [img_patch_raw, img_with_gt, img_with_pd, img_with_init, img_with_mesh], # grid_ids = [[0, 1, 2, 3, 4]], grid_ids = [[1, 2, 3, 4]], ) imgs_spliced.append(img_spliced) img_final = splice_img(imgs_spliced, grid_ids=[[i] for i in range(len(log_data['img_patch']))]) img_final = to_tensor(img_final, device=None).permute(2, 0, 1) # (3, H, W) summary_writer.add_image('skelify/visualization', img_final, step_cnt) self.render_frames.append(img_final) except Exception as e: get_logger().error(f'Failed to visualize the optimization process: {e}') # traceback.print_exc() @rank_zero_only def _tb_log_for_report(self, step_cnt:int, log_data:Dict, *args, **kwargs): ''' Write the logging information to the TensorBoard. ''' get_logger().warning(f'This logging functions is just for presentation.') if len(self.render_frames) < 1: self.init_v = log_data['pd_verts'] self.init_ct = log_data['cam_t'] if step_cnt != 0 and (step_cnt + 1) % self.cfg.logger.interval != 0: return summary_writer = self.tb_logger.experiment # Save losses. for loss_name, loss_val in log_data['losses'].items(): summary_writer.add_scalar(f'losses/{loss_name}', loss_val.detach().item(), step_cnt) # Visualization of the optimization process. TODO: Maybe we can make this more elegant. if log_data['img_patch'] is None: log_data['img_patch'] = [np.zeros((self.cfg.img_patch_size, self.cfg.img_patch_size, 3), dtype=np.uint8)] \ * len(log_data['gt_kp2d']) # Overlay the skin mesh of the results on the original image. try: imgs_spliced = [] for i, img_patch in enumerate(log_data['img_patch']): img_with_init = render_mesh_overlay_img( faces = self.skel_model.skin_f, verts = self.init_v[i], K4 = [self.cfg.focal_length, self.cfg.focal_length, 0, 0], img = img_patch, Rt = [torch.eye(3), self.init_ct[i]], mesh_color = 'pink', ) img_with_init = annotate_img(img_with_init, 'init') img_with_mesh = render_mesh_overlay_img( faces = self.skel_model.skin_f, verts = log_data['pd_verts'][i], K4 = [self.cfg.focal_length, self.cfg.focal_length, 0, 0], img = img_patch, Rt = [torch.eye(3), log_data['cam_t'][i]], mesh_color = 'pink', ) img_with_mesh = annotate_img(img_with_mesh, 'pd_mesh') img_patch_raw = annotate_img(img_patch, 'raw') log_data['gt_kp2d'][i][..., :2] = (log_data['gt_kp2d'][i][..., :2] + 0.5) * self.cfg.img_patch_size img_with_gt = annotate_img(img_patch, 'gt_kp2d') img_with_gt = draw_kp2d_on_img( img_with_gt, log_data['gt_kp2d'][i], Skeleton_OpenPose25.bones, Skeleton_OpenPose25.bone_colors, ) img_spliced = splice_img([img_patch_raw, img_with_gt, img_with_init, img_with_mesh], grid_ids=[[0, 1, 2, 3]]) imgs_spliced.append(img_spliced) img_final = splice_img(imgs_spliced, grid_ids=[[i] for i in range(len(log_data['img_patch']))]) img_final = to_tensor(img_final, device=None).permute(2, 0, 1) summary_writer.add_image('visualization', img_final, step_cnt) self.render_frames.append(img_final) except Exception as e: get_logger().error(f'Failed to visualize the optimization process: {e}') traceback.print_exc()