lib/modeling/optim/skelify/skelify.py

from lib.kits.basic import *

import cv2
import traceback
from tqdm import tqdm

from lib.body_models.common import make_SKEL
from lib.body_models.abstract_skeletons import Skeleton_OpenPose25
from lib.utils.vis import render_mesh_overlay_img
from lib.utils.data import to_tensor
from lib.utils.media import draw_kp2d_on_img, annotate_img, splice_img
from lib.utils.camera import perspective_projection

from .utils import (
    compute_rel_change,
    gmof,
)

from .closure import build_closure

class SKELify():

    def __init__(self, cfg, tb_logger=None, device='cuda:0', name='SKELify'):
        self.cfg = cfg
        self.name = name
        self.eq_thre = cfg.early_quit_thresholds

        self.tb_logger = tb_logger

        self.device = device
        # self.skel_model = make_SKEL(device=device)
        self.skel_model = instantiate(cfg.skel_model).to(device)

        # Shortcuts.
        self.n_samples = cfg.logger.samples_per_record

        # 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 it is None, the visualization will simply use a black background.

        ### Returns:
        - dict, containing the optimized parameters.
            - poses: torch.Tensor, (B, 46)
            - betas: torch.Tensor, (B, 10)
            - cam_t: torch.Tensor, (B, 3)
        '''
        with PM.time_monitor('input preparation'):
            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 = float(self.cfg.focal_length / self.cfg.img_patch_size)  # float

        # ⛩️ Optimization phases, controlled by config file.
        with PM.time_monitor('optim') as tm:
            prev_steps = 0  # accumulate the steps are *supposed* to be done in the previous phases
            n_phases = len(self.cfg.phases)
            for phase_id, phase_name in enumerate(self.cfg.phases):
                phase_cfg = self.cfg.phases[phase_name]
                # 📦 Data 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
                log_data = {}
                tm.tick(f'Data preparation')

                # ⚙️ Optimization preparation.
                optimizer = instantiate(phase_cfg.optimizer, optim_params, _recursive_=True)
                closure = self._build_closure(
                        cfg=phase_cfg, optimizer=optimizer,  # basic
                        inputs=inputs, focal_length=focal_length, gt_kp2d=gt_kp2d,  # data reference
                        log_data=log_data,  # monitoring
                    )
                tm.tick(f'Optimizer * closure prepared.')

                # 🚀 Optimization loop.
                with tqdm(range(phase_cfg.max_loop)) as bar:
                    prev_loss = None
                    bar.set_description(f'[{phase_name}] Loss: ???')
                    for i in bar:
                        # 1. Main part of the optimization loop.
                        log_data.clear()
                        curr_loss = optimizer.step(closure)

                        # 2. Log.
                        if self.tb_logger is not None:
                            log_data.update({
                                'img_patch' : img_patch[:self.n_samples] if img_patch is not None else None,
                                'gt_kp2d'   : gt_kp2d[:self.n_samples].detach().clone(),
                            })
                            self._tb_log(prev_steps + i, phase_name, log_data)

                        # 3. The end of one optimization loop.
                        bar.set_description(f'[{phase_id+1}/{n_phases}] @ {phase_name} - Loss: {curr_loss:.4f}')
                        if self._can_early_quit(optim_params, prev_loss, curr_loss):
                            break
                        prev_loss = curr_loss

                    prev_steps += phase_cfg.max_loop
                    tm.tick(f'{phase_name} finished.')

        with PM.time_monitor('last infer'):
            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)
            skel_outputs = self.skel_model(poses=poses, betas=betas, skelmesh=False)  # (B, 44, 3)
            optim_kp3d = skel_outputs.joints  # (B, 44, 3)
            # Evaluate the confidence of the results.
            focal_length_xy = np.ones((len(poses), 2)) * focal_length  # (B, 2)
            optim_kp2d = perspective_projection(
                    points       = optim_kp3d,
                    translation  = cam_t,
                    focal_length = to_tensor(focal_length_xy, device=self.device),
                )
            kp2d_err = SKELify.eval_kp2d_err(gt_kp2d, optim_kp2d)  # (B,)

        # ⛩️ Prepare the output data.
        outputs = {
                'poses'    : poses,     # (B, 46)
                'betas'    : betas,     # (B, 10)
                'cam_t'    : cam_t,     # (B, 3)
                'kp2d_err' : kp2d_err,  # (B,)
            }
        return outputs


    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} = 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


    def _build_closure(self, *args, **kwargs):
        # Using this way to hide the very details and simplify the code.
        return build_closure(self, *args, **kwargs)


    @staticmethod
    def eval_kp2d_err(gt_kp2d_with_conf:torch.Tensor, pd_kp2d:torch.Tensor):
        '''  Evaluate the mean 2D keypoints L2 error. The formula is: ∑(gt - pd)^2 * conf / ∑conf. '''
        assert len(gt_kp2d_with_conf.shape) == len(gt_kp2d_with_conf.shape), f'gt_kp2d_with_conf.shape={gt_kp2d_with_conf.shape}, pd_kp2d.shape={pd_kp2d.shape} but they should both be ((B,) J, D).'
        if len(gt_kp2d_with_conf.shape) == 2:
            gt_kp2d_with_conf, pd_kp2d = gt_kp2d_with_conf[None], pd_kp2d[None]
        assert len(gt_kp2d_with_conf.shape) == 3, f'gt_kp2d_with_conf.shape={gt_kp2d_with_conf.shape}, pd_kp2d.shape={pd_kp2d.shape} but they should both be ((B,) J, D).'
        B, J, _ = gt_kp2d_with_conf.shape
        assert gt_kp2d_with_conf.shape == (B, J, 3), f'gt_kp2d_with_conf.shape={gt_kp2d_with_conf.shape} but it should be ((B,) J, 3).'
        assert pd_kp2d.shape == (B, J, 2), f'pd_kp2d.shape={pd_kp2d.shape} but it should be ((B,) J, 2).'

        conf = gt_kp2d_with_conf[..., 2]  # (B, J)
        gt_kp2d = gt_kp2d_with_conf[..., :2]  # (B, J, 2)
        kp2d_err = torch.sum((gt_kp2d - pd_kp2d) ** 2, dim=-1) * conf  # (B, J)
        kp2d_err = kp2d_err.sum(dim=-1) / (torch.sum(conf, dim=-1) + 1e-6)  # (B,)
        return kp2d_err


    @rank_zero_only
    def _tb_log(self, step_cnt:int, phase_name:str, log_data:Dict, *args, **kwargs):
        ''' Write the logging information to the TensorBoard. '''
        if step_cnt != 0 and (step_cnt + 1) % self.cfg.logger.interval_skelify != 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, 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_kp2d_err = log_data['kp2d_err']
            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']):
                kp2d_err = log_data['kp2d_err'][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, 128, 128],
                        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_kp2d_err[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, 128, 128],
                        img        = img_patch,
                        Rt         = [torch.eye(3), log_data['cam_t'][i]],
                        mesh_color = 'pink',
                    )
                betas_max = log_data['optim_betas'][i].abs().max().item()
                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 = cv2.addWeighted(img_with_mesh, 0.7, img_patch, 0.3, 0)
                img_with_pd = draw_kp2d_on_img(
                        img_with_pd,
                        log_data['pd_kp2d'][i],
                        Skeleton_OpenPose25.bones,
                        Skeleton_OpenPose25.bone_colors,
                    )

                img_with_pd = annotate_img(img_with_pd, 'pd')
                img_with_pd = annotate_img(img_with_pd, f'Quality: {kp2d_err*1000:.3f}/1e3\nbetas_max: {betas_max:.3f}', pos='tl')
                img_with_mesh = annotate_img(img_with_mesh, f'Quality: {kp2d_err*1000:.3f}/1e3\nbetas_max: {betas_max:.3f}', pos='tl')
                img_with_mesh = annotate_img(img_with_mesh, 'pd_mesh')

                img_spliced = splice_img(
                        img_grids = [img_patch_raw, img_with_gt, img_with_pd, img_with_mesh, img_with_init],
                        # grid_ids  = [[0, 1, 2, 3, 4]],
                        grid_ids  = [[1, 2, 3, 4]],
                    )
                img_spliced = annotate_img(img_spliced, f'{phase_name}/{step_cnt}', pos='tl')
                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()