SHOW/stage2_main.py

#  -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright2023 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Authors: paper author. 
# Special Acknowlegement:  Wojciech Zielonka and Justus Thies
# Contact: [email protected]

from pathlib import Path
import numpy as np
from tqdm import tqdm
import mmcv
import numpy as np
import os
from tqdm import tqdm
import cv2
import os.path
from functools import reduce
from pathlib import Path
from loguru import logger
import face_alignment
import mmcv
from pathlib import Path
import numpy as np
from tqdm import tqdm
import mmcv
import numpy as np
import os
import os.path as osp
from tqdm import tqdm
import cv2
import glob
import os.path
from functools import reduce
from pathlib import Path
from loguru import logger
import face_alignment
import mmcv

from SHOW.utils.video import images_to_video
from torchvision.transforms.functional import gaussian_blur
from pytorch3d.transforms import axis_angle_to_matrix
from pytorch3d.renderer import RasterizationSettings, PointLights, MeshRenderer, MeshRasterizer, TexturesVertex, SoftPhongShader, look_at_view_transform, PerspectiveCameras
from pytorch3d.transforms import axis_angle_to_matrix
from pytorch3d.io import load_obj

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torch.backends.cudnn
import torch.nn.functional as F

import SHOW
from SHOW.utils import default_timers
from SHOW.datasets import op_base
from SHOW.detector.face_detector import FaceDetector
from SHOW.load_models import load_smplx_model, load_vposer_model
from SHOW.save_results import save_one_results
from SHOW.load_models import load_save_pkl
from SHOW.flame.FLAME import FLAMETex
from SHOW.smplx_dataset import ImagesDataset
from SHOW.renderer import Renderer
from SHOW.load_assets import load_assets
from SHOW.loggers.logger import setup_logger
from SHOW.save_tracker import save_tracker
from SHOW.utils import is_valid_json
from configs.cfg_ins import condor_cfg


@logger.catch
def SHOW_stage2(*args, **kwargs):

    machine_info = SHOW.get_machine_info()
    import pprint
    pprint.pprint(f'machine_info: {machine_info}')

    loggers = kwargs.get('loggers', None)

    tracker_cfg = SHOW.from_rela_path(__file__,
                                      './configs/mmcv_tracker_config.py')
    tracker_cfg.update(**kwargs)
    tracker_cfg.merge_from_dict(condor_cfg)

    if tracker_cfg.get('over_write_cfg', None):
        tracker_cfg.update(tracker_cfg.over_write_cfg)
    

    mmcv.dump(tracker_cfg, tracker_cfg.tracker_cfg_path)
    
    
    try:
        gpu_mem = machine_info['gpu_info']['gpu_Total']

        import platform
        if platform.system() == 'Linux':
            # 50.0 * 24220 / (65.0*1024)
            tracker_cfg.bs_at_a_time = int(50.0 * gpu_mem / (80.0 * 1024))
            logger.warning(f'bs_at_a_time: {tracker_cfg.bs_at_a_time}')
    except:
        import traceback
        traceback.print_exc()

    Path(tracker_cfg.mica_save_path).mkdir(exist_ok=True, parents=True)
    Path(tracker_cfg.mica_org_out_path).mkdir(exist_ok=True, parents=True)

    iters = tracker_cfg.iters
    sampling = tracker_cfg.sampling
    device = tracker_cfg.device
    tracker_cfg.dtype = dtype = SHOW.str_to_torch_dtype(tracker_cfg.dtype)

    face_ider = SHOW.build_ider(tracker_cfg.ider_cfg)
    img_folder = tracker_cfg.img_folder
    template_im = os.listdir(img_folder)[0]
    template_im = os.path.join(img_folder, template_im)

    assets = load_assets(
        tracker_cfg,
        face_ider=face_ider,
        template_im=template_im,
    )
    if assets is None:
        return

    setup_logger(tracker_cfg.mica_all_dir, filename='mica.log', mode='o')

    
    if not Path(tracker_cfg.ours_pkl_file_path).exists():
        logger.warning(
            f'ours_pkl_file_path not exists: {tracker_cfg.ours_pkl_file_path}')
        return False

    if not is_valid_json(tracker_cfg.final_losses_json_path):
        logger.warning(
            f'final_losses_json_path not valid: {tracker_cfg.final_losses_json_path}'
        )
        return False

    
    with default_timers['build_vars_stage']:
        face_ider = SHOW.build_ider(tracker_cfg.ider_cfg)
        person_face_emb = assets.person_face_emb
        face_detector_mediapipe = FaceDetector('google', device=device)
        face_detector = face_alignment.FaceAlignment(
            face_alignment.LandmarksType._2D, device=device)

        body_model = load_smplx_model(dtype=dtype, **tracker_cfg.smplx_cfg)
        body_params_dict = load_save_pkl(tracker_cfg.ours_pkl_file_path,
                                         device)

        width = body_params_dict['width']
        height = body_params_dict['height']
        center = body_params_dict['center']
        camera_transl = body_params_dict['camera_transl']
        focal_length = body_params_dict['focal_length']
        total_batch_size = body_params_dict['batch_size']

        opt_bs = tracker_cfg.bs_at_a_time
        opt_iters = total_batch_size // opt_bs
        st_et_list = []
        for i in range(opt_iters):
            st = i * opt_bs
            et = (i + 1) * opt_bs
            if et > total_batch_size - 1:
                et = total_batch_size - 1
            st_et_list.append((st, et))

        op = op_base()
        smplx2flame_idx = assets.smplx2flame_idx

        mesh_file = Path(__file__).parent.joinpath(
            '../data/head_template_mesh.obj')

        diff_renderer = Renderer(torch.Tensor([[512, 512]]),
                                 obj_filename=mesh_file)

        flame_faces = load_obj(mesh_file)[1]
        flametex = FLAMETex(tracker_cfg.flame_cfg).to(device)

        mesh_rasterizer = MeshRasterizer(
            raster_settings=RasterizationSettings(image_size=[512, 512],
                                                  faces_per_pixel=1,
                                                  cull_backfaces=True,
                                                  perspective_correct=True))

        debug_renderer = MeshRenderer(
            rasterizer=mesh_rasterizer,
            shader=SoftPhongShader(device=device,
                                   lights=PointLights(
                                       device=device,
                                       location=((0.0, 0.0, -5.0), ),
                                       ambient_color=((0.5, 0.5, 0.5), ),
                                       diffuse_color=((0.5, 0.5, 0.5), ))))

    pre_frame_exp = None
    for opt_idx, (start_frame, end_frame) in enumerate(st_et_list):
        if assets.person_face_emb is not None:
            mica_part_file_path = f'w_mica_part_{start_frame}_{end_frame}_{opt_idx}_{opt_iters}.pkl'
            mica_part_pkl_path = os.path.join(tracker_cfg.mica_all_dir,
                                              mica_part_file_path)

            if Path(mica_part_pkl_path).exists():
                logger.info(
                    f'mica_part_pkl_path exists,skipping: {mica_part_pkl_path}'
                )
                pre_con = mmcv.load(mica_part_pkl_path)
                pre_frame_exp = pre_con['expression'][-1]
                pre_frame_exp = torch.Tensor(pre_frame_exp).to(device)
                continue

            opt_bs = end_frame - start_frame

            com_tex = torch.zeros(1, 150).to('cuda')
            com_sh = torch.zeros(1, 9, 3).to('cuda')

            use_shared_tex = 1
            if not use_shared_tex:
                opt_bs_tex = nn.Parameter(com_tex).expand(opt_bs, -1).detach()
            else:
                opt_bs_tex = nn.Parameter(com_tex).expand(1, -1).detach()

            opt_bs_sh = nn.Parameter(com_sh).expand(opt_bs, -1, -1).detach()

            logger.info(f'origin input data frame batchsize:{opt_bs}')
            
            with default_timers['load_dataset_stage']:

                debug = 0
                if debug: opt_bs = 30

                dataset = ImagesDataset(
                    tracker_cfg,
                    start_frame=start_frame,
                    face_ider=face_ider,
                    person_face_emb=person_face_emb,
                    face_detector_mediapipe=face_detector_mediapipe,
                    face_detector=face_detector)
                dataloader = DataLoader(dataset,
                                        batch_size=opt_bs,
                                        num_workers=0,
                                        shuffle=False,
                                        pin_memory=True,
                                        drop_last=False)
                iterator = iter(dataloader)
                batch = next(iterator)
                if not debug:
                    batch = SHOW.utils.to_cuda(batch)
                valid_bool = batch['is_person_deted'].bool()
                valid_bs = valid_bool.count_nonzero()
                logger.info(f'valid input data frame batchsize:{valid_bs}')
                logger.info(f'valid_bool: {valid_bool}')

                if valid_bs == 0:
                    logger.warning('valid bs == 0, skipping')
                    open(mica_part_pkl_path + '.empty', 'a').close()
                    continue

                bbox = batch['bbox']
                images = batch['cropped_image']
                landmarks = batch['cropped_lmk']
                h = batch['h']
                w = batch['w']
                py = batch['py']
                px = batch['px']

                diff_renderer.masking.set_bs(valid_bs)
                diff_renderer = diff_renderer.to(device)

            debug = 0
            report_wandb = 0
            use_opt_pose = 1
            save_traing_img = 0
            observe_idx_list = [4, 8]
            
            with default_timers['optimize_stage']:
                model_output = None

                def get_pose_opt(start_frame, end_frame):
                    tmp = body_params_dict['body_pose_axis'][
                        start_frame:end_frame, ...].clone().detach()
                    tmp = tmp.reshape(tmp.shape[0], -1, 3)
                    return torch.stack([tmp[:, 12 - 1, :], 
                                        tmp[:, 15 - 1, :]],
                                       dim=1)

                def clone_params_color(start_frame, end_frame):
                    opt_var_clone_detach = [
                        {
                            'params': [
                                nn.Parameter(
                                    body_params_dict['expression']
                                    [start_frame:end_frame].clone().detach())
                            ],
                            'lr':
                            0.025,
                            'name': ['exp']
                        },
                        {
                            'params': [
                                nn.Parameter(body_params_dict['leye_pose']
                                             [start_frame:end_frame].clone(
                                             ).clone().detach())
                            ],
                            'lr':
                            0.001,
                            'name': ['leyes']
                        },
                        {
                            'params': [
                                nn.Parameter(
                                    body_params_dict['reye_pose']
                                    [start_frame:end_frame].clone().detach())
                            ],
                            'lr':
                            0.001,
                            'name': ['reyes']
                        },
                        {
                            'params': [
                                nn.Parameter(
                                    body_params_dict['jaw_pose']
                                    [start_frame:end_frame].clone().detach())
                            ],
                            'lr':
                            0.001,
                            'name': ['jaw']
                        },
                        {
                            'params':
                            [nn.Parameter(opt_bs_sh.clone().detach())],
                            'lr': 0.01,
                            'name': ['sh']
                        },
                        {
                            'params':
                            [nn.Parameter(opt_bs_tex.clone().detach())],
                            'lr': 0.005,
                            'name': ['tex']
                        },
                    ]
                    if use_opt_pose:
                        opt_var_clone_detach.append({
                            'params': [
                                nn.Parameter(
                                    get_pose_opt(start_frame, end_frame))
                            ],
                            'lr':
                            0.005,
                            'name': ['body_pose']
                        })
                    return opt_var_clone_detach

                save_traing_img_dir = tracker_cfg.mica_process_path + f'_{start_frame}_{end_frame}'

                if save_traing_img:
                    Path(save_traing_img_dir).mkdir(parents=True,
                                                    exist_ok=True)

                with tqdm(total=iters * 3,
                          position=0,
                          leave=True,
                          bar_format="{percentage:3.0f}%|{bar}{r_bar}{desc}"
                          ) as pbar:
                    for k, scale in enumerate(sampling):

                        size = [int(512 * scale), int(512 * scale)]
                        img = F.interpolate(images.float().clone(),
                                            size,
                                            mode='bilinear',
                                            align_corners=False)

                        if k > 0:
                            img = gaussian_blur(img, [9, 9]).detach()

                        flipped = torch.flip(img, [2, 3])
                        flipped = flipped[valid_bool.bool(), ...]

                        best_loss = np.inf
                        prev_loss = np.inf

                        xb_min, xb_max, yb_min, yb_max = bbox.values()
                        box_w = xb_max - xb_min
                        box_h = yb_max - yb_min
                        box_w = box_w.int()
                        box_h = box_h.int()

                        image_size = size

                        diff_renderer.rasterizer.reset()
                        diff_renderer.set_size(image_size)
                        debug_renderer.rasterizer.raster_settings.image_size = size

                        image_lmks = landmarks * size[0] / 512
                        image_lmks = image_lmks[valid_bool.bool(), ...]

                        optimizer = torch.optim.Adam(
                            clone_params_color(start_frame, end_frame))
                        params = optimizer.param_groups
                        get_param = SHOW.utils.get_param

                        cur_tex = get_param('tex', params)
                        cur_sh = get_param('sh', params)

                        cur_exp = get_param('exp', params)
                        cur_leyes = get_param('leyes', params)
                        cur_reyes = get_param('reyes', params)
                        cur_jaw = get_param('jaw', params)

                        if use_opt_pose:
                            two_opt = get_param('body_pose', params)
                            frame_pose = body_params_dict['body_pose_axis'][
                                start_frame:end_frame]
                            bs = frame_pose.shape[0]
                            frame_pose = frame_pose.reshape(bs, -1, 3)
                            cur_pose = torch.cat(
                                [
                                    frame_pose[:, :11, :],
                                    two_opt[:, 0:1],  #11
                                    frame_pose[:, 12:14, :],
                                    two_opt[:, 1:2],  #14
                                    frame_pose[:, 15:, :]
                                ],
                                dim=1).reshape(bs, 1, -1)
                        else:
                            frame_pose = body_params_dict['body_pose_axis'][
                                start_frame:end_frame]
                            bs = frame_pose.shape[0]
                            cur_pose = frame_pose.reshape(bs, 1, -1)

                        cur_transl = body_params_dict['transl'][
                            start_frame:end_frame]
                        cur_global_orient = body_params_dict['global_orient'][
                            start_frame:end_frame]
                        cur_left_hand_pose = body_params_dict[
                            'left_hand_pose'][start_frame:end_frame]
                        cur_right_hand_pose = body_params_dict[
                            'right_hand_pose'][start_frame:end_frame]

                        R = torch.Tensor([[[-1, 0, 0], [0, -1, 0], [0, 0, 1]]])
                        bs_image_size = torch.Tensor(image_size).repeat(
                            opt_bs, 1).to(device)
                        bs_camera_transl = (camera_transl).repeat(opt_bs,
                                                                  1).to(device)

                        # bs_center: torch.Size([10, 2])
                        bs_center = torch.Tensor(center).repeat(opt_bs,
                                                                1).to(device)
                        bs_box_min = torch.stack([xb_min, yb_min],
                                                 dim=-1).to(device)
                        bs_R = R.repeat(opt_bs, 1, 1).to(device)

                        s_w = size[0] / box_w
                        s_h = size[1] / box_h
                        s_1to2 = torch.stack([s_w, s_h], dim=1)
                        s_1to2 = s_1to2.to(device)
                        bs_pp = (bs_center - bs_box_min) * (s_1to2)

                        bs_pp[:, 0] = (bs_pp[:, 0] + px) * 512 / (w + 2 * px)
                        bs_pp[:, 1] = (bs_pp[:, 1] + py) * 512 / (h + 2 * py)

                        s_1to2[:, 0] = s_1to2[:, 0] * 512 / (w + 2 * px)
                        s_1to2[:, 1] = s_1to2[:, 1] * 512 / (h + 2 * py)

                        cam_cfg = dict(
                            principal_point=bs_pp,
                            focal_length=focal_length * s_1to2,
                            R=bs_R,
                            T=bs_camera_transl,
                            image_size=bs_image_size,
                        )

                        for key, val in cam_cfg.items():
                            cam_cfg[key] = val[valid_bool.bool(), ...]

                        cameras = PerspectiveCameras(**cam_cfg,
                                                     device=device,
                                                     in_ndc=False)

                        if True:
                            for p in range(iters):  # p=0
                                if (p + 1) % 32 == 0:
                                    diff_renderer.rasterizer.reset()
                                losses = {}
                                model_output = body_model(
                                    return_verts=True,
                                    jaw_pose=cur_jaw,
                                    leye_pose=cur_leyes,
                                    reye_pose=cur_reyes,
                                    expression=cur_exp,
                                    betas=body_params_dict['betas'],
                                    transl=cur_transl,
                                    body_pose=cur_pose,
                                    global_orient=cur_global_orient,
                                    left_hand_pose=cur_left_hand_pose,
                                    right_hand_pose=cur_right_hand_pose,
                                )

                                vertices = model_output.vertices[:,
                                                                 smplx2flame_idx
                                                                 .long(), :]
                                vertices = vertices[valid_bool.bool(), ...]

                                lmk68_all = model_output.joints[:, 67:67 + 51 +
                                                                17, :]
                                lmk68 = lmk68_all[valid_bool.bool(), ...]

                                proj_lmks = cameras.transform_points_screen(
                                    lmk68)[:, :, :2]
                                proj_lmks = torch.cat([
                                    proj_lmks[:, -17:, :],
                                    proj_lmks[:, :-17, :]
                                ],
                                                      dim=1)

                                I = torch.eye(3)[None].to(device)

                                if pre_frame_exp is not None and start_frame != 0:
                                    losses['pre_exp'] = 0.001 * torch.sum(
                                        (pre_frame_exp - cur_exp[0])**2)

                                if False:
                                    linear_rot_left = (axis_angle_to_matrix(
                                        cur_leyes[valid_bool.bool(), ...]))
                                    linear_rot_right = (axis_angle_to_matrix(
                                        cur_reyes[valid_bool.bool(), ...]))
                                    losses['eyes_sym_reg'] = torch.sum(
                                        (linear_rot_right - linear_rot_left)**
                                        2) / opt_bs
                                    losses['eyes_left_reg'] = torch.sum(
                                        (I - linear_rot_left)**2) / opt_bs
                                    losses['eyes_right_reg'] = torch.sum(
                                        (I - linear_rot_right)**2) / opt_bs

                                w_lmks = tracker_cfg.w_lmks
                                losses['lmk'] = SHOW.utils.lmk_loss(
                                    proj_lmks, image_lmks,
                                    image_size) * w_lmks * 8.0
                                losses[
                                    'lmk_mount'] = SHOW.utils.mouth_loss(  #(49, 68)
                                        proj_lmks, image_lmks,
                                        image_size) * w_lmks * 4.0 * 4
                                losses['lmk_oval'] = SHOW.utils.lmk_loss(
                                    proj_lmks[:, :17, ...], image_lmks[:, :17,
                                                                       ...],
                                    image_size) * w_lmks

                                losses['jaw_reg'] = torch.sum(
                                    (I - axis_angle_to_matrix(
                                        cur_jaw[valid_bool.bool(), ...]))**
                                    2) * 16.0 / opt_bs
                                losses['exp_reg'] = torch.sum(
                                    cur_exp[valid_bool.bool(),
                                            ...]**2) * 0.01 / opt_bs

                                if use_shared_tex:
                                    losses['tex_reg'] = torch.sum(cur_tex**
                                                                  2) * 0.02
                                else:
                                    losses['tex_reg'] = torch.sum(
                                        cur_tex[valid_bool.bool(),
                                                ...]**2) * 0.02 / opt_bs

                                def temporary_loss(o_w, i_w, gmof, param):
                                    assert param.shape[
                                        0] > 2, f'optimize batchsize must > 2 to enable temporary smooth'
                                    return (o_w**2) * (gmof(
                                        i_w *
                                        (param[2:, ...] + param[:-2, ...] -
                                         2 * param[1:-1, ...]))).mean()

                                def pow(x):
                                    return x.pow(2)

                                if cur_exp.shape[0] > 2:
                                    losses['loss_sexp'] = temporary_loss(
                                        1.0, 2.0, pow, cur_exp)
                                    losses['loss_sjaw'] = temporary_loss(
                                        1.0, 2.0, pow, cur_jaw)
                                    
                                def k_fun(k):
                                    return tracker_cfg.w_pho * 32.0 if k > 0 else tracker_cfg.w_pho

                                albedos = flametex(cur_tex) / 255.

                                if use_shared_tex:
                                    albedos = albedos.expand(
                                        valid_bs, -1, -1, -1)
                                else:
                                    albedos = albedos[valid_bool.bool(), ...]

                                ops = diff_renderer(
                                    vertices, albedos,
                                    cur_sh[valid_bool.bool(), ...], cameras)

                                grid = ops['position_images'].permute(
                                    0, 2, 3, 1)[:, :, :, :2]
                                sampled_image = F.grid_sample(
                                    flipped, grid, align_corners=False)
                                ops_mask = SHOW.utils.parse_mask(ops)
                                tmp_img = ops['images']

                                losses['pho'] = SHOW.utils.pixel_loss(
                                    tmp_img, sampled_image,
                                    ops_mask) * k_fun(k)

                                all_loss = 0.
                                for key in losses.keys():
                                    all_loss = all_loss + losses[key]
                                losses['all_loss'] = all_loss

                                log_str = SHOW.print_dict_losses(losses)

                                if report_wandb:
                                    if globals().get('wandb', None) is None:
                                        os.environ[
                                            'WANDB_API_KEY'] = 'xxx'
                                        os.environ['WANDB_NAME'] = 'tracker'
                                        import wandb
                                        wandb.init(
                                            reinit=True,
                                            resume='allow',
                                            project='tracker',
                                        )
                                        globals()['wandb'] = wandb

                                    if globals().get('wandb',
                                                     None) is not None:
                                        globals()['wandb'].log(losses)

                                if save_traing_img:

                                    def save_callback(frame, final_views):
                                        cur_idx = (frame + opt_bs * opt_idx)
                                        if cur_idx in observe_idx_list:

                                            observe_idx_frame_dir = os.path.join(
                                                save_traing_img_dir,
                                                f'{cur_idx:03d}')
                                            Path(observe_idx_frame_dir).mkdir(
                                                parents=True, exist_ok=True)

                                            cv2.imwrite(
                                                os.path.join(
                                                    observe_idx_frame_dir,
                                                    f'{k}_{p}.jpg'),
                                                final_views)

                                    save_tracker(
                                        img,
                                        valid_bool,
                                        valid_bs,
                                        ops,
                                        vertices,
                                        cameras,
                                        image_lmks,
                                        proj_lmks,
                                        flame_faces,
                                        mesh_rasterizer,
                                        debug_renderer,
                                        save_callback,
                                    )

                                if loggers is not None:
                                    loggers.log_bs(losses)
                                    if torch.isnan(all_loss).sum():
                                        loggers.alert(
                                            title='Nan error',
                                            msg=
                                            f'tracker nan in: {tracker_cfg.ours_output_folder}'
                                        )
                                        open(
                                            tracker_cfg.ours_output_folder +
                                            '/mica_opt_nan.info', 'a').close()
                                        break

                                else:
                                    pbar.set_description(log_str)
                                    pbar.update(1)

                                optimizer.zero_grad()
                                all_loss.backward()
                                optimizer.step()

                                if all_loss.item() < best_loss:
                                    best_loss = all_loss.item()
                                    opt_bs_tex = cur_tex.clone().detach()
                                    opt_bs_sh = cur_sh.clone().detach()
                                    body_params_dict['expression'][
                                        start_frame:end_frame] = cur_exp.clone(
                                        ).detach()
                                    body_params_dict['leye_pose'][
                                        start_frame:
                                        end_frame] = cur_leyes.clone().detach(
                                        )
                                    body_params_dict['reye_pose'][
                                        start_frame:
                                        end_frame] = cur_reyes.clone().detach(
                                        )
                                    body_params_dict['jaw_pose'][
                                        start_frame:end_frame] = cur_jaw.clone(
                                        ).detach()
                                    body_params_dict['body_pose_axis'][
                                        start_frame:
                                        end_frame] = cur_pose.clone().detach(
                                        ).squeeze()

        
        with default_timers['saving_stage']:
            if save_traing_img:
                for idx in observe_idx_list:
                    observe_idx_frame_dir = os.path.join(
                        save_traing_img_dir, f'{idx:03d}')
                    Path(observe_idx_frame_dir).mkdir(parents=True,
                                                      exist_ok=True)

                    if not SHOW.is_empty_dir(observe_idx_frame_dir):
                        images_to_video(
                            input_folder=observe_idx_frame_dir,
                            output_path=observe_idx_frame_dir + '.mp4',
                            img_format=None,
                            fps=30,
                        )

            dict_to_save = dict(
                expression=body_params_dict['expression']
                [start_frame:end_frame].clone().detach().cpu().numpy(),
                leye_pose=body_params_dict['leye_pose']
                [start_frame:end_frame].clone().detach().cpu().numpy(),
                reye_pose=body_params_dict['reye_pose']
                [start_frame:end_frame].clone().detach().cpu().numpy(),
                jaw_pose=body_params_dict['jaw_pose']
                [start_frame:end_frame].clone().detach().cpu().numpy(),
                body_pose_axis=body_params_dict['body_pose_axis']
                [start_frame:end_frame].clone().detach().cpu().numpy(),
                tex=opt_bs_tex.clone().detach().cpu().numpy(),
                sh=opt_bs_sh.clone().detach().cpu().numpy(),
            )
            mmcv.dump(dict_to_save, mica_part_pkl_path)
            logger.info(f'mica pkl part path: {mica_part_pkl_path}')
            pre_frame_exp = dict_to_save['expression'][-1]
            pre_frame_exp = torch.Tensor(pre_frame_exp).to(device)
            vertices_ = model_output.vertices.clone().detach().cpu()
            logger.info(
                f'mica render to origin path: {tracker_cfg.mica_org_out_path}')

            import platform
            if platform.system() == "Linux":
                os.environ['PYOPENGL_PLATFORM'] = 'egl'
            else:
                if 'PYOPENGL_PLATFORM' in os.environ:
                    os.environ.__delitem__('PYOPENGL_PLATFORM')

            import pyrender
            input_renderer = pyrender.OffscreenRenderer(viewport_width=width,
                                                        viewport_height=height,
                                                        point_size=1.0)

            for idx in tqdm(range(vertices_.shape[0]),
                            desc='saving ours final pyrender images'):  # idx=0

                cur_idx = idx + start_frame + 1

                input_img = SHOW.find_full_impath_by_name(
                    root=tracker_cfg.img_folder, name=f'{cur_idx:06d}')
                output_name = os.path.join(
                    tracker_cfg.mica_org_out_path,
                    f"{cur_idx:06}.{tracker_cfg.output_img_ext}")

                camera_pose = op.get_smplx_to_pyrender_K(camera_transl)

                meta_data = dict(
                    input_img=input_img,
                    output_name=output_name,
                )

                save_one_results(
                    vertices_[idx],
                    body_model.faces,
                    img_size=(height, width),
                    center=center,
                    focal_length=[focal_length, focal_length],
                    camera_pose=camera_pose,
                    meta_data=meta_data,
                    input_renderer=input_renderer,
                )
            input_renderer.delete()

            
            if tracker_cfg.save_final_vis:

                def save_callback(frame, final_views):
                    cur_idx = (frame + opt_bs * opt_idx)

                    if loggers is not None:
                        loggers.log_image(f"final_mica_img/{cur_idx:03d}",
                                          final_views / 255.0)

                    cv2.imwrite(
                        os.path.join(tracker_cfg.mica_save_path,
                                     f'{cur_idx:03d}.jpg'), final_views)

                if True:
                    save_tracker(
                        img,
                        valid_bool,
                        valid_bs,
                        ops,
                        vertices,
                        cameras,
                        image_lmks,
                        proj_lmks,
                        flame_faces,
                        mesh_rasterizer,
                        debug_renderer,
                        save_callback,
                    )

    load_data = mmcv.load(tracker_cfg.ours_pkl_file_path)[0]
    load_data = SHOW.replace_mica_exp(tracker_cfg.mica_all_dir, load_data)
    mmcv.dump([load_data], tracker_cfg.mica_merge_pkl)

    if not Path(tracker_cfg.mica_org_out_video).exists():
        if not SHOW.is_empty_dir(tracker_cfg.mica_org_out_path):
            images_to_video(
                input_folder=tracker_cfg.mica_org_out_path,
                output_path=tracker_cfg.mica_org_out_video,
                img_format=None,
                fps=30,
            )
    if not Path(tracker_cfg.mica_grid_video).exists():
        if not SHOW.is_empty_dir(tracker_cfg.mica_save_path):
            images_to_video(
                input_folder=tracker_cfg.mica_save_path,
                output_path=tracker_cfg.mica_grid_video,
                img_format=None,
                fps=30,
            )