# modified from Multi-HMR
# Copyright (c) 2024-present NAVER Corp.
# CC BY-NC-SA 4.0 license

import torch
import numpy as np
import smplx
from smplx.joint_names import JOINT_NAMES
from dust3r.utils.geometry import (
    perspective_projection,
    resize_camera_intrinsics,
    get_camera_parameters
)
from dust3r.utils.image import pad_image
import roma
import pickle
import os

current_dir = os.path.dirname(os.path.abspath(__file__))
src_dir = os.path.dirname(current_dir)
SMPLX_DIR = os.path.join(src_dir, 'models')
MEAN_PARAMS = os.path.join(src_dir, 'models', 'smpl_mean_params.npz')
SMPLX2SMPL = os.path.join(src_dir, 'models', 'smplx', 'smplx2smpl.pkl')

class SMPLModel(object):
    def __init__(self, device, model_args={}, eval_args={}):
        self.device = device
        self.person_center = 'head'

        self.patch_size = model_args.get('patch_size', 16)
        self.mhmr_img_res = model_args.get('mhmr_img_res', 896)
        self.bb_patch_size = model_args.get('bb_patch_size', 14)

        # Parametric 3D human models
        self.smplx_neutral_11 = smplx.create(
            SMPLX_DIR, 'smplx', gender='neutral', use_pca=False, flat_hand_mean=True, num_betas=11).to(self.device)
        self.smplx_neutral_10 = smplx.create(
            SMPLX_DIR, 'smplx', gender='neutral', use_pca=False, flat_hand_mean=True, num_betas=10).to(self.device)

        # Evaluation
        self.use_fake_K = eval_args.get('use_fake_K', False)
        dataset = eval_args.get('dataset', None)
        if dataset is not None:
            self.smpl = [
                smplx.create(SMPLX_DIR, 'smpl', gender=g).to(self.device) for g in ['neutral', 'male', 'female']]
            self.smpl_faces = {'smpl': self.smpl[0].faces, 'smplx': self.smplx_neutral_11.faces}
            with open(SMPLX2SMPL, 'rb') as f:
                self.smplx2smpl = torch.from_numpy(pickle.load(f)['matrix'].astype(np.float32)).to(self.device)

            if dataset in ['rich']:
                self.smplx = {
                    g: smplx.create(SMPLX_DIR, 'smplx', gender=g, num_pca_comps=12
                                    ).to(self.device) for g in ['male', 'female']}
            self._setup_dataset_config(dataset)

    def _setup_dataset_config(self, dataset):
        self.j_smpl = self.smpl[0].J_regressor[:24]
        if dataset in ['3dpw']:
            h36m_to_14 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9][:14]
            self.j_h36m = torch.Tensor(np.load('src/models/smpl/J_regressor_h36m.npy'))
            self.j_regressor = self.j_h36m[h36m_to_14]
            self.pelvis_idx = [2, 3]
            self.params_type = 'smpl'
        elif dataset in ['bedlam', 'rich']:
            self.j_regressor = self.j_smpl
            self.pelvis_idx = [1, 2]
            self.params_type = 'smplx'
        else:
            self.j_regressor = self.j_smpl
            self.pelvis_idx = [1, 2]
            self.params_type = 'smpl'

    def forward_smpl(self, dataset, smpl_dict, smpl_mask):
        nhv = int(smpl_mask.sum())

        if dataset in ['bedlam']:
            out = self.smplx_neutral_11(
                global_orient=smpl_dict['smplx_root_pose'][smpl_mask].reshape(-1, 3),
                body_pose=smpl_dict['smplx_body_pose'][smpl_mask].reshape(-1, 21*3),
                jaw_pose=smpl_dict['smplx_jaw_pose'][smpl_mask].reshape(-1, 3),
                leye_pose=smpl_dict['smplx_leye_pose'][smpl_mask].reshape(-1, 3),
                reye_pose=smpl_dict['smplx_reye_pose'][smpl_mask].reshape(-1, 3),
                left_hand_pose=smpl_dict['smplx_left_hand_pose'][smpl_mask].reshape(-1, 15*3),
                right_hand_pose=smpl_dict['smplx_right_hand_pose'][smpl_mask].reshape(-1, 15*3),
                betas=smpl_dict['smplx_shape'][smpl_mask].reshape(-1, 11),
                transl=smpl_dict['smplx_transl'][smpl_mask].reshape(-1, 3),
                expression=self.smplx_neutral_11.expression.repeat(nhv, 1),
            )
            verts = out.vertices.reshape(nhv, -1, 3)

        elif dataset in ['3dpw']:
            smpl_params = {
                'global_orient': smpl_dict['smpl_root_pose'][smpl_mask].reshape(-1,3),
                'body_pose': smpl_dict['smpl_body_pose'][smpl_mask].reshape(-1,23*3),
                'betas': smpl_dict['smpl_shape'][smpl_mask].reshape(-1,10),
                'transl': smpl_dict['smpl_transl'][smpl_mask].reshape(-1,3),
                }
            out = self.smpl[1](**smpl_params)
            verts = out.vertices.reshape(nhv, -1, 3)

            # update verts/joints if this is not the right gender
            if int(smpl_dict['smpl_gender_id'].max()) == 2:
                out_female = self.smpl[2](**smpl_params)
                idx = torch.where(smpl_dict['smpl_gender_id'] == 2)[1]
                verts[idx] = out_female.vertices.reshape(nhv, -1, 3)[idx]

        elif dataset in ['emdb', 'emdb1', 'emdb2']:
            gender = smpl_dict['smpl_gender_id'].max()
            out = self.smpl[gender](
                global_orient=smpl_dict['smpl_root_pose_w'][smpl_mask].reshape(-1,3),
                body_pose=smpl_dict['smpl_body_pose'][smpl_mask].reshape(-1,23*3),
                betas=smpl_dict['smpl_shape'][smpl_mask].reshape(-1,10),
                transl=smpl_dict['smpl_transl_w'][smpl_mask].reshape(-1,3),
            )
            verts = out.vertices.reshape(nhv, -1, 3) # world space

        elif dataset in ['rich']:
            gender = {1: 'male', 2: 'female'}[int(smpl_dict['smplx_gender_id'].max())]
            out = self.smplx[gender](
                global_orient=smpl_dict['smplx_global_orient'][smpl_mask].reshape(-1,3),
                body_pose=smpl_dict['smplx_body_pose'][smpl_mask].reshape(-1,21*3),
                jaw_pose=torch.zeros([nhv, 3]),
                leye_pose=torch.zeros([nhv, 3]),
                reye_pose=torch.zeros([nhv, 3]),
                left_hand_pose=torch.zeros([nhv, 12]),
                right_hand_pose=torch.zeros([nhv, 12]),
                betas=smpl_dict['smplx_betas'][smpl_mask].reshape(-1,10),
                transl=smpl_dict['smplx_transl'][smpl_mask].reshape(-1,3),
                expression=torch.zeros([nhv, 10]),
            )
            verts = out.vertices.reshape(nhv, -1, 3)

        if self.params_type == 'smplx':
            verts = self.smplx2smpl @ verts
        jts = self.j_regressor @ verts

        return verts, jts

    def update_smpl_gt(self, views):
        target = {}

        batch_size = views[0]["img"].shape[0]

        smpl_keys = [k for k in views[0].keys() if 'smpl' in k]
        smpl_dict = {
            k: (stacked := torch.stack(
                [view.pop(k) for view in views], dim=0)).view(-1, *stacked.shape[2:])
            for k in smpl_keys
        }   # Shape: (num_views * batch_size, 10, ...)
        smpl_mask = smpl_dict['smpl_mask']
        idx_h = torch.where(smpl_mask) # frame_idx, batch_idx, human_idx
        K = torch.stack(
            [view['camera_intrinsics'] for view in views], dim=0
        )
        K = K.view(-1, *K.shape[2:])
        nhv = int(smpl_mask.sum())

        # Get MHMR input image (high-res, square)
        imgs = torch.stack([view["img"] for view in views], dim=0)
        imgs = imgs.view(-1, *imgs.shape[2:])
        K_mhmr = resize_camera_intrinsics(K, *imgs.shape[2:], self.mhmr_img_res)
        imgs_mhmr = pad_image(imgs, self.mhmr_img_res)

        # SMPLX forward - BEDLAM
        has_smplx_params = 1
        out = self.smplx_neutral_11(
            global_orient=smpl_dict['smplx_root_pose'][smpl_mask].reshape(-1, 3),
            body_pose=smpl_dict['smplx_body_pose'][smpl_mask].reshape(-1, 21*3),
            jaw_pose=smpl_dict['smplx_jaw_pose'][smpl_mask].reshape(-1, 3),
            leye_pose=smpl_dict['smplx_leye_pose'][smpl_mask].reshape(-1, 3),
            reye_pose=smpl_dict['smplx_reye_pose'][smpl_mask].reshape(-1, 3),
            left_hand_pose=smpl_dict['smplx_left_hand_pose'][smpl_mask].reshape(-1, 15*3),
            right_hand_pose=smpl_dict['smplx_right_hand_pose'][smpl_mask].reshape(-1, 15*3),
            betas=smpl_dict['smplx_shape'][smpl_mask].reshape(-1, 11),
            transl=smpl_dict['smplx_transl'][smpl_mask].reshape(-1, 3),
            expression=self.smplx_neutral_11.expression.repeat(nhv, 1),
        )
        verts, jts = out.vertices.reshape(nhv, -1, 3), out.joints.reshape(nhv, -1, 3)

        j2d = perspective_projection(jts, K[idx_h[0]])
        v2d = perspective_projection(verts, K[idx_h[0]])

        # Translation of the primary keypoint
        root_joint_idx = JOINT_NAMES.index(self.person_center)
        target['smpl_transl'] = jts[:,root_joint_idx] # [nhv,3]
        target['smpl_transl_pelvis'] = jts[:,0] # [nhv,3]

        # Fill in target
        target['smpl_v3d'] = verts
        target['smpl_j3d'] = jts
        target['smpl_j2d'] = j2d
        target['smpl_v2d'] = v2d

        if has_smplx_params:
            target['smpl_rotvec'] = torch.cat([smpl_dict['smplx_root_pose'],
                                        smpl_dict['smplx_body_pose'],
                                        smpl_dict['smplx_left_hand_pose'],
                                        smpl_dict['smplx_right_hand_pose'],
                                        smpl_dict['smplx_jaw_pose']],2)[smpl_mask] # [bs,nhmax]
            target['smpl_rotmat'] = roma.rotvec_to_rotmat(target['smpl_rotvec'])
            target['smpl_shape'] = smpl_dict['smplx_shape'][smpl_mask]

        true_shapes = torch.stack([view["true_shape"] for view in views], dim=0)
        if len(torch.unique(true_shapes, dim=0)) != 1:
            raise NotImplementedError

        # Creating the target heatmap for the primary keypoint
        pk = target['smpl_transl'].unsqueeze(1) # (nhv,3)

        # For 512 res (CUT3R, patch_size=16)
        pk_loc = perspective_projection(pk, K[idx_h[0]]).squeeze(1) # original pixel uv coordinates (nhv,2): W, H
        n_patch_16, pk_idx_16 = get_patch_uv(true_shapes[0][0], self.patch_size, pk_loc)
        target['smpl_uv_16'] = pk_idx_16[:, [1, 0]]

        # For 896 res (MHMR, patch_size=14)
        pk_loc_mhmr = perspective_projection(pk, K_mhmr[idx_h[0]]).squeeze(1) # original pixel uv coordinates (nhv,2): W, H
        n_patch_14, pk_idx_14 = get_patch_uv(self.mhmr_img_res, self.bb_patch_size, pk_loc_mhmr)
        smpl_mask_14, visible_humans_14, scores_14 = get_score(n_patch_14, pk_idx_14, smpl_mask.clone())
        target['smpl_uv'] = pk_idx_14[:, [1, 0]]

        # Rebatch and Update with visibility indice
        _target = {}
        num_view = len(views)
        max_humans = smpl_mask_14.shape[1]
        idx_vis = torch.where(visible_humans_14)[0]

        for k, v in target.items():
            full_out = torch.zeros(
                num_view * batch_size, max_humans, *v.shape[1:],
                device=v.device, dtype=v.dtype,
            )
            full_out[smpl_mask_14] = v[idx_vis] # discard unvisible humans due to olccusion
            _target[k] = full_out.chunk(num_view, dim=0) # .view(num_view, batch_size, *full_out.shape[1:])

        _target['smpl_scores'] = scores_14.chunk(num_view, dim=0)
        _target['smpl_mask'] = smpl_mask_14.chunk(num_view, dim=0)
        _target['K_mhmr'] = K_mhmr.chunk(num_view, dim=0)
        _target['img_mhmr'] = imgs_mhmr.chunk(num_view, dim=0)

        if "msk" in views[0]:
            msks = torch.stack([view["msk"] for view in views], dim=0)
            msks = msks.view(-1, *msks.shape[2:])
            msks_mhmr = pad_image(msks, self.mhmr_img_res, pad_value=0.0)  # bs,288,512->bs,896,896
            msks_mhmr = (msks_mhmr > 0.1).float()
            _target['msk_mhmr'] = msks_mhmr.chunk(num_view, dim=0)

        for i, v in enumerate(zip(*_target.values())):
            views[i].update(dict(zip(_target.keys(), v)))

        torch.cuda.empty_cache()

    def update_smpl_gt_eval(self, views, dataset):
        from dust3r.utils.geometry import geotrf

        target = {}
        batch_size = views[0]["img"].shape[0]

        smpl_keys = [k for k in views[0].keys() if 'smpl' in k]
        smpl_dict = {
            k: (stacked := torch.stack(
                [view.pop(k) for view in views], dim=0)).view(-1, *stacked.shape[2:])
            for k in smpl_keys
        }   # Shape: (num_views * batch_size, 10, ...)
        smpl_mask = smpl_dict['smpl_mask']
        idx_h = torch.where(smpl_mask) # frame_idx, batch_idx, human_idx
        K = torch.stack([view['camera_intrinsics'] for view in views], dim=0)
        K = K.view(-1, *K.shape[2:])

        # Get MHMR input image (high-res, square)
        imgs = torch.stack([view["img"] for view in views], dim=0)
        imgs = imgs.view(-1, *imgs.shape[2:])
        K_mhmr = resize_camera_intrinsics(K, *imgs.shape[2:], self.mhmr_img_res)
        imgs_mhmr = pad_image(imgs, self.mhmr_img_res)

        verts, jts = self.forward_smpl(dataset, smpl_dict, smpl_mask)

        if dataset in ['emdb', 'emdb1', 'emdb2', 'rich']:
            target['smpl_v3d_w'] = verts
            target['smpl_j3d_w'] = jts
            T_w2c = torch.stack([view['T_w2c'] for view in views], dim=0)
            T_w2c = T_w2c.view(-1, *T_w2c.shape[2:])
            target['smpl_v3d_c'] = geotrf(T_w2c[idx_h[0]], verts)
            target['smpl_j3d_c'] = geotrf(T_w2c[idx_h[0]], jts)

        else:
            target['smpl_v3d_c'] = verts
            target['smpl_j3d_c'] = jts
            T_c2w = torch.stack([view['camera_pose'] for view in views], dim=0)
            T_c2w = T_c2w.view(-1, *T_c2w.shape[2:])
            target['smpl_v3d_w'] = geotrf(T_c2w[idx_h[0]], verts)
            target['smpl_j3d_w'] = geotrf(T_c2w[idx_h[0]], jts)

        target['smpl_j2d'] = perspective_projection(target['smpl_j3d_c'], K[idx_h[0]])
        target['smpl_v2d'] = perspective_projection(target['smpl_v3d_c'], K[idx_h[0]])

        # Rebatch and Update with visibility indice
        _target = {}
        num_view = len(views)
        max_humans = smpl_mask.shape[1]
        for k, v in target.items():
            full_out = torch.zeros(
                num_view * batch_size, max_humans, *v.shape[1:],
                device=v.device, dtype=v.dtype,
            )
            full_out[smpl_mask] = v # discard unvisible humans due to olccusion
            _target[k] = full_out.chunk(num_view, dim=0) # .view(num_view, batch_size, *full_out.shape[1:])

        if self.use_fake_K:
            K_mhmr = get_camera_parameters(self.mhmr_img_res, device=K.device) # if use pseudo K
            K_mhmr = K_mhmr.expand(K.shape[0], -1, -1)

        _target['smpl_mask'] = smpl_mask.chunk(num_view, dim=0)
        _target['K_mhmr'] = K_mhmr.chunk(num_view, dim=0)
        _target['img_mhmr'] = imgs_mhmr.chunk(num_view, dim=0)

        if "msk" in views[0]:
            msks = torch.stack([view["msk"] for view in views], dim=0)
            msks = msks.view(-1, *msks.shape[2:])
            msks_mhmr = pad_image(msks, self.mhmr_img_res, pad_value=0.0)  # bs,288,512->bs,896,896
            msks_mhmr = (msks_mhmr > 0.1).float()
            _target['msk_mhmr'] = msks_mhmr.chunk(num_view, dim=0)

        for i, v in enumerate(zip(*_target.values())):
            views[i].update(dict(zip(_target.keys(), v)))

        torch.cuda.empty_cache()

def get_patch_uv(imgshape, patch_size, pk_loc):
    n_patch = imgshape // patch_size  # H, W
    pk_idx = (pk_loc // patch_size).int()
    return n_patch, pk_idx

def get_score(n_patch, pk_idx, smpl_mask):
    # Scores & updating valid_humans according to occlusion - wap X and Y for scores only
    idx_h = torch.where(smpl_mask)
    nhv = int(smpl_mask.sum())
    bs = smpl_mask.shape[0]
    device = smpl_mask.device

    if isinstance(n_patch, (int, float)):
        patch_h, patch_w = int(n_patch), int(n_patch)
    else:
        patch_h, patch_w = n_patch[0], n_patch[1]

    scores = torch.zeros((bs, patch_h, patch_w)).to(device)
    visible_humans = torch.ones(nhv).to(device) # by default no occlusion so all visible

    for k in range(nhv):
        i = int(idx_h[0][k]) # index of the image
        j = int(idx_h[1][k]) # index of the human in this image
        _x = pk_idx[k,1] # patch center H
        _y = pk_idx[k,0] # patch center W
        # filter out heads out of cropping bounds
        if _x >= 0 and _x < patch_h and _y >= 0 and _y < patch_w:
            if scores[i,_x,_y] == 1:
                smpl_mask[i,j] = 0
                visible_humans[k] = 0
            else:
                scores[i,_x,_y] = 1
        else:
            smpl_mask[i,j] = 0
            visible_humans[k] = 0

    return smpl_mask, visible_humans, scores

import torch.nn as nn
from croco.models.blocks import Mlp_flex

class SMPLDecoder(nn.Module):
    def __init__(
        self,
        hidden_size=768,
        target_dim=1,
        mlp_ratio=1,
        num_layers=2,
    ):
        super().__init__()
        self.mlp = Mlp_flex(
            in_features=hidden_size,
            hidden_features=int(hidden_size * mlp_ratio),
            out_features=target_dim,
            num_layers=num_layers,
            drop=0,
        )

    def forward(
        self,
        feat,
    ):
        """
        feat: BxC
        """

        pred = self.mlp(feat)
        return pred

def regression_mlp(layers_sizes):
    """
    Return a fully connected network.
    """
    assert len(layers_sizes) >= 2
    in_features = layers_sizes[0]
    layers = []
    for i in range(1, len(layers_sizes)-1):
        out_features = layers_sizes[i]
        layers.append(torch.nn.Linear(in_features, out_features))
        layers.append(torch.nn.ReLU())
        in_features = out_features
    layers.append(torch.nn.Linear(in_features, layers_sizes[-1]))
    return torch.nn.Sequential(*layers)

def apply_threshold(det_thresh, _scores):
    """ Apply thresholding to detection scores; if stack_K is used and det_thresh is a list, apply to each channel separately """
    if isinstance(det_thresh, list):
        det_thresh = det_thresh[0]
    idx = torch.where(_scores >= det_thresh)
    return idx

def nms(heat, kernel=3):
    """ easy non maximal supression (as in CenterNet) """

    if kernel not in [2, 4]:
        pad = (kernel - 1) // 2
    else:
        if kernel == 2:
            pad = 1
        else:
            pad = 2

    hmax = nn.functional.max_pool2d( heat, (kernel, kernel), stride=1, padding=pad)

    if hmax.shape[2] > heat.shape[2]:
        hmax = hmax[:, :, :heat.shape[2], :heat.shape[3]]

    keep = (hmax == heat).float()

    return heat * keep


SMPLModel

import torch
import numpy as np
import smplx
from smplx.joint_names import JOINT_NAMES
from dust3r.utils.geometry import (
    perspective_projection,
    resize_camera_intrinsics,
    get_camera_parameters
)
from dust3r.utils.image import pad_image
import roma
import pickle
import os

current_dir = os.path.dirname(os.path.abspath(__file__))
src_dir = os.path.dirname(current_dir)
SMPLX_DIR = os.path.join(src_dir, 'models')
MEAN_PARAMS = os.path.join(src_dir, 'models', 'smpl_mean_params.npz')
SMPLX2SMPL = os.path.join(src_dir, 'models', 'smplx', 'smplx2smpl.pkl')

class SMPLModel(object):
    def __init__(self, device, model_args={}, eval_args={}):
        self.device = device
        self.person_center = 'head'

        self.patch_size = model_args.get('patch_size', 16)
        self.mhmr_img_res = model_args.get('mhmr_img_res', 896)
        self.bb_patch_size = model_args.get('bb_patch_size', 14)

        # Parametric 3D human models
        self.smplx_neutral_11 = smplx.create(
            SMPLX_DIR, 'smplx', gender='neutral', use_pca=False, flat_hand_mean=True, num_betas=11).to(self.device)
        self.smplx_neutral_10 = smplx.create(
            SMPLX_DIR, 'smplx', gender='neutral', use_pca=False, flat_hand_mean=True, num_betas=10).to(self.device)

        # Evaluation
        self.use_fake_K = eval_args.get('use_fake_K', False)
        dataset = eval_args.get('dataset', None)
        if dataset is not None:
            self.smpl = [
                smplx.create(SMPLX_DIR, 'smpl', gender=g).to(self.device) for g in ['neutral', 'male', 'female']]
            self.smpl_faces = {'smpl': self.smpl[0].faces, 'smplx': self.smplx_neutral_11.faces}
            with open(SMPLX2SMPL, 'rb') as f:
                self.smplx2smpl = torch.from_numpy(pickle.load(f)['matrix'].astype(np.float32)).to(self.device)

            if dataset in ['rich']:
                self.smplx = {
                    g: smplx.create(SMPLX_DIR, 'smplx', gender=g, num_pca_comps=12
                                    ).to(self.device) for g in ['male', 'female']}
            self._setup_dataset_config(dataset)

    def _setup_dataset_config(self, dataset):
        self.j_smpl = self.smpl[0].J_regressor[:24]
        if dataset in ['3dpw']:
            h36m_to_14 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9][:14]
            self.j_h36m = torch.Tensor(np.load('src/models/smpl/J_regressor_h36m.npy'))
            self.j_regressor = self.j_h36m[h36m_to_14]
            self.pelvis_idx = [2, 3]
            self.params_type = 'smpl'
        elif dataset in ['bedlam', 'rich']:
            self.j_regressor = self.j_smpl
            self.pelvis_idx = [1, 2]
            self.params_type = 'smplx'
        else:
            self.j_regressor = self.j_smpl
            self.pelvis_idx = [1, 2]
            self.params_type = 'smpl'

    def forward_smpl(self, dataset, smpl_dict, smpl_mask):
        nhv = int(smpl_mask.sum())

        if dataset in ['bedlam']:
            out = self.smplx_neutral_11(
                global_orient=smpl_dict['smplx_root_pose'][smpl_mask].reshape(-1, 3),
                body_pose=smpl_dict['smplx_body_pose'][smpl_mask].reshape(-1, 21*3),
                jaw_pose=smpl_dict['smplx_jaw_pose'][smpl_mask].reshape(-1, 3),
                leye_pose=smpl_dict['smplx_leye_pose'][smpl_mask].reshape(-1, 3),
                reye_pose=smpl_dict['smplx_reye_pose'][smpl_mask].reshape(-1, 3),
                left_hand_pose=smpl_dict['smplx_left_hand_pose'][smpl_mask].reshape(-1, 15*3),
                right_hand_pose=smpl_dict['smplx_right_hand_pose'][smpl_mask].reshape(-1, 15*3),
                betas=smpl_dict['smplx_shape'][smpl_mask].reshape(-1, 11),
                transl=smpl_dict['smplx_transl'][smpl_mask].reshape(-1, 3),
                expression=self.smplx_neutral_11.expression.repeat(nhv, 1),
            )
            verts = out.vertices.reshape(nhv, -1, 3)

        elif dataset in ['3dpw']:
            smpl_params = {
                'global_orient': smpl_dict['smpl_root_pose'][smpl_mask].reshape(-1,3),
                'body_pose': smpl_dict['smpl_body_pose'][smpl_mask].reshape(-1,23*3),
                'betas': smpl_dict['smpl_shape'][smpl_mask].reshape(-1,10),
                'transl': smpl_dict['smpl_transl'][smpl_mask].reshape(-1,3),
                }
            out = self.smpl[1](**smpl_params)
            verts = out.vertices.reshape(nhv, -1, 3)

            # update verts/joints if this is not the right gender
            if int(smpl_dict['smpl_gender_id'].max()) == 2:
                out_female = self.smpl[2](**smpl_params)
                idx = torch.where(smpl_dict['smpl_gender_id'] == 2)[1]
                verts[idx] = out_female.vertices.reshape(nhv, -1, 3)[idx]

        elif dataset in ['emdb', 'emdb1', 'emdb2']:
            gender = smpl_dict['smpl_gender_id'].max()
            out = self.smpl[gender](
                global_orient=smpl_dict['smpl_root_pose_w'][smpl_mask].reshape(-1,3),
                body_pose=smpl_dict['smpl_body_pose'][smpl_mask].reshape(-1,23*3),
                betas=smpl_dict['smpl_shape'][smpl_mask].reshape(-1,10),
                transl=smpl_dict['smpl_transl_w'][smpl_mask].reshape(-1,3),
            )
            verts = out.vertices.reshape(nhv, -1, 3) # world space

        elif dataset in ['rich']:
            gender = {1: 'male', 2: 'female'}[int(smpl_dict['smplx_gender_id'].max())]
            out = self.smplx[gender](
                global_orient=smpl_dict['smplx_global_orient'][smpl_mask].reshape(-1,3),
                body_pose=smpl_dict['smplx_body_pose'][smpl_mask].reshape(-1,21*3),
                jaw_pose=torch.zeros([nhv, 3]),
                leye_pose=torch.zeros([nhv, 3]),
                reye_pose=torch.zeros([nhv, 3]),
                left_hand_pose=torch.zeros([nhv, 12]),
                right_hand_pose=torch.zeros([nhv, 12]),
                betas=smpl_dict['smplx_betas'][smpl_mask].reshape(-1,10),
                transl=smpl_dict['smplx_transl'][smpl_mask].reshape(-1,3),
                expression=torch.zeros([nhv, 10]),
            )
            verts = out.vertices.reshape(nhv, -1, 3)

        if self.params_type == 'smplx':
            verts = self.smplx2smpl @ verts
        jts = self.j_regressor @ verts

        return verts, jts

    def update_smpl_gt(self, views):
        target = {}

        batch_size = views[0]["img"].shape[0]

        smpl_keys = [k for k in views[0].keys() if 'smpl' in k]
        smpl_dict = {
            k: (stacked := torch.stack(
                [view.pop(k) for view in views], dim=0)).view(-1, *stacked.shape[2:])
            for k in smpl_keys
        }   # Shape: (num_views * batch_size, 10, ...)
        smpl_mask = smpl_dict['smpl_mask']
        idx_h = torch.where(smpl_mask) # frame_idx, batch_idx, human_idx
        K = torch.stack(
            [view['camera_intrinsics'] for view in views], dim=0
        )
        K = K.view(-1, *K.shape[2:])
        nhv = int(smpl_mask.sum())

        # Get MHMR input image (high-res, square)
        imgs = torch.stack([view["img"] for view in views], dim=0)
        imgs = imgs.view(-1, *imgs.shape[2:])
        K_mhmr = resize_camera_intrinsics(K, *imgs.shape[2:], self.mhmr_img_res)
        imgs_mhmr = pad_image(imgs, self.mhmr_img_res)

        # SMPLX forward - BEDLAM
        has_smplx_params = 1
        out = self.smplx_neutral_11(
            global_orient=smpl_dict['smplx_root_pose'][smpl_mask].reshape(-1, 3),
            body_pose=smpl_dict['smplx_body_pose'][smpl_mask].reshape(-1, 21*3),
            jaw_pose=smpl_dict['smplx_jaw_pose'][smpl_mask].reshape(-1, 3),
            leye_pose=smpl_dict['smplx_leye_pose'][smpl_mask].reshape(-1, 3),
            reye_pose=smpl_dict['smplx_reye_pose'][smpl_mask].reshape(-1, 3),
            left_hand_pose=smpl_dict['smplx_left_hand_pose'][smpl_mask].reshape(-1, 15*3),
            right_hand_pose=smpl_dict['smplx_right_hand_pose'][smpl_mask].reshape(-1, 15*3),
            betas=smpl_dict['smplx_shape'][smpl_mask].reshape(-1, 11),
            transl=smpl_dict['smplx_transl'][smpl_mask].reshape(-1, 3),
            expression=self.smplx_neutral_11.expression.repeat(nhv, 1),
        )
        verts, jts = out.vertices.reshape(nhv, -1, 3), out.joints.reshape(nhv, -1, 3)

        j2d = perspective_projection(jts, K[idx_h[0]])
        v2d = perspective_projection(verts, K[idx_h[0]])

        # Translation of the primary keypoint
        root_joint_idx = JOINT_NAMES.index(self.person_center)
        target['smpl_transl'] = jts[:,root_joint_idx] # [nhv,3]
        target['smpl_transl_pelvis'] = jts[:,0] # [nhv,3]

        # Fill in target
        target['smpl_v3d'] = verts
        target['smpl_j3d'] = jts
        target['smpl_j2d'] = j2d
        target['smpl_v2d'] = v2d

        if has_smplx_params:
            target['smpl_rotvec'] = torch.cat([smpl_dict['smplx_root_pose'],
                                        smpl_dict['smplx_body_pose'],
                                        smpl_dict['smplx_left_hand_pose'],
                                        smpl_dict['smplx_right_hand_pose'],
                                        smpl_dict['smplx_jaw_pose']],2)[smpl_mask] # [bs,nhmax]
            target['smpl_rotmat'] = roma.rotvec_to_rotmat(target['smpl_rotvec'])
            target['smpl_shape'] = smpl_dict['smplx_shape'][smpl_mask]

        true_shapes = torch.stack([view["true_shape"] for view in views], dim=0)
        if len(torch.unique(true_shapes, dim=0)) != 1:
            raise NotImplementedError

        # Creating the target heatmap for the primary keypoint
        pk = target['smpl_transl'].unsqueeze(1) # (nhv,3)

        # For 512 res (CUT3R, patch_size=16)
        pk_loc = perspective_projection(pk, K[idx_h[0]]).squeeze(1) # original pixel uv coordinates (nhv,2): W, H
        n_patch_16, pk_idx_16 = get_patch_uv(true_shapes[0][0], self.patch_size, pk_loc)
        target['smpl_uv_16'] = pk_idx_16[:, [1, 0]]

        # For 896 res (MHMR, patch_size=14)
        pk_loc_mhmr = perspective_projection(pk, K_mhmr[idx_h[0]]).squeeze(1) # original pixel uv coordinates (nhv,2): W, H
        n_patch_14, pk_idx_14 = get_patch_uv(self.mhmr_img_res, self.bb_patch_size, pk_loc_mhmr)
        smpl_mask_14, visible_humans_14, scores_14 = get_score(n_patch_14, pk_idx_14, smpl_mask.clone())
        target['smpl_uv'] = pk_idx_14[:, [1, 0]]

        # Rebatch and Update with visibility indice
        _target = {}
        num_view = len(views)
        max_humans = smpl_mask_14.shape[1]
        idx_vis = torch.where(visible_humans_14)[0]

        for k, v in target.items():
            full_out = torch.zeros(
                num_view * batch_size, max_humans, *v.shape[1:],
                device=v.device, dtype=v.dtype,
            )
            full_out[smpl_mask_14] = v[idx_vis] # discard unvisible humans due to olccusion
            _target[k] = full_out.chunk(num_view, dim=0) # .view(num_view, batch_size, *full_out.shape[1:])

        _target['smpl_scores'] = scores_14.chunk(num_view, dim=0)
        _target['smpl_mask'] = smpl_mask_14.chunk(num_view, dim=0)
        _target['K_mhmr'] = K_mhmr.chunk(num_view, dim=0)
        _target['img_mhmr'] = imgs_mhmr.chunk(num_view, dim=0)

        if "msk" in views[0]:
            msks = torch.stack([view["msk"] for view in views], dim=0)
            msks = msks.view(-1, *msks.shape[2:])
            msks_mhmr = pad_image(msks, self.mhmr_img_res, pad_value=0.0)  # bs,288,512->bs,896,896
            msks_mhmr = (msks_mhmr > 0.1).float()
            _target['msk_mhmr'] = msks_mhmr.chunk(num_view, dim=0)

        for i, v in enumerate(zip(*_target.values())):
            views[i].update(dict(zip(_target.keys(), v)))

        torch.cuda.empty_cache()

    def update_smpl_gt_eval(self, views, dataset):
        from dust3r.utils.geometry import geotrf

        target = {}
        batch_size = views[0]["img"].shape[0]

        smpl_keys = [k for k in views[0].keys() if 'smpl' in k]
        smpl_dict = {
            k: (stacked := torch.stack(
                [view.pop(k) for view in views], dim=0)).view(-1, *stacked.shape[2:])
            for k in smpl_keys
        }   # Shape: (num_views * batch_size, 10, ...)
        smpl_mask = smpl_dict['smpl_mask']
        idx_h = torch.where(smpl_mask) # frame_idx, batch_idx, human_idx
        K = torch.stack([view['camera_intrinsics'] for view in views], dim=0)
        K = K.view(-1, *K.shape[2:])

        # Get MHMR input image (high-res, square)
        imgs = torch.stack([view["img"] for view in views], dim=0)
        imgs = imgs.view(-1, *imgs.shape[2:])
        K_mhmr = resize_camera_intrinsics(K, *imgs.shape[2:], self.mhmr_img_res)
        imgs_mhmr = pad_image(imgs, self.mhmr_img_res)

        verts, jts = self.forward_smpl(dataset, smpl_dict, smpl_mask)

        if dataset in ['emdb', 'emdb1', 'emdb2', 'rich']:
            target['smpl_v3d_w'] = verts
            target['smpl_j3d_w'] = jts
            T_w2c = torch.stack([view['T_w2c'] for view in views], dim=0)
            T_w2c = T_w2c.view(-1, *T_w2c.shape[2:])
            target['smpl_v3d_c'] = geotrf(T_w2c[idx_h[0]], verts)
            target['smpl_j3d_c'] = geotrf(T_w2c[idx_h[0]], jts)

        else:
            target['smpl_v3d_c'] = verts
            target['smpl_j3d_c'] = jts
            T_c2w = torch.stack([view['camera_pose'] for view in views], dim=0)
            T_c2w = T_c2w.view(-1, *T_c2w.shape[2:])
            target['smpl_v3d_w'] = geotrf(T_c2w[idx_h[0]], verts)
            target['smpl_j3d_w'] = geotrf(T_c2w[idx_h[0]], jts)

        target['smpl_j2d'] = perspective_projection(target['smpl_j3d_c'], K[idx_h[0]])
        target['smpl_v2d'] = perspective_projection(target['smpl_v3d_c'], K[idx_h[0]])

        # Rebatch and Update with visibility indice
        _target = {}
        num_view = len(views)
        max_humans = smpl_mask.shape[1]
        for k, v in target.items():
            full_out = torch.zeros(
                num_view * batch_size, max_humans, *v.shape[1:],
                device=v.device, dtype=v.dtype,
            )
            full_out[smpl_mask] = v # discard unvisible humans due to olccusion
            _target[k] = full_out.chunk(num_view, dim=0) # .view(num_view, batch_size, *full_out.shape[1:])

        if self.use_fake_K:
            K_mhmr = get_camera_parameters(self.mhmr_img_res, device=K.device) # if use pseudo K
            K_mhmr = K_mhmr.expand(K.shape[0], -1, -1)

        _target['smpl_mask'] = smpl_mask.chunk(num_view, dim=0)
        _target['K_mhmr'] = K_mhmr.chunk(num_view, dim=0)
        _target['img_mhmr'] = imgs_mhmr.chunk(num_view, dim=0)

        if "msk" in views[0]:
            msks = torch.stack([view["msk"] for view in views], dim=0)
            msks = msks.view(-1, *msks.shape[2:])
            msks_mhmr = pad_image(msks, self.mhmr_img_res, pad_value=0.0)  # bs,288,512->bs,896,896
            msks_mhmr = (msks_mhmr > 0.1).float()
            _target['msk_mhmr'] = msks_mhmr.chunk(num_view, dim=0)

        for i, v in enumerate(zip(*_target.values())):
            views[i].update(dict(zip(_target.keys(), v)))

        torch.cuda.empty_cache()

def get_patch_uv(imgshape, patch_size, pk_loc):
    n_patch = imgshape // patch_size  # H, W
    pk_idx = (pk_loc // patch_size).int()
    return n_patch, pk_idx

def get_score(n_patch, pk_idx, smpl_mask):
    # Scores & updating valid_humans according to occlusion - wap X and Y for scores only
    idx_h = torch.where(smpl_mask)
    nhv = int(smpl_mask.sum())
    bs = smpl_mask.shape[0]
    device = smpl_mask.device

    if isinstance(n_patch, (int, float)):
        patch_h, patch_w = int(n_patch), int(n_patch)
    else:
        patch_h, patch_w = n_patch[0], n_patch[1]

    scores = torch.zeros((bs, patch_h, patch_w)).to(device)


            K_mhmr = get_camera_parameters(self.mhmr_img_res, device=K.device) # if use pseudo K
            K_mhmr = K_mhmr.expand(K.shape[0], -1, -1)

        _target['smpl_mask'] = smpl_mask.chunk(num_view, dim=0)
        _target['K_mhmr'] = K_mhmr.chunk(num_view, dim=0)
        _target['img_mhmr'] = imgs_mhmr.chunk(num_view, dim=0)

        if "msk" in views[0]:
            msks = torch.stack([view["msk"] for view in views], dim=0)
            msks = msks.view(-1, *msks.shape[2:])
            msks_mhmr = pad_image(msks, self.mhmr_img_res, pad_value=0.0)  # bs,288,512->bs,896,896
            msks_mhmr = (msks_mhmr > 0.1).float()
            _target['msk_mhmr'] = msks_mhmr.chunk(num_view, dim=0)

        for i, v in enumerate(zip(*_target.values())):
            views[i].update(dict(zip(_target.keys(), v)))

        torch.cuda.empty_cache()

def get_patch_uv(imgshape, patch_size, pk_loc):
    n_patch = imgshape // patch_size  # H, W
    pk_idx = (pk_loc // patch_size).int()
    return n_patch, pk_idx

def get_score(n_patch, pk_idx, smpl_mask):
    # Scores & updating valid_humans according to occlusion - wap X and Y for scores only
    idx_h = torch.where(smpl_mask)
    nhv = int(smpl_mask.sum())
    bs = smpl_mask.shape[0]
    device = smpl_mask.device

    if isinstance(n_patch, (int, float)):
        patch_h, patch_w = int(n_patch), int(n_patch)
    else:
        patch_h, patch_w = n_patch[0], n_patch[1]

    scores = torch.zeros((bs, patch_h, patch_w)).to(device)
    visible_humans = torch.ones(nhv).to(device) # by default no occlusion so all visible

    for k in range(nhv):
        i = int(idx_h[0][k]) # index of the image
        j = int(idx_h[1][k]) # index of the human in this image
        _x = pk_idx[k,1] # patch center H


        _target['img_mhmr'] = imgs_mhmr.chunk(num_view, dim=0)

        if "msk" in views[0]:
            msks = torch.stack([view["msk"] for view in views], dim=0)
            msks = msks.view(-1, *msks.shape[2:])
            msks_mhmr = pad_image(msks, self.mhmr_img_res, pad_value=0.0)  # bs,288,512->bs,896,896
            msks_mhmr = (msks_mhmr > 0.1).float()
            _target['msk_mhmr'] = msks_mhmr.chunk(num_view, dim=0)

        for i, v in enumerate(zip(*_target.values())):
            views[i].update(dict(zip(_target.keys(), v)))

        torch.cuda.empty_cache()

def get_patch_uv(imgshape, patch_size, pk_loc):
    n_patch = imgshape // patch_size  # H, W
    pk_idx = (pk_loc // patch_size).int()
    return n_patch, pk_idx

def get_score(n_patch, pk_idx, smpl_mask):
    # Scores & updating valid_humans according to occlusion - wap X and Y for scores only
    idx_h = torch.where(smpl_mask)
    nhv = int(smpl_mask.sum())
    bs = smpl_mask.shape[0]
    device = smpl_mask.device

    if isinstance(n_patch, (int, float)):
        patch_h, patch_w = int(n_patch), int(n_patch)
    else:
        patch_h, patch_w = n_patch[0], n_patch[1]

    scores = torch.zeros((bs, patch_h, patch_w)).to(device)
    visible_humans = torch.ones(nhv).to(device) # by default no occlusion so all visible

    for k in range(nhv):
        i = int(idx_h[0][k]) # index of the image
        j = int(idx_h[1][k]) # index of the human in this image
        _x = pk_idx[k,1] # patch center H
        _y = pk_idx[k,0] # patch center W
        # filter out heads out of cropping bounds
        if _x >= 0 and _x < patch_h and _y >= 0 and _y < patch_w:
            if scores[i,_x,_y] == 1:
                smpl_mask[i,j] = 0
                visible_humans[k] = 0
            else:
                scores[i,_x,_y] = 1
        else:
            smpl_mask[i,j] = 0
            visible_humans[k] = 0

    return smpl_mask, visible_humans, scores

import torch.nn as nn
from croco.models.blocks import Mlp_flex

class SMPLDecoder(nn.Module):
    def __init__(
        self,
        hidden_size=768,
        target_dim=1,
        mlp_ratio=1,
        num_layers=2,
    ):
        super().__init__()
        self.mlp = Mlp_flex(
            in_features=hidden_size,
            hidden_features=int(hidden_size * mlp_ratio),
            out_features=target_dim,
            num_layers=num_layers,
            drop=0,


SMPLDecoder

        j = int(idx_h[1][k]) # index of the human in this image
        _x = pk_idx[k,1] # patch center H
        _y = pk_idx[k,0] # patch center W
        # filter out heads out of cropping bounds
        if _x >= 0 and _x < patch_h and _y >= 0 and _y < patch_w:
            if scores[i,_x,_y] == 1:
                smpl_mask[i,j] = 0
                visible_humans[k] = 0
            else:
                scores[i,_x,_y] = 1
        else:
            smpl_mask[i,j] = 0
            visible_humans[k] = 0

    return smpl_mask, visible_humans, scores

import torch.nn as nn
from croco.models.blocks import Mlp_flex

class SMPLDecoder(nn.Module):
    def __init__(
        self,
        hidden_size=768,
        target_dim=1,
        mlp_ratio=1,
        num_layers=2,
    ):
        super().__init__()
        self.mlp = Mlp_flex(
            in_features=hidden_size,
            hidden_features=int(hidden_size * mlp_ratio),
            out_features=target_dim,
            num_layers=num_layers,
            drop=0,
        )

    def forward(
        self,
        feat,
    ):
        """
        feat: BxC
        """

        pred = self.mlp(feat)
        return pred

def regression_mlp(layers_sizes):
    """
    Return a fully connected network.
    """
    assert len(layers_sizes) >= 2
    in_features = layers_sizes[0]
    layers = []
    for i in range(1, len(layers_sizes)-1):
        out_features = layers_sizes[i]
        layers.append(torch.nn.Linear(in_features, out_features))
        layers.append(torch.nn.ReLU())
        in_features = out_features
    layers.append(torch.nn.Linear(in_features, layers_sizes[-1]))
    return torch.nn.Sequential(*layers)

def apply_threshold(det_thresh, _scores):
    """ Apply thresholding to detection scores; if stack_K is used and det_thresh is a list, apply to each channel separately """
    if isinstance(det_thresh, list):


        self.mlp = Mlp_flex(
            in_features=hidden_size,
            hidden_features=int(hidden_size * mlp_ratio),
            out_features=target_dim,
            num_layers=num_layers,
            drop=0,
        )

    def forward(
        self,
        feat,
    ):
        """
        feat: BxC
        """

        pred = self.mlp(feat)
        return pred

def regression_mlp(layers_sizes):
    """
    Return a fully connected network.
    """
    assert len(layers_sizes) >= 2
    in_features = layers_sizes[0]
    layers = []
    for i in range(1, len(layers_sizes)-1):
        out_features = layers_sizes[i]
        layers.append(torch.nn.Linear(in_features, out_features))
        layers.append(torch.nn.ReLU())
        in_features = out_features
    layers.append(torch.nn.Linear(in_features, layers_sizes[-1]))
    return torch.nn.Sequential(*layers)

def apply_threshold(det_thresh, _scores):
    """ Apply thresholding to detection scores; if stack_K is used and det_thresh is a list, apply to each channel separately """
    if isinstance(det_thresh, list):
        det_thresh = det_thresh[0]
    idx = torch.where(_scores >= det_thresh)
    return idx

def nms(heat, kernel=3):
    """ easy non maximal supression (as in CenterNet) """

    if kernel not in [2, 4]:
        pad = (kernel - 1) // 2
    else:
        if kernel == 2:
            pad = 1
        else:
            pad = 2

    hmax = nn.functional.max_pool2d( heat, (kernel, kernel), stride=1, padding=pad)


        pred = self.mlp(feat)
        return pred

def regression_mlp(layers_sizes):
    """
    Return a fully connected network.
    """
    assert len(layers_sizes) >= 2
    in_features = layers_sizes[0]
    layers = []
    for i in range(1, len(layers_sizes)-1):
        out_features = layers_sizes[i]
        layers.append(torch.nn.Linear(in_features, out_features))
        layers.append(torch.nn.ReLU())
        in_features = out_features
    layers.append(torch.nn.Linear(in_features, layers_sizes[-1]))
    return torch.nn.Sequential(*layers)

def apply_threshold(det_thresh, _scores):
    """ Apply thresholding to detection scores; if stack_K is used and det_thresh is a list, apply to each channel separately """
    if isinstance(det_thresh, list):
        det_thresh = det_thresh[0]
    idx = torch.where(_scores >= det_thresh)
    return idx

def nms(heat, kernel=3):
    """ easy non maximal supression (as in CenterNet) """

    if kernel not in [2, 4]:
        pad = (kernel - 1) // 2
    else:
        if kernel == 2:
            pad = 1
        else:
            pad = 2

    hmax = nn.functional.max_pool2d( heat, (kernel, kernel), stride=1, padding=pad)

    if hmax.shape[2] > heat.shape[2]:
        hmax = hmax[:, :, :heat.shape[2], :heat.shape[3]]

    keep = (hmax == heat).float()

    return heat * keep


nms

    Return a fully connected network.
    """
    assert len(layers_sizes) >= 2
    in_features = layers_sizes[0]
    layers = []
    for i in range(1, len(layers_sizes)-1):
        out_features = layers_sizes[i]
        layers.append(torch.nn.Linear(in_features, out_features))
        layers.append(torch.nn.ReLU())
        in_features = out_features
    layers.append(torch.nn.Linear(in_features, layers_sizes[-1]))
    return torch.nn.Sequential(*layers)

def apply_threshold(det_thresh, _scores):
    """ Apply thresholding to detection scores; if stack_K is used and det_thresh is a list, apply to each channel separately """
    if isinstance(det_thresh, list):
        det_thresh = det_thresh[0]
    idx = torch.where(_scores >= det_thresh)
    return idx

def nms(heat, kernel=3):
    """ easy non maximal supression (as in CenterNet) """

    if kernel not in [2, 4]:
        pad = (kernel - 1) // 2
    else:
        if kernel == 2:
            pad = 1
        else:
            pad = 2

    hmax = nn.functional.max_pool2d( heat, (kernel, kernel), stride=1, padding=pad)

    if hmax.shape[2] > heat.shape[2]:
        hmax = hmax[:, :, :heat.shape[2], :heat.shape[3]]

    keep = (hmax == heat).float()

    return heat * keep


        _x = pk_idx[k,1] # patch center H
        _y = pk_idx[k,0] # patch center W
        # filter out heads out of cropping bounds
        if _x >= 0 and _x < patch_h and _y >= 0 and _y < patch_w:
            if scores[i,_x,_y] == 1:
                smpl_mask[i,j] = 0
                visible_humans[k] = 0
            else:
                scores[i,_x,_y] = 1
        else:
            smpl_mask[i,j] = 0
            visible_humans[k] = 0

    return smpl_mask, visible_humans, scores

import torch.nn as nn
from croco.models.blocks import Mlp_flex

class SMPLDecoder(nn.Module):
    def __init__(
        self,
        hidden_size=768,
        target_dim=1,
        mlp_ratio=1,
        num_layers=2,
    ):
        super().__init__()
        self.mlp = Mlp_flex(
            in_features=hidden_size,
            hidden_features=int(hidden_size * mlp_ratio),
            out_features=target_dim,
            num_layers=num_layers,
            drop=0,
        )

    def forward(
        self,
        feat,
    ):
        """
        feat: BxC
        """

        pred = self.mlp(feat)
        return pred

def regression_mlp(layers_sizes):
    """
    Return a fully connected network.
    """
    assert len(layers_sizes) >= 2
    in_features = layers_sizes[0]
    layers = []


            SMPLX_DIR, 'smplx', gender='neutral', use_pca=False, flat_hand_mean=True, num_betas=11).to(self.device)
        self.smplx_neutral_10 = smplx.create(
            SMPLX_DIR, 'smplx', gender='neutral', use_pca=False, flat_hand_mean=True, num_betas=10).to(self.device)

        # Evaluation
        self.use_fake_K = eval_args.get('use_fake_K', False)
        dataset = eval_args.get('dataset', None)
        if dataset is not None:
            self.smpl = [
                smplx.create(SMPLX_DIR, 'smpl', gender=g).to(self.device) for g in ['neutral', 'male', 'female']]
            self.smpl_faces = {'smpl': self.smpl[0].faces, 'smplx': self.smplx_neutral_11.faces}
            with open(SMPLX2SMPL, 'rb') as f:
                self.smplx2smpl = torch.from_numpy(pickle.load(f)['matrix'].astype(np.float32)).to(self.device)

            if dataset in ['rich']:
                self.smplx = {
                    g: smplx.create(SMPLX_DIR, 'smplx', gender=g, num_pca_comps=12
                                    ).to(self.device) for g in ['male', 'female']}
            self._setup_dataset_config(dataset)

    def _setup_dataset_config(self, dataset):
        self.j_smpl = self.smpl[0].J_regressor[:24]
        if dataset in ['3dpw']:
            h36m_to_14 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9][:14]
            self.j_h36m = torch.Tensor(np.load('src/models/smpl/J_regressor_h36m.npy'))
            self.j_regressor = self.j_h36m[h36m_to_14]
            self.pelvis_idx = [2, 3]
            self.params_type = 'smpl'
        elif dataset in ['bedlam', 'rich']:
            self.j_regressor = self.j_smpl
            self.pelvis_idx = [1, 2]
            self.params_type = 'smplx'
        else:
            self.j_regressor = self.j_smpl
            self.pelvis_idx = [1, 2]
            self.params_type = 'smpl'

    def forward_smpl(self, dataset, smpl_dict, smpl_mask):
        nhv = int(smpl_mask.sum())

        if dataset in ['bedlam']:
            out = self.smplx_neutral_11(
                global_orient=smpl_dict['smplx_root_pose'][smpl_mask].reshape(-1, 3),
                body_pose=smpl_dict['smplx_body_pose'][smpl_mask].reshape(-1, 21*3),
                jaw_pose=smpl_dict['smplx_jaw_pose'][smpl_mask].reshape(-1, 3),
                leye_pose=smpl_dict['smplx_leye_pose'][smpl_mask].reshape(-1, 3),
                reye_pose=smpl_dict['smplx_reye_pose'][smpl_mask].reshape(-1, 3),
                left_hand_pose=smpl_dict['smplx_left_hand_pose'][smpl_mask].reshape(-1, 15*3),
                right_hand_pose=smpl_dict['smplx_right_hand_pose'][smpl_mask].reshape(-1, 15*3),
                betas=smpl_dict['smplx_shape'][smpl_mask].reshape(-1, 11),
                transl=smpl_dict['smplx_transl'][smpl_mask].reshape(-1, 3),
                expression=self.smplx_neutral_11.expression.repeat(nhv, 1),
            )
            verts = out.vertices.reshape(nhv, -1, 3)

        elif dataset in ['3dpw']:


                                    ).to(self.device) for g in ['male', 'female']}
            self._setup_dataset_config(dataset)

    def _setup_dataset_config(self, dataset):
        self.j_smpl = self.smpl[0].J_regressor[:24]
        if dataset in ['3dpw']:
            h36m_to_14 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9][:14]
            self.j_h36m = torch.Tensor(np.load('src/models/smpl/J_regressor_h36m.npy'))
            self.j_regressor = self.j_h36m[h36m_to_14]
            self.pelvis_idx = [2, 3]
            self.params_type = 'smpl'
        elif dataset in ['bedlam', 'rich']:
            self.j_regressor = self.j_smpl
            self.pelvis_idx = [1, 2]
            self.params_type = 'smplx'
        else:
            self.j_regressor = self.j_smpl
            self.pelvis_idx = [1, 2]
            self.params_type = 'smpl'

    def forward_smpl(self, dataset, smpl_dict, smpl_mask):
        nhv = int(smpl_mask.sum())

        if dataset in ['bedlam']:
            out = self.smplx_neutral_11(
                global_orient=smpl_dict['smplx_root_pose'][smpl_mask].reshape(-1, 3),
                body_pose=smpl_dict['smplx_body_pose'][smpl_mask].reshape(-1, 21*3),
                jaw_pose=smpl_dict['smplx_jaw_pose'][smpl_mask].reshape(-1, 3),
                leye_pose=smpl_dict['smplx_leye_pose'][smpl_mask].reshape(-1, 3),
                reye_pose=smpl_dict['smplx_reye_pose'][smpl_mask].reshape(-1, 3),
                left_hand_pose=smpl_dict['smplx_left_hand_pose'][smpl_mask].reshape(-1, 15*3),
                right_hand_pose=smpl_dict['smplx_right_hand_pose'][smpl_mask].reshape(-1, 15*3),
                betas=smpl_dict['smplx_shape'][smpl_mask].reshape(-1, 11),
                transl=smpl_dict['smplx_transl'][smpl_mask].reshape(-1, 3),
                expression=self.smplx_neutral_11.expression.repeat(nhv, 1),
            )
            verts = out.vertices.reshape(nhv, -1, 3)

        elif dataset in ['3dpw']:
            smpl_params = {
                'global_orient': smpl_dict['smpl_root_pose'][smpl_mask].reshape(-1,3),
                'body_pose': smpl_dict['smpl_body_pose'][smpl_mask].reshape(-1,23*3),
                'betas': smpl_dict['smpl_shape'][smpl_mask].reshape(-1,10),
                'transl': smpl_dict['smpl_transl'][smpl_mask].reshape(-1,3),
                }
            out = self.smpl[1](**smpl_params)
            verts = out.vertices.reshape(nhv, -1, 3)

            # update verts/joints if this is not the right gender
            if int(smpl_dict['smpl_gender_id'].max()) == 2:
                out_female = self.smpl[2](**smpl_params)
                idx = torch.where(smpl_dict['smpl_gender_id'] == 2)[1]
                verts[idx] = out_female.vertices.reshape(nhv, -1, 3)[idx]

        elif dataset in ['emdb', 'emdb1', 'emdb2']:
            gender = smpl_dict['smpl_gender_id'].max()
            out = self.smpl[gender](
                global_orient=smpl_dict['smpl_root_pose_w'][smpl_mask].reshape(-1,3),
                body_pose=smpl_dict['smpl_body_pose'][smpl_mask].reshape(-1,23*3),
                betas=smpl_dict['smpl_shape'][smpl_mask].reshape(-1,10),
                transl=smpl_dict['smpl_transl_w'][smpl_mask].reshape(-1,3),
            )
            verts = out.vertices.reshape(nhv, -1, 3) # world space

        elif dataset in ['rich']:
            gender = {1: 'male', 2: 'female'}[int(smpl_dict['smplx_gender_id'].max())]
            out = self.smplx[gender](
                global_orient=smpl_dict['smplx_global_orient'][smpl_mask].reshape(-1,3),
                body_pose=smpl_dict['smplx_body_pose'][smpl_mask].reshape(-1,21*3),
                jaw_pose=torch.zeros([nhv, 3]),
                leye_pose=torch.zeros([nhv, 3]),
                reye_pose=torch.zeros([nhv, 3]),
                left_hand_pose=torch.zeros([nhv, 12]),
                right_hand_pose=torch.zeros([nhv, 12]),
                betas=smpl_dict['smplx_betas'][smpl_mask].reshape(-1,10),
                transl=smpl_dict['smplx_transl'][smpl_mask].reshape(-1,3),
                expression=torch.zeros([nhv, 10]),
            )
            verts = out.vertices.reshape(nhv, -1, 3)

        if self.params_type == 'smplx':
            verts = self.smplx2smpl @ verts
        jts = self.j_regressor @ verts

        return verts, jts

    def update_smpl_gt(self, views):
        target = {}

        batch_size = views[0]["img"].shape[0]

        smpl_keys = [k for k in views[0].keys() if 'smpl' in k]
        smpl_dict = {
            k: (stacked := torch.stack(
                [view.pop(k) for view in views], dim=0)).view(-1, *stacked.shape[2:])
            for k in smpl_keys
        }   # Shape: (num_views * batch_size, 10, ...)
        smpl_mask = smpl_dict['smpl_mask']
        idx_h = torch.where(smpl_mask) # frame_idx, batch_idx, human_idx
        K = torch.stack(
            [view['camera_intrinsics'] for view in views], dim=0
        )
        K = K.view(-1, *K.shape[2:])
        nhv = int(smpl_mask.sum())


            out = self.smplx[gender](
                global_orient=smpl_dict['smplx_global_orient'][smpl_mask].reshape(-1,3),
                body_pose=smpl_dict['smplx_body_pose'][smpl_mask].reshape(-1,21*3),
                jaw_pose=torch.zeros([nhv, 3]),
                leye_pose=torch.zeros([nhv, 3]),
                reye_pose=torch.zeros([nhv, 3]),
                left_hand_pose=torch.zeros([nhv, 12]),
                right_hand_pose=torch.zeros([nhv, 12]),
                betas=smpl_dict['smplx_betas'][smpl_mask].reshape(-1,10),
                transl=smpl_dict['smplx_transl'][smpl_mask].reshape(-1,3),
                expression=torch.zeros([nhv, 10]),
            )
            verts = out.vertices.reshape(nhv, -1, 3)

        if self.params_type == 'smplx':
            verts = self.smplx2smpl @ verts
        jts = self.j_regressor @ verts

        return verts, jts

    def update_smpl_gt(self, views):
        target = {}

        batch_size = views[0]["img"].shape[0]

        smpl_keys = [k for k in views[0].keys() if 'smpl' in k]
        smpl_dict = {
            k: (stacked := torch.stack(
                [view.pop(k) for view in views], dim=0)).view(-1, *stacked.shape[2:])
            for k in smpl_keys
        }   # Shape: (num_views * batch_size, 10, ...)
        smpl_mask = smpl_dict['smpl_mask']
        idx_h = torch.where(smpl_mask) # frame_idx, batch_idx, human_idx
        K = torch.stack(
            [view['camera_intrinsics'] for view in views], dim=0
        )
        K = K.view(-1, *K.shape[2:])
        nhv = int(smpl_mask.sum())

        # Get MHMR input image (high-res, square)
        imgs = torch.stack([view["img"] for view in views], dim=0)
        imgs = imgs.view(-1, *imgs.shape[2:])
        K_mhmr = resize_camera_intrinsics(K, *imgs.shape[2:], self.mhmr_img_res)
        imgs_mhmr = pad_image(imgs, self.mhmr_img_res)

        # SMPLX forward - BEDLAM
        has_smplx_params = 1
        out = self.smplx_neutral_11(
            global_orient=smpl_dict['smplx_root_pose'][smpl_mask].reshape(-1, 3),
            body_pose=smpl_dict['smplx_body_pose'][smpl_mask].reshape(-1, 21*3),
            jaw_pose=smpl_dict['smplx_jaw_pose'][smpl_mask].reshape(-1, 3),
            leye_pose=smpl_dict['smplx_leye_pose'][smpl_mask].reshape(-1, 3),
            reye_pose=smpl_dict['smplx_reye_pose'][smpl_mask].reshape(-1, 3),
            left_hand_pose=smpl_dict['smplx_left_hand_pose'][smpl_mask].reshape(-1, 15*3),
            right_hand_pose=smpl_dict['smplx_right_hand_pose'][smpl_mask].reshape(-1, 15*3),
            betas=smpl_dict['smplx_shape'][smpl_mask].reshape(-1, 11),
            transl=smpl_dict['smplx_transl'][smpl_mask].reshape(-1, 3),
            expression=self.smplx_neutral_11.expression.repeat(nhv, 1),
        )
        verts, jts = out.vertices.reshape(nhv, -1, 3), out.joints.reshape(nhv, -1, 3)

        j2d = perspective_projection(jts, K[idx_h[0]])
        v2d = perspective_projection(verts, K[idx_h[0]])

        # Translation of the primary keypoint
        root_joint_idx = JOINT_NAMES.index(self.person_center)
        target['smpl_transl'] = jts[:,root_joint_idx] # [nhv,3]
        target['smpl_transl_pelvis'] = jts[:,0] # [nhv,3]

        # Fill in target
        target['smpl_v3d'] = verts
        target['smpl_j3d'] = jts
        target['smpl_j2d'] = j2d
        target['smpl_v2d'] = v2d

        if has_smplx_params:
            target['smpl_rotvec'] = torch.cat([smpl_dict['smplx_root_pose'],
                                        smpl_dict['smplx_body_pose'],
                                        smpl_dict['smplx_left_hand_pose'],
                                        smpl_dict['smplx_right_hand_pose'],
                                        smpl_dict['smplx_jaw_pose']],2)[smpl_mask] # [bs,nhmax]
            target['smpl_rotmat'] = roma.rotvec_to_rotmat(target['smpl_rotvec'])
            target['smpl_shape'] = smpl_dict['smplx_shape'][smpl_mask]

        true_shapes = torch.stack([view["true_shape"] for view in views], dim=0)
        if len(torch.unique(true_shapes, dim=0)) != 1:
            raise NotImplementedError

        # Creating the target heatmap for the primary keypoint
        pk = target['smpl_transl'].unsqueeze(1) # (nhv,3)

        # For 512 res (CUT3R, patch_size=16)
        pk_loc = perspective_projection(pk, K[idx_h[0]]).squeeze(1) # original pixel uv coordinates (nhv,2): W, H
        n_patch_16, pk_idx_16 = get_patch_uv(true_shapes[0][0], self.patch_size, pk_loc)
        target['smpl_uv_16'] = pk_idx_16[:, [1, 0]]

        # For 896 res (MHMR, patch_size=14)
        pk_loc_mhmr = perspective_projection(pk, K_mhmr[idx_h[0]]).squeeze(1) # original pixel uv coordinates (nhv,2): W, H
        n_patch_14, pk_idx_14 = get_patch_uv(self.mhmr_img_res, self.bb_patch_size, pk_loc_mhmr)
        smpl_mask_14, visible_humans_14, scores_14 = get_score(n_patch_14, pk_idx_14, smpl_mask.clone())
        target['smpl_uv'] = pk_idx_14[:, [1, 0]]

        # Rebatch and Update with visibility indice
        _target = {}
        num_view = len(views)
        max_humans = smpl_mask_14.shape[1]
        idx_vis = torch.where(visible_humans_14)[0]

        for k, v in target.items():
            full_out = torch.zeros(
                num_view * batch_size, max_humans, *v.shape[1:],
                device=v.device, dtype=v.dtype,
            )
            full_out[smpl_mask_14] = v[idx_vis] # discard unvisible humans due to olccusion
            _target[k] = full_out.chunk(num_view, dim=0) # .view(num_view, batch_size, *full_out.shape[1:])

        _target['smpl_scores'] = scores_14.chunk(num_view, dim=0)
        _target['smpl_mask'] = smpl_mask_14.chunk(num_view, dim=0)
        _target['K_mhmr'] = K_mhmr.chunk(num_view, dim=0)
        _target['img_mhmr'] = imgs_mhmr.chunk(num_view, dim=0)

        if "msk" in views[0]:
            msks = torch.stack([view["msk"] for view in views], dim=0)
            msks = msks.view(-1, *msks.shape[2:])
            msks_mhmr = pad_image(msks, self.mhmr_img_res, pad_value=0.0)  # bs,288,512->bs,896,896
            msks_mhmr = (msks_mhmr > 0.1).float()
            _target['msk_mhmr'] = msks_mhmr.chunk(num_view, dim=0)

        for i, v in enumerate(zip(*_target.values())):
            views[i].update(dict(zip(_target.keys(), v)))

        torch.cuda.empty_cache()

    def update_smpl_gt_eval(self, views, dataset):
        from dust3r.utils.geometry import geotrf

        target = {}
        batch_size = views[0]["img"].shape[0]

        smpl_keys = [k for k in views[0].keys() if 'smpl' in k]
        smpl_dict = {
            k: (stacked := torch.stack(
                [view.pop(k) for view in views], dim=0)).view(-1, *stacked.shape[2:])
            for k in smpl_keys
        }   # Shape: (num_views * batch_size, 10, ...)
        smpl_mask = smpl_dict['smpl_mask']
        idx_h = torch.where(smpl_mask) # frame_idx, batch_idx, human_idx
        K = torch.stack([view['camera_intrinsics'] for view in views], dim=0)
        K = K.view(-1, *K.shape[2:])

        # Get MHMR input image (high-res, square)
        imgs = torch.stack([view["img"] for view in views], dim=0)


            full_out[smpl_mask_14] = v[idx_vis] # discard unvisible humans due to olccusion
            _target[k] = full_out.chunk(num_view, dim=0) # .view(num_view, batch_size, *full_out.shape[1:])

        _target['smpl_scores'] = scores_14.chunk(num_view, dim=0)
        _target['smpl_mask'] = smpl_mask_14.chunk(num_view, dim=0)
        _target['K_mhmr'] = K_mhmr.chunk(num_view, dim=0)
        _target['img_mhmr'] = imgs_mhmr.chunk(num_view, dim=0)

        if "msk" in views[0]:
            msks = torch.stack([view["msk"] for view in views], dim=0)
            msks = msks.view(-1, *msks.shape[2:])
            msks_mhmr = pad_image(msks, self.mhmr_img_res, pad_value=0.0)  # bs,288,512->bs,896,896
            msks_mhmr = (msks_mhmr > 0.1).float()
            _target['msk_mhmr'] = msks_mhmr.chunk(num_view, dim=0)

        for i, v in enumerate(zip(*_target.values())):
            views[i].update(dict(zip(_target.keys(), v)))

        torch.cuda.empty_cache()

    def update_smpl_gt_eval(self, views, dataset):
        from dust3r.utils.geometry import geotrf

        target = {}
        batch_size = views[0]["img"].shape[0]

        smpl_keys = [k for k in views[0].keys() if 'smpl' in k]
        smpl_dict = {
            k: (stacked := torch.stack(
                [view.pop(k) for view in views], dim=0)).view(-1, *stacked.shape[2:])
            for k in smpl_keys
        }   # Shape: (num_views * batch_size, 10, ...)
        smpl_mask = smpl_dict['smpl_mask']
        idx_h = torch.where(smpl_mask) # frame_idx, batch_idx, human_idx
        K = torch.stack([view['camera_intrinsics'] for view in views], dim=0)
        K = K.view(-1, *K.shape[2:])

        # Get MHMR input image (high-res, square)
        imgs = torch.stack([view["img"] for view in views], dim=0)
        imgs = imgs.view(-1, *imgs.shape[2:])
        K_mhmr = resize_camera_intrinsics(K, *imgs.shape[2:], self.mhmr_img_res)
        imgs_mhmr = pad_image(imgs, self.mhmr_img_res)

        verts, jts = self.forward_smpl(dataset, smpl_dict, smpl_mask)

        if dataset in ['emdb', 'emdb1', 'emdb2', 'rich']:
            target['smpl_v3d_w'] = verts
            target['smpl_j3d_w'] = jts
            T_w2c = torch.stack([view['T_w2c'] for view in views], dim=0)
            T_w2c = T_w2c.view(-1, *T_w2c.shape[2:])
            target['smpl_v3d_c'] = geotrf(T_w2c[idx_h[0]], verts)
            target['smpl_j3d_c'] = geotrf(T_w2c[idx_h[0]], jts)

        else:
            target['smpl_v3d_c'] = verts
            target['smpl_j3d_c'] = jts
            T_c2w = torch.stack([view['camera_pose'] for view in views], dim=0)
            T_c2w = T_c2w.view(-1, *T_c2w.shape[2:])
            target['smpl_v3d_w'] = geotrf(T_c2w[idx_h[0]], verts)
            target['smpl_j3d_w'] = geotrf(T_c2w[idx_h[0]], jts)

        target['smpl_j2d'] = perspective_projection(target['smpl_j3d_c'], K[idx_h[0]])
        target['smpl_v2d'] = perspective_projection(target['smpl_v3d_c'], K[idx_h[0]])

        # Rebatch and Update with visibility indice
        _target = {}
        num_view = len(views)
        max_humans = smpl_mask.shape[1]
        for k, v in target.items():
            full_out = torch.zeros(
                num_view * batch_size, max_humans, *v.shape[1:],
                device=v.device, dtype=v.dtype,
            )
            full_out[smpl_mask] = v # discard unvisible humans due to olccusion
            _target[k] = full_out.chunk(num_view, dim=0) # .view(num_view, batch_size, *full_out.shape[1:])

        if self.use_fake_K:
            K_mhmr = get_camera_parameters(self.mhmr_img_res, device=K.device) # if use pseudo K
            K_mhmr = K_mhmr.expand(K.shape[0], -1, -1)

        _target['smpl_mask'] = smpl_mask.chunk(num_view, dim=0)
        _target['K_mhmr'] = K_mhmr.chunk(num_view, dim=0)
        _target['img_mhmr'] = imgs_mhmr.chunk(num_view, dim=0)

        if "msk" in views[0]:
            msks = torch.stack([view["msk"] for view in views], dim=0)
            msks = msks.view(-1, *msks.shape[2:])
            msks_mhmr = pad_image(msks, self.mhmr_img_res, pad_value=0.0)  # bs,288,512->bs,896,896
            msks_mhmr = (msks_mhmr > 0.1).float()
            _target['msk_mhmr'] = msks_mhmr.chunk(num_view, dim=0)

        for i, v in enumerate(zip(*_target.values())):
            views[i].update(dict(zip(_target.keys(), v)))

        torch.cuda.empty_cache()

def get_patch_uv(imgshape, patch_size, pk_loc):
    n_patch = imgshape // patch_size  # H, W
    pk_idx = (pk_loc // patch_size).int()
    return n_patch, pk_idx

def get_score(n_patch, pk_idx, smpl_mask):
    # Scores & updating valid_humans according to occlusion - wap X and Y for scores only
    idx_h = torch.where(smpl_mask)
    nhv = int(smpl_mask.sum())
    bs = smpl_mask.shape[0]
    device = smpl_mask.device

    if isinstance(n_patch, (int, float)):
        patch_h, patch_w = int(n_patch), int(n_patch)
    else:
        patch_h, patch_w = n_patch[0], n_patch[1]

    scores = torch.zeros((bs, patch_h, patch_w)).to(device)


forward

import torch.nn as nn
from croco.models.blocks import Mlp_flex

class SMPLDecoder(nn.Module):
    def __init__(
        self,
        hidden_size=768,
        target_dim=1,
        mlp_ratio=1,
        num_layers=2,
    ):
        super().__init__()
        self.mlp = Mlp_flex(
            in_features=hidden_size,
            hidden_features=int(hidden_size * mlp_ratio),
            out_features=target_dim,
            num_layers=num_layers,
            drop=0,
        )

    def forward(
        self,
        feat,
    ):
        """
        feat: BxC
        """

        pred = self.mlp(feat)
        return pred

def regression_mlp(layers_sizes):
    """
    Return a fully connected network.
    """
    assert len(layers_sizes) >= 2
    in_features = layers_sizes[0]
    layers = []
    for i in range(1, len(layers_sizes)-1):
        out_features = layers_sizes[i]
        layers.append(torch.nn.Linear(in_features, out_features))
        layers.append(torch.nn.ReLU())
        in_features = out_features
    layers.append(torch.nn.Linear(in_features, layers_sizes[-1]))
    return torch.nn.Sequential(*layers)

def apply_threshold(det_thresh, _scores):
    """ Apply thresholding to detection scores; if stack_K is used and det_thresh is a list, apply to each channel separately """
    if isinstance(det_thresh, list):


import tqdm
import torch
from dust3r.utils.device import to_cpu, collate_with_cat
from dust3r.utils.misc import invalid_to_nans
from dust3r.utils.geometry import depthmap_to_pts3d, geotrf
from dust3r.model import ARCroco3DStereo
from dust3r.smpl_model import SMPLModel
from accelerate import Accelerator
import re

def custom_sort_key(key):
    text = key.split("/")
    if len(text) > 1:
        text, num = text[0], text[-1]
        return (text, int(num))
    else:
        return (key, -1)

def merge_chunk_dict(old_dict, curr_dict, add_number):
    new_dict = {}
    for key, value in curr_dict.items():

        match = re.search(r"(\d+)$", key)
        if match:

            num_part = int(match.group()) + add_number

            new_key = re.sub(r"(\d+)$", str(num_part), key, 1)
            new_dict[new_key] = value
        else:
            new_dict[key] = value
    new_dict = old_dict | new_dict
    return {k: new_dict[k] for k in sorted(new_dict.keys(), key=custom_sort_key)}

def _interleave_imgs(img1, img2):
    res = {}
    for key, value1 in img1.items():
        value2 = img2[key]
        if isinstance(value1, torch.Tensor):
            value = torch.stack((value1, value2), dim=1).flatten(0, 1)
        else:
            value = [x for pair in zip(value1, value2) for x in pair]
        res[key] = value
    return res

def make_batch_symmetric(batch):
    view1, view2 = batch
    view1, view2 = (_interleave_imgs(view1, view2), _interleave_imgs(view2, view1))
    return view1, view2

def loss_of_one_batch(
    batch,
    model,
    criterion,
    accelerator: Accelerator,
    symmetrize_batch=False,
    use_amp=False,
    ret=None,
    img_mask=None,
    inference=False,
    smpl_model: SMPLModel = None
):
    if len(batch) > 2:
        assert (
            symmetrize_batch is False
        ), "cannot symmetrize batch with more than 2 views"
    if symmetrize_batch:
        batch = make_batch_symmetric(batch)

    with torch.cuda.amp.autocast(enabled=not inference):
        if inference:
            output, state_args = model(batch, ret_state=True, inference=True)
            preds, batch = output.ress, output.views
            result = dict(views=batch, pred=preds)
            return result[ret] if ret else result, state_args
        else:
            smpl_model.update_smpl_gt(batch)
            output = model(batch)
            preds, batch = output.ress, output.views

        with torch.cuda.amp.autocast(enabled=False):
            loss = criterion(batch, preds) if criterion is not None else None

    result = dict(views=batch, pred=preds, loss=loss)
    return result[ret] if ret else result

@torch.no_grad()
def inference(groups, model, device, verbose=True):
    ignore_keys = set(
        ["depthmap", "dataset", "label", "instance", "idx", "true_shape", "rng"]
    )
    for view in groups:
        for name in view.keys():  # pseudo_focal
            if name in ignore_keys:
                continue
            if isinstance(view[name], tuple) or isinstance(view[name], list):
                view[name] = [x.to(device, non_blocking=True) for x in view[name]]
            else:
                view[name] = view[name].to(device, non_blocking=True)

    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    res, state_args = loss_of_one_batch(groups, model, None, None, inference=True)
    result = to_cpu(res)
    return result, state_args

@torch.no_grad()
def inference_step(view, state_args, model, device, verbose=True):
    ignore_keys = set(
        ["depthmap", "dataset", "label", "instance", "idx", "true_shape", "rng"]
    )
    for name in view.keys():  # pseudo_focal
        if name in ignore_keys:
            continue
        if isinstance(view[name], tuple) or isinstance(view[name], list):
            view[name] = [x.to(device, non_blocking=True) for x in view[name]]
        else:
            view[name] = view[name].to(device, non_blocking=True)

    with torch.cuda.amp.autocast(enabled=False):
        state_feat, state_pos, init_state_feat, mem, init_mem = state_args
        pred, _ = model.inference_step(
            view, state_feat, state_pos, init_state_feat, mem, init_mem
        )

    res = dict(pred=pred)
    result = to_cpu(res)
    return result

@torch.no_grad()
def inference_recurrent(groups, model, device, verbose=True):
    ignore_keys = set(
        ["depthmap", "dataset", "label", "instance", "idx", "true_shape", "rng"]
    )
    for view in groups:
        for name in view.keys():  # pseudo_focal
            if name in ignore_keys:
                continue
            if isinstance(view[name], tuple) or isinstance(view[name], list):
                view[name] = [x.to(device, non_blocking=True) for x in view[name]]
            else:
                view[name] = view[name].to(device, non_blocking=True)

    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    with torch.cuda.amp.autocast(enabled=False):
        preds, batch, state_args = model.forward_recurrent(
            groups, device, ret_state=True
        )
        res = dict(views=batch, pred=preds)
    result = to_cpu(res)
    return result, state_args

@torch.no_grad()
def inference_recurrent_lighter(groups, model, device, verbose=True, is_naive=False, use_ttt3r=False):
    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    with torch.cuda.amp.autocast(enabled=False):
        if is_naive:
            preds, batch, state_args = model.forward_recurrent_lighter_naive(
            groups, device, ret_state=True, use_ttt3r=use_ttt3r
        )
        else:
            preds, batch, state_args = model.forward_recurrent_lighter(
                groups, device, ret_state=True, use_ttt3r=use_ttt3r
            )
        res = dict(views=batch, pred=preds)
    return res, state_args

def check_if_same_size(pairs):
    shapes1 = [img1["img"].shape[-2:] for img1, img2 in pairs]
    shapes2 = [img2["img"].shape[-2:] for img1, img2 in pairs]
    return all(shapes1[0] == s for s in shapes1) and all(
        shapes2[0] == s for s in shapes2
    )

def get_pred_pts3d(gt, pred, use_pose=False, inplace=False):
    if "depth" in pred and "pseudo_focal" in pred:
        try:
            pp = gt["camera_intrinsics"][..., :2, 2]
        except KeyError:
            pp = None
        pts3d = depthmap_to_pts3d(**pred, pp=pp)

    elif "pts3d" in pred:

        pts3d = pred["pts3d"]

    elif "pts3d_in_other_view" in pred:

        assert use_pose is True
        return (
            pred["pts3d_in_other_view"]
            if inplace
            else pred["pts3d_in_other_view"].clone()
        )

    if use_pose:
        camera_pose = pred.get("camera_pose")
        assert camera_pose is not None
        pts3d = geotrf(camera_pose, pts3d)

    return pts3d

def find_opt_scaling(
    gt_pts1,
    gt_pts2,
    pr_pts1,
    pr_pts2=None,
    fit_mode="weiszfeld_stop_grad",
    valid1=None,
    valid2=None,
):
    assert gt_pts1.ndim == pr_pts1.ndim == 4
    assert gt_pts1.shape == pr_pts1.shape
    if gt_pts2 is not None:
        assert gt_pts2.ndim == pr_pts2.ndim == 4
        assert gt_pts2.shape == pr_pts2.shape

    nan_gt_pts1 = invalid_to_nans(gt_pts1, valid1).flatten(1, 2)
    nan_gt_pts2 = (
        invalid_to_nans(gt_pts2, valid2).flatten(1, 2) if gt_pts2 is not None else None
    )

    pr_pts1 = invalid_to_nans(pr_pts1, valid1).flatten(1, 2)
    pr_pts2 = (
        invalid_to_nans(pr_pts2, valid2).flatten(1, 2) if pr_pts2 is not None else None
    )

    all_gt = (
        torch.cat((nan_gt_pts1, nan_gt_pts2), dim=1)
        if gt_pts2 is not None
        else nan_gt_pts1
    )
    all_pr = torch.cat((pr_pts1, pr_pts2), dim=1) if pr_pts2 is not None else pr_pts1

    dot_gt_pr = (all_pr * all_gt).sum(dim=-1)
    dot_gt_gt = all_gt.square().sum(dim=-1)

    if fit_mode.startswith("avg"):

        scaling = dot_gt_pr.nanmean(dim=1) / dot_gt_gt.nanmean(dim=1)
    elif fit_mode.startswith("median"):
        scaling = (dot_gt_pr / dot_gt_gt).nanmedian(dim=1).values
    elif fit_mode.startswith("weiszfeld"):

        scaling = dot_gt_pr.nanmean(dim=1) / dot_gt_gt.nanmean(dim=1)

        for iter in range(10):

            dis = (all_pr - scaling.view(-1, 1, 1) * all_gt).norm(dim=-1)

            w = dis.clip_(min=1e-8).reciprocal()

            scaling = (w * dot_gt_pr).nanmean(dim=1) / (w * dot_gt_gt).nanmean(dim=1)
    else:
        raise ValueError(f"bad {fit_mode=}")

    if fit_mode.endswith("stop_grad"):
        scaling = scaling.detach()

    scaling = scaling.clip(min=1e-3)

    return scaling


import tqdm
import torch
from dust3r.utils.device import to_cpu, collate_with_cat
from dust3r.utils.misc import invalid_to_nans
from dust3r.utils.geometry import depthmap_to_pts3d, geotrf
from dust3r.model import ARCroco3DStereo
from dust3r.smpl_model import SMPLModel
from accelerate import Accelerator
import re

def custom_sort_key(key):
    text = key.split("/")
    if len(text) > 1:
        text, num = text[0], text[-1]
        return (text, int(num))
    else:
        return (key, -1)

def merge_chunk_dict(old_dict, curr_dict, add_number):
    new_dict = {}
    for key, value in curr_dict.items():

        match = re.search(r"(\d+)$", key)
        if match:

            num_part = int(match.group()) + add_number

            new_key = re.sub(r"(\d+)$", str(num_part), key, 1)
            new_dict[new_key] = value
        else:
            new_dict[key] = value
    new_dict = old_dict | new_dict
    return {k: new_dict[k] for k in sorted(new_dict.keys(), key=custom_sort_key)}

def _interleave_imgs(img1, img2):


import tqdm
import torch
from dust3r.utils.device import to_cpu, collate_with_cat
from dust3r.utils.misc import invalid_to_nans
from dust3r.utils.geometry import depthmap_to_pts3d, geotrf
from dust3r.model import ARCroco3DStereo
from dust3r.smpl_model import SMPLModel
from accelerate import Accelerator
import re

def custom_sort_key(key):
    text = key.split("/")
    if len(text) > 1:
        text, num = text[0], text[-1]
        return (text, int(num))
    else:
        return (key, -1)

def merge_chunk_dict(old_dict, curr_dict, add_number):
    new_dict = {}
    for key, value in curr_dict.items():

        match = re.search(r"(\d+)$", key)
        if match:

            num_part = int(match.group()) + add_number

            new_key = re.sub(r"(\d+)$", str(num_part), key, 1)
            new_dict[new_key] = value
        else:
            new_dict[key] = value
    new_dict = old_dict | new_dict
    return {k: new_dict[k] for k in sorted(new_dict.keys(), key=custom_sort_key)}

def _interleave_imgs(img1, img2):
    res = {}
    for key, value1 in img1.items():
        value2 = img2[key]
        if isinstance(value1, torch.Tensor):
            value = torch.stack((value1, value2), dim=1).flatten(0, 1)
        else:
            value = [x for pair in zip(value1, value2) for x in pair]
        res[key] = value
    return res

def make_batch_symmetric(batch):
    view1, view2 = batch
    view1, view2 = (_interleave_imgs(view1, view2), _interleave_imgs(view2, view1))
    return view1, view2


        return (key, -1)

def merge_chunk_dict(old_dict, curr_dict, add_number):
    new_dict = {}
    for key, value in curr_dict.items():

        match = re.search(r"(\d+)$", key)
        if match:

            num_part = int(match.group()) + add_number

            new_key = re.sub(r"(\d+)$", str(num_part), key, 1)
            new_dict[new_key] = value
        else:
            new_dict[key] = value
    new_dict = old_dict | new_dict
    return {k: new_dict[k] for k in sorted(new_dict.keys(), key=custom_sort_key)}

def _interleave_imgs(img1, img2):
    res = {}
    for key, value1 in img1.items():
        value2 = img2[key]
        if isinstance(value1, torch.Tensor):
            value = torch.stack((value1, value2), dim=1).flatten(0, 1)
        else:
            value = [x for pair in zip(value1, value2) for x in pair]
        res[key] = value
    return res

def make_batch_symmetric(batch):
    view1, view2 = batch
    view1, view2 = (_interleave_imgs(view1, view2), _interleave_imgs(view2, view1))
    return view1, view2

def loss_of_one_batch(
    batch,
    model,
    criterion,
    accelerator: Accelerator,
    symmetrize_batch=False,
    use_amp=False,
    ret=None,
    img_mask=None,
    inference=False,
    smpl_model: SMPLModel = None
):


            new_key = re.sub(r"(\d+)$", str(num_part), key, 1)
            new_dict[new_key] = value
        else:
            new_dict[key] = value
    new_dict = old_dict | new_dict
    return {k: new_dict[k] for k in sorted(new_dict.keys(), key=custom_sort_key)}

def _interleave_imgs(img1, img2):
    res = {}
    for key, value1 in img1.items():
        value2 = img2[key]
        if isinstance(value1, torch.Tensor):
            value = torch.stack((value1, value2), dim=1).flatten(0, 1)
        else:
            value = [x for pair in zip(value1, value2) for x in pair]
        res[key] = value
    return res

def make_batch_symmetric(batch):
    view1, view2 = batch
    view1, view2 = (_interleave_imgs(view1, view2), _interleave_imgs(view2, view1))
    return view1, view2

def loss_of_one_batch(
    batch,
    model,
    criterion,
    accelerator: Accelerator,
    symmetrize_batch=False,
    use_amp=False,
    ret=None,
    img_mask=None,
    inference=False,
    smpl_model: SMPLModel = None
):
    if len(batch) > 2:
        assert (
            symmetrize_batch is False
        ), "cannot symmetrize batch with more than 2 views"
    if symmetrize_batch:
        batch = make_batch_symmetric(batch)


def _interleave_imgs(img1, img2):
    res = {}
    for key, value1 in img1.items():
        value2 = img2[key]
        if isinstance(value1, torch.Tensor):
            value = torch.stack((value1, value2), dim=1).flatten(0, 1)
        else:
            value = [x for pair in zip(value1, value2) for x in pair]
        res[key] = value
    return res

def make_batch_symmetric(batch):
    view1, view2 = batch
    view1, view2 = (_interleave_imgs(view1, view2), _interleave_imgs(view2, view1))
    return view1, view2

def loss_of_one_batch(
    batch,
    model,
    criterion,
    accelerator: Accelerator,
    symmetrize_batch=False,
    use_amp=False,
    ret=None,
    img_mask=None,
    inference=False,
    smpl_model: SMPLModel = None
):
    if len(batch) > 2:
        assert (
            symmetrize_batch is False
        ), "cannot symmetrize batch with more than 2 views"
    if symmetrize_batch:
        batch = make_batch_symmetric(batch)

    with torch.cuda.amp.autocast(enabled=not inference):
        if inference:
            output, state_args = model(batch, ret_state=True, inference=True)
            preds, batch = output.ress, output.views
            result = dict(views=batch, pred=preds)
            return result[ret] if ret else result, state_args
        else:
            smpl_model.update_smpl_gt(batch)
            output = model(batch)
            preds, batch = output.ress, output.views

        with torch.cuda.amp.autocast(enabled=False):
            loss = criterion(batch, preds) if criterion is not None else None

    result = dict(views=batch, pred=preds, loss=loss)
    return result[ret] if ret else result

@torch.no_grad()
def inference(groups, model, device, verbose=True):
    ignore_keys = set(
        ["depthmap", "dataset", "label", "instance", "idx", "true_shape", "rng"]
    )
    for view in groups:
        for name in view.keys():  # pseudo_focal
            if name in ignore_keys:
                continue
            if isinstance(view[name], tuple) or isinstance(view[name], list):
                view[name] = [x.to(device, non_blocking=True) for x in view[name]]
            else:
                view[name] = view[name].to(device, non_blocking=True)

    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    res, state_args = loss_of_one_batch(groups, model, None, None, inference=True)
    result = to_cpu(res)


inference

        batch = make_batch_symmetric(batch)

    with torch.cuda.amp.autocast(enabled=not inference):
        if inference:
            output, state_args = model(batch, ret_state=True, inference=True)
            preds, batch = output.ress, output.views
            result = dict(views=batch, pred=preds)
            return result[ret] if ret else result, state_args
        else:
            smpl_model.update_smpl_gt(batch)
            output = model(batch)
            preds, batch = output.ress, output.views

        with torch.cuda.amp.autocast(enabled=False):
            loss = criterion(batch, preds) if criterion is not None else None

    result = dict(views=batch, pred=preds, loss=loss)
    return result[ret] if ret else result

@torch.no_grad()
def inference(groups, model, device, verbose=True):
    ignore_keys = set(
        ["depthmap", "dataset", "label", "instance", "idx", "true_shape", "rng"]
    )
    for view in groups:
        for name in view.keys():  # pseudo_focal
            if name in ignore_keys:
                continue
            if isinstance(view[name], tuple) or isinstance(view[name], list):
                view[name] = [x.to(device, non_blocking=True) for x in view[name]]
            else:
                view[name] = view[name].to(device, non_blocking=True)

    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    res, state_args = loss_of_one_batch(groups, model, None, None, inference=True)
    result = to_cpu(res)
    return result, state_args

@torch.no_grad()
def inference_step(view, state_args, model, device, verbose=True):
    ignore_keys = set(
        ["depthmap", "dataset", "label", "instance", "idx", "true_shape", "rng"]
    )
    for name in view.keys():  # pseudo_focal
        if name in ignore_keys:
            continue
        if isinstance(view[name], tuple) or isinstance(view[name], list):
            view[name] = [x.to(device, non_blocking=True) for x in view[name]]
        else:
            view[name] = view[name].to(device, non_blocking=True)

    with torch.cuda.amp.autocast(enabled=False):
        state_feat, state_pos, init_state_feat, mem, init_mem = state_args
        pred, _ = model.inference_step(
            view, state_feat, state_pos, init_state_feat, mem, init_mem
        )


        ["depthmap", "dataset", "label", "instance", "idx", "true_shape", "rng"]
    )
    for view in groups:
        for name in view.keys():  # pseudo_focal
            if name in ignore_keys:
                continue
            if isinstance(view[name], tuple) or isinstance(view[name], list):
                view[name] = [x.to(device, non_blocking=True) for x in view[name]]
            else:
                view[name] = view[name].to(device, non_blocking=True)

    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    res, state_args = loss_of_one_batch(groups, model, None, None, inference=True)
    result = to_cpu(res)
    return result, state_args

@torch.no_grad()
def inference_step(view, state_args, model, device, verbose=True):
    ignore_keys = set(
        ["depthmap", "dataset", "label", "instance", "idx", "true_shape", "rng"]
    )
    for name in view.keys():  # pseudo_focal
        if name in ignore_keys:
            continue
        if isinstance(view[name], tuple) or isinstance(view[name], list):
            view[name] = [x.to(device, non_blocking=True) for x in view[name]]
        else:
            view[name] = view[name].to(device, non_blocking=True)

    with torch.cuda.amp.autocast(enabled=False):
        state_feat, state_pos, init_state_feat, mem, init_mem = state_args
        pred, _ = model.inference_step(
            view, state_feat, state_pos, init_state_feat, mem, init_mem
        )

    res = dict(pred=pred)
    result = to_cpu(res)
    return result

@torch.no_grad()
def inference_recurrent(groups, model, device, verbose=True):
    ignore_keys = set(
        ["depthmap", "dataset", "label", "instance", "idx", "true_shape", "rng"]
    )
    for view in groups:
        for name in view.keys():  # pseudo_focal
            if name in ignore_keys:
                continue
            if isinstance(view[name], tuple) or isinstance(view[name], list):
                view[name] = [x.to(device, non_blocking=True) for x in view[name]]
            else:
                view[name] = view[name].to(device, non_blocking=True)

    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    with torch.cuda.amp.autocast(enabled=False):


    for name in view.keys():  # pseudo_focal
        if name in ignore_keys:
            continue
        if isinstance(view[name], tuple) or isinstance(view[name], list):
            view[name] = [x.to(device, non_blocking=True) for x in view[name]]
        else:
            view[name] = view[name].to(device, non_blocking=True)

    with torch.cuda.amp.autocast(enabled=False):
        state_feat, state_pos, init_state_feat, mem, init_mem = state_args
        pred, _ = model.inference_step(
            view, state_feat, state_pos, init_state_feat, mem, init_mem
        )

    res = dict(pred=pred)
    result = to_cpu(res)
    return result

@torch.no_grad()
def inference_recurrent(groups, model, device, verbose=True):
    ignore_keys = set(
        ["depthmap", "dataset", "label", "instance", "idx", "true_shape", "rng"]
    )
    for view in groups:
        for name in view.keys():  # pseudo_focal
            if name in ignore_keys:
                continue
            if isinstance(view[name], tuple) or isinstance(view[name], list):
                view[name] = [x.to(device, non_blocking=True) for x in view[name]]
            else:
                view[name] = view[name].to(device, non_blocking=True)

    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    with torch.cuda.amp.autocast(enabled=False):
        preds, batch, state_args = model.forward_recurrent(
            groups, device, ret_state=True
        )
        res = dict(views=batch, pred=preds)
    result = to_cpu(res)
    return result, state_args

@torch.no_grad()
def inference_recurrent_lighter(groups, model, device, verbose=True, is_naive=False, use_ttt3r=False):
    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    with torch.cuda.amp.autocast(enabled=False):
        if is_naive:
            preds, batch, state_args = model.forward_recurrent_lighter_naive(
            groups, device, ret_state=True, use_ttt3r=use_ttt3r
        )
        else:
            preds, batch, state_args = model.forward_recurrent_lighter(
                groups, device, ret_state=True, use_ttt3r=use_ttt3r
            )
        res = dict(views=batch, pred=preds)
    return res, state_args

def check_if_same_size(pairs):
    shapes1 = [img1["img"].shape[-2:] for img1, img2 in pairs]


        for name in view.keys():  # pseudo_focal
            if name in ignore_keys:
                continue
            if isinstance(view[name], tuple) or isinstance(view[name], list):
                view[name] = [x.to(device, non_blocking=True) for x in view[name]]
            else:
                view[name] = view[name].to(device, non_blocking=True)

    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    with torch.cuda.amp.autocast(enabled=False):
        preds, batch, state_args = model.forward_recurrent(
            groups, device, ret_state=True
        )
        res = dict(views=batch, pred=preds)
    result = to_cpu(res)
    return result, state_args

@torch.no_grad()
def inference_recurrent_lighter(groups, model, device, verbose=True, is_naive=False, use_ttt3r=False):
    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    with torch.cuda.amp.autocast(enabled=False):
        if is_naive:
            preds, batch, state_args = model.forward_recurrent_lighter_naive(
            groups, device, ret_state=True, use_ttt3r=use_ttt3r
        )
        else:
            preds, batch, state_args = model.forward_recurrent_lighter(
                groups, device, ret_state=True, use_ttt3r=use_ttt3r
            )
        res = dict(views=batch, pred=preds)
    return res, state_args

def check_if_same_size(pairs):
    shapes1 = [img1["img"].shape[-2:] for img1, img2 in pairs]
    shapes2 = [img2["img"].shape[-2:] for img1, img2 in pairs]
    return all(shapes1[0] == s for s in shapes1) and all(
        shapes2[0] == s for s in shapes2
    )

def get_pred_pts3d(gt, pred, use_pose=False, inplace=False):
    if "depth" in pred and "pseudo_focal" in pred:
        try:
            pp = gt["camera_intrinsics"][..., :2, 2]
        except KeyError:
            pp = None
        pts3d = depthmap_to_pts3d(**pred, pp=pp)

    elif "pts3d" in pred:

        pts3d = pred["pts3d"]


    result = to_cpu(res)
    return result, state_args

@torch.no_grad()
def inference_recurrent_lighter(groups, model, device, verbose=True, is_naive=False, use_ttt3r=False):
    if verbose:
        print(f">> Inference with model on {len(groups)} image/raymaps")

    with torch.cuda.amp.autocast(enabled=False):
        if is_naive:
            preds, batch, state_args = model.forward_recurrent_lighter_naive(
            groups, device, ret_state=True, use_ttt3r=use_ttt3r
        )
        else:
            preds, batch, state_args = model.forward_recurrent_lighter(
                groups, device, ret_state=True, use_ttt3r=use_ttt3r
            )
        res = dict(views=batch, pred=preds)
    return res, state_args

def check_if_same_size(pairs):
    shapes1 = [img1["img"].shape[-2:] for img1, img2 in pairs]
    shapes2 = [img2["img"].shape[-2:] for img1, img2 in pairs]
    return all(shapes1[0] == s for s in shapes1) and all(
        shapes2[0] == s for s in shapes2
    )

def get_pred_pts3d(gt, pred, use_pose=False, inplace=False):
    if "depth" in pred and "pseudo_focal" in pred:
        try:
            pp = gt["camera_intrinsics"][..., :2, 2]
        except KeyError:
            pp = None
        pts3d = depthmap_to_pts3d(**pred, pp=pp)

    elif "pts3d" in pred:

        pts3d = pred["pts3d"]

    elif "pts3d_in_other_view" in pred:

        assert use_pose is True
        return (
            pred["pts3d_in_other_view"]
            if inplace


    with torch.cuda.amp.autocast(enabled=False):
        if is_naive:
            preds, batch, state_args = model.forward_recurrent_lighter_naive(
            groups, device, ret_state=True, use_ttt3r=use_ttt3r
        )
        else:
            preds, batch, state_args = model.forward_recurrent_lighter(
                groups, device, ret_state=True, use_ttt3r=use_ttt3r
            )
        res = dict(views=batch, pred=preds)
    return res, state_args

def check_if_same_size(pairs):
    shapes1 = [img1["img"].shape[-2:] for img1, img2 in pairs]
    shapes2 = [img2["img"].shape[-2:] for img1, img2 in pairs]
    return all(shapes1[0] == s for s in shapes1) and all(
        shapes2[0] == s for s in shapes2
    )

def get_pred_pts3d(gt, pred, use_pose=False, inplace=False):
    if "depth" in pred and "pseudo_focal" in pred:
        try:
            pp = gt["camera_intrinsics"][..., :2, 2]
        except KeyError:
            pp = None
        pts3d = depthmap_to_pts3d(**pred, pp=pp)

    elif "pts3d" in pred:

        pts3d = pred["pts3d"]

    elif "pts3d_in_other_view" in pred:

        assert use_pose is True
        return (
            pred["pts3d_in_other_view"]
            if inplace
            else pred["pts3d_in_other_view"].clone()
        )

    if use_pose:
        camera_pose = pred.get("camera_pose")
        assert camera_pose is not None
        pts3d = geotrf(camera_pose, pts3d)

    return pts3d

def find_opt_scaling(
    gt_pts1,
    gt_pts2,
    pr_pts1,
    pr_pts2=None,
    fit_mode="weiszfeld_stop_grad",
    valid1=None,
    valid2=None,
):
    assert gt_pts1.ndim == pr_pts1.ndim == 4
    assert gt_pts1.shape == pr_pts1.shape
    if gt_pts2 is not None:
        assert gt_pts2.ndim == pr_pts2.ndim == 4
        assert gt_pts2.shape == pr_pts2.shape

    nan_gt_pts1 = invalid_to_nans(gt_pts1, valid1).flatten(1, 2)
    nan_gt_pts2 = (
        invalid_to_nans(gt_pts2, valid2).flatten(1, 2) if gt_pts2 is not None else None


        pts3d = pred["pts3d"]

    elif "pts3d_in_other_view" in pred:

        assert use_pose is True
        return (
            pred["pts3d_in_other_view"]
            if inplace
            else pred["pts3d_in_other_view"].clone()
        )

    if use_pose:
        camera_pose = pred.get("camera_pose")
        assert camera_pose is not None
        pts3d = geotrf(camera_pose, pts3d)

    return pts3d

def find_opt_scaling(
    gt_pts1,
    gt_pts2,
    pr_pts1,
    pr_pts2=None,
    fit_mode="weiszfeld_stop_grad",
    valid1=None,
    valid2=None,
):
    assert gt_pts1.ndim == pr_pts1.ndim == 4
    assert gt_pts1.shape == pr_pts1.shape
    if gt_pts2 is not None:
        assert gt_pts2.ndim == pr_pts2.ndim == 4
        assert gt_pts2.shape == pr_pts2.shape

    nan_gt_pts1 = invalid_to_nans(gt_pts1, valid1).flatten(1, 2)
    nan_gt_pts2 = (
        invalid_to_nans(gt_pts2, valid2).flatten(1, 2) if gt_pts2 is not None else None
    )

    pr_pts1 = invalid_to_nans(pr_pts1, valid1).flatten(1, 2)
    pr_pts2 = (
        invalid_to_nans(pr_pts2, valid2).flatten(1, 2) if pr_pts2 is not None else None
    )

    all_gt = (
        torch.cat((nan_gt_pts1, nan_gt_pts2), dim=1)
        if gt_pts2 is not None
        else nan_gt_pts1
    )
    all_pr = torch.cat((pr_pts1, pr_pts2), dim=1) if pr_pts2 is not None else pr_pts1

    dot_gt_pr = (all_pr * all_gt).sum(dim=-1)
    dot_gt_gt = all_gt.square().sum(dim=-1)

    if fit_mode.startswith("avg"):

        scaling = dot_gt_pr.nanmean(dim=1) / dot_gt_gt.nanmean(dim=1)
    elif fit_mode.startswith("median"):
        scaling = (dot_gt_pr / dot_gt_gt).nanmedian(dim=1).values
    elif fit_mode.startswith("weiszfeld"):

        scaling = dot_gt_pr.nanmean(dim=1) / dot_gt_gt.nanmean(dim=1)

        for iter in range(10):

            dis = (all_pr - scaling.view(-1, 1, 1) * all_gt).norm(dim=-1)

            w = dis.clip_(min=1e-8).reciprocal()

            scaling = (w * dot_gt_pr).nanmean(dim=1) / (w * dot_gt_gt).nanmean(dim=1)
    else:
        raise ValueError(f"bad {fit_mode=}")

    if fit_mode.endswith("stop_grad"):
        scaling = scaling.detach()

    scaling = scaling.clip(min=1e-3)

    return scaling