code/SMPLer_X.py

import torch
import torch.nn as nn
from torch.nn import functional as F
from nets.smpler_x import PositionNet, HandRotationNet, FaceRegressor, BoxNet, HandRoI, BodyRotationNet, BodyTokenNet, GlobalHandPositionNet, HandTokenRegressor
from tokenhmr.lib.models.heads.token_head import SMPLTokenDecoderHead
from nets.loss import CoordLoss, ParamLoss, CELoss
from utils.human_models import smpl_x, smpl
from utils.transforms import rot6d_to_axis_angle, restore_bbox
from config import cfg
import math
import copy
from mmpose.models import build_posenet
from mmcv import Config
from FishEyeCalibrated import FishEyeCameraCalibrated
import numpy as np
import cv2
import matplotlib.pyplot as plt
from pdb import set_trace
# from common.utils.vis import render_mesh_fisheye
from scipy.spatial import cKDTree
from data import humandata
import os
from omegaconf import OmegaConf
from DPoserX.run.tester.wholebody.smplify import DPoser
from torch.autograd import Function


        
# from main.transformer_utils.mmpose.models.backbones.ViT_DINO import vit_large
# from main.transformer_utils.mmpose.models.heads.RAFTDepthNormalDPTDecoder5 import RAFTDepthNormalDPT5
# from mmengine.config import Config

class GradReverse(Function):
    @staticmethod
    def forward(ctx, x, lambd):
        ctx.lambd = float(lambd)
        return x.view_as(x)
    @staticmethod
    def backward(ctx, grad_output):
        return grad_output.neg() * ctx.lambd, None
    

bone_pairs = [
    [0, 1], [1, 2], [2, 3],         # right arm
    [0, 4], [4, 5], [5, 6],         # left arm
    [1, 7], [7, 8], [8, 9], [9,10], # right leg
    [4,11], [11,12], [12,13], [13,14], # left leg
]

mo2cap2_to_smplx = {
    "neck": 7,
    "right_shoulder": 9,
    "right_elbow": 11,
    "right_wrist": 13,
    "left_shoulder": 8,
    "left_elbow": 10,
    "left_wrist": 12,
    "right_hip": 2,
    "right_knee": 4,
    "right_ankle": 6,
    "right_foot": 19,
    "left_hip": 1,
    "left_knee": 3,
    "left_ankle": 5,
    "left_foot": 16,
}

mo2cap2_joint_names = [
    "neck","right_shoulder","right_elbow","right_wrist",
    "left_shoulder","left_elbow","left_wrist",
    "right_hip","right_knee","right_ankle","right_foot",
    "left_hip","left_knee","left_ankle","left_foot",
]
mo2cap2_chain = [
    [0, 1, 2, 3],
    [0, 4, 5, 6],
    [1, 7, 8, 9, 10],
    [4, 11, 12, 13, 14],
    [7, 11],
]


class Model(nn.Module):
    def __init__(self, encoder, token_decoder, body_position_net, body_rotation_net, box_net, hand_position_net, hand_roi_net,
                 hand_rotation_net, face_regressor, mode):
        super(Model, self).__init__()

        if getattr(cfg, 'fisheye_camera_path', False):
            fisheye_camera_path = cfg.fisheye_camera_path
        else:
            fisheye_camera_path = '/home/cv1/works/SMPLer-X/main/fisheye.calibration_01_12.json' 
        self.fisheye_camera = FishEyeCameraCalibrated(fisheye_camera_path)
           
        # body
        self.encoder = encoder

        self.box_net = box_net
        
        # hand
        self.hand_roi_net = hand_roi_net
        self.hand_position_net = hand_position_net
        self.hand_regressor = hand_rotation_net
        self.neck = [self.box_net, self.hand_roi_net]
        # face
        self.face_regressor = face_regressor

        self.smplx_layer = copy.deepcopy(smpl_x.layer['neutral']).cuda()
        self.smpl_layer  = copy.deepcopy(smpl.layer['neutral']).cuda()  # SMPL(6890 verts)

        self.coord_loss = CoordLoss()
        self.param_loss = ParamLoss()
        self.ce_loss = CELoss()

        self.body_num_joints = len(smpl_x.pos_joint_part['body'])
        self.hand_joint_num = len(smpl_x.pos_joint_part['rhand'])

        self.body_position_net = body_position_net
        if getattr(cfg, 'use_token_decoder', False):
            if getattr(cfg, 'use_smpl', False):
                self.token_decoder = token_decoder
                

                self.head = [self.token_decoder, self.body_position_net,
                        self.hand_position_net, self.hand_regressor, 
                        self.face_regressor]

                self.trainable_modules = [self.encoder, self.token_decoder, self.body_position_net]
            else:
                self.token_decoder = token_decoder
                

                self.head = [self.token_decoder, self.body_position_net,
                        self.hand_position_net, self.hand_regressor, 
                        self.face_regressor]

                self.trainable_modules = [self.encoder, self.token_decoder, self.body_position_net,
                                        self.box_net, self.hand_position_net,
                                        self.hand_roi_net, self.hand_regressor, self.face_regressor]

        else:
            if getattr(cfg, 'use_smpl', False):
                
                self.body_regressor = body_rotation_net
                self.head = [self.body_position_net, self.body_regressor,
                        self.hand_position_net, self.hand_regressor, 
                        self.face_regressor]
                
                if getattr(cfg, 'use_dposerx', False):
                    if mode == 'train':
                        self.dposer_x = DPoser(batch_size=cfg.train_batch_size, config_path=cfg.dposerx_cfg_path).cuda()
                        self.trainable_modules = [self.encoder, self.dposer_x, self.body_position_net, self.body_regressor]
                        #self.trainable_modules = [self.encoder, self.body_position_net, self.body_regressor]
                    else:
                        self.dposer_x = DPoser(batch_size=cfg.test_batch_size, config_path=cfg.dposerx_cfg_path).cuda()
                
                else:
                    self.trainable_modules = [self.encoder, self.body_position_net, self.body_regressor]

            else:
                
                self.body_regressor = body_rotation_net
                self.head = [self.body_position_net, self.body_regressor,
                        self.hand_position_net, self.hand_regressor, 
                        self.face_regressor]

                if getattr(cfg, 'use_dposerx', False):
                    if mode == 'train':
                        self.dposer_x = DPoser(batch_size=cfg.train_batch_size, config_path=cfg.dposerx_cfg_path).cuda()

                        self.trainable_modules = [self.encoder, self.dposer_x, self.body_position_net, self.body_regressor, 
                                        self.box_net, self.hand_position_net,
                                        self.hand_roi_net, self.hand_regressor]
                    
                    else:
                        self.dposer_x = DPoser(batch_size=cfg.test_batch_size, config_path=cfg.dposerx_cfg_path).cuda()
                
                else:
                    self.trainable_modules = [self.encoder, self.body_position_net, self.body_regressor, 
                                        self.box_net, self.hand_position_net,
                                        self.hand_roi_net, self.hand_regressor, self.face_regressor]
                    
                           
        self.special_trainable_modules = []
        
        # --- ID adversarial head (optional) ---
        self.id_adv_lambda = float(getattr(cfg, 'id_adv_lambda', 0.0))   # e.g., 0.5
        self.num_subjects  = int(getattr(cfg, 'num_subjects', 0))        # ID 클래스 수
        if self.id_adv_lambda > 0 and self.num_subjects > 0:
            # img_feat의 채널 차원 = encoder 출력 채널(보통 cfg.feat_dim과 일치)
            self.id_head = nn.Linear(cfg.feat_dim, self.num_subjects)
            self.id_loss_fn = nn.BCEWithLogitsLoss()
            # 학습 모듈에 포함(선택)
            self.trainable_modules.append(self.id_head)
            self.head.append(self.id_head) 

        # backbone:
        param_bb = sum(p.numel() for p in self.encoder.parameters() if p.requires_grad)
        # neck 
        param_neck = 0
        for module in self.neck:
            param_neck += sum(p.numel() for p in module.parameters() if p.requires_grad)
        # head
        param_head = 0
        for module in self.head:
            param_head += sum(p.numel() for p in module.parameters() if p.requires_grad)

        param_net = param_bb + param_neck + param_head

                # ---- 파라미터 개수 출력 ----
        total_params = sum(p.numel() for p in self.parameters())
        total_trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)

        def _m(x):
            return x / 1e6  # million 단위

        print("========== SMPLer-X Parameter Counts ==========")
        print(f"Backbone (encoder, trainable) : {param_bb:,}  ({_m(param_bb):.3f} M)")
        print(f"Neck (box + hand_roi)        : {param_neck:,}  ({_m(param_neck):.3f} M)")
        print(f"Head (body/hand/face 등)     : {param_head:,}  ({_m(param_head):.3f} M)")
        print("----------------------------------------------")
        print(f"Total trainable              : {param_net:,}  ({_m(param_net):.3f} M)")
        print(f"Total (all params)           : {total_params:,}  ({_m(total_params):.3f} M)")
        print("==============================================")


        

        self.undist_w, self.undist_h = 1280, 1024
        fov_deg = 120
        fov_rad = np.deg2rad(fov_deg)

        x = np.linspace(-1, 1, self.undist_w)
        y = np.linspace(-1, 1, self.undist_h)
        xx, yy = np.meshgrid(x, y)
        z = 1.0 / np.tan(fov_rad / 2)
        rays = np.stack([xx, yy, np.full_like(xx, z)], axis=-1)
        rays /= np.linalg.norm(rays, axis=-1, keepdims=True)
        rays_flat = rays.reshape(-1, 3)

        mapped_pixels = []
        for r in rays_flat:
            try:
                uv = self.fisheye_camera.world2camera(np.array([r]))
                mapped_pixels.append(uv[0])
            except:
                mapped_pixels.append([-1, -1])

        self.mapped_pixels = np.array(mapped_pixels, dtype=np.float32).reshape(self.undist_h, self.undist_w, 2)

    def _get_subject_scene_dir(self, img_path):
        """
        img_path: .../{subject}/{scene}/imgs/filename 형태라고 가정하고
        마지막 'imgs' 기준으로 바로 앞의 두 디렉토리(subject, scene)를 리턴.
        실패하면 빈 문자열 반환.
        """
        if not isinstance(img_path, str):
            return ""
        path = img_path.replace("\\", "/")
        parts = path.split("/")

        # 마지막 'imgs' 위치 찾기
        idx = None
        for j in range(len(parts) - 1, -1, -1):
            if parts[j] == "imgs":
                idx = j
                break

        if idx is None or idx < 2:
            return ""

        subject = parts[idx - 2]
        scene = parts[idx - 1]
        return os.path.join(subject, scene)

        
    def _save_activation_heatmap(
        self,
        img_3chw: torch.Tensor,
        feat_3chw: torch.Tensor,
        save_path: str,
        alpha: float = 0.45,
        # 새 인자들: 원래 토큰 그리드와 잘려나간 칸 수
        orig_grid_hw=(16, 16),           # (H0, W0) = 원래 16x16
        crop=(0, 2, 0, 2),               # (top, right, bottom, left) = 좌우 2칸
        upsample: str = "nearest"        # 토큰 경계 또렷: 'nearest' 권장
    ):
        """
        img_3chw : [3,H,W]
        feat_3chw: [C,h,w]  (현재 16x12 토큰 그리드로 재배치된 feature)
        """
        import os
        import matplotlib.pyplot as plt
        import torch.nn.functional as F

        # 1) 채널 평균 → (h,w)
        hm = feat_3chw.detach().float().mean(0)  # [h,w]  (※ .values 쓰지 마세요)

        # 2) (선택) 간단 정규화
        hm = (hm - hm.min()) / (hm.max() - hm.min() + 1e-6)

        # 3) 잘려나간 칸만큼 패딩해서 원래 16x16 위치로 복원 + per-pixel alpha 마스크
        H0, W0 = orig_grid_hw         # ex) 16, 16
        t, r, b, l = crop             # ex) (0,2,0,2)
        cam_pad   = torch.zeros((H0, W0), dtype=hm.dtype, device=hm.device)
        alpha_pad = torch.zeros_like(cam_pad)  # 패딩 영역은 투명(0)
        cam_pad[t:H0-b, l:W0-r]   = hm
        alpha_pad[t:H0-b, l:W0-r] = 1.0

        # 4) 이미지 해상도로 업샘플 (히트맵/알파)
        H_img, W_img = img_3chw.shape[-2], img_3chw.shape[-1]
        cam_up = F.interpolate(cam_pad[None, None], size=(H_img, W_img),
                            mode=upsample, align_corners=False if upsample!="nearest" else None)[0, 0].cpu().numpy()
        a_up   = F.interpolate(alpha_pad[None, None], size=(H_img, W_img),
                            mode="nearest")[0, 0].cpu().numpy()  # 알파는 nearest 권장

        # 5) 원본 오버레이 (잘려나간 영역은 자동으로 투명)
        img = img_3chw.detach().cpu().float().clamp(0, 1).permute(1, 2, 0).numpy()
        plt.figure(figsize=(6, 4.5))
        plt.imshow(img)
        plt.imshow(cam_up, cmap="viridis", alpha=alpha * a_up)  # per-pixel alpha
        plt.axis("off"); plt.tight_layout()
        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
        plt.savefig(save_path, dpi=200, bbox_inches="tight")
        plt.close()

    def visualize_hand_bboxes_on_input(self, img_batch, lhand_bbox, rhand_bbox,
                                       save_dir, meta_info=None):
        """
        img_batch : [B, 3, H, W] (inputs['img_ori'])
        lhand_bbox, rhand_bbox : [B, 4] (xyxy, input_img_shape 기준)
        """
        os.makedirs(save_dir, exist_ok=True)

        # [B, H, W, C] 로 변환
        img_np = img_batch.detach().cpu().numpy().transpose(0, 2, 3, 1)

        B = img_np.shape[0]
        for i in range(B):
            img = img_np[i]

            # 이미지 값 범위에 따라 0~255 uint8 로 변환
            if img.max() <= 1.0:
                img = (img * 255.0).astype(np.uint8)
            else:
                img = img.astype(np.uint8)
            vis = img.copy()

            h, w = vis.shape[:2]

            def _draw_box(box, color, label_text=None):
                box = box.detach().cpu().numpy()
                x1, y1, x2, y2 = box

                x1 = int(np.clip(x1, 0, w - 1))
                x2 = int(np.clip(x2, 0, w - 1))
                y1 = int(np.clip(y1, 0, h - 1))
                y2 = int(np.clip(y2, 0, h - 1))

                cv2.rectangle(vis, (x1, y1), (x2, y2), color, 2)
                if label_text is not None:
                    cv2.putText(vis, label_text, (x1, max(0, y1 - 5)),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1,
                                lineType=cv2.LINE_AA)

            # 왼손: 초록, 오른손: 파랑
            _draw_box(lhand_bbox[i], (0, 255, 0), 'L-hand')
            _draw_box(rhand_bbox[i], (255, 0, 0), 'R-hand')

            # # 파일 이름 결정 (img_path 있으면 거기서 베이스 이름 사용)
            # if meta_info is not None and 'img_path' in meta_info:
            #     base = meta_info['img_path'][i]
            #     if not isinstance(base, str):
            #         base = f"{i:06d}.png"
            #     name = os.path.splitext(os.path.basename(base))[0]
            # else:
            #     name = f"{i:06d}"

            # out_path = os.path.join(save_dir, f"{name}_hand_bbox.png")
            # # OpenCV는 BGR 을 기대하므로 RGB -> BGR 로 뒤집어서 저장
            # cv2.imwrite(out_path, vis[:, :, ::-1])

            img_path_i = None
            sub_rel = ""
            if meta_info is not None and 'img_path' in meta_info:
                img_path_i = meta_info['img_path'][i]
                if isinstance(img_path_i, str):
                    sub_rel = self._get_subject_scene_dir(img_path_i)
                    name = os.path.splitext(os.path.basename(img_path_i))[0]
                else:
                    name = f"{i:06d}"
            else:
                name = f"{i:06d}"

            if sub_rel:
                cur_dir = os.path.join(save_dir, sub_rel)
            else:
                cur_dir = save_dir

            os.makedirs(cur_dir, exist_ok=True)

            out_path = os.path.join(cur_dir, f"{name}_hand_bbox.png")
            cv2.imwrite(out_path, vis[:, :, ::-1])


    def get_camera_trans(self, cam_param):
        # camera translation
        t_xy = cam_param[:, :2]
        gamma = torch.sigmoid(cam_param[:, 2])  # apply sigmoid to make it positive
        k_value = torch.FloatTensor([math.sqrt(cfg.focal[0] * cfg.focal[1] * cfg.camera_3d_size * cfg.camera_3d_size / (
                cfg.input_img_shape[0] * cfg.input_img_shape[1]))]).cuda().view(-1)
        t_z = k_value * gamma
        cam_trans = torch.cat((t_xy, t_z[:, None]), 1)
        return cam_trans


    def undistort_tensor_image(self, tensor_img, fisheye_camera: FishEyeCameraCalibrated):
        """
        tensor_img: [B, C, H, W] - torch.Tensor (range 0~1 or 0~255)
        return: undistorted image in torch.Tensor [B, C, H, W]
        """
        device = tensor_img.device
        B, C, H, W = tensor_img.shape
        tensor_img_np = tensor_img.detach().cpu().numpy().transpose(0, 2, 3, 1)  # [B, H, W, C]

        undistorted_np = []
        for i in range(B):
            img = (tensor_img_np[i] * 255).astype(np.uint8) if tensor_img_np[i].max() <= 1 else tensor_img_np[i].astype(np.uint8)

            # grid 생성
            x = np.linspace(-1, 1, W)
            y = np.linspace(-1, 1, H)
            xx, yy = np.meshgrid(x, y)
            z = 1.0  # FOV에 해당하는 focal length 기준
            rays = np.stack([xx, yy, np.full_like(xx, z)], axis=-1)
            rays /= np.linalg.norm(rays, axis=-1, keepdims=True)
            rays_flat = rays.reshape(-1, 3)

            mapped_pixels = []
            for r in rays_flat:
                try:
                    uv = self.fisheye_camera.world2camera(np.array([r]))  # (1, 2)
                    mapped_pixels.append(uv[0])
                except:
                    mapped_pixels.append([-1, -1])
            mapped_pixels = np.array(mapped_pixels, dtype=np.float32).reshape(H, W, 2)

            undistorted_img = cv2.remap(
                img,
                mapped_pixels[:, :, 0],
                mapped_pixels[:, :, 1],
                interpolation=cv2.INTER_LINEAR,
                borderMode=cv2.BORDER_CONSTANT,
                borderValue=0
            )
            undistorted_np.append(undistorted_img)

        undistorted_np = np.stack(undistorted_np).astype(np.float32) / 255.0
        undistorted_tensor = torch.from_numpy(undistorted_np).permute(0, 3, 1, 2).to(device)
        return undistorted_tensor
    



    def visualize_body_img_comparison(self, original, undistorted):
        """
        original, undistorted: [B, C, H, W] torch.Tensor
        """
        import matplotlib.pyplot as plt

        B = original.shape[0]
        for i in range(B):
            orig_np = original[i].detach().cpu().permute(1, 2, 0).numpy()
            undist_np = undistorted[i].detach().cpu().permute(1, 2, 0).numpy()

            plt.figure(figsize=(10, 5))
            plt.subplot(1, 2, 1)
            plt.imshow(np.clip(orig_np, 0, 1))
            plt.title('Original (Fisheye)')
            plt.axis('off')

            plt.subplot(1, 2, 2)
            plt.imshow(np.clip(undist_np, 0, 1))
            plt.title('Undistorted')
            plt.axis('off')

            plt.tight_layout()
            plt.show()

    def get_coord(self, root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose, shape, expr, cam_trans, mode):
        batch_size = root_pose.shape[0]
        # set_trace()
        zero_pose = torch.zeros((1, 3)).float().cuda().repeat(batch_size, 1)  # eye poses
        
        # if getattr(cfg, 'use_token_decoder', False):
        #     pad = torch.zeros(body_pose.size(0), 6, dtype=body_pose.dtype, device=body_pose.device) # (B, 2, 3)
        #     body_pose = torch.cat((body_pose, pad), 1)
        #     output = self.smpl_layer(betas=shape, body_pose=body_pose, global_orient=root_pose)
        
        # else:
        output = self.smplx_layer(betas=shape, body_pose=body_pose, global_orient=root_pose, right_hand_pose=rhand_pose,
                            left_hand_pose=lhand_pose, jaw_pose=jaw_pose, leye_pose=zero_pose,
                            reye_pose=zero_pose, expression=expr)

        mesh_cam = output.vertices + cam_trans[:, None, :]  # ← 최종 mesh_cam은 SMPL 기준

        # smplx decoder 사용
        
        # camera-centered 3D coordinate

        mesh_cam = output.vertices
        if mode == 'test' and cfg.testset == 'AGORA':  # use 144 joints for AGORA evaluation
            joint_cam = output.joints
        else:
            joint_cam = output.joints[:, smpl_x.joint_idx, :]
            # set_trace()
            


        B, J, _ = joint_cam.shape
        # set_trace()
        joint_cam_translation =joint_cam.detach() + cam_trans[:, None, :]

        joint_cam_flat = (joint_cam_translation).view(-1, 3)  # (B*J, 3)

        # (B*J, 3) → (B*J, 2): fisheye projection (image 좌표계)
        joint_proj_flat = self.fisheye_camera.world2camera_pytorch(joint_cam_flat)  # (B*J, 2)

        # (B*J, 2) → (B, J, 2)
        joint_proj = joint_proj_flat.view(B, J, 2)
        
        # root-relative 3D coordinates
        root_cam = joint_cam[:, smpl_x.root_joint_idx, None, :]
        joint_cam = joint_cam - root_cam
        
        mesh_cam = mesh_cam + cam_trans[:, None, :]  # for rendering
        joint_cam_wo_ra = joint_cam.clone()
        
        # 2. heatmap 좌표계로 변환
        joint_proj[..., 0] = joint_proj[..., 0] / cfg.input_img_shape[1] * cfg.output_hm_shape[2] # x: w 기준
        joint_proj[..., 1] = joint_proj[..., 1] / cfg.input_img_shape[0] * cfg.output_hm_shape[1] # y: h 기준
        
   
        # left hand root (left wrist)-relative 3D coordinatese
        lhand_idx = smpl_x.joint_part['lhand']
        lhand_cam = joint_cam[:, lhand_idx, :]
        lwrist_cam = joint_cam[:, smpl_x.lwrist_idx, None, :]
        lhand_cam = lhand_cam - lwrist_cam
        joint_cam = torch.cat((joint_cam[:, :lhand_idx[0], :], lhand_cam, joint_cam[:, lhand_idx[-1] + 1:, :]), 1)

        # right hand root (right wrist)-relative 3D coordinatese
        rhand_idx = smpl_x.joint_part['rhand']
        rhand_cam = joint_cam[:, rhand_idx, :]
        rwrist_cam = joint_cam[:, smpl_x.rwrist_idx, None, :]
        rhand_cam = rhand_cam - rwrist_cam
        joint_cam = torch.cat((joint_cam[:, :rhand_idx[0], :], rhand_cam, joint_cam[:, rhand_idx[-1] + 1:, :]), 1)

        # face root (neck)-relative 3D coordinates
        face_idx = smpl_x.joint_part['face']
        face_cam = joint_cam[:, face_idx, :]
        neck_cam = joint_cam[:, smpl_x.neck_idx, None, :]
        face_cam = face_cam - neck_cam
        joint_cam = torch.cat((joint_cam[:, :face_idx[0], :], face_cam, joint_cam[:, face_idx[-1] + 1:, :]), 1)


        return joint_proj, joint_cam, joint_cam_wo_ra, mesh_cam

    def generate_mesh_gt(self, targets, mode):
        if 'smplx_mesh_cam' in targets:
            return targets['smplx_mesh_cam']
        nums = [3, 63, 45, 45, 3]
        accu = []
        temp = 0
        for num in nums:
            temp += num
            accu.append(temp)
        pose = targets['smplx_pose']
        root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose = \
            pose[:, :accu[0]], pose[:, accu[0]:accu[1]], pose[:, accu[1]:accu[2]], pose[:, accu[2]:accu[3]], pose[:,
                                                                                                             accu[3]:
                                                                                                             accu[4]]
        # print(lhand_pose)
        shape = targets['smplx_shape']
        expr = targets['smplx_expr']
        cam_trans = targets['smplx_cam_trans']

        # final output
        joint_proj, joint_cam, joint_cam_wo_ra, mesh_cam = self.get_coord(root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose, shape,
                                                         expr, cam_trans, mode)

        return mesh_cam

    def bbox_split(self, bbox):
        # bbox:[bs, 3, 3]
        lhand_bbox_center, rhand_bbox_center, face_bbox_center = \
            bbox[:, 0, :2], bbox[:, 1, :2], bbox[:, 2, :2]
        return lhand_bbox_center, rhand_bbox_center, face_bbox_center
    
    def build_forward_map(self,mapped_pixels):
        """
        mapped_pixels: (H, W, 2) — dst(i,j) 위치에서 원본 좌표를 가져오는 역방향 맵
        반환: KDTree와 보정 이미지 좌표 (dst_x, dst_y)
        """
        H, W = mapped_pixels.shape[:2]
        src_points = mapped_pixels.reshape(-1, 2)  # (H*W, 2): 원본 좌표계 기준
        dst_coords = np.stack(np.meshgrid(np.arange(W), np.arange(H)), axis=-1).reshape(-1, 2)  # (H*W, 2)

        # 유효한 원본 좌표만 사용
        valid_mask = np.all((src_points >= 0) & (src_points < np.array([W, H])), axis=1)
        src_points = src_points[valid_mask]
        dst_coords = dst_coords[valid_mask]

        tree = cKDTree(src_points)
        return tree, dst_coords

    def remap_joints_forward(self, joint_orig, tree, dst_coords):
        """
        joint_orig: (N, 2) — 원본 이미지 기준 joint
        tree: KDTree of mapped_pixels (src → dst)
        dst_coords: 보정된 이미지 좌표

        반환: 보정된 joint 위치 (N, 2)
        """
        dist, idx = tree.query(joint_orig, k=1)
        joint_remapped = dst_coords[idx]
        return joint_remapped.astype(np.float32)
    
   
    def draw_joint_lines(self, img_np, joints_2d, joint_names, chains, color=(0,255,0)):
        """
        img_np: (H,W,3) numpy array (uint8)
        joints_2d: (N,2) float (xy)
        joint_names: 이름 리스트
        chains: 각 limb의 joint 인덱스 chain(순서대로 연결)
        color: 선 색
        """
        # joint 찍기
        for x, y in joints_2d:
            x_int, y_int = int(round(x)), int(round(y))
            if 0 <= x_int < img_np.shape[1] and 0 <= y_int < img_np.shape[0]:
                cv2.circle(img_np, (x_int, y_int), 3, color, -1)

        # 선 그리기
        for chain in chains:
            for i in range(len(chain)-1):
                idx1, idx2 = chain[i], chain[i+1]
                x1, y1 = joints_2d[idx1]
                x2, y2 = joints_2d[idx2]
                pt1 = (int(round(x1)), int(round(y1)))
                pt2 = (int(round(x2)), int(round(y2)))
                if all(0 <= pt < img_np.shape[1] for pt in pt1) and all(0 <= pt < img_np.shape[0] for pt in pt1) and \
                all(0 <= pt < img_np.shape[1] for pt in pt2) and all(0 <= pt < img_np.shape[0] for pt in pt2):
                    cv2.line(img_np, pt1, pt2, (255, 128, 0), 2)  # 선 색은 예시 (주황)

        return img_np
    
    def visualize_mesh_on_image(self, img_np, mesh_cam, color=(0, 255, 0)):
        """
        img_np: (H, W, 3) numpy array (uint8)
        mesh_cam: (N, 3) torch.Tensor - 카메라 좌표계의 mesh vertices
        """
        # numpy로 변환
        mesh_cam_np = mesh_cam.detach().cpu().numpy()

        # 3D mesh를 fisheye projection을 통해 2D로 투영
        mesh_proj_2d = self.fisheye_camera.world2camera_pytorch(
            torch.from_numpy(mesh_cam_np).float().cuda()
        ).detach().cpu().numpy()  # (N, 2)

        # 각 vertex를 이미지 위에 표시
        for x, y in mesh_proj_2d:
            x_int, y_int = int(round(x)), int(round(y))
            if 0 <= x_int < img_np.shape[1] and 0 <= y_int < img_np.shape[0]:
                cv2.circle(img_np, (x_int, y_int), 1, color, -1)

        return img_np

    def forward(self, inputs, targets, meta_info, mode):
        # body_img = F.interpolate(inputs['img_ori'], cfg.input_body_shape)
        # body_img = inputs['img']
        #     # 0. Fisheye 보정 수행
        

        # # 0.5 시각화 (선택)
        # self.visualize_body_img_comparison(body_img, body_img_undistorted)
        body_img = inputs['img_ori']
        # set_trace()
        # 1. Encoder
        ## img_feat: [bs, 1280, 16, 12], task_tokens: [bs, 31, 1280]
        img_feat, task_tokens = self.encoder(body_img)  # task_token:[bs, N, c]
        shape_token, cam_token, expr_token, jaw_pose_token, hand_token, body_pose_token = \
            task_tokens[:, 0], task_tokens[:, 1], task_tokens[:, 2], task_tokens[:, 3], task_tokens[:, 4:6], task_tokens[:, 6:]


        # 1-1. ID Regressor

        # person_id = self.id_regressor(id_token)
        
        # set_trace()
        # 2. Body Regressor
        if not getattr(cfg, 'use_token_decoder', False):
            body_joint_hm, body_joint_img = self.body_position_net(img_feat)
            root_pose, body_pose, shape, cam_param, = self.body_regressor(body_pose_token, shape_token, cam_token, body_joint_img.detach())
            

        else:
            body_joint_hm, body_joint_img = self.body_position_net(img_feat)
            # set_trace()
            smpl_param = self.token_decoder(img_feat)
            # set_trace()
            root_pose = smpl_param['root_pose']
            body_pose = smpl_param['body_pose']
            shape = smpl_param['betas']
            cam_param = smpl_param['cam_param']


        # set_trace()
        root_pose = rot6d_to_axis_angle(root_pose)
        body_pose = rot6d_to_axis_angle(body_pose.reshape(-1, 6)).reshape(body_pose.shape[0], -1)  # (N, J_R*3)
        # set_trace()
    
        # 3. Hand and Face BBox Estimation
        lhand_bbox_center, lhand_bbox_size, rhand_bbox_center, rhand_bbox_size, face_bbox_center, face_bbox_size = self.box_net(img_feat, body_joint_hm.detach())
        lhand_bbox = restore_bbox(lhand_bbox_center, lhand_bbox_size, cfg.input_hand_shape[1] / cfg.input_hand_shape[0], 2.0).detach()  # xyxy in (cfg.input_body_shape[1], cfg.input_body_shape[0]) space
        rhand_bbox = restore_bbox(rhand_bbox_center, rhand_bbox_size, cfg.input_hand_shape[1] / cfg.input_hand_shape[0], 2.0).detach()  # xyxy in (cfg.input_body_shape[1], cfg.input_body_shape[0]) space
        face_bbox = restore_bbox(face_bbox_center, face_bbox_size, cfg.input_face_shape[1] / cfg.input_face_shape[0], 1.5).detach()  # xyxy in (cfg.input_body_shape[1], cfg.input_body_shape[0]) space

        # 4. Differentiable Feature-level Hand Crop-Upsample
        # hand_feat: list, [bsx2, c, cfg.output_hm_shape[1]*scale, cfg.output_hm_shape[2]*scale]
        hand_feat = self.hand_roi_net(img_feat, lhand_bbox, rhand_bbox)  # hand_feat: flipped left hand + right hand

        # 5. Hand/Face Regressor
        # hand regressor
        _, hand_joint_img = self.hand_position_net(hand_feat)  # (2N, J_P, 3)
        hand_pose = self.hand_regressor(hand_feat, hand_joint_img.detach())
        hand_pose = rot6d_to_axis_angle(hand_pose.reshape(-1, 6)).reshape(hand_feat.shape[0], -1)  # (2N, J_R*3)

        # restore flipped left hand joint coordinates
        batch_size = hand_joint_img.shape[0] // 2
        lhand_joint_img = hand_joint_img[:batch_size, :, :]
        lhand_joint_img = torch.cat((cfg.output_hand_hm_shape[2] - 1 - lhand_joint_img[:, :, 0:1], lhand_joint_img[:, :, 1:]), 2)
        rhand_joint_img = hand_joint_img[batch_size:, :, :]
        # restore flipped left hand joint rotations
        batch_size = hand_pose.shape[0] // 2
        lhand_pose = hand_pose[:batch_size, :].reshape(-1, len(smpl_x.orig_joint_part['lhand']), 3)
        lhand_pose = torch.cat((lhand_pose[:, :, 0:1], -lhand_pose[:, :, 1:3]), 2).view(batch_size, -1)
        rhand_pose = hand_pose[batch_size:, :]

        

        # # 5. Hand regressor without ROI
        # hand_joint_hm, hand_joint_coord_global = self.hand_position_net(img_feat)  # (B, 2J_all, 3)

        # B = img_feat.size(0)
        # J_all = self.hand_joint_num                      # 20 (keypoints/joints)
        # J_p = len(smpl_x.orig_joint_part['lhand'])       # 보통 15 (pose parameter joints)

        # # split left/right (each is (B, J_all, 3))
        # lhand_joint_img = hand_joint_coord_global[:, :J_all, :]
        # rhand_joint_img = hand_joint_coord_global[:, J_all:, :]

        # # --- 중요: pose에 쓰는 joint만 선택 ---
        # # 대부분 ordering이 [wrist(0), finger 15개(1~15), tips 4개(16~19)] 라서
        # # wrist/tips 제외하고 finger 15개만 쓰면 됨.
        # lhand_joint_img_pose = lhand_joint_img[:, 1:1+J_p, :]   # (B, J_p, 3)
        # rhand_joint_img_pose = rhand_joint_img[:, 1:1+J_p, :]   # (B, J_p, 3)

        # hand_joint_coord = torch.stack([lhand_joint_img_pose, rhand_joint_img_pose], dim=1)  # (B, 2, J_p, 3)

        # token+joint -> 6D  (B, 2, J_p, 6)
        # hand_pose_6d = self.hand_regressor(hand_token, hand_joint_coord.detach())

        # 6D -> axis-angle  (B, 2, J_p, 3)
        # hand_pose_aa = rot6d_to_axis_angle(hand_pose_6d.reshape(-1, 6)).view(B, 2, J_p, 3)

        # (B, J_p*3) == (B, 45) 로 맞춰짐
        # lhand_pose = hand_pose_aa[:, 0].reshape(B, -1)
        # rhand_pose = hand_pose_aa[:, 1].reshape(B, -1)


        # face regressor
        expr, jaw_pose = self.face_regressor(expr_token, jaw_pose_token)
        jaw_pose = rot6d_to_axis_angle(jaw_pose)

        # final output
        joint_proj, joint_cam, joint_cam_wo_ra, mesh_cam = self.get_coord(root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose, shape, expr, cam_param, mode)
       
    
        pose = torch.cat((root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose), 1)
        # set_trace()


        # set_trace()
        B      = body_pose.size(0)
        device = body_pose.device
        dtype  = body_pose.dtype
        expr_zeros = torch.zeros(B, 100, device=device, dtype=dtype) 

        pose_experssion = torch.cat((body_pose, lhand_pose, rhand_pose, jaw_pose, expr_zeros), 1)
        joint_img = torch.cat((body_joint_img, lhand_joint_img, rhand_joint_img), 1)

        # if mode == 'test' and 'smplx_pose' in targets:
        #     mesh_pseudo_gt = self.generate_mesh_gt(targets, mode)

        

        if mode == 'train':
            loss = {}

            if mode == 'train' and getattr(self, 'id_head', None) is not None:
                id_loss_weight = getattr(cfg, 'id_loss_weight', 1.0)
                # img_feat: [B, C, H, W]  -> GAP로 [B, C]
                g = F.adaptive_avg_pool2d(img_feat, 1).flatten(1)


                # GRL 통과(부호 반전 + 스케일링)
                # g_rev = GradReverse.apply(g, self.id_adv_lambda)

                # set_trace()
                id_logits = self.id_head(g)   # [B, num_subjects]

                # 타깃: 정수 인덱스 형태 권장 (DataLoader에서 meta_info['id_idx']로 전달)
                id_idx = meta_info['id_idx'].long()                 # [B]
                id_target = F.one_hot(id_idx, num_classes=self.num_subjects).float()  # [B, K]
                loss['id_adv'] = self.id_loss_fn(id_logits, id_target) * id_loss_weight

            # loss functions
            

            smplx_kps_3d_weight = getattr(cfg, 'smplx_kps_3d_weight', 1.0)
            smplx_kps_3d_weight = getattr(cfg, 'smplx_kps_weight', smplx_kps_3d_weight) # old config

            smplx_kps_2d_weight = getattr(cfg, 'smplx_kps_2d_weight', 1.0)
            net_kps_2d_weight = getattr(cfg, 'net_kps_2d_weight', 1.0)

            smplx_pose_weight = getattr(cfg, 'smplx_pose_weight', 1.0)
            smplx_shape_weight = getattr(cfg, 'smplx_loss_weight', 1.0)

            dposer_x_weight = getattr(cfg, 'dposer_x_weight', 1.0)
            # smplx_orient_weight = getattr(cfg, 'smplx_orient_weight', smplx_pose_weight) # if not specified, use the same weight as pose
    

            # do not supervise root pose if original agora json is used


            ###  SMPL parameter regression ###
            if getattr(cfg, 'agora_fix_global_orient_transl', False):
                # loss['smplx_pose'] = self.param_loss(pose, targets['smplx_pose'], meta_info['smplx_pose_valid'])[:, 3:] * smplx_pose_weight
                if hasattr(cfg, 'smplx_orient_weight'):
                    smplx_orient_weight = getattr(cfg, 'smplx_orient_weight')
                    loss['smplx_orient'] = self.param_loss(pose, targets['smplx_pose'], meta_info['smplx_pose_valid'])[:, :3] * smplx_orient_weight

                loss['smplx_pose'] = self.param_loss(pose, targets['smplx_pose'], meta_info['smplx_pose_valid']) * smplx_pose_weight
                # loss['smplx_lhand_pose'] = self.param_loss(lhand_pose, targets['smplx_lhand_pose'], meta_info['smplx_pose_valid'][:, 70:116]) * smplx_pose_weight
                # loss['smplx_lhand_pose'] = self.param_loss(lhand_pose, targets['smplx_lhand_pose'], meta_info['smplx_pose_valid'][:, 116:162]) * smplx_pose_weight

            else:
                loss['smplx_pose'] = self.param_loss(pose, targets['smplx_pose'], meta_info['smplx_pose_valid']) * smplx_pose_weight
                # loss['smplx_lhand_pose'] = self.param_loss(lhand_pose, targets['smplx_lhand_pose'], meta_info['smplx_pose_valid'][:, 70:116]) * smplx_pose_weight
                # loss['smplx_lhand_pose'] = self.param_loss(lhand_pose, targets['smplx_lhand_pose'], meta_info['smplx_pose_valid'][:, 116:162]) * smplx_pose_weight


            loss['smplx_shape'] = self.param_loss(shape, targets['smplx_shape'],
                                                  meta_info['smplx_shape_valid'][:, None]) * smplx_shape_weight 
            loss['smplx_expr'] = self.param_loss(expr, targets['smplx_expr'], meta_info['smplx_expr_valid'][:, None])

            loss['dposerx'] = self.dposer_x(pose_experssion) * dposer_x_weight
            # supervision for keypoints3d wo/ ra
            # loss['joint_cam'] = self.coord_loss(joint_cam_wo_ra, targets['joint_cam'], meta_info['joint_valid'] * meta_info['is_3D'][:, None, None]) * smplx_kps_3d_weight
            # supervision for keypoints3d w/ ra


            loss['smplx_joint_cam'] = self.coord_loss(joint_cam, targets['smplx_joint_cam'], meta_info['smplx_joint_valid']) * smplx_kps_3d_weight


            # loss['bone_length'] = self.compute_bone_length_loss(
            #     joint_cam[:, [mo2cap2_to_smplx[name] for name in mo2cap2_joint_names], :],
            #     targets['smplx_joint_cam'][:, [mo2cap2_to_smplx[name] for name in mo2cap2_joint_names], :],
            #     bone_pairs,
            #     weight=bone_loss_weight
            # )
            if not (meta_info['lhand_bbox_valid'] == 0).all():
                loss['lhand_bbox'] = (self.coord_loss(lhand_bbox_center, targets['lhand_bbox_center'], meta_info['lhand_bbox_valid'][:, None]) +
                                    self.coord_loss(lhand_bbox_size, targets['lhand_bbox_size'], meta_info['lhand_bbox_valid'][:, None]))
            if not (meta_info['rhand_bbox_valid'] == 0).all():
                loss['rhand_bbox'] = (self.coord_loss(rhand_bbox_center, targets['rhand_bbox_center'], meta_info['rhand_bbox_valid'][:, None]) +
                                    self.coord_loss(rhand_bbox_size, targets['rhand_bbox_size'], meta_info['rhand_bbox_valid'][:, None]))
            if not (meta_info['face_bbox_valid'] == 0).all():
                loss['face_bbox'] = (self.coord_loss(face_bbox_center, targets['face_bbox_center'], meta_info['face_bbox_valid'][:, None]) +
                                 self.coord_loss(face_bbox_size, targets['face_bbox_size'], meta_info['face_bbox_valid'][:, None]))
            
            # if (meta_info['face_bbox_valid'] == 0).all():
            #     out = {}
            targets['original_joint_img'] = targets['joint_img'].clone()
            targets['original_smplx_joint_img'] = targets['smplx_joint_img'].clone()
            # out['original_joint_proj'] = joint_proj.clone()
            if not (meta_info['lhand_bbox_valid'] + meta_info['rhand_bbox_valid'] == 0).all():

                # change hand target joint_img and joint_trunc according to hand bbox (cfg.output_hm_shape -> downsampled hand bbox space)
                for part_name, bbox in (('lhand', lhand_bbox), ('rhand', rhand_bbox)):
                    for coord_name, trunc_name in (('joint_img', 'joint_trunc'), ('smplx_joint_img', 'smplx_joint_trunc')):
                        x = targets[coord_name][:, smpl_x.joint_part[part_name], 0]
                        y = targets[coord_name][:, smpl_x.joint_part[part_name], 1]
                        z = targets[coord_name][:, smpl_x.joint_part[part_name], 2]
                        trunc = meta_info[trunc_name][:, smpl_x.joint_part[part_name], 0]

                        x -= (bbox[:, None, 0] / cfg.input_body_shape[1] * cfg.output_hm_shape[2])
                        x *= (cfg.output_hand_hm_shape[2] / (
                                (bbox[:, None, 2] - bbox[:, None, 0]) / cfg.input_body_shape[1] * cfg.output_hm_shape[
                            2]))
                        y -= (bbox[:, None, 1] / cfg.input_body_shape[0] * cfg.output_hm_shape[1])
                        y *= (cfg.output_hand_hm_shape[1] / (
                                (bbox[:, None, 3] - bbox[:, None, 1]) / cfg.input_body_shape[0] * cfg.output_hm_shape[
                            1]))
                        z *= cfg.output_hand_hm_shape[0] / cfg.output_hm_shape[0]
                        trunc *= ((x >= 0) * (x < cfg.output_hand_hm_shape[2]) * (y >= 0) * (
                                y < cfg.output_hand_hm_shape[1]))

                        coord = torch.stack((x, y, z), 2)
                        trunc = trunc[:, :, None]
                        targets[coord_name] = torch.cat((targets[coord_name][:, :smpl_x.joint_part[part_name][0], :], coord,
                                                        targets[coord_name][:, smpl_x.joint_part[part_name][-1] + 1:, :]),
                                                        1)
                        meta_info[trunc_name] = torch.cat((meta_info[trunc_name][:, :smpl_x.joint_part[part_name][0], :],
                                                        trunc,
                                                        meta_info[trunc_name][:, smpl_x.joint_part[part_name][-1] + 1:,
                                                        :]), 1)

                # change hand projected joint coordinates according to hand bbox (cfg.output_hm_shape -> hand bbox space)
                for part_name, bbox in (('lhand', lhand_bbox), ('rhand', rhand_bbox)):
                    x = joint_proj[:, smpl_x.joint_part[part_name], 0]
                    y = joint_proj[:, smpl_x.joint_part[part_name], 1]

                    x -= (bbox[:, None, 0] / cfg.input_body_shape[1] * cfg.output_hm_shape[2])
                    x *= (cfg.output_hand_hm_shape[2] / (
                            (bbox[:, None, 2] - bbox[:, None, 0]) / cfg.input_body_shape[1] * cfg.output_hm_shape[2]))
                    y -= (bbox[:, None, 1] / cfg.input_body_shape[0] * cfg.output_hm_shape[1])
                    y *= (cfg.output_hand_hm_shape[1] / (
                            (bbox[:, None, 3] - bbox[:, None, 1]) / cfg.input_body_shape[0] * cfg.output_hm_shape[1]))

                    coord = torch.stack((x, y), 2)
                    trans = []
                    for bid in range(coord.shape[0]):
                        mask = meta_info['joint_trunc'][bid, smpl_x.joint_part[part_name], 0] == 1
                        if torch.sum(mask) == 0:
                            trans.append(torch.zeros((2)).float().cuda())
                        else:
                            trans.append((-coord[bid, mask, :2] + targets['joint_img'][:, smpl_x.joint_part[part_name], :][
                                                                bid, mask, :2]).mean(0))
                    trans = torch.stack(trans)[:, None, :]
                    coord = coord + trans  # global translation alignment
                    joint_proj = torch.cat((joint_proj[:, :smpl_x.joint_part[part_name][0], :], coord,
                                            joint_proj[:, smpl_x.joint_part[part_name][-1] + 1:, :]), 1)

            if not (meta_info['face_bbox_valid'] == 0).all():
                # change face projected joint coordinates according to face bbox (cfg.output_hm_shape -> face bbox space)
                coord = joint_proj[:, smpl_x.joint_part['face'], :]
                trans = []
                for bid in range(coord.shape[0]):
                    mask = meta_info['joint_trunc'][bid, smpl_x.joint_part['face'], 0] == 1
                    if torch.sum(mask) == 0:
                        trans.append(torch.zeros((2)).float().cuda())
                    else:
                        trans.append((-coord[bid, mask, :2] + targets['joint_img'][:, smpl_x.joint_part['face'], :][bid,
                                                            mask, :2]).mean(0))
                trans = torch.stack(trans)[:, None, :]
                coord = coord + trans  # global translation alignment
                joint_proj = torch.cat((joint_proj[:, :smpl_x.joint_part['face'][0], :], coord,
                                        joint_proj[:, smpl_x.joint_part['face'][-1] + 1:, :]), 1)
           # print("joint_proj:", joint_proj.shape)
           # print("coord_gt  :", targets['joint_img'][:, :, :2].shape)
           # print("valid     :", meta_info['joint_trunc'].shape)
            # set_trace()


            loss['joint_proj'] = self.coord_loss(joint_proj, targets['smplx_joint_img'][:, :, :2], meta_info['smplx_joint_trunc']) * smplx_kps_2d_weight
            


            # loss['joint_img'] = self.coord_loss(joint_img, smpl_x.reduce_joint_set(targets['joint_img']),
            #                                     smpl_x.reduce_joint_set(meta_info['joint_trunc']), meta_info['is_3D']) * net_kps_2d_weight
            # set_trace()

            # PositionNet
            loss['smplx_joint_img'] = self.coord_loss(joint_img, smpl_x.reduce_joint_set(targets['smplx_joint_img']),
                                                     smpl_x.reduce_joint_set(meta_info['smplx_joint_trunc'])) * net_kps_2d_weight

            return loss
        else:

            save_dir = os.path.join(cfg.result_dir, 'vis_test_result')
            os.makedirs(save_dir, exist_ok=True)
            B = body_img.shape[0]

            if getattr(cfg, 'vis_feature', False):
                base_feat_dir = os.path.join(cfg.result_dir, 'vis_test_result', 'feature')
                B = body_img.shape[0]

                for i in range(B):
                    img_path_i = None
                    if 'img_path' in meta_info:
                        img_path_i = meta_info['img_path'][i]

                    # dataset/.../{subject}/{scene}/imgs/... 에서 subject/scene 뽑기
                    sub_rel = self._get_subject_scene_dir(img_path_i) if isinstance(img_path_i, str) else ""

                    if sub_rel:
                        feat_dir = os.path.join(base_feat_dir, sub_rel)
                    else:
                        feat_dir = base_feat_dir

                    os.makedirs(feat_dir, exist_ok=True)

                    if isinstance(img_path_i, str):
                        base = os.path.basename(img_path_i)
                        name = os.path.splitext(base)[0]
                    else:
                        name = f"{i:06d}"

                    save_path = os.path.join(feat_dir, f"{name}_feat.png")
                    # body_img: [B,3,H,W], img_feat: [B,C,h,w]
                    self._save_activation_heatmap(body_img[i], img_feat[i], save_path)


            # change hand output joint_img according to hand bbox
            for part_name, bbox in (('lhand', lhand_bbox), ('rhand', rhand_bbox)):
                joint_img[:, smpl_x.pos_joint_part[part_name], 0] *= (
                        ((bbox[:, None, 2] - bbox[:, None, 0]) / cfg.input_body_shape[1] * cfg.output_hm_shape[2]) /
                        cfg.output_hand_hm_shape[2])
                joint_img[:, smpl_x.pos_joint_part[part_name], 0] += (
                        bbox[:, None, 0] / cfg.input_body_shape[1] * cfg.output_hm_shape[2])
                joint_img[:, smpl_x.pos_joint_part[part_name], 1] *= (
                        ((bbox[:, None, 3] - bbox[:, None, 1]) / cfg.input_body_shape[0] * cfg.output_hm_shape[1]) /
                        cfg.output_hand_hm_shape[1])
                joint_img[:, smpl_x.pos_joint_part[part_name], 1] += (
                        bbox[:, None, 1] / cfg.input_body_shape[0] * cfg.output_hm_shape[1])

            # change input_body_shape to input_img_shape
            for bbox in (lhand_bbox, rhand_bbox, face_bbox):
                bbox[:, 0] *= cfg.input_img_shape[1] / cfg.input_body_shape[1]
                bbox[:, 1] *= cfg.input_img_shape[0] / cfg.input_body_shape[0]
                bbox[:, 2] *= cfg.input_img_shape[1] / cfg.input_body_shape[1]
                bbox[:, 3] *= cfg.input_img_shape[0] / cfg.input_body_shape[0]


            if getattr(cfg, 'vis_hand_bbox', False):
                hand_bbox_dir = os.path.join(save_dir, 'hand_bbox')
                os.makedirs(hand_bbox_dir, exist_ok=True)
                self.visualize_hand_bboxes_on_input(
                    body_img, lhand_bbox, rhand_bbox,
                    hand_bbox_dir, meta_info
                )

           #  for i in range(B):
                # set_trace()
                # img_np = inputs['img_vis'][i].detach().cpu().numpy()  # [H, W, C]
                # img_np = np.clip(img_np, 0, 1) * 255.0
                # img_np = img_np.astype(np.uint8).copy()

                # selected_joint_indices = list(mo2cap2_to_smplx.values()) 
                # joints_2d = joint_proj[i][selected_joint_indices].detach().cpu().numpy()  # [N_valid, 2]
                # for x, y in joints_2d:
                #     x_int, y_int = int(round(x)), int(round(y))
                #     if 0 <= x_int < img_np.shape[1] and 0 <= y_int < img_np.shape[0]:
                #         cv2.circle(img_np, (x_int, y_int), 3, (0, 255, 0), -1)

                # joint 시각화
                # img_np = self.draw_joint_lines(img_np, joints_2d, mo2cap2_joint_names, mo2cap2_chain)

                # mesh 시각화
                # img_np = self.visualize_mesh_on_image(img_np, mesh_cam[i])


                # save_path = os.path.join(save_dir, f'proj_{i:03d}.png')
                # cv2.imwrite(save_path, img_np[:, :, ::-1])

            # test output
            out = {}
            out['img'] = inputs['img_ori']
            # out['img_vis'] = inputs['img_vis']
            out['joint_img'] = joint_img

            out['smplx_joint_proj'] = joint_proj
            out['smplx_mesh_cam'] = mesh_cam
            out['smplx_root_pose'] = root_pose
            out['smplx_body_pose'] = body_pose
            out['smplx_lhand_pose'] = lhand_pose
            out['smplx_rhand_pose'] = rhand_pose
            out['smplx_jaw_pose'] = jaw_pose
            out['smplx_shape'] = shape
            out['smplx_expr'] = expr
            # out['cam_trans'] = cam_trans
            # out['lhand_bbox'] = lhand_bbox
            # out['rhand_bbox'] = rhand_bbox
            # out['face_bbox'] = face_bbox
            if 'smplx_shape' in targets:
                out['smplx_shape_target'] = targets['smplx_shape']
            if 'img_path' in meta_info:
                out['img_path'] = meta_info['img_path']
            if 'joint_cam' in meta_info:
                out['gt_joint'] = meta_info['joint_cam']
            # if 'smplx_pose' in targets:
            #     out['smplx_mesh_cam_pseudo_gt'] = mesh_pseudo_gt
            if 'smplx_mesh_cam' in targets:
                out['smplx_mesh_cam_target'] = targets['smplx_mesh_cam']
            if 'smpl_mesh_cam' in targets:
                out['smpl_mesh_cam_target'] = targets['smpl_mesh_cam']
            if 'bb2img_trans' in meta_info:
                out['bb2img_trans'] = meta_info['bb2img_trans']
            if 'gt_smplx_transl' in meta_info:
                out['gt_smplx_transl'] = meta_info['gt_smplx_transl']

            return out

def init_weights(m):
    try:
        if type(m) == nn.ConvTranspose2d:
            nn.init.normal_(m.weight, std=0.001)
        elif type(m) == nn.Conv2d:
            nn.init.normal_(m.weight, std=0.001)
            nn.init.constant_(m.bias, 0)
        elif type(m) == nn.BatchNorm2d:
            nn.init.constant_(m.weight, 1)
            nn.init.constant_(m.bias, 0)
        elif type(m) == nn.Linear:
            nn.init.normal_(m.weight, std=0.01)
            nn.init.constant_(m.bias, 0)
    except AttributeError:
        pass

def reinit_cam_out(m):
    """cam_out만 랜덤 초기화"""
    if isinstance(m, nn.Linear):
        nn.init.normal_(m.weight, std=0.01)
        nn.init.constant_(m.bias, 0)


def get_model(mode):

    # body
    vit_cfg = Config.fromfile(cfg.encoder_config_file)
    vit = build_posenet(vit_cfg.model)

    token_decoder = None
    if getattr(cfg, 'use_token_decoder', False):
        token_cfg_path = cfg.token_cfg
        token_cfg = OmegaConf.load(token_cfg_path)
        token_decoder = SMPLTokenDecoderHead(token_cfg)

    body_position_net = PositionNet('body', feat_dim=cfg.feat_dim)
    body_rotation_net = BodyRotationNet(feat_dim=cfg.feat_dim)
    box_net = BoxNet(feat_dim=cfg.feat_dim)

    # hand

    ## no-ROI
    # hand_position_net = GlobalHandPositionNet(feat_dim = cfg.feat_dim)
    # hand_rotation_net = HandTokenRegressor()

    hand_position_net = PositionNet('hand', feat_dim=cfg.feat_dim)
    hand_roi_net = HandRoI(feat_dim=cfg.feat_dim, upscale=cfg.upscale)
    hand_rotation_net = HandRotationNet('hand', feat_dim=cfg.feat_dim)

    # face
    face_regressor = FaceRegressor(feat_dim=cfg.feat_dim)

    if mode == 'train':
        # body
        if not getattr(cfg, 'random_init', False):
            encoder_pretrained_model = torch.load(cfg.encoder_pretrained_model_path)['state_dict']
            vit.load_state_dict(encoder_pretrained_model, strict=False)
            print(f"Initialize encoder from {cfg.encoder_pretrained_model_path}")
        else:
            print('Random init!!!!!!!')

        if getattr(cfg, 'use_smpl', False):
            body_position_net.apply(init_weights)
            body_rotation_net.apply(init_weights)

        else:

            body_position_net.apply(init_weights)
            body_rotation_net.apply(init_weights)
            
            box_net.apply(init_weights)

            # if getattr(cfg, 'use_token_decoder', False):
            #     set_trace()
            #     for name, child in token_decoder.named_children():
            #         if name == 'tokenizer':
            #             continue
            #         child.apply(init_weights)

            # hand
            hand_position_net.apply(init_weights)
            hand_roi_net.apply(init_weights)
            hand_rotation_net.apply(init_weights)

            # face
            face_regressor.apply(init_weights)

        

    encoder = vit.backbone

    model = Model(encoder, token_decoder, body_position_net, body_rotation_net, box_net, hand_position_net, hand_roi_net, hand_rotation_net,
                  face_regressor, mode)
  
    return model