main

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--output_dir",
        type=str,
        default="",
        help="value for outdir",
    )
    parser.add_argument(
        "--eval_dataset", type=str, default="nyu", choices=list(dataset_metadata.keys())
    )
    parser.add_argument(
        "--align",
        type=str,
        default="scale&shift",
        choices=["scale&shift", "scale", "metric"],
    )
    return parser

def main(args):
    if args.eval_dataset == "sintel":
        TAG_FLOAT = 202021.25

        def depth_read(filename):
            """Read depth data from file, return as numpy array."""
            f = open(filename, "rb")
            check = np.fromfile(f, dtype=np.float32, count=1)[0]
            assert (
                check == TAG_FLOAT
            ), " depth_read:: Wrong tag in flow file (should be: {0}, is: {1}). Big-endian machine? ".format(
                TAG_FLOAT, check
            )
            width = np.fromfile(f, dtype=np.int32, count=1)[0]
            height = np.fromfile(f, dtype=np.int32, count=1)[0]
            size = width * height
            assert (
                width > 0 and height > 0 and size > 1 and size < 100000000
            ), " depth_read:: Wrong input size (width = {0}, height = {1}).".format(
                width, height
            )
            depth = np.fromfile(f, dtype=np.float32, count=-1).reshape((height, width))
            return depth

        pred_pathes = glob.glob(
            f"{args.output_dir}/*/frame_*.npy"
        )  # TODO: update the path to your prediction
        pred_pathes = sorted(pred_pathes)

        if len(pred_pathes) > 643:
            full = True
        else:
            full = False

        if full:
            depth_pathes = glob.glob(f"/path/to/sintel/training/depth/*/*.dpt")
            depth_pathes = sorted(depth_pathes)
        else:
            seq_list = [
                "alley_2",
                "ambush_4",
                "ambush_5",
                "ambush_6",
                "cave_2",
                "cave_4",
                "market_2",
                "market_5",
                "market_6",
                "shaman_3",
                "sleeping_1",
                "sleeping_2",
                "temple_2",
                "temple_3",
            ]
            depth_pathes_folder = [
                f"/path/to/sintel/training/depth/{seq}" for seq in seq_list
            ]
            depth_pathes = []
            for depth_pathes_folder_i in depth_pathes_folder:
                depth_pathes += glob.glob(depth_pathes_folder_i + "/*.dpt")
            depth_pathes = sorted(depth_pathes)

        def get_video_results():
            grouped_pred_depth = group_by_directory(pred_pathes)

            grouped_gt_depth = group_by_directory(depth_pathes)
            gathered_depth_metrics = []

            for key in tqdm(grouped_pred_depth.keys()):
                pd_pathes = grouped_pred_depth[key]
                gt_pathes = grouped_gt_depth[key.replace("_pred_depth", "")]

                gt_depth = np.stack(
                    [depth_read(gt_path) for gt_path in gt_pathes], axis=0
                )
                pr_depth = np.stack(
                    [
                        cv2.resize(
                            np.load(pd_path),
                            (gt_depth.shape[2], gt_depth.shape[1]),
                            interpolation=cv2.INTER_CUBIC,
                        )
                        for pd_path in pd_pathes
                    ],
                    axis=0,
                )
                # for depth eval, set align_with_lad2=False to use median alignment; set align_with_lad2=True to use scale&shift alignment
                if args.align == "scale&shift":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=70,
                            align_with_lad2=True,
                            use_gpu=True,
                            post_clip_max=70,
                        )
                    )
                elif args.align == "scale":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=70,
                            align_with_scale=True,
                            use_gpu=True,
                            post_clip_max=70,
                        )
                    )
                elif args.align == "metric":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=70,
                            metric_scale=True,
                            use_gpu=True,
                            post_clip_max=70,
                        )
                    )
                gathered_depth_metrics.append(depth_results)

            depth_log_path = f"{args.output_dir}/result_{args.align}.json"
            average_metrics = {
                key: np.average(
                    [metrics[key] for metrics in gathered_depth_metrics],
                    weights=[
                        metrics["valid_pixels"] for metrics in gathered_depth_metrics
                    ],
                )
                for key in gathered_depth_metrics[0].keys()
                if key != "valid_pixels"
            }
            print("Average depth evaluation metrics:", average_metrics)
            with open(depth_log_path, "w") as f:
                f.write(json.dumps(average_metrics))

        get_video_results()
    elif args.eval_dataset.startswith("bonn"):

        def depth_read(filename):
            # loads depth map D from png file
            # and returns it as a numpy array
            depth_png = np.asarray(Image.open(filename))
            # make sure we have a proper 16bit depth map here.. not 8bit!
            assert np.max(depth_png) > 255
            depth = depth_png.astype(np.float64) / 5000.0
            depth[depth_png == 0] = -1.0
            return depth

        seq_list = ["balloon2", "crowd2", "crowd3", "person_tracking2", "synchronous"]

        if "_" in args.eval_dataset:
            bonn_number = args.eval_dataset.split("_")[-1]
        else:
            bonn_number = "110"  # default value

        img_pathes_folder = [
            f"/path/to/long_bonn/rgbd_bonn_{seq}/rgb_{bonn_number}/*.png"
            for seq in seq_list
        ]
        img_pathes = []
        for img_pathes_folder_i in img_pathes_folder:
            img_pathes += glob.glob(img_pathes_folder_i)
        img_pathes = sorted(img_pathes)
        depth_pathes_folder = [
            f"/path/to/long_bonn/rgbd_bonn_{seq}/depth_{bonn_number}/*.png"
            for seq in seq_list
        ]
        depth_pathes = []
        for depth_pathes_folder_i in depth_pathes_folder:
            depth_pathes += glob.glob(depth_pathes_folder_i)
        depth_pathes = sorted(depth_pathes)
        pred_pathes = glob.glob(
            f"{args.output_dir}/*/frame*.npy"
        )  # TODO: update the path to your prediction
        pred_pathes = sorted(pred_pathes)

        def get_video_results():
            grouped_pred_depth = group_by_directory(pred_pathes)
            grouped_gt_depth = group_by_directory(depth_pathes, idx=-2)
            gathered_depth_metrics = []
            for key in tqdm(grouped_gt_depth.keys()):
                pd_pathes = grouped_pred_depth[key[10:]]
                gt_pathes = grouped_gt_depth[key]
                gt_depth = np.stack(
                    [depth_read(gt_path) for gt_path in gt_pathes], axis=0
                )
                pr_depth = np.stack(
                    [
                        cv2.resize(
                            np.load(pd_path),
                            (gt_depth.shape[2], gt_depth.shape[1]),
                            interpolation=cv2.INTER_CUBIC,
                        )
                        for pd_path in pd_pathes
                    ],
                    axis=0,
                )
                # for depth eval, set align_with_lad2=False to use median alignment; set align_with_lad2=True to use scale&shift alignment
                if args.align == "scale&shift":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=70,
                            align_with_lad2=True,
                            use_gpu=True,
                        )
                    )
                elif args.align == "scale":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=70,
                            align_with_scale=True,
                            use_gpu=True,
                        )
                    )
                elif args.align == "metric":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=70,
                            metric_scale=True,
                            use_gpu=True,
                        )
                    )
                gathered_depth_metrics.append(depth_results)

                # seq_len = gt_depth.shape[0]
                # error_map = error_map.reshape(seq_len, -1, error_map.shape[-1]).cpu()
                # error_map_colored = colorize(error_map, range=(error_map.min(), error_map.max()), append_cbar=True)
                # ImageSequenceClip([x for x in (error_map_colored.numpy()*255).astype(np.uint8)], fps=10).write_videofile(f'{args.output_dir}/errormap_{key}_{args.align}.mp4', fps=10)

            depth_log_path = f"{args.output_dir}/result_{args.align}.json"
            average_metrics = {
                key: np.average(
                    [metrics[key] for metrics in gathered_depth_metrics],
                    weights=[
                        metrics["valid_pixels"] for metrics in gathered_depth_metrics
                    ],
                )
                for key in gathered_depth_metrics[0].keys()
                if key != "valid_pixels"
            }
            print("Average depth evaluation metrics:", average_metrics)
            with open(depth_log_path, "w") as f:
                f.write(json.dumps(average_metrics))

        get_video_results()
    elif args.eval_dataset == "kitti":

        def depth_read(filename):
            # loads depth map D from png file
            # and returns it as a numpy array,
            # for details see readme.txt
            img_pil = Image.open(filename)
            depth_png = np.array(img_pil, dtype=int)
            # make sure we have a proper 16bit depth map here.. not 8bit!
            assert np.max(depth_png) > 255

            depth = depth_png.astype(float) / 256.0
            depth[depth_png == 0] = -1.0
            return depth

        depth_pathes = glob.glob(
            "data/kitti/depth_selection/val_selection_cropped/groundtruth_depth_gathered/*/*.png"
        )
        depth_pathes = sorted(depth_pathes)
        pred_pathes = glob.glob(
            f"{args.output_dir}/*/frame_*.npy"
        )  # TODO: update the path to your prediction
        pred_pathes = sorted(pred_pathes)

        def get_video_results():
            grouped_pred_depth = group_by_directory(pred_pathes)
            grouped_gt_depth = group_by_directory(depth_pathes)
            gathered_depth_metrics = []
            for key in tqdm(grouped_pred_depth.keys()):
                pd_pathes = grouped_pred_depth[key]
                gt_pathes = grouped_gt_depth[key]
                gt_depth = np.stack(
                    [depth_read(gt_path) for gt_path in gt_pathes], axis=0
                )
                pr_depth = np.stack(
                    [
                        cv2.resize(
                            np.load(pd_path),
                            (gt_depth.shape[2], gt_depth.shape[1]),
                            interpolation=cv2.INTER_CUBIC,
                        )
                        for pd_path in pd_pathes
                    ],
                    axis=0,
                )

                # for depth eval, set align_with_lad2=False to use median alignment; set align_with_lad2=True to use scale&shift alignment
                if args.align == "scale&shift":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=None,
                            align_with_lad2=True,
                            use_gpu=True,
                        )
                    )
                elif args.align == "scale":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=None,
                            align_with_scale=True,
                            use_gpu=True,
                        )
                    )
                elif args.align == "metric":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=None,
                            metric_scale=True,
                            use_gpu=True,
                        )
                    )
                gathered_depth_metrics.append(depth_results)

            depth_log_path = f"{args.output_dir}/result_{args.align}.json"
            average_metrics = {
                key: np.average(
                    [metrics[key] for metrics in gathered_depth_metrics],
                    weights=[
                        metrics["valid_pixels"] for metrics in gathered_depth_metrics
                    ],
                )
                for key in gathered_depth_metrics[0].keys()
                if key != "valid_pixels"
            }
            print("Average depth evaluation metrics:", average_metrics)
            with open(depth_log_path, "w") as f:
                f.write(json.dumps(average_metrics))

        get_video_results()

if __name__ == "__main__":
    args = get_args_parser()
    args = args.parse_args()
    main(args)


                # ImageSequenceClip([x for x in (error_map_colored.numpy()*255).astype(np.uint8)], fps=10).write_videofile(f'{args.output_dir}/errormap_{key}_{args.align}.mp4', fps=10)

            depth_log_path = f"{args.output_dir}/result_{args.align}.json"
            average_metrics = {
                key: np.average(
                    [metrics[key] for metrics in gathered_depth_metrics],
                    weights=[
                        metrics["valid_pixels"] for metrics in gathered_depth_metrics
                    ],
                )
                for key in gathered_depth_metrics[0].keys()
                if key != "valid_pixels"
            }
            print("Average depth evaluation metrics:", average_metrics)
            with open(depth_log_path, "w") as f:
                f.write(json.dumps(average_metrics))

        get_video_results()
    elif args.eval_dataset == "kitti":

        def depth_read(filename):
            # loads depth map D from png file
            # and returns it as a numpy array,
            # for details see readme.txt
            img_pil = Image.open(filename)
            depth_png = np.array(img_pil, dtype=int)
            # make sure we have a proper 16bit depth map here.. not 8bit!
            assert np.max(depth_png) > 255

            depth = depth_png.astype(float) / 256.0
            depth[depth_png == 0] = -1.0
            return depth

        depth_pathes = glob.glob(
            "data/kitti/depth_selection/val_selection_cropped/groundtruth_depth_gathered/*/*.png"
        )
        depth_pathes = sorted(depth_pathes)
        pred_pathes = glob.glob(
            f"{args.output_dir}/*/frame_*.npy"
        )  # TODO: update the path to your prediction
        pred_pathes = sorted(pred_pathes)

        def get_video_results():
            grouped_pred_depth = group_by_directory(pred_pathes)
            grouped_gt_depth = group_by_directory(depth_pathes)
            gathered_depth_metrics = []
            for key in tqdm(grouped_pred_depth.keys()):
                pd_pathes = grouped_pred_depth[key]
                gt_pathes = grouped_gt_depth[key]
                gt_depth = np.stack(
                    [depth_read(gt_path) for gt_path in gt_pathes], axis=0
                )


            # and returns it as a numpy array,
            # for details see readme.txt
            img_pil = Image.open(filename)
            depth_png = np.array(img_pil, dtype=int)
            # make sure we have a proper 16bit depth map here.. not 8bit!
            assert np.max(depth_png) > 255

            depth = depth_png.astype(float) / 256.0
            depth[depth_png == 0] = -1.0
            return depth

        depth_pathes = glob.glob(
            "data/kitti/depth_selection/val_selection_cropped/groundtruth_depth_gathered/*/*.png"
        )
        depth_pathes = sorted(depth_pathes)
        pred_pathes = glob.glob(
            f"{args.output_dir}/*/frame_*.npy"
        )  # TODO: update the path to your prediction
        pred_pathes = sorted(pred_pathes)

        def get_video_results():
            grouped_pred_depth = group_by_directory(pred_pathes)
            grouped_gt_depth = group_by_directory(depth_pathes)
            gathered_depth_metrics = []
            for key in tqdm(grouped_pred_depth.keys()):
                pd_pathes = grouped_pred_depth[key]
                gt_pathes = grouped_gt_depth[key]
                gt_depth = np.stack(
                    [depth_read(gt_path) for gt_path in gt_pathes], axis=0
                )
                pr_depth = np.stack(
                    [
                        cv2.resize(
                            np.load(pd_path),
                            (gt_depth.shape[2], gt_depth.shape[1]),
                            interpolation=cv2.INTER_CUBIC,
                        )
                        for pd_path in pd_pathes
                    ],
                    axis=0,
                )

                # for depth eval, set align_with_lad2=False to use median alignment; set align_with_lad2=True to use scale&shift alignment
                if args.align == "scale&shift":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=None,
                            align_with_lad2=True,
                            use_gpu=True,
                        )
                    )
                elif args.align == "scale":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=None,
                            align_with_scale=True,
                            use_gpu=True,
                        )
                    )
                elif args.align == "metric":
                    depth_results, error_map, depth_predict, depth_gt = (
                        depth_evaluation(
                            pr_depth,
                            gt_depth,
                            max_depth=None,
                            metric_scale=True,
                            use_gpu=True,
                        )
                    )
                gathered_depth_metrics.append(depth_results)

            depth_log_path = f"{args.output_dir}/result_{args.align}.json"
            average_metrics = {
                key: np.average(
                    [metrics[key] for metrics in gathered_depth_metrics],
                    weights=[
                        metrics["valid_pixels"] for metrics in gathered_depth_metrics
                    ],
                )
                for key in gathered_depth_metrics[0].keys()
                if key != "valid_pixels"
            }
            print("Average depth evaluation metrics:", average_metrics)
            with open(depth_log_path, "w") as f:
                f.write(json.dumps(average_metrics))

        get_video_results()

if __name__ == "__main__":
    args = get_args_parser()
    args = args.parse_args()
    main(args)


import os
import torch
import numpy as np
import pickle
from eval.global_human.data_utils import *

# Define the merged dataset metadata dictionary

def create_emdb(split):
    return {
        "img_path": "/path/to/EMDB",
        "dir_path_func": lambda img_path, seq: os.path.join(f"{img_path}/{seq}/images"),
        "seq_list": None,
        "full_seq": True,
        "mask_path_func": lambda filelist: [],
        "split": split,
        "subsample": 1,
        "get_view_func": lambda inputs: load_view_emdb(*inputs),
        "get_seq_func": lambda img_path, split, annots: get_seq_emdb(img_path, split, annots),
        "get_annot_func": lambda img_path, split: get_annot_emdb(img_path, split),
        "is_global": lambda split: {1: False, 2: True}[split],
    }

dataset_metadata = {
    "bedlam": {
        "img_path": "/path/to/processed_bedlam",
        "dir_path_func": lambda img_path, seq: os.path.join(f"{img_path}/Test/{seq}/rgb"),
        "seq_list": None,
        "full_seq": True,
        "mask_path_func": lambda filelist: [f.replace("/rgb/", "/mask/") for f in filelist],
        "split": "Test",
        "subsample": 25, # 25 is used in Multi-HMR's code
        "get_view_func": lambda inputs: load_view_bedlam(*inputs[:2]),
        "get_seq_func": None,
        "is_global": lambda split: False,
    },
    "3dpw": {
        "img_path": "/path/to/3DPW",
        "dir_path_func": lambda img_path, seq: os.path.join(f"{img_path}/imageFiles/{seq}"),
        "seq_list": None,
        "full_seq": True,
        "mask_path_func": lambda filelist: [],
        "split": "test",
        "subsample": 1,
        "get_view_func": lambda inputs: load_view_3dpw(*inputs[:3]),
        "get_seq_func": lambda img_path, split, annots: get_seq_3dpw(img_path, split, annots),
        "get_annot_func": lambda img_path, split: get_annot(img_path, split, "3dpw"),
        "is_global": lambda split: False,
    },
    "emdb1": create_emdb(1),
    "emdb2": create_emdb(2),
    "rich": {
        "img_path": "/path/to/RICH",
        "dir_path_func": lambda img_path, seq: os.path.join(f"{img_path}/{seq}"),
        "seq_list": None,
        "full_seq": True,
        "mask_path_func": lambda filelist: [],
        "split": "test",
        "subsample": 1,
        "get_view_func": lambda inputs: load_view_rich(*inputs),
        "get_seq_func": lambda img_path, split, annots: get_seq_rich(img_path, split, annots),
        "get_annot_func": lambda img_path, split: get_annot_rich(),
        "is_global": lambda split: True,
    }
}

def get_annot(img_path, split, dataset):
    annot_file = os.path.join(f"eval/global_human/annots/{dataset}_{split}.pkl")
    with open(annot_file, 'rb') as f:
        annots = pickle.load(f)
    return annots

def get_annot_emdb(img_path, split):
    annots = {}
    for pkl_name in EMDB_LIST[split]:
        data = load_pkl(os.path.join(img_path, pkl_name))
        annots[data["name"]] = data
    return annots

def get_annot_rich():
    """
    For annotations of RICH dataset,
    please download from https://github.com/zju3dv/GVHMR/blob/main/docs/INSTALL.md#inputs--outputs
    """
    annot_path = os.path.join(os.path.dirname(__file__), "annots/RICH")
    annots = torch.load(os.path.join(annot_path, "hmr4d_support/rich_test_labels.pt"))
    cam_params = torch.load(os.path.join(annot_path, "resource/cam2params.pt"))

    for vid in list(annots.keys()):
        _, sname, cname = vid.split("/")
        scene = sname.split("_")[0]
        cid = int(cname.split("_")[1])
        cam_key = f"{scene}_{cid}"
        annots[vid]["T_w2c"], annots[vid]["K"] = cam_params[cam_key]

    return annots

def get_seq_emdb(img_path, split, annots=None):
    if annots is None:
        annots = get_annot_emdb(img_path, split)

    seq_list = list(annots.keys())
    seq_to_images = {}
    for seq in seq_list:
        mask = annots[seq]["mask"]
        seq_images_dir = os.path.join(img_path, seq, "images")
        all_images = np.array(sorted(os.listdir(seq_images_dir)))
        seq_to_images[seq] = all_images[mask].tolist()

    return seq_list, seq_to_images

def get_seq_rich(img_path, split, annots=None):
    if annots is None:
        annots = get_annot_rich()

    seq_list = list(annots.keys())

    seq_to_images = {}
    for seq in seq_list:
        mask = annots[seq]["frame_id"]
        seq_images_dir = os.path.join(img_path, seq)
        all_images = np.array(sorted(os.listdir(seq_images_dir)))
        seq_to_images[seq] = all_images[mask].tolist()

    return seq_list, seq_to_images

def get_seq_3dpw(img_path, split, annots=None):
    if annots is None:
        annots = get_annot(img_path, split, "3dpw")

    imagenames = sorted(annots.keys())

    seq_to_images = {}
    for imgname in imagenames:
        seq, image = imgname.split('/')
        seq_to_images.setdefault(seq, []).append(image)

    return list(seq_to_images.keys()), seq_to_images

def load_view_bedlam(img_paths, images):
    max_humans = 10
    smpl_key2shape = {
        'smplx_root_pose': (1, 3),
        'smplx_body_pose': (21, 3),
        'smplx_jaw_pose': (1, 3),
        'smplx_leye_pose': (1, 3),
        'smplx_reye_pose': (1, 3),
        'smplx_left_hand_pose': (15, 3),
        'smplx_right_hand_pose': (15, 3),
        'smplx_shape': (11,),
        'smplx_transl': (3,),
        'smplx_gender_id': (),
        }
    for img_path, image in zip(img_paths, images):
        # load camera
        cam_path = img_path.replace("/rgb/", "/cam/").replace(".png", ".npz")
        cam = np.load(cam_path)
        camera_pose = cam["pose"]
        intrinsics = cam["intrinsics"]

        # calculate scale factor
        fy_scale, fx_scale = (image['true_shape'] / image['ori_shape'])[0]

        # update intrinsics
        intrinsics[0, 0] *= fx_scale  # fx
        intrinsics[1, 1] *= fy_scale  # fy
        intrinsics[0, 2] *= fx_scale  # cx
        intrinsics[1, 2] *= fy_scale  # cy

        # load smpl
        annot_file = img_path.replace("/rgb/", "/smpl/").replace(".png", ".pkl")
        annots = []
        smpl_mask = np.zeros(max_humans, dtype=np.bool_)

        if os.path.isfile(annot_file):
            with open(annot_file, 'rb') as f:
                annots = pickle.load(f)
            humans = [hum for hum in annots]
            # humans = [hum for hum in annots if hum['smplx_transl'][-1] > 0.01] # the person should be in front of the camera
            if len(humans) > 0:
                smpl_mask[:len(humans)] = 1.
                l_dist = [hum['smplx_transl'][-1] for hum in humans]
                indexed_lst = list(enumerate(l_dist))
                sorted_indexed = sorted(indexed_lst, key=lambda x: x[1], reverse=False)
                sorted_indices = [index for index, _ in sorted_indexed]
                annots = [humans[h_idx] for h_idx in sorted_indices]

                # Update smplx_gender - 0=neutral - 1=male - 2=female - kids?
                for hum in annots:
                    hum['smplx_gender_id'] = np.asarray({'neutral': 0}[hum['smplx_gender']])

        smpl_dict = {}
        for k, shape in smpl_key2shape.items():
            smpl_dict[k] = np.zeros((max_humans, *shape), dtype=np.float32)
            if len(humans) > 0:
                for h in range(len(humans)):
                    smpl_dict[k][h] = annots[h][k].astype(np.float32)

        image.update({
            'camera_pose': camera_pose.astype(np.float32),
            'intrinsics': intrinsics.astype(np.float32),
            'smpl_mask': smpl_mask,
            **smpl_dict,
        })

    return images

def load_view_3dpw(img_paths, images, annots):
    max_humans = 2
    smpl_key2shape = {
        'smpl_root_pose': (1, 3),
        'smpl_body_pose': (23, 3),
        'smpl_shape': (10,),
        'smpl_transl': (3,),
        'smpl_gender_id': (),
        }

    for img_path, image in zip(img_paths, images):
        img_key = '/'.join(img_path.split('/')[-2:])

        # load camera
        T_w2c = annots[img_key]["cam_poses"]
        camera_pose = np.linalg.inv(T_w2c) # T_c2w
        focal = annots[img_key]["focal"]
        princpt = annots[img_key]["princpt"]
        intrinsics = np.array([
            [focal[0], 0, princpt[0]],
            [0, focal[1], princpt[1]],
            [0, 0, 1]
        ], dtype=np.float32)

        # calculate scale factor
        fy_scale, fx_scale = (image['true_shape'] / image['ori_shape'])[0]

        # update intrinsics
        intrinsics[0, 0] *= fx_scale  # fx
        intrinsics[1, 1] *= fy_scale  # fy
        intrinsics[0, 2] *= fx_scale  # cx
        intrinsics[1, 2] *= fy_scale  # cy

        # load smpl
        humans = annots[img_key]["humans"]
        smpl_mask = np.zeros(max_humans, dtype=np.bool_)

        if len(humans) > 0:
            smpl_mask[:len(humans)] = 1.
            l_dist = [hum['smpl_transl'][-1] for hum in humans]
            indexed_lst = list(enumerate(l_dist))
            sorted_indexed = sorted(indexed_lst, key=lambda x: x[1], reverse=False)
            sorted_indices = [index for index, _ in sorted_indexed]
            annots_humans = [humans[h_idx] for h_idx in sorted_indices]

            # Update smplx_gender - 0=neutral - 1=male - 2=female - kids?
            for hum in annots_humans:
                hum['smpl_gender_id'] = np.asarray({'male': 1, 'female': 2}[hum['smpl_gender']])

        smpl_dict = {}
        for k, shape in smpl_key2shape.items():
            smpl_dict[k] = np.zeros((max_humans, *shape), dtype=np.float32)
            if len(humans) > 0:
                for h in range(len(humans)):
                    smpl_dict[k][h] = annots_humans[h][k].astype(np.float32)

        image.update({
            'camera_pose': camera_pose.astype(np.float32),
            'T_w2c': T_w2c.astype(np.float32),
            'intrinsics': intrinsics.astype(np.float32),
            'smpl_mask': smpl_mask,
            **smpl_dict,
        })

    return images

def load_view_emdb(img_paths, images, annots, indices):
    seq = '/'.join(img_paths[0].split('/')[-4:-2])
    gender = annots[seq]['gender']
    masks = annots[seq]['mask']

    # load camera
    T_w2c = annots[seq]['T_w2c'][masks] # (n_frame,4,4)
    camera_poses = np.linalg.inv(T_w2c) # T_c2w: (n_frame,4,4)
    intrinsics = annots[seq]['K_fullimg'].copy() # (3,3)

    # update intrinsics
    fy_scale, fx_scale = (
        images[0]['true_shape'] / images[0]['ori_shape'])[0]
    intrinsics[0, 0] *= fx_scale  # fx
    intrinsics[1, 1] *= fy_scale  # fy
    intrinsics[0, 2] *= fx_scale  # cx
    intrinsics[1, 2] *= fy_scale  # cy

    # load smpl
    smpl_params = annots[seq]['smpl_params']    # world space

    max_humans = 1
    smpl_key2shape = {
        'smpl_root_pose_w': (1, 3),
        'smpl_body_pose': (23, 3),
        'smpl_shape': (10,),
        'smpl_transl_w': (3,),
        }

    for img_id, image in zip(indices, images):
        smpl_mask = np.ones(max_humans, dtype=np.bool_)
        smpl_dict = {}
        for k, shape in smpl_key2shape.items():
            smpl_dict[k] = np.zeros((max_humans, *shape), dtype=np.float32)
            smpl_dict[k][0] = smpl_params[k][masks][img_id].reshape(*shape).astype(np.float32)

        image.update({
            'T_w2c': T_w2c[img_id].astype(np.float32),
            'camera_pose': camera_poses[img_id].astype(np.float32),
            'intrinsics': intrinsics.astype(np.float32),
            'smpl_mask': smpl_mask,
            'smpl_gender_id': np.asarray({'male': 1, 'female': 2}[gender]),
            **smpl_dict,
        })

    return images

def load_view_rich(img_paths, images, annots, indices):
    seq = '/'.join(img_paths[0].split('/')[-4:-1])
    gender = annots[seq]['gender']

    # load camera
    T_w2c = annots[seq]['T_w2c'].numpy() # (4,4)
    camera_poses = np.linalg.inv(T_w2c) # T_c2w: (4,4)
    intrinsics = annots[seq]['K'].numpy().copy() # (3,3)

    # update intrinsics
    fy_scale, fx_scale = (
        images[0]['true_shape'] / images[0]['ori_shape'])[0]
    intrinsics[0, 0] *= fx_scale  # fx
    intrinsics[1, 1] *= fy_scale  # fy
    intrinsics[0, 2] *= fx_scale  # cx
    intrinsics[1, 2] *= fy_scale  # cy

    # load smpl
    smpl_params = annots[seq]['gt_smplx_params']    # world space

    max_humans = 1
    smpl_key2shape = {
        'global_orient': (1, 3),
        'body_pose': (21, 3),
        'betas': (10,),
        'transl': (3,),
        }

    for img_id, image in zip(indices, images):
        smpl_mask = np.ones(max_humans, dtype=np.bool_)
        smpl_dict = {}
        for k, shape in smpl_key2shape.items():
            smpl_dict['smplx_'+k] = np.zeros((max_humans, *shape), dtype=np.float32)
            smpl_dict['smplx_'+k][0] = smpl_params[k][img_id].reshape(*shape).numpy().astype(np.float32)

        image.update({
            'T_w2c': T_w2c.astype(np.float32),
            'camera_pose': camera_poses.astype(np.float32),
            'intrinsics': intrinsics.astype(np.float32),
            'smpl_mask': smpl_mask,
            'smplx_gender_id': np.asarray({'male': 1, 'female': 2}[gender]),
            **smpl_dict,
        })

    return images


import os
import torch
import numpy as np
import pickle
from eval.global_human.data_utils import *

# Define the merged dataset metadata dictionary

def create_emdb(split):
    return {
        "img_path": "/path/to/EMDB",
        "dir_path_func": lambda img_path, seq: os.path.join(f"{img_path}/{seq}/images"),
        "seq_list": None,
        "full_seq": True,
        "mask_path_func": lambda filelist: [],
        "split": split,
        "subsample": 1,
        "get_view_func": lambda inputs: load_view_emdb(*inputs),
        "get_seq_func": lambda img_path, split, annots: get_seq_emdb(img_path, split, annots),
        "get_annot_func": lambda img_path, split: get_annot_emdb(img_path, split),
        "is_global": lambda split: {1: False, 2: True}[split],
    }

dataset_metadata = {
    "bedlam": {
        "img_path": "/path/to/processed_bedlam",
        "dir_path_func": lambda img_path, seq: os.path.join(f"{img_path}/Test/{seq}/rgb"),
        "seq_list": None,
        "full_seq": True,
        "mask_path_func": lambda filelist: [f.replace("/rgb/", "/mask/") for f in filelist],
        "split": "Test",
        "subsample": 25, # 25 is used in Multi-HMR's code
        "get_view_func": lambda inputs: load_view_bedlam(*inputs[:2]),
        "get_seq_func": None,
        "is_global": lambda split: False,
    },
    "3dpw": {
        "img_path": "/path/to/3DPW",
        "dir_path_func": lambda img_path, seq: os.path.join(f"{img_path}/imageFiles/{seq}"),
        "seq_list": None,
        "full_seq": True,
        "mask_path_func": lambda filelist: [],


        "is_global": lambda split: False,
    },
    "emdb1": create_emdb(1),
    "emdb2": create_emdb(2),
    "rich": {
        "img_path": "/path/to/RICH",
        "dir_path_func": lambda img_path, seq: os.path.join(f"{img_path}/{seq}"),
        "seq_list": None,
        "full_seq": True,
        "mask_path_func": lambda filelist: [],
        "split": "test",
        "subsample": 1,
        "get_view_func": lambda inputs: load_view_rich(*inputs),
        "get_seq_func": lambda img_path, split, annots: get_seq_rich(img_path, split, annots),
        "get_annot_func": lambda img_path, split: get_annot_rich(),
        "is_global": lambda split: True,
    }
}

def get_annot(img_path, split, dataset):
    annot_file = os.path.join(f"eval/global_human/annots/{dataset}_{split}.pkl")
    with open(annot_file, 'rb') as f:
        annots = pickle.load(f)
    return annots

def get_annot_emdb(img_path, split):
    annots = {}
    for pkl_name in EMDB_LIST[split]:
        data = load_pkl(os.path.join(img_path, pkl_name))
        annots[data["name"]] = data
    return annots

def get_annot_rich():
    """
    For annotations of RICH dataset,
    please download from https://github.com/zju3dv/GVHMR/blob/main/docs/INSTALL.md#inputs--outputs
    """
    annot_path = os.path.join(os.path.dirname(__file__), "annots/RICH")
    annots = torch.load(os.path.join(annot_path, "hmr4d_support/rich_test_labels.pt"))
    cam_params = torch.load(os.path.join(annot_path, "resource/cam2params.pt"))

    for vid in list(annots.keys()):
        _, sname, cname = vid.split("/")
        scene = sname.split("_")[0]


        "dir_path_func": lambda img_path, seq: os.path.join(f"{img_path}/{seq}"),
        "seq_list": None,
        "full_seq": True,
        "mask_path_func": lambda filelist: [],
        "split": "test",
        "subsample": 1,
        "get_view_func": lambda inputs: load_view_rich(*inputs),
        "get_seq_func": lambda img_path, split, annots: get_seq_rich(img_path, split, annots),
        "get_annot_func": lambda img_path, split: get_annot_rich(),
        "is_global": lambda split: True,
    }
}

def get_annot(img_path, split, dataset):
    annot_file = os.path.join(f"eval/global_human/annots/{dataset}_{split}.pkl")
    with open(annot_file, 'rb') as f:
        annots = pickle.load(f)
    return annots

def get_annot_emdb(img_path, split):
    annots = {}
    for pkl_name in EMDB_LIST[split]:
        data = load_pkl(os.path.join(img_path, pkl_name))
        annots[data["name"]] = data
    return annots

def get_annot_rich():
    """
    For annotations of RICH dataset,
    please download from https://github.com/zju3dv/GVHMR/blob/main/docs/INSTALL.md#inputs--outputs
    """
    annot_path = os.path.join(os.path.dirname(__file__), "annots/RICH")
    annots = torch.load(os.path.join(annot_path, "hmr4d_support/rich_test_labels.pt"))
    cam_params = torch.load(os.path.join(annot_path, "resource/cam2params.pt"))

    for vid in list(annots.keys()):
        _, sname, cname = vid.split("/")
        scene = sname.split("_")[0]
        cid = int(cname.split("_")[1])
        cam_key = f"{scene}_{cid}"
        annots[vid]["T_w2c"], annots[vid]["K"] = cam_params[cam_key]

    return annots

def get_seq_emdb(img_path, split, annots=None):


        "get_seq_func": lambda img_path, split, annots: get_seq_rich(img_path, split, annots),
        "get_annot_func": lambda img_path, split: get_annot_rich(),
        "is_global": lambda split: True,
    }
}

def get_annot(img_path, split, dataset):
    annot_file = os.path.join(f"eval/global_human/annots/{dataset}_{split}.pkl")
    with open(annot_file, 'rb') as f:
        annots = pickle.load(f)
    return annots

def get_annot_emdb(img_path, split):
    annots = {}
    for pkl_name in EMDB_LIST[split]:
        data = load_pkl(os.path.join(img_path, pkl_name))
        annots[data["name"]] = data
    return annots

def get_annot_rich():
    """
    For annotations of RICH dataset,
    please download from https://github.com/zju3dv/GVHMR/blob/main/docs/INSTALL.md#inputs--outputs
    """
    annot_path = os.path.join(os.path.dirname(__file__), "annots/RICH")
    annots = torch.load(os.path.join(annot_path, "hmr4d_support/rich_test_labels.pt"))
    cam_params = torch.load(os.path.join(annot_path, "resource/cam2params.pt"))

    for vid in list(annots.keys()):
        _, sname, cname = vid.split("/")
        scene = sname.split("_")[0]
        cid = int(cname.split("_")[1])
        cam_key = f"{scene}_{cid}"
        annots[vid]["T_w2c"], annots[vid]["K"] = cam_params[cam_key]

    return annots

def get_seq_emdb(img_path, split, annots=None):
    if annots is None:
        annots = get_annot_emdb(img_path, split)

    seq_list = list(annots.keys())
    seq_to_images = {}
    for seq in seq_list:
        mask = annots[seq]["mask"]
        seq_images_dir = os.path.join(img_path, seq, "images")
        all_images = np.array(sorted(os.listdir(seq_images_dir)))
        seq_to_images[seq] = all_images[mask].tolist()

    return seq_list, seq_to_images

def get_seq_rich(img_path, split, annots=None):
    if annots is None:
        annots = get_annot_rich()

    seq_list = list(annots.keys())


    return annots

def get_annot_rich():
    """
    For annotations of RICH dataset,
    please download from https://github.com/zju3dv/GVHMR/blob/main/docs/INSTALL.md#inputs--outputs
    """
    annot_path = os.path.join(os.path.dirname(__file__), "annots/RICH")
    annots = torch.load(os.path.join(annot_path, "hmr4d_support/rich_test_labels.pt"))
    cam_params = torch.load(os.path.join(annot_path, "resource/cam2params.pt"))

    for vid in list(annots.keys()):
        _, sname, cname = vid.split("/")
        scene = sname.split("_")[0]
        cid = int(cname.split("_")[1])
        cam_key = f"{scene}_{cid}"
        annots[vid]["T_w2c"], annots[vid]["K"] = cam_params[cam_key]

    return annots

def get_seq_emdb(img_path, split, annots=None):
    if annots is None:
        annots = get_annot_emdb(img_path, split)

    seq_list = list(annots.keys())
    seq_to_images = {}
    for seq in seq_list:
        mask = annots[seq]["mask"]
        seq_images_dir = os.path.join(img_path, seq, "images")
        all_images = np.array(sorted(os.listdir(seq_images_dir)))
        seq_to_images[seq] = all_images[mask].tolist()

    return seq_list, seq_to_images

def get_seq_rich(img_path, split, annots=None):
    if annots is None:
        annots = get_annot_rich()

    seq_list = list(annots.keys())

    seq_to_images = {}
    for seq in seq_list:
        mask = annots[seq]["frame_id"]
        seq_images_dir = os.path.join(img_path, seq)
        all_images = np.array(sorted(os.listdir(seq_images_dir)))
        seq_to_images[seq] = all_images[mask].tolist()

    return seq_list, seq_to_images

def get_seq_3dpw(img_path, split, annots=None):
    if annots is None:
        annots = get_annot(img_path, split, "3dpw")


        cid = int(cname.split("_")[1])
        cam_key = f"{scene}_{cid}"
        annots[vid]["T_w2c"], annots[vid]["K"] = cam_params[cam_key]

    return annots

def get_seq_emdb(img_path, split, annots=None):
    if annots is None:
        annots = get_annot_emdb(img_path, split)

    seq_list = list(annots.keys())
    seq_to_images = {}
    for seq in seq_list:
        mask = annots[seq]["mask"]
        seq_images_dir = os.path.join(img_path, seq, "images")
        all_images = np.array(sorted(os.listdir(seq_images_dir)))
        seq_to_images[seq] = all_images[mask].tolist()

    return seq_list, seq_to_images

def get_seq_rich(img_path, split, annots=None):
    if annots is None:
        annots = get_annot_rich()

    seq_list = list(annots.keys())

    seq_to_images = {}
    for seq in seq_list:
        mask = annots[seq]["frame_id"]
        seq_images_dir = os.path.join(img_path, seq)
        all_images = np.array(sorted(os.listdir(seq_images_dir)))
        seq_to_images[seq] = all_images[mask].tolist()

    return seq_list, seq_to_images

def get_seq_3dpw(img_path, split, annots=None):
    if annots is None:
        annots = get_annot(img_path, split, "3dpw")

    imagenames = sorted(annots.keys())

    seq_to_images = {}
    for imgname in imagenames:
        seq, image = imgname.split('/')
        seq_to_images.setdefault(seq, []).append(image)

    return list(seq_to_images.keys()), seq_to_images

def load_view_bedlam(img_paths, images):
    max_humans = 10
    smpl_key2shape = {
        'smplx_root_pose': (1, 3),
        'smplx_body_pose': (21, 3),
        'smplx_jaw_pose': (1, 3),


        all_images = np.array(sorted(os.listdir(seq_images_dir)))
        seq_to_images[seq] = all_images[mask].tolist()

    return seq_list, seq_to_images

def get_seq_rich(img_path, split, annots=None):
    if annots is None:
        annots = get_annot_rich()

    seq_list = list(annots.keys())

    seq_to_images = {}
    for seq in seq_list:
        mask = annots[seq]["frame_id"]
        seq_images_dir = os.path.join(img_path, seq)
        all_images = np.array(sorted(os.listdir(seq_images_dir)))
        seq_to_images[seq] = all_images[mask].tolist()

    return seq_list, seq_to_images

def get_seq_3dpw(img_path, split, annots=None):
    if annots is None:
        annots = get_annot(img_path, split, "3dpw")

    imagenames = sorted(annots.keys())

    seq_to_images = {}
    for imgname in imagenames:
        seq, image = imgname.split('/')
        seq_to_images.setdefault(seq, []).append(image)

    return list(seq_to_images.keys()), seq_to_images

def load_view_bedlam(img_paths, images):
    max_humans = 10
    smpl_key2shape = {
        'smplx_root_pose': (1, 3),
        'smplx_body_pose': (21, 3),
        'smplx_jaw_pose': (1, 3),
        'smplx_leye_pose': (1, 3),
        'smplx_reye_pose': (1, 3),
        'smplx_left_hand_pose': (15, 3),
        'smplx_right_hand_pose': (15, 3),
        'smplx_shape': (11,),
        'smplx_transl': (3,),
        'smplx_gender_id': (),
        }
    for img_path, image in zip(img_paths, images):
        # load camera
        cam_path = img_path.replace("/rgb/", "/cam/").replace(".png", ".npz")
        cam = np.load(cam_path)
        camera_pose = cam["pose"]


        mask = annots[seq]["frame_id"]
        seq_images_dir = os.path.join(img_path, seq)
        all_images = np.array(sorted(os.listdir(seq_images_dir)))
        seq_to_images[seq] = all_images[mask].tolist()

    return seq_list, seq_to_images

def get_seq_3dpw(img_path, split, annots=None):
    if annots is None:
        annots = get_annot(img_path, split, "3dpw")

    imagenames = sorted(annots.keys())

    seq_to_images = {}
    for imgname in imagenames:
        seq, image = imgname.split('/')
        seq_to_images.setdefault(seq, []).append(image)

    return list(seq_to_images.keys()), seq_to_images

def load_view_bedlam(img_paths, images):
    max_humans = 10
    smpl_key2shape = {
        'smplx_root_pose': (1, 3),
        'smplx_body_pose': (21, 3),
        'smplx_jaw_pose': (1, 3),
        'smplx_leye_pose': (1, 3),
        'smplx_reye_pose': (1, 3),
        'smplx_left_hand_pose': (15, 3),
        'smplx_right_hand_pose': (15, 3),
        'smplx_shape': (11,),
        'smplx_transl': (3,),
        'smplx_gender_id': (),
        }
    for img_path, image in zip(img_paths, images):
        # load camera
        cam_path = img_path.replace("/rgb/", "/cam/").replace(".png", ".npz")
        cam = np.load(cam_path)
        camera_pose = cam["pose"]
        intrinsics = cam["intrinsics"]

        # calculate scale factor
        fy_scale, fx_scale = (image['true_shape'] / image['ori_shape'])[0]

        # update intrinsics
        intrinsics[0, 0] *= fx_scale  # fx
        intrinsics[1, 1] *= fy_scale  # fy
        intrinsics[0, 2] *= fx_scale  # cx
        intrinsics[1, 2] *= fy_scale  # cy

        # load smpl
        annot_file = img_path.replace("/rgb/", "/smpl/").replace(".png", ".pkl")
        annots = []
        smpl_mask = np.zeros(max_humans, dtype=np.bool_)

        if os.path.isfile(annot_file):
            with open(annot_file, 'rb') as f:
                annots = pickle.load(f)
            humans = [hum for hum in annots]
            # humans = [hum for hum in annots if hum['smplx_transl'][-1] > 0.01] # the person should be in front of the camera
            if len(humans) > 0:
                smpl_mask[:len(humans)] = 1.
                l_dist = [hum['smplx_transl'][-1] for hum in humans]
                indexed_lst = list(enumerate(l_dist))
                sorted_indexed = sorted(indexed_lst, key=lambda x: x[1], reverse=False)
                sorted_indices = [index for index, _ in sorted_indexed]
                annots = [humans[h_idx] for h_idx in sorted_indices]

                # Update smplx_gender - 0=neutral - 1=male - 2=female - kids?
                for hum in annots:
                    hum['smplx_gender_id'] = np.asarray({'neutral': 0}[hum['smplx_gender']])

        smpl_dict = {}
        for k, shape in smpl_key2shape.items():
            smpl_dict[k] = np.zeros((max_humans, *shape), dtype=np.float32)
            if len(humans) > 0:
                for h in range(len(humans)):
                    smpl_dict[k][h] = annots[h][k].astype(np.float32)

        image.update({
            'camera_pose': camera_pose.astype(np.float32),
            'intrinsics': intrinsics.astype(np.float32),
            'smpl_mask': smpl_mask,
            **smpl_dict,
        })

    return images

def load_view_3dpw(img_paths, images, annots):
    max_humans = 2
    smpl_key2shape = {
        'smpl_root_pose': (1, 3),
        'smpl_body_pose': (23, 3),
        'smpl_shape': (10,),
        'smpl_transl': (3,),
        'smpl_gender_id': (),
        }

    for img_path, image in zip(img_paths, images):
        img_key = '/'.join(img_path.split('/')[-2:])

        # load camera
        T_w2c = annots[img_key]["cam_poses"]
        camera_pose = np.linalg.inv(T_w2c) # T_c2w
        focal = annots[img_key]["focal"]
        princpt = annots[img_key]["princpt"]
        intrinsics = np.array([


                # Update smplx_gender - 0=neutral - 1=male - 2=female - kids?
                for hum in annots:
                    hum['smplx_gender_id'] = np.asarray({'neutral': 0}[hum['smplx_gender']])

        smpl_dict = {}
        for k, shape in smpl_key2shape.items():
            smpl_dict[k] = np.zeros((max_humans, *shape), dtype=np.float32)
            if len(humans) > 0:
                for h in range(len(humans)):
                    smpl_dict[k][h] = annots[h][k].astype(np.float32)

        image.update({
            'camera_pose': camera_pose.astype(np.float32),
            'intrinsics': intrinsics.astype(np.float32),
            'smpl_mask': smpl_mask,
            **smpl_dict,
        })

    return images

def load_view_3dpw(img_paths, images, annots):
    max_humans = 2
    smpl_key2shape = {
        'smpl_root_pose': (1, 3),
        'smpl_body_pose': (23, 3),
        'smpl_shape': (10,),
        'smpl_transl': (3,),
        'smpl_gender_id': (),
        }

    for img_path, image in zip(img_paths, images):
        img_key = '/'.join(img_path.split('/')[-2:])

        # load camera
        T_w2c = annots[img_key]["cam_poses"]
        camera_pose = np.linalg.inv(T_w2c) # T_c2w
        focal = annots[img_key]["focal"]
        princpt = annots[img_key]["princpt"]
        intrinsics = np.array([
            [focal[0], 0, princpt[0]],
            [0, focal[1], princpt[1]],
            [0, 0, 1]
        ], dtype=np.float32)

        # calculate scale factor
        fy_scale, fx_scale = (image['true_shape'] / image['ori_shape'])[0]

        # update intrinsics
        intrinsics[0, 0] *= fx_scale  # fx
        intrinsics[1, 1] *= fy_scale  # fy
        intrinsics[0, 2] *= fx_scale  # cx
        intrinsics[1, 2] *= fy_scale  # cy

        # load smpl
        humans = annots[img_key]["humans"]
        smpl_mask = np.zeros(max_humans, dtype=np.bool_)

        if len(humans) > 0:
            smpl_mask[:len(humans)] = 1.
            l_dist = [hum['smpl_transl'][-1] for hum in humans]
            indexed_lst = list(enumerate(l_dist))
            sorted_indexed = sorted(indexed_lst, key=lambda x: x[1], reverse=False)
            sorted_indices = [index for index, _ in sorted_indexed]
            annots_humans = [humans[h_idx] for h_idx in sorted_indices]

            # Update smplx_gender - 0=neutral - 1=male - 2=female - kids?
            for hum in annots_humans:
                hum['smpl_gender_id'] = np.asarray({'male': 1, 'female': 2}[hum['smpl_gender']])

        smpl_dict = {}
        for k, shape in smpl_key2shape.items():
            smpl_dict[k] = np.zeros((max_humans, *shape), dtype=np.float32)
            if len(humans) > 0:
                for h in range(len(humans)):
                    smpl_dict[k][h] = annots_humans[h][k].astype(np.float32)

        image.update({
            'camera_pose': camera_pose.astype(np.float32),
            'T_w2c': T_w2c.astype(np.float32),
            'intrinsics': intrinsics.astype(np.float32),
            'smpl_mask': smpl_mask,
            **smpl_dict,
        })

    return images

def load_view_emdb(img_paths, images, annots, indices):
    seq = '/'.join(img_paths[0].split('/')[-4:-2])
    gender = annots[seq]['gender']
    masks = annots[seq]['mask']

    # load camera
    T_w2c = annots[seq]['T_w2c'][masks] # (n_frame,4,4)
    camera_poses = np.linalg.inv(T_w2c) # T_c2w: (n_frame,4,4)
    intrinsics = annots[seq]['K_fullimg'].copy() # (3,3)

    # update intrinsics
    fy_scale, fx_scale = (
        images[0]['true_shape'] / images[0]['ori_shape'])[0]
    intrinsics[0, 0] *= fx_scale  # fx
    intrinsics[1, 1] *= fy_scale  # fy
    intrinsics[0, 2] *= fx_scale  # cx
    intrinsics[1, 2] *= fy_scale  # cy

    # load smpl


            for hum in annots_humans:
                hum['smpl_gender_id'] = np.asarray({'male': 1, 'female': 2}[hum['smpl_gender']])

        smpl_dict = {}
        for k, shape in smpl_key2shape.items():
            smpl_dict[k] = np.zeros((max_humans, *shape), dtype=np.float32)
            if len(humans) > 0:
                for h in range(len(humans)):
                    smpl_dict[k][h] = annots_humans[h][k].astype(np.float32)

        image.update({
            'camera_pose': camera_pose.astype(np.float32),
            'T_w2c': T_w2c.astype(np.float32),
            'intrinsics': intrinsics.astype(np.float32),
            'smpl_mask': smpl_mask,
            **smpl_dict,
        })

    return images

def load_view_emdb(img_paths, images, annots, indices):
    seq = '/'.join(img_paths[0].split('/')[-4:-2])
    gender = annots[seq]['gender']
    masks = annots[seq]['mask']

    # load camera
    T_w2c = annots[seq]['T_w2c'][masks] # (n_frame,4,4)
    camera_poses = np.linalg.inv(T_w2c) # T_c2w: (n_frame,4,4)
    intrinsics = annots[seq]['K_fullimg'].copy() # (3,3)

    # update intrinsics
    fy_scale, fx_scale = (
        images[0]['true_shape'] / images[0]['ori_shape'])[0]
    intrinsics[0, 0] *= fx_scale  # fx
    intrinsics[1, 1] *= fy_scale  # fy
    intrinsics[0, 2] *= fx_scale  # cx
    intrinsics[1, 2] *= fy_scale  # cy

    # load smpl
    smpl_params = annots[seq]['smpl_params']    # world space

    max_humans = 1
    smpl_key2shape = {
        'smpl_root_pose_w': (1, 3),
        'smpl_body_pose': (23, 3),
        'smpl_shape': (10,),
        'smpl_transl_w': (3,),
        }

    for img_id, image in zip(indices, images):
        smpl_mask = np.ones(max_humans, dtype=np.bool_)
        smpl_dict = {}
        for k, shape in smpl_key2shape.items():
            smpl_dict[k] = np.zeros((max_humans, *shape), dtype=np.float32)
            smpl_dict[k][0] = smpl_params[k][masks][img_id].reshape(*shape).astype(np.float32)

        image.update({
            'T_w2c': T_w2c[img_id].astype(np.float32),
            'camera_pose': camera_poses[img_id].astype(np.float32),
            'intrinsics': intrinsics.astype(np.float32),
            'smpl_mask': smpl_mask,
            'smpl_gender_id': np.asarray({'male': 1, 'female': 2}[gender]),
            **smpl_dict,
        })

    return images

def load_view_rich(img_paths, images, annots, indices):
    seq = '/'.join(img_paths[0].split('/')[-4:-1])
    gender = annots[seq]['gender']

    # load camera
    T_w2c = annots[seq]['T_w2c'].numpy() # (4,4)
    camera_poses = np.linalg.inv(T_w2c) # T_c2w: (4,4)
    intrinsics = annots[seq]['K'].numpy().copy() # (3,3)

    # update intrinsics
    fy_scale, fx_scale = (
        images[0]['true_shape'] / images[0]['ori_shape'])[0]
    intrinsics[0, 0] *= fx_scale  # fx
    intrinsics[1, 1] *= fy_scale  # fy
    intrinsics[0, 2] *= fx_scale  # cx
    intrinsics[1, 2] *= fy_scale  # cy

    # load smpl
    smpl_params = annots[seq]['gt_smplx_params']    # world space


        }

    for img_id, image in zip(indices, images):
        smpl_mask = np.ones(max_humans, dtype=np.bool_)
        smpl_dict = {}
        for k, shape in smpl_key2shape.items():
            smpl_dict[k] = np.zeros((max_humans, *shape), dtype=np.float32)
            smpl_dict[k][0] = smpl_params[k][masks][img_id].reshape(*shape).astype(np.float32)

        image.update({
            'T_w2c': T_w2c[img_id].astype(np.float32),
            'camera_pose': camera_poses[img_id].astype(np.float32),
            'intrinsics': intrinsics.astype(np.float32),
            'smpl_mask': smpl_mask,
            'smpl_gender_id': np.asarray({'male': 1, 'female': 2}[gender]),
            **smpl_dict,
        })

    return images

def load_view_rich(img_paths, images, annots, indices):
    seq = '/'.join(img_paths[0].split('/')[-4:-1])
    gender = annots[seq]['gender']

    # load camera
    T_w2c = annots[seq]['T_w2c'].numpy() # (4,4)
    camera_poses = np.linalg.inv(T_w2c) # T_c2w: (4,4)
    intrinsics = annots[seq]['K'].numpy().copy() # (3,3)

    # update intrinsics
    fy_scale, fx_scale = (
        images[0]['true_shape'] / images[0]['ori_shape'])[0]
    intrinsics[0, 0] *= fx_scale  # fx
    intrinsics[1, 1] *= fy_scale  # fy
    intrinsics[0, 2] *= fx_scale  # cx
    intrinsics[1, 2] *= fy_scale  # cy

    # load smpl
    smpl_params = annots[seq]['gt_smplx_params']    # world space

    max_humans = 1
    smpl_key2shape = {
        'global_orient': (1, 3),
        'body_pose': (21, 3),
        'betas': (10,),
        'transl': (3,),
        }

    for img_id, image in zip(indices, images):
        smpl_mask = np.ones(max_humans, dtype=np.bool_)
        smpl_dict = {}
        for k, shape in smpl_key2shape.items():
            smpl_dict['smplx_'+k] = np.zeros((max_humans, *shape), dtype=np.float32)
            smpl_dict['smplx_'+k][0] = smpl_params[k][img_id].reshape(*shape).numpy().astype(np.float32)

        image.update({
            'T_w2c': T_w2c.astype(np.float32),
            'camera_pose': camera_poses.astype(np.float32),
            'intrinsics': intrinsics.astype(np.float32),
            'smpl_mask': smpl_mask,
            'smplx_gender_id': np.asarray({'male': 1, 'female': 2}[gender]),
            **smpl_dict,
        })

    return images


# Modified from GVHMR [https://github.com/zju3dv/GVHMR].
# Load EMDB data

import pickle
import torch

EMDB1_LIST = [
    'P8/69_outdoor_cartwheel/P8_69_outdoor_cartwheel_data.pkl', # 656
    'P5/42_indoor_dancing/P5_42_indoor_dancing_data.pkl', # 1291
    'P6/51_outdoor_dancing/P6_51_outdoor_dancing_data.pkl',  # 1427
    'P2/23_outdoor_hug_tree/P2_23_outdoor_hug_tree_data.pkl',  # 1086
    'P6/49_outdoor_big_stairs_down/P6_49_outdoor_big_stairs_down_data.pkl',  # DUPLICATE 1559

    'P7/59_outdoor_rom/P7_59_outdoor_rom_data.pkl', # 1839
    'P3/31_outdoor_workout/P3_31_outdoor_workout_data.pkl',  # 1216
    'P3/33_outdoor_soccer_warmup_b/P3_33_outdoor_soccer_warmup_b_data.pkl',  # 1433
    'P7/57_outdoor_rock_chair/P7_57_outdoor_rock_chair_data.pkl', # DUPLICATE 1558

    'P3/32_outdoor_soccer_warmup_a/P3_32_outdoor_soccer_warmup_a_data.pkl',  # 1084
    'P8/64_outdoor_skateboard/P8_64_outdoor_skateboard_data.pkl',   # DUPLICATE 1704
    'P7/60_outdoor_workout/P7_60_outdoor_workout_data.pkl',  # 1693
    'P6/50_outdoor_workout/P6_50_outdoor_workout_data.pkl', # 1532

    'P8/68_outdoor_handstand/P8_68_outdoor_handstand_data.pkl',  # 1606
    'P9/76_outdoor_sitting/P9_76_outdoor_sitting_data.pkl',  # 1768
    'P1/14_outdoor_climb/P1_14_outdoor_climb_data.pkl',  # 1284
    'P5/44_indoor_rom/P5_44_indoor_rom_data.pkl', # 1381
]
EMDB1_NAMES = ["_".join(p.split("/")[:2]) for p in EMDB1_LIST]

EMDB2_LIST = [
    'P2/19_indoor_walk_off_mvs/P2_19_indoor_walk_off_mvs_data.pkl',  # 1299
    'P3/29_outdoor_stairs_up/P3_29_outdoor_stairs_up_data.pkl', # 1205
    'P4/35_indoor_walk/P4_35_indoor_walk_data.pkl',  # 1226
    'P7/55_outdoor_walk/P7_55_outdoor_walk_data.pkl',  # 2179
    'P9/80_outdoor_walk_big_circle/P9_80_outdoor_walk_big_circle_data.pkl',  # 2240
    'P9/77_outdoor_stairs_up/P9_77_outdoor_stairs_up_data.pkl',  # DUPLICATE 728
    'P9/79_outdoor_walk_rectangle/P9_79_outdoor_walk_rectangle_data.pkl',  # 1917

    'P7/57_outdoor_rock_chair/P7_57_outdoor_rock_chair_data.pkl',  # DUPLICATE 1558
    'P2/24_outdoor_long_walk/P2_24_outdoor_long_walk_data.pkl',  # 3280
    'P3/30_outdoor_stairs_down/P3_30_outdoor_stairs_down_data.pkl',  #1137
    'P4/36_outdoor_long_walk/P4_36_outdoor_long_walk_data.pkl',  # 2160
    'P6/49_outdoor_big_stairs_down/P6_49_outdoor_big_stairs_down_data.pkl',   # DUPLICATE 1559
    'P9/78_outdoor_stairs_up_down/P9_78_outdoor_stairs_up_down_data.pkl',  # 1083

    'P7/56_outdoor_stairs_up_down/P7_56_outdoor_stairs_up_down_data.pkl', # 1120
    'P2/20_outdoor_walk/P2_20_outdoor_walk_data.pkl',  # 2713
    'P3/27_indoor_walk_off_mvs/P3_27_indoor_walk_off_mvs_data.pkl',  # 1448
    'P4/37_outdoor_run_circle/P4_37_outdoor_run_circle_data.pkl', # 881
    'P5/40_indoor_walk_big_circle/P5_40_indoor_walk_big_circle_data.pkl', # 2661
    'P6/48_outdoor_walk_downhill/P6_48_outdoor_walk_downhill_data.pkl',  # 1959

    'P0/09_outdoor_walk/P0_09_outdoor_walk_data.pkl',  # 2009
    'P3/28_outdoor_walk_lunges/P3_28_outdoor_walk_lunges_data.pkl',  # 1836
    'P7/58_outdoor_parcours/P7_58_outdoor_parcours_data.pkl',  # 1332
    'P7/61_outdoor_sit_lie_walk/P7_61_outdoor_sit_lie_walk_data.pkl', # 1914
    'P8/64_outdoor_skateboard/P8_64_outdoor_skateboard_data.pkl',   # DUPLICATE 1704
    'P8/65_outdoor_walk_straight/P8_65_outdoor_walk_straight_data.pkl',  # 1981
]

EMDB2_NAMES = ["_".join(p.split("/")[:2]) for p in EMDB2_LIST]
EMDB_NAMES = {1: EMDB1_NAMES, 2: EMDB2_NAMES}
EMDB_LIST = {1: EMDB1_LIST, 2: EMDB2_LIST}

def load_pkl(fp):
    annot = pickle.load(open(fp, "rb"))
    # ['gender', 'name', 'emdb1', 'emdb2', 'n_frames', 'good_frames_mask', 'camera', 'smpl', 'kp2d', 'bboxes', 'subfolder']
    data = {}

    F = annot["n_frames"]
    data["smpl_params"] = {
        "smpl_body_pose": annot["smpl"]["poses_body"],  # (F, 69)
        "smpl_shape": annot["smpl"]["betas"][None].repeat(F, axis=0),  # (F, 10)
        "smpl_root_pose_w": annot["smpl"]["poses_root"],  # (F, 3)
        "smpl_transl_w": annot["smpl"]["trans"],  # (F, 3)
    }

    data["name"] = annot["name"].replace('_', '/', 1)
    data["gender"] = annot["gender"]
    data["mask"] = annot["good_frames_mask"]   # (L,)
    data["K_fullimg"] = annot["camera"]["intrinsics"]  # (3, 3)
    data["T_w2c"] = annot["camera"]["extrinsics"]  # (L, 4, 4)

    return data

def to_tensor(ndarray):
    tensor = torch.from_numpy(ndarray).float()
    return tensor


    'P7/56_outdoor_stairs_up_down/P7_56_outdoor_stairs_up_down_data.pkl', # 1120
    'P2/20_outdoor_walk/P2_20_outdoor_walk_data.pkl',  # 2713
    'P3/27_indoor_walk_off_mvs/P3_27_indoor_walk_off_mvs_data.pkl',  # 1448
    'P4/37_outdoor_run_circle/P4_37_outdoor_run_circle_data.pkl', # 881
    'P5/40_indoor_walk_big_circle/P5_40_indoor_walk_big_circle_data.pkl', # 2661
    'P6/48_outdoor_walk_downhill/P6_48_outdoor_walk_downhill_data.pkl',  # 1959

    'P0/09_outdoor_walk/P0_09_outdoor_walk_data.pkl',  # 2009
    'P3/28_outdoor_walk_lunges/P3_28_outdoor_walk_lunges_data.pkl',  # 1836
    'P7/58_outdoor_parcours/P7_58_outdoor_parcours_data.pkl',  # 1332
    'P7/61_outdoor_sit_lie_walk/P7_61_outdoor_sit_lie_walk_data.pkl', # 1914
    'P8/64_outdoor_skateboard/P8_64_outdoor_skateboard_data.pkl',   # DUPLICATE 1704
    'P8/65_outdoor_walk_straight/P8_65_outdoor_walk_straight_data.pkl',  # 1981
]

EMDB2_NAMES = ["_".join(p.split("/")[:2]) for p in EMDB2_LIST]
EMDB_NAMES = {1: EMDB1_NAMES, 2: EMDB2_NAMES}
EMDB_LIST = {1: EMDB1_LIST, 2: EMDB2_LIST}

def load_pkl(fp):
    annot = pickle.load(open(fp, "rb"))
    # ['gender', 'name', 'emdb1', 'emdb2', 'n_frames', 'good_frames_mask', 'camera', 'smpl', 'kp2d', 'bboxes', 'subfolder']
    data = {}

    F = annot["n_frames"]
    data["smpl_params"] = {
        "smpl_body_pose": annot["smpl"]["poses_body"],  # (F, 69)
        "smpl_shape": annot["smpl"]["betas"][None].repeat(F, axis=0),  # (F, 10)
        "smpl_root_pose_w": annot["smpl"]["poses_root"],  # (F, 3)
        "smpl_transl_w": annot["smpl"]["trans"],  # (F, 3)
    }

    data["name"] = annot["name"].replace('_', '/', 1)
    data["gender"] = annot["gender"]
    data["mask"] = annot["good_frames_mask"]   # (L,)
    data["K_fullimg"] = annot["camera"]["intrinsics"]  # (3, 3)
    data["T_w2c"] = annot["camera"]["extrinsics"]  # (L, 4, 4)

    return data

def to_tensor(ndarray):
    tensor = torch.from_numpy(ndarray).float()
    return tensor


    annot = pickle.load(open(fp, "rb"))
    # ['gender', 'name', 'emdb1', 'emdb2', 'n_frames', 'good_frames_mask', 'camera', 'smpl', 'kp2d', 'bboxes', 'subfolder']
    data = {}

    F = annot["n_frames"]
    data["smpl_params"] = {
        "smpl_body_pose": annot["smpl"]["poses_body"],  # (F, 69)
        "smpl_shape": annot["smpl"]["betas"][None].repeat(F, axis=0),  # (F, 10)
        "smpl_root_pose_w": annot["smpl"]["poses_root"],  # (F, 3)
        "smpl_transl_w": annot["smpl"]["trans"],  # (F, 3)
    }

    data["name"] = annot["name"].replace('_', '/', 1)
    data["gender"] = annot["gender"]
    data["mask"] = annot["good_frames_mask"]   # (L,)
    data["K_fullimg"] = annot["camera"]["intrinsics"]  # (3, 3)
    data["T_w2c"] = annot["camera"]["extrinsics"]  # (L, 4, 4)

    return data

def to_tensor(ndarray):
    tensor = torch.from_numpy(ndarray).float()
    return tensor


from copy import deepcopy
import cv2

import os
import re
import numpy as np
import torch
import roma
import math
import tqdm
import matplotlib as mpl
import matplotlib.cm as cm
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg
from scipy.spatial.transform import Rotation
from eval.relpose.evo_utils import *
from PIL import Image
import imageio.v2 as iio
from matplotlib.figure import Figure
from itertools import product

def todevice(batch, device, callback=None, non_blocking=False):
    """Transfer some variables to another device (i.e. GPU, CPU:torch, CPU:numpy).

    batch: list, tuple, dict of tensors or other things
    device: pytorch device or 'numpy'
    callback: function that would be called on every sub-elements.
    """
    if callback:
        batch = callback(batch)

    if isinstance(batch, dict):
        return {k: todevice(v, device) for k, v in batch.items()}

    if isinstance(batch, (tuple, list)):
        return type(batch)(todevice(x, device) for x in batch)

    x = batch
    if device == "numpy":
        if isinstance(x, torch.Tensor):
            x = x.detach().cpu().numpy()
    elif x is not None:
        if isinstance(x, np.ndarray):
            x = torch.from_numpy(x)
        if torch.is_tensor(x):
            x = x.to(device, non_blocking=non_blocking)
    return x

to_device = todevice  # alias

def to_numpy(x):
    return todevice(x, "numpy")

def compute_prf1(count, miss, fp):
    """
    Code modified from https://github.com/Arthur151/ROMP/blob/4eebd3647f57d291d26423e51f0d514ff7197cb3/simple_romp/evaluation/RH_evaluation/evaluation.py#L90
    """
    if count == 0:
        return 0, 0, 0
    all_tp = count - miss
    all_fp = fp
    all_fn = miss
    if all_tp == 0:
        return 0., 0., 0.
    all_f1_score = round(all_tp / (all_tp + 0.5 * (all_fp + all_fn)), 2)
    all_recall = round(all_tp / (all_tp + all_fn), 2)
    all_precision = round(all_tp / (all_tp + all_fp), 2)
    return 100. * all_precision, 100.* all_recall, 100. * all_f1_score

def match_2d_greedy(
        pred_kps,
        gtkp,
        valid_mask,
        imgPath=None,
        baseline=None,
        iou_thresh=0.05,
        valid=None,
        ind=-1):
    '''
    Code modified from: https://github.com/Arthur151/ROMP/blob/4eebd3647f57d291d26423e51f0d514ff7197cb3/simple_romp/trace2/evaluation/eval_3DPW.py#L232
    matches groundtruth keypoints to the detection by considering all possible matchings.
    :return: best possible matching, a list of tuples, where each tuple corresponds to one match of pred_person.to gt_person.
            the order within one tuple is as follows (idx_pred_kps, idx_gt_kps)
    '''
    predList = np.arange(len(pred_kps))
    gtList = np.arange(len(gtkp))
    # get all pairs of elements in pred_kps, gtkp
    # all combinations of 2 elements from l1 and l2
    combs = list(product(predList, gtList))

    errors_per_pair = {}
    errors_per_pair_list = []
    for comb in combs:
        vmask = valid_mask[comb[1]]
        assert vmask.sum()>0, print('no valid points')
        errors_per_pair[str(comb)] = np.linalg.norm(pred_kps[comb[0]][vmask, :2] - gtkp[comb[1]][vmask, :2], 2)
        errors_per_pair_list.append(errors_per_pair[str(comb)])

    gtAssigned = np.zeros((len(gtkp),), dtype=bool)
    opAssigned = np.zeros((len(pred_kps),), dtype=bool)
    errors_per_pair_list = np.array(errors_per_pair_list)

    bestMatch = []
    excludedGtBecauseInvalid = []
    falsePositiveCounter = 0
    while np.sum(gtAssigned) < len(gtAssigned) and np.sum(
            opAssigned) + falsePositiveCounter < len(pred_kps):
        found = False
        falsePositive = False
        while not(found):
            if sum(np.inf == errors_per_pair_list) == len(
                    errors_per_pair_list):
                print('something went wrong here')

            minIdx = np.argmin(errors_per_pair_list)
            minComb = combs[minIdx]
            # compute IOU
            iou = get_bbx_overlap(
                pred_kps[minComb[0]], gtkp[minComb[1]]) #, imgPath, baseline)
            # if neither prediction nor ground truth has been matched before and iou
            # is larger than threshold
            if not(opAssigned[minComb[0]]) and not(
                    gtAssigned[minComb[1]]) and iou >= iou_thresh:
                #print(imgPath + ': found matching')
                found = True
                errors_per_pair_list[minIdx] = np.inf
            else:
                errors_per_pair_list[minIdx] = np.inf
                # if errors_per_pair_list[minIdx] >
                # matching_threshold*headBboxs[combs[minIdx][1]]:
                if iou < iou_thresh:
                    #print(
                    #   imgPath + ': false positive detected using threshold')
                    found = True
                    falsePositive = True
                    falsePositiveCounter += 1

        # if ground truth of combination is valid keep the match, else exclude
        # gt from matching
        if not(valid is None):
            if valid[minComb[1]]:
                if not falsePositive:
                    bestMatch.append(minComb)
                    opAssigned[minComb[0]] = True
                    gtAssigned[minComb[1]] = True
            else:
                gtAssigned[minComb[1]] = True
                excludedGtBecauseInvalid.append(minComb[1])

        elif not falsePositive:
            # same as above but without checking for valid
            bestMatch.append(minComb)
            opAssigned[minComb[0]] = True
            gtAssigned[minComb[1]] = True

    bestMatch = np.array(bestMatch)
    # add false positives and false negatives to the matching
    # find which elements have been successfully assigned
    opAssigned = []
    gtAssigned = []
    for pair in bestMatch:
        opAssigned.append(pair[0])
        gtAssigned.append(pair[1])
    opAssigned.sort()
    gtAssigned.sort()

    falsePositives = []
    misses = []

    # handle false positives
    opIds = np.arange(len(pred_kps))
    # returns values of oIds that are not in opAssigned
    notAssignedIds = np.setdiff1d(opIds, opAssigned)
    for notAssignedId in notAssignedIds:
        falsePositives.append(notAssignedId)
    gtIds = np.arange(len(gtList))
    # returns values of gtIds that are not in gtAssigned
    notAssignedIdsGt = np.setdiff1d(gtIds, gtAssigned)

    # handle false negatives/misses
    for notAssignedIdGt in notAssignedIdsGt:
        if not(valid is None):  # if using the new matching
            if valid[notAssignedIdGt]:
                #print(imgPath + ': miss')
                misses.append(notAssignedIdGt)
            else:
                excludedGtBecauseInvalid.append(notAssignedIdGt)
        else:
            #print(imgPath + ': miss')
            misses.append(notAssignedIdGt)

    return bestMatch, falsePositives, misses  # tuples are (idx_pred_kps, idx_gt_kps)

def get_bbx_overlap(p1, p2):
    """
    Code modifed from https://github.com/Arthur151/ROMP/blob/4eebd3647f57d291d26423e51f0d514ff7197cb3/simple_romp/trace2/evaluation/eval_3DPW.py#L185
    """
    min_p1 = np.min(p1, axis=0)
    min_p2 = np.min(p2, axis=0)
    max_p1 = np.max(p1, axis=0)
    max_p2 = np.max(p2, axis=0)

    bb1 = {}
    bb2 = {}

    bb1['x1'] = min_p1[0]
    bb1['x2'] = max_p1[0]
    bb1['y1'] = min_p1[1]
    bb1['y2'] = max_p1[1]
    bb2['x1'] = min_p2[0]
    bb2['x2'] = max_p2[0]
    bb2['y1'] = min_p2[1]
    bb2['y2'] = max_p2[1]

    assert bb1['x1'] < bb1['x2']
    assert bb1['y1'] < bb1['y2']
    assert bb2['x1'] < bb2['x2']
    assert bb2['y1'] < bb2['y2']
    # determine the coordinates of the intersection rectangle
    x_left = max(bb1['x1'], bb2['x1'])
    y_top = max(bb1['y1'], bb2['y1'])
    x_right = min(bb1['x2'], bb2['x2'])
    y_bottom = min(bb1['y2'], bb2['y2'])

    # The intersection of two axis-aligned bounding boxes is always an
    # axis-aligned bounding box
    intersection_area = max(0, x_right - x_left + 1) * \
        max(0, y_bottom - y_top + 1)

    # compute the area of both AABBs
    bb1_area = (bb1['x2'] - bb1['x1'] + 1) * (bb1['y2'] - bb1['y1'] + 1)
    bb2_area = (bb2['x2'] - bb2['x1'] + 1) * (bb2['y2'] - bb2['y1'] + 1)

    # compute the intersection over union by taking the intersection
    # area and dividing it by the sum of prediction + ground-truth
    # areas - the interesection area
    iou = intersection_area / float(bb1_area + bb2_area - intersection_area)

    return iou

def avg_per_human(lst, default=0.0):
    values = [x for arr in lst for x in np.array(arr).flatten()]
    return np.mean(values) if values else default

def extract_metrics(file_path):
    with open(file_path, "r") as file:
        content = file.read()

    matches = re.findall(
        r'n_human: (\d+).*?PVE: ([\d.]+), PA-PVE: ([\d.]+), Metric-PVE: ([\d.]+), MPJPE: ([\d.]+), PA-MPJPE: ([\d.]+), Metric-MPJPE: ([\d.]+), RootError: ([\d.]+), W-MPJPE: ([\d.]+), WA-MPJPE: ([\d.]+), RTE: ([\d.]+), Scaled-RTE: ([\d.]+), Jitter: ([\d.]+), Foot-Sliding: ([\d.]+), Precision: ([\d.]+), Recall: ([\d.]+), F1-Score: ([\d.]+)',
        content
    )

    if not matches:
        return 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0

    total_humans = 0
    weighted_sums = [0.0] * 16

    for match in matches:
        n_human = int(match[0])
        metrics = [float(x) for x in match[1:]]

        total_humans += n_human
        for i, metric in enumerate(metrics):
            weighted_sums[i] += metric * n_human

    if total_humans == 0:
        return 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0

    averages = [ws / total_humans for ws in weighted_sums]

    return total_humans, *averages

def process_directory(directory):
    results = []
    for root, _, files in os.walk(directory):
        for file in sorted(files):
            if file.endswith(".txt"):
                file_path = os.path.join(root, file)
                result = extract_metrics(file_path)
                if result[0] > 0:
                    results.append(result)
    return results

def calculate_averages(results):
    if not results:
        return 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0

    total_humans = sum(r[0] for r in results)
    if total_humans == 0:
        return 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0

    weighted_sums = [sum(r[i] * r[0] for r in results) for i in range(1, 17)]
    averages = [ws / total_humans for ws in weighted_sums]

    return total_humans, *averages

def visualize(
    save_dir, img_path, view, gt_v3d_c, pred_v3d_c,
    K_to_proj, gt_K, bestMatch, smpl_face
):
    from dust3r.utils import render_meshes, denormalize_rgb
    from viser_utils import get_color
    from PIL import Image

    os.makedirs(save_dir, exist_ok=True)
    n_humans_i = pred_v3d_c.shape[0]

    # image
    img_array = denormalize_rgb(view.cpu().numpy())

    focal = K_to_proj[[0,1],[0,1]].cpu().numpy()
    princpt = K_to_proj[[0,1],[-1,-1]].cpu().numpy()
    gt_focal = gt_K[[0,1],[0,1]].cpu().numpy()
    gt_princpt = gt_K[[0,1],[-1,-1]].cpu().numpy()

    gt_color_indices = list(range(gt_v3d_c.shape[0]))
    pred_color_indices = [-1] * n_humans_i
    for (pid, gid) in bestMatch:
        pred_color_indices[pid] = gt_color_indices[gid]

    next_color_idx = max(gt_color_indices) + 1 if gt_color_indices else 0
    for i in range(len(pred_color_indices)):
        if pred_color_indices[i] == -1:
            pred_color_indices[i] = next_color_idx
            next_color_idx += 1

    # gt
    gt_verts, gt_faces = [], []
    for j in range(gt_v3d_c.shape[0]):
        gt_verts.append(gt_v3d_c[j].cpu().numpy().reshape(-1,3))
        gt_faces.append(smpl_face)

    gt_colors = [get_color(gt_color_indices[j])/255 for j in range(len(gt_verts))]
    gt_rend_array = render_meshes(img_array.copy(),
                                    gt_verts,
                                    gt_faces,
                                    {'focal': gt_focal, 'princpt': gt_princpt},
                                    color=gt_colors)

    # pred
    pred_verts, pred_faces = [], []
    for j in range(n_humans_i):
        pred_verts.append(pred_v3d_c[j].cpu().numpy().reshape(-1,3))
        pred_faces.append(smpl_face)

    pred_colors = [get_color(pred_color_indices[j])/255 for j in range(len(pred_verts))]
    pred_rend_array = render_meshes(img_array.copy(),
                                    pred_verts,
                                    pred_faces,
                                    {'focal': focal, 'princpt': princpt},
                                    color=pred_colors)

    img = np.concatenate([img_array, pred_rend_array, gt_rend_array], 1)
    Image.fromarray(img).save(
        os.path.join(
            f"{save_dir}/{os.path.splitext(os.path.basename(img_path))[0]}.jpg"
            ))

def write_log(log_path, dataset, seq, counter, metrics):
    with open(log_path, "a") as f:
        f.write(
            f"{dataset}-{seq: <16} |  "
            f"n_human: {counter['n_human']:06d} | "
            f"PVE: {avg_per_human(metrics['ca_pve']):.1f}, "
            f"PA-PVE: {avg_per_human(metrics['pa_pve']):.1f}, "
            f"Metric-PVE: {avg_per_human(metrics['me_pve']):.1f}, "
            f"MPJPE: {avg_per_human(metrics['ca_mpjpe']):.1f}, "
            f"PA-MPJPE: {avg_per_human(metrics['pa_mpjpe']):.1f}, "
            f"Metric-MPJPE: {avg_per_human(metrics['me_mpjpe']):.1f}, "
            f"RootError: {avg_per_human(metrics['rt_error']):.1f}, "
            f"W-MPJPE: {avg_per_human(metrics['wa2_mpjpe']):.1f}, "
            f"WA-MPJPE: {avg_per_human(metrics['waa_mpjpe']):.1f}, "
            f"RTE: {avg_per_human(metrics['rte']):.1f}, "
            f"Scaled-RTE: {avg_per_human(metrics['rte_scaled']):.1f}, "
            f"Jitter: {avg_per_human(metrics['jitter']):.1f}, "
            f"Foot-Sliding: {avg_per_human(metrics['fs']):.1f}, "
            f"Precision: {metrics['precision']:.1f}, "
            f"Recall: {metrics['recall']:.1f}, "
            f"F1-Score: {metrics['f1_score']:.1f}\n"
        )

def get_summary_log(summary):
    """Generate summary log for evaluation results"""
    return (
        f"EVALUATION SUMMARY\n"
        f"{'='*7}EVALUATION SUMMARY{'='*7}\n"
        f"Total Humans: {summary[0]}\n"
        f"\n"
        f"Camera Coordinate Metrics (mm):\n"
        f"  PVE:         {summary[1]:6.1f}\n"
        f"  PA-PVE:      {summary[2]:6.1f}\n"
        f"  Metric-PVE:  {summary[3]:6.1f}\n"
        f"\n"
        f"  MPJPE:       {summary[4]:6.1f}\n"
        f"  PA-MPJPE:    {summary[5]:6.1f}\n"
        f"  Metric-MPJPE: {summary[6]:6.1f}\n"
        f"  Root-Error:  {summary[7]:6.1f}\n"
        f"\n"
        f"Global Coordinate Metrics (cm):\n"
        f"  W-MPJPE:     {summary[8]:6.1f}\n"
        f"  WA-MPJPE:    {summary[9]:6.1f}\n"
        f"  RTE:         {summary[10]:6.1f}\n"
        f"  Scaled-RTE:  {summary[11]:6.1f}\n"
        f"  Jitter:      {summary[12]:6.1f}\n"
        f"  Foot-Sliding: {summary[13]:6.1f}\n"
        f"\n"
        f"Detection Metrics (%):\n"
        f"  Precision:   {summary[14]:6.1f}\n"
        f"  Recall:      {summary[15]:6.1f}\n"
        f"  F1-Score:    {summary[16]:6.1f}\n"
        f"{'='*32}\n"
    )

# Evaluation metrics
# Code modifed from https://github.com/zju3dv/GVHMR/blob/088caff492aa38c2d82cea363b78a3c65a83118f/hmr4d/utils/eval/eval_utils.py

def compute_jpe(S1, S2):
    return torch.sqrt(((S1 - S2) ** 2).sum(dim=-1)).mean(dim=-1).numpy()

def compute_perjoint_jpe(S1, S2):
    return torch.sqrt(((S1 - S2) ** 2).sum(dim=-1)).numpy()

def batch_compute_similarity_transform_torch(S1, S2):
    '''
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    '''
    S1 = S1.permute(0,2,1)
    S2 = S2.permute(0,2,1)
    transposed = True

    # 1. Remove mean.
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=True)

    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1**2, dim=1).sum(dim=1)

    # 3. The outer product of X1 and X2.
    K = X1.bmm(X2.permute(0,2,1))

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, s, V = torch.svd(K)

    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0],1,1)
    Z[:,-1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0,2,1))))

    # Construct R.
    R = V.bmm(Z.bmm(U.permute(0,2,1)))

    # 5. Recover scale.
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1

    # 6. Recover translation.
    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))

    # 7. Error:
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t

    if transposed:
        S1_hat = S1_hat.permute(0,2,1)

    return S1_hat

def batch_align_by_pelvis(data_list, pelvis_idxs):
    """
    Assumes data is given as [pred_j3d, target_j3d, pred_verts, target_verts].
    Each data is in shape of (batch, num_points, 3)
    Pelvis is notated as one / two joints indices.
    Align all data to the corresponding pelvis location.
    """

    pred_j3d, target_j3d, pred_verts, target_verts = data_list

    pred_pelvis = pred_j3d[:, pelvis_idxs].mean(dim=1, keepdims=True).clone()
    target_pelvis = target_j3d[:, pelvis_idxs].mean(dim=1, keepdims=True).clone()

    # Align to the pelvis
    pred_j3d = pred_j3d - pred_pelvis
    target_j3d = target_j3d - target_pelvis
    pred_verts = pred_verts - pred_pelvis
    target_verts = target_verts - target_pelvis

    return (pred_j3d, target_j3d, pred_verts, target_verts, pred_pelvis, target_pelvis)

def align_pcl(Y, X, weight=None, fixed_scale=False):
    """
    align similarity transform to align X with Y using umeyama method
    X' = s * R * X + t is aligned with Y
    :param Y (*, N, 3) first trajectory
    :param X (*, N, 3) second trajectory
    :param weight (*, N, 1) optional weight of valid correspondences
    :returns s (*, 1), R (*, 3, 3), t (*, 3)
    """
    *dims, N, _ = Y.shape
    N = torch.ones(*dims, 1, 1) * N

    if weight is not None:
        Y = Y * weight
        X = X * weight
        N = weight.sum(dim=-2, keepdim=True)  # (*, 1, 1)

    # subtract mean
    my = Y.sum(dim=-2) / N[..., 0]  # (*, 3)
    mx = X.sum(dim=-2) / N[..., 0]
    y0 = Y - my[..., None, :]  # (*, N, 3)
    x0 = X - mx[..., None, :]

    if weight is not None:
        y0 = y0 * weight
        x0 = x0 * weight

    # correlation
    C = torch.matmul(y0.transpose(-1, -2), x0) / N  # (*, 3, 3)
    U, D, Vh = torch.linalg.svd(C)  # (*, 3, 3), (*, 3), (*, 3, 3)

    S = torch.eye(3).reshape(*(1,) * (len(dims)), 3, 3).repeat(*dims, 1, 1)
    neg = torch.det(U) * torch.det(Vh.transpose(-1, -2)) < 0
    S[neg, 2, 2] = -1

    R = torch.matmul(U, torch.matmul(S, Vh))  # (*, 3, 3)

    D = torch.diag_embed(D)  # (*, 3, 3)
    if fixed_scale:
        s = torch.ones(*dims, 1, device=Y.device, dtype=torch.float32)
    else:
        var = torch.sum(torch.square(x0), dim=(-1, -2), keepdim=True) / N  # (*, 1, 1)
        s = torch.diagonal(torch.matmul(D, S), dim1=-2, dim2=-1).sum(dim=-1, keepdim=True) / var[..., 0]  # (*, 1)

    t = my - s * torch.matmul(R, mx[..., None])[..., 0]  # (*, 3)

    return s, R, t

def global_align_joints(gt_joints, pred_joints):
    """
    :param gt_joints (T, J, 3)
    :param pred_joints (T, J, 3)
    """
    s_glob, R_glob, t_glob = align_pcl(gt_joints.reshape(-1, 3), pred_joints.reshape(-1, 3))
    pred_glob = s_glob * torch.einsum("ij,tnj->tn
# ... truncated ...


todevice

import cv2

import os
import re
import numpy as np
import torch
import roma
import math
import tqdm
import matplotlib as mpl
import matplotlib.cm as cm
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg
from scipy.spatial.transform import Rotation
from eval.relpose.evo_utils import *
from PIL import Image
import imageio.v2 as iio
from matplotlib.figure import Figure
from itertools import product

def todevice(batch, device, callback=None, non_blocking=False):
    """Transfer some variables to another device (i.e. GPU, CPU:torch, CPU:numpy).

    batch: list, tuple, dict of tensors or other things
    device: pytorch device or 'numpy'
    callback: function that would be called on every sub-elements.
    """
    if callback:
        batch = callback(batch)

    if isinstance(batch, dict):
        return {k: todevice(v, device) for k, v in batch.items()}

    if isinstance(batch, (tuple, list)):
        return type(batch)(todevice(x, device) for x in batch)

    x = batch
    if device == "numpy":
        if isinstance(x, torch.Tensor):
            x = x.detach().cpu().numpy()
    elif x is not None:
        if isinstance(x, np.ndarray):
            x = torch.from_numpy(x)
        if torch.is_tensor(x):
            x = x.to(device, non_blocking=non_blocking)
    return x

to_device = todevice  # alias

def to_numpy(x):
    return todevice(x, "numpy")

def compute_prf1(count, miss, fp):
    """
    Code modified from https://github.com/Arthur151/ROMP/blob/4eebd3647f57d291d26423e51f0d514ff7197cb3/simple_romp/evaluation/RH_evaluation/evaluation.py#L90
    """
    if count == 0:
        return 0, 0, 0
    all_tp = count - miss
    all_fp = fp
    all_fn = miss
    if all_tp == 0:
        return 0., 0., 0.


        return {k: todevice(v, device) for k, v in batch.items()}

    if isinstance(batch, (tuple, list)):
        return type(batch)(todevice(x, device) for x in batch)

    x = batch
    if device == "numpy":
        if isinstance(x, torch.Tensor):
            x = x.detach().cpu().numpy()
    elif x is not None:
        if isinstance(x, np.ndarray):
            x = torch.from_numpy(x)
        if torch.is_tensor(x):
            x = x.to(device, non_blocking=non_blocking)
    return x

to_device = todevice  # alias

def to_numpy(x):
    return todevice(x, "numpy")

def compute_prf1(count, miss, fp):
    """
    Code modified from https://github.com/Arthur151/ROMP/blob/4eebd3647f57d291d26423e51f0d514ff7197cb3/simple_romp/evaluation/RH_evaluation/evaluation.py#L90
    """
    if count == 0:
        return 0, 0, 0
    all_tp = count - miss
    all_fp = fp
    all_fn = miss
    if all_tp == 0:
        return 0., 0., 0.
    all_f1_score = round(all_tp / (all_tp + 0.5 * (all_fp + all_fn)), 2)
    all_recall = round(all_tp / (all_tp + all_fn), 2)
    all_precision = round(all_tp / (all_tp + all_fp), 2)
    return 100. * all_precision, 100.* all_recall, 100. * all_f1_score

def match_2d_greedy(


    x = batch
    if device == "numpy":
        if isinstance(x, torch.Tensor):
            x = x.detach().cpu().numpy()
    elif x is not None:
        if isinstance(x, np.ndarray):
            x = torch.from_numpy(x)
        if torch.is_tensor(x):
            x = x.to(device, non_blocking=non_blocking)
    return x

to_device = todevice  # alias

def to_numpy(x):
    return todevice(x, "numpy")

def compute_prf1(count, miss, fp):
    """
    Code modified from https://github.com/Arthur151/ROMP/blob/4eebd3647f57d291d26423e51f0d514ff7197cb3/simple_romp/evaluation/RH_evaluation/evaluation.py#L90
    """
    if count == 0:
        return 0, 0, 0
    all_tp = count - miss
    all_fp = fp
    all_fn = miss
    if all_tp == 0:
        return 0., 0., 0.
    all_f1_score = round(all_tp / (all_tp + 0.5 * (all_fp + all_fn)), 2)
    all_recall = round(all_tp / (all_tp + all_fn), 2)
    all_precision = round(all_tp / (all_tp + all_fp), 2)
    return 100. * all_precision, 100.* all_recall, 100. * all_f1_score

def match_2d_greedy(
        pred_kps,
        gtkp,
        valid_mask,
        imgPath=None,
        baseline=None,
        iou_thresh=0.05,
        valid=None,
        ind=-1):
    '''
    Code modified from: https://github.com/Arthur151/ROMP/blob/4eebd3647f57d291d26423e51f0d514ff7197cb3/simple_romp/trace2/evaluation/eval_3DPW.py#L232
    matches groundtruth keypoints to the detection by considering all possible matchings.
    :return: best possible matching, a list of tuples, where each tuple corresponds to one match of pred_person.to gt_person.
            the order within one tuple is as follows (idx_pred_kps, idx_gt_kps)
    '''
    predList = np.arange(len(pred_kps))
    gtList = np.arange(len(gtkp))
    # get all pairs of elements in pred_kps, gtkp


    return todevice(x, "numpy")

def compute_prf1(count, miss, fp):
    """
    Code modified from https://github.com/Arthur151/ROMP/blob/4eebd3647f57d291d26423e51f0d514ff7197cb3/simple_romp/evaluation/RH_evaluation/evaluation.py#L90
    """
    if count == 0:
        return 0, 0, 0
    all_tp = count - miss
    all_fp = fp
    all_fn = miss
    if all_tp == 0:
        return 0., 0., 0.
    all_f1_score = round(all_tp / (all_tp + 0.5 * (all_fp + all_fn)), 2)
    all_recall = round(all_tp / (all_tp + all_fn), 2)
    all_precision = round(all_tp / (all_tp + all_fp), 2)
    return 100. * all_precision, 100.* all_recall, 100. * all_f1_score

def match_2d_greedy(
        pred_kps,
        gtkp,
        valid_mask,
        imgPath=None,
        baseline=None,
        iou_thresh=0.05,
        valid=None,
        ind=-1):
    '''
    Code modified from: https://github.com/Arthur151/ROMP/blob/4eebd3647f57d291d26423e51f0d514ff7197cb3/simple_romp/trace2/evaluation/eval_3DPW.py#L232
    matches groundtruth keypoints to the detection by considering all possible matchings.
    :return: best possible matching, a list of tuples, where each tuple corresponds to one match of pred_person.to gt_person.
            the order within one tuple is as follows (idx_pred_kps, idx_gt_kps)
    '''
    predList = np.arange(len(pred_kps))
    gtList = np.arange(len(gtkp))
    # get all pairs of elements in pred_kps, gtkp
    # all combinations of 2 elements from l1 and l2
    combs = list(product(predList, gtList))

    errors_per_pair = {}
    errors_per_pair_list = []
    for comb in combs:
        vmask = valid_mask[comb[1]]
        assert vmask.sum()>0, print('no valid points')
        errors_per_pair[str(comb)] = np.linalg.norm(pred_kps[comb[0]][vmask, :2] - gtkp[comb[1]][vmask, :2], 2)
        errors_per_pair_list.append(errors_per_pair[str(comb)])

    gtAssigned = np.zeros((len(gtkp),), dtype=bool)
    opAssigned = np.zeros((len(pred_kps),), dtype=bool)
    errors_per_pair_list = np.array(errors_per_pair_list)

    bestMatch = []
    excludedGtBecauseInvalid = []
    falsePositiveCounter = 0
    while np.sum(gtAssigned) < len(gtAssigned) and np.sum(
            opAssigned) + falsePositiveCounter < len(pred_kps):
        found = False
        falsePositive = False
        while not(found):
            if sum(np.inf == errors_per_pair_list) == len(
                    errors_per_pair_list):
                print('something went wrong here')

            minIdx = np.argmin(errors_per_pair_list)
            minComb = combs[minIdx]
            # compute IOU
            iou = get_bbx_overlap(
                pred_kps[minComb[0]], gtkp[minComb[1]]) #, imgPath, baseline)
            # if neither prediction nor ground truth has been matched before and iou
            # is larger than threshold
            if not(opAssigned[minComb[0]]) and not(
                    gtAssigned[minComb[1]]) and iou >= iou_thresh:
                #print(imgPath + ': found matching')
                found = True
                errors_per_pair_list[minIdx] = np.inf
            else:
                errors_per_pair_list[minIdx] = np.inf
                # if errors_per_pair_list[minIdx] >
                # matching_threshold*headBboxs[combs[minIdx][1]]:
                if iou < iou_thresh:
                    #print(
                    #   imgPath + ': false positive detected using threshold')
                    found = True
                    falsePositive = True
                    falsePositiveCounter += 1

        # if ground truth of combination is valid keep the match, else exclude
        # gt from matching
        if not(valid is None):
            if valid[minComb[1]]:
                if not falsePositive:
                    bestMatch.append(minComb)
                    opAssigned[minComb[0]] = True
                    gtAssigned[minComb[1]] = True
            else:
                gtAssigned[minComb[1]] = True
                excludedGtBecauseInvalid.append(minComb[1])

        elif not falsePositive:
            # same as above but without checking for valid
            bestMatch.append(minComb)
            opAssigned[minComb[0]] = True
            gtAssigned[minComb[1]] = True

    bestMatch = np.array(bestMatch)
    # add false positives and false negatives to the matching
    # find which elements have been successfully assigned
    opAssigned = []
    gtAssigned = []
    for pair in bestMatch:
        opAssigned.append(pair[0])
        gtAssigned.append(pair[1])
    opAssigned.sort()
    gtAssigned.sort()

    falsePositives = []
    misses = []

    # handle false positives
    opIds = np.arange(len(pred_kps))
    # returns values of oIds that are not in opAssigned
    notAssignedIds = np.setdiff1d(opIds, opAssigned)
    for notAssignedId in notAssignedIds:
        falsePositives.append(notAssignedId)
    gtIds = np.arange(len(gtList))
    # returns values of gtIds that are not in gtAssigned
    notAssignedIdsGt = np.setdiff1d(gtIds, gtAssigned)

    # handle false negatives/misses
    for notAssignedIdGt in notAssignedIdsGt:
        if not(valid is None):  # if using the new matching
            if valid[notAssignedIdGt]:
                #print(imgPath + ': miss')
                misses.append(notAssignedIdGt)
            else:
                excludedGtBecauseInvalid.append(notAssignedIdGt)
        else:
            #print(imgPath + ': miss')
            misses.append(notAssignedIdGt)

    return bestMatch, falsePositives, misses  # tuples are (idx_pred_kps, idx_gt_kps)

def get_bbx_overlap(p1, p2):
    """
    Code modifed from https://github.com/Arthur151/ROMP/blob/4eebd3647f57d291d26423e51f0d514ff7197cb3/simple_romp/trace2/evaluation/eval_3DPW.py#L185
    """
    min_p1 = np.min(p1, axis=0)
    min_p2 = np.min(p2, axis=0)
    max_p1 = np.max(p1, axis=0)
    max_p2 = np.max(p2, axis=0)

    bb1 = {}
    bb2 = {}

    bb1['x1'] = min_p1[0]
    bb1['x2'] = max_p1[0]
    bb1['y1'] = min_p1[1]
    bb1['y2'] = max_p1[1]
    bb2['x1'] = min_p2[0]
    bb2['x2'] = max_p2[0]
    bb2['y1'] = min_p2[1]