import numpy as np
import torch

DEFAULT_AGGR_SETTINGS = dict(
    USE_SMPL=True,
    GRID_RESOLUTION=48,
    GLOBAL=dict(
        USE_STAR_POSE_TEMPLATE=True,
        VISUALIZE_POSE_NUM=3,
        AUXILIARY_EXIST_OK=False,
        TAG_AUXILIARY=False,
        ADD_HUMAN_FOR_VIS=False,
        USE_TORCH=True,
        DEVICE='cuda',
        INTERACTION_REGIONS_PRECOMPUTE_SETTINGS=dict(
            PART_INTERACTION_REGION_PRECOMPUTE_DEVICE='cpu',
            PART_INTERACTION_REGION_PRECOMPUTE_CHUNK_SIZE=50000,
        ),
        ZEROBETA=True,
        LEARNED_K=True,
        SYMMETRIC_AUGMENTATION=False,
        MAX_CAM_SAVE_NUM=5,
    ),
    SMPL=dict(
        GRID_LENGTH=3.0,
        CENTER=np.array([0,0,0],dtype=np.float32),
        O3D_VOXEL_COLOR_FILLED=[1,0,0],
        O3D_VOXEL_COLOR_UNFILLED=[0.8,0.8,0.8],
        O3D_VOXEL_SPHERE_RESOLUTION=2,
        GRID_COLOR=[0,0,0],
        FRUSTUM_SIZE=0.25,
        CAMERA_SAMPLING='azimuth-uniform-fine',
        USE_NORMALIZED_K=False,
        MANUAL_SELECTION=None,
        CAM_COLOR=[0,0,0],
        CANONCAM2SMPL_TRANS=np.zeros([3], dtype=np.float32),
        V_TEMPLATE_COLOR=np.array([0.5,0.5,0.5]),
        BODY_PART_DEFINE_METHOD='quant',
        BODY_SEGMENTATION_JSON_PTH="./imports/meshcapade_wiki/assets/SMPL_body_segmentation/smpl/smpl_vert_segmentation.json",
        LBS_PRECOMPUTE_SETTINGS=dict(
            PRECOMPUTE_DEVICE="cpu",
            CHUNK_SIZE=50000,
            LBS_NUM_NEIGHBORS=30,
            LBS_MERGE_METHOD='nn_inv_distance_mean+identity_mixing_linear_ratio_with_constant',
            LBS_SMOOTHING_METHOD='laplace-smoothing+identity_thresh_preserve',
            LBS_NUM_SMOOTH_TIMES=30,
            LBS_Y_SCALE=0.5,
            LBS_DIST_ORIGIN_THRESH=0.8,
            LBS_IDENTITY_CENTER=np.zeros([3], dtype=np.float32),
            LBS_IDENTITY_DIST_THRESHOLD=1.0,
        ),
    )
)
constants/aggr_settings/default_aggr_core.py
jellyheadandrew/CHORUS
QUAL_AGGR_SETTINGS = {
    "qual:demo":
    dict(
        USE_SMPL = True,
        USE_STAR_POSE_TEMPLATE = True,
        SYMMETRIC_AUGMENTATION = False,
        GRID_LENGTH_SMPL = 3.0,
        GRID_RESOLUTION = 48,
        BODY_PART_DEFINE_METHOD_SMPL = "qual:demo",
        CAMERA_SAMPLING_SMPL = "azimuth-uniform-fine",
        LBS_MERGE_METHOD_SMPL =  "nn_inv_distance_mean+identity_mixing_linear_ratio_with_constant",
        LBS_SMOOTHING_METHOD_SMPL = "laplace-smoothing+identity_thresh_preserve",
        LBS_NUM_NEIGHBORS_SMPL = 30,
        LBS_NUM_SMOOTH_TIMES_SMPL = 30, 
        LBS_Y_SCALE_SMPL = 0.5,
        LBS_DIST_ORIGIN_THRESH_SMPL = 0.8,
        LBS_IDENTITY_CENTER_SMPL = None,
        LBS_IDENTITY_DIST_THRESHOLD_SMPL = 0.8,
        DATA_RATIO = 1.,
        NAME="qual:demo",
    ),
}
constants/aggr_settings/qual.py
jellyheadandrew/CHORUS
QUANT_AGGR_SETTINGS = {
    "quant:full": dict(
        USE_SMPL = True,
        USE_STAR_POSE_TEMPLATE = True,
        SYMMETRIC_AUGMENTATION = False,
        GRID_LENGTH_SMPL = 3.0,  
        GRID_RESOLUTION = 48,
        BODY_PART_DEFINE_METHOD_SMPL = "quant",
        CAMERA_SAMPLING_SMPL = "azimuth-uniform-fine",
        LBS_MERGE_METHOD_SMPL =  "nn_inv_distance_mean+identity_mixing_linear_ratio_with_constant",
        LBS_SMOOTHING_METHOD_SMPL = "laplace-smoothing+identity_thresh_preserve",
        LBS_NUM_NEIGHBORS_SMPL = 30,
        LBS_NUM_SMOOTH_TIMES_SMPL = 30, 
        LBS_Y_SCALE_SMPL = 0.5,
        LBS_DIST_ORIGIN_THRESH_SMPL = 0.8,
        LBS_IDENTITY_CENTER_SMPL = None,
        LBS_IDENTITY_DIST_THRESHOLD_SMPL = 0.8,
        DATA_RATIO = 1.,
        NAME="quant:full",
    ),
}
constants/aggr_settings/quant.py
jellyheadandrew/CHORUS
EPSILON_DEFAULT = 0.13 # interaction-distance threshold
ALL_BODY_SEGMENTATION_NAMES = [ # 24 different body parts
    'rightHand', 
    'rightUpLeg', 
    'leftArm', 
    'leftLeg', 
    'leftToeBase', 
    'leftFoot', 
    'spine1', 
    'spine2', 
    'leftShoulder', 
    'rightShoulder', 
    'rightFoot', 
    'head', 
    'rightArm', 
    'leftHandIndex1', 
    'rightLeg', 
    'rightHandIndex1', 
    'leftForeArm', 
    'rightForeArm', 
    'neck', 
    'rightToeBase', 
    'spine', 
    'leftUpLeg', 
    'leftHand', 
    'hips'
]
BODY_PART_CORRESPONDENCE = { # key: 'body_part_define_method'
    # quant (dummy)
    'quant': {
        'fullbody': {
            'correspondence_map': {
                'pass': [
                    ALL_BODY_SEGMENTATION_NAMES,
                ],
                'block': [
                ]
            },
            'interaction_threshold_per_world': {'SMPL': EPSILON_DEFAULT}
        },        
    },
    # demo
    'qual:demo': {
        'rightArm': {
            'correspondence_map': {
                'pass': [
                    ['rightArm', 'rightForeArm'],
                ],
                'block': [
                ]
            },
            'interaction_threshold_per_world': {'SMPL': EPSILON_DEFAULT}
        },
        'leftFoot': {
            'correspondence_map': {
                'pass': [
                    ['leftFoot', 'leftToeBase'],
                ],
                'block': [
                ]
            },
            'interaction_threshold_per_world': {'SMPL': EPSILON_DEFAULT}
        },
        'torso': {
            'correspondence_map': {
                'pass': [
                    ['spine', 'spine1', 'spine2', 'leftShoulder', 'rightShoulder'],
                ],
                'block': [
                ]
            },
            'interaction_threshold_per_world': {'SMPL': EPSILON_DEFAULT}
        },
    },
}
constants/aggr_settings/semantic_clustering.py
jellyheadandrew/CHORUS
import json

# coco paths
COCO_METADATA_PTH = "constants/coco_metadata.pickle"
COCO_THING_CLASSES_PTH="constants/coco_thing_classes.json"
COCO_VAL_ANNOTATION_PTH = './imports/COCO/annotations/instances_val2014.json'
COCO_VAL_IMAGE_DIR = './imports/COCO/images/val2014'

# lvis paths
LVIS_METADATA_PTH="constants/lvis_metadata.pickle"
LVIS_THING_CLASSES_PTH="constants/lvis_thing_classes.json"

# coco-eft paths
COCO_EFT_VAL_PTH = './imports/eft/eft_fit/COCO2014-Val-ver10.json'
EXTENDED_COCO_EFT_VAL_SAVE_PTH = './imports/eft/eft_fit/extended_coco2014_eft_val.pickle'

# coco categories: if it exists, load. if it doesn't exist, create & save.
with open(COCO_THING_CLASSES_PTH, "r") as rf:
    coco_thing_classes = json.load(rf)
COCO_CLASS_ID2NAME = {idx: name for idx, name in enumerate(coco_thing_classes)}
COCO_CLASS_NAME2ID = {v: k for k, v in COCO_CLASS_ID2NAME.items()}

# lvis categories
with open(LVIS_THING_CLASSES_PTH, "r") as rf:
    lvis_thing_classes = json.load(rf)
LVIS_CLASS_ID2NAME = {idx: name for idx, name in enumerate(lvis_thing_classes)}
LVIS_CLASS_NAME2ID = {v: k for k, v in LVIS_CLASS_ID2NAME.items()}

# category exceptions
CATEGORY_EXCEPTIONS = {
    'surfboard-demo': COCO_CLASS_NAME2ID['surfboard']
}
constants/datasets.py
jellyheadandrew/CHORUS
DEFAULT_NUM_ENFORCE_MULTIVIEW = 30
DEFAULT_DOWNSCALE_LENGTH = 32
DEFAULT_PRECISON_THRESHOLDS = [round(i * 0.01, 2) for i in range(1,100)]
constants/evaluation.py
jellyheadandrew/CHORUS
DEFAULT_BBOX_OVERLAP_THRES = 0.1
DEFAULT_KEYPOINT_CONFIDENCE_THRES = 0.7
DEFAULT_FILTER_MULTI_PERSON_OBJECT_METHODS = [
    "only-one-person-in-image", 
    "only-one-object-in-image", 
    "1-bbox-per-person"
]
DEFAULT_KEYPOINT_FILTER_NUM=3
DEFAULT_FILTER_SETTING_NUM=3
KEYPOINT_FILTERS = {
        # 1. [face (hand allowed)]
        1: "nose/left_eye/right_eye/left_ear/right_ear | left_knee/right_knee/left_ankle/right_ankle/left_hip_extra/right_hip_extra",
        # 2. [face (strict)]
        2: "nose/left_eye/right_eye/left_ear/right_ear | left_knee/right_knee/left_ankle/right_ankle/left_hip_extra/right_hip_extra/left_wrist/right_wrist/left_elbow/right_elbow",
        # 3. [full body (one of shoulder, one of hip)]
        3: "left_shoulder,right_shoulder / left_hip_extra,right_hip_extra | ",
        # 4. [full body (one of shoulder, one of hip, one of elbow or hand)]
        4: "left_shoulder,right_shoulder / left_hip_extra,right_hip_extra / left_elbow, right_elbow | ",
        # 5. [full body (one of shoulder, one of hip, one of knees or ankle)]
        5: "left_shoulder,right_shoulder / left_hip_extra,right_hip_extra / left_knee, right_knee, left_ankle, right_ankle | ",
        # 6. [hand (elbow, shoulder allowed: both sides available)]
        6: "left_wrist, right_wrist | left_knee/right_knee/left_ankle/right_ankle/left_hip_extra/right_hip_extra/nose/left_eye/right_eye/left_ear/right_ear",
        # 7. [hand (elbow allowed: both sides available)]
        7: "left_wrist, right_wrist | left_knee/right_knee/left_ankle/right_ankle/left_hip_extra/right_hip_extra/nose/left_eye/right_eye/left_ear/right_ear/left_shoulder/right_shoulder",
        # 8. [hand (strict: both sides available)]
        8: "left_wrist, right_wrist | left_knee/right_knee/left_ankle/right_ankle/left_hip_extra/right_hip_extra/nose/left_eye/right_eye/left_ear/right_ear/left_shoulder/right_shoulder/left_elbow/right_elbow",
        # 9. [hand (elbow allowed: only left)]
        9: "left_wrist | left_knee/right_knee/left_ankle/right_ankle/left_hip_extra/right_hip_extra/nose/left_eye/right_eye/left_ear/right_ear/left_shoulder/right_shoulder/right_wrist/right_elbow",
        # 10. [hand (elbow allowed: only right)]
        10: "right_wrist | left_knee/right_knee/left_ankle/right_ankle/left_hip_extra/right_hip_extra/nose/left_eye/right_eye/left_ear/right_ear/left_shoulder/right_shoulder/left_wrist/left_elbow",
        # 11. [full body (one of shoulder, one of hip, one of knees, one of ankle)]
        11: "left_shoulder,right_shoulder / left_hip_extra,right_hip_extra / left_knee, right_knee / left_ankle, right_ankle | ",
        # 12. [face (hand allowed, side view allowed)]
        12: "left_eye,right_eye/left_ear,right_ear,nose | left_knee/right_knee/left_ankle/right_ankle/left_hip_extra/right_hip_extra",
        # 13. [face (strict, side view allowed)]
        13: "left_eye,right_eye/left_ear,right_ear,nose | left_knee/right_knee/left_ankle/right_ankle/left_hip_extra/right_hip_extra/left_wrist/right_wrist/left_elbow/right_elbow",
        # 14. [face+body]
        14: "nose/left_eye,right_eye/left_ear,right_ear /left_shoulder,right_shoulder | ",
        # 15. [no keypoint filter]
        15: "|"
}
constants/filtering.py
jellyheadandrew/CHORUS
FRANKMOCAP_DIR="./imports/frankmocap"
BODY_MOCAP_REGRESSOR_CKPT=f"{FRANKMOCAP_DIR}/extra_data/body_module/pretrained_weights/2020_05_31-00_50_43-best-51.749683916568756.pt"
BODY_MOCAP_SMPL_DIR=f"{FRANKMOCAP_DIR}/extra_data/smpl"
BODY_MOCAP_SMPL_PTH = f"{BODY_MOCAP_SMPL_DIR}/basicmodel_neutral_lbs_10_207_0_v1.1.0.pkl"
VISUALIZER_MODE="opengl"
HMR_INPUT_SIZE = 224
DEFAULT_EXCONF_FRANKMOCAP = 0.98
DEFAULT_MINOVERLAP_FRANKMOCAP = 0.8
SMPL_KEYPOINTS_IDX2NAME = {
    0:	'pelvis',
    1:	'left_hip',
    2:	'right_hip',
    3:	'spine1',
    4:	'left_knee',
    5:	'right_knee',
    6:	'spine2',
    7:	'left_ankle',
    8:	'right_ankle',
    9:	'spine3',
    10:	'left_foot',
    11:	'right_foot',
    12:	'neck',
    13:	'left_collar',
    14:	'right_collar',
    15:	'head',
    16:	'left_shoulder',
    17:	'right_shoulder',
    18:	'left_elbow',
    19:	'right_elbow',
    20:	"left_wrist",
    21:	"right_wrist",
    22: "left_hand",
    23: "right_hand",
}
constants/frankmocap.py
jellyheadandrew/CHORUS
DEFAULT_MMPOSE_DETECTOR_KEY = "HRNet+Dark384x288_COCO"
DEFAULT_KEYPOINT_THRESHOLD = 0.7
DEFAULT_KEYPOINT_BBOX_THRESHOLD = 0.0
MMPOSE_DETECTOR_SETTINGS = {
    # HRNet+Dark:384x288, COCO-keypoint
    "HRNet+Dark384x288_COCO":
        [
            f'./imports/mmpose/configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/coco/hrnet_w48_coco_384x288_dark.py',
            'https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w48_coco_384x288_dark-e881a4b6_20210203.pth',
        ]
}

# keypoint names (coco)
COCO_2DKEYPOINT_IDX2NAME = {
    i: name for i, name in 
    enumerate([
        "nose","left_eye","right_eye","left_ear","right_ear","left_shoulder",
        "right_shoulder","left_elbow","right_elbow","left_wrist","right_wrist",
        "left_hip_extra","right_hip_extra","left_knee","right_knee","left_ankle","right_ankle"
        ])
    }
COCO_2DKEYPOINT_NAME2IDX = {v: k for k, v in COCO_2DKEYPOINT_IDX2NAME.items()}
constants/keypoint_detection.py
jellyheadandrew/CHORUS
from easydict import EasyDict

AUGMENTATION_PROMPT_PTH = "constants/augmentations.txt"

DEFAULT_LDM_MODEL_KEY = "stable-v1-4"
DEFAULT_BATCH_SIZE = 6

MODELS = {
    "stable-v1-4": dict(
        config="imports/ldm/configs/stable-diffusion/v1-inference.yaml", 
        ckpt="imports/ldm/models/ldm/stable-diffusion/sd-v1-4.ckpt",
        plms=True
    )
}

# Classifier-Guidance Scale: 7.5 / Image Size: 512 / PLMS Sampling: On
def GetLDMOptions(model_key: str, models: dict):
    # Basic Options
    opt = EasyDict()
    opt.H = 512
    opt.W = 512
    opt.C = 4
    opt.f = 8
    opt.ddim_steps = 50
    opt.fixed_code = False
    opt.ddim_eta = 0.0
    # opt.seed = 42
    opt.scale = 7.5
    opt.precision = "autocast"
    opt.model_type = model_key
    opt.plms = models[opt.model_type]["plms"]
    opt.config = models[opt.model_type]["config"]
    opt.ckpt = models[opt.model_type]["ckpt"]
    if "C" in list(models[opt.model_type].keys()):
        print("C changed.")
        opt.C = models[opt.model_type]["C"]
    return opt
constants/ldm.py
jellyheadandrew/CHORUS
from easydict import EasyDict
import numpy as np

QUAL_CAMS = []

QUAL_CAMS.append(EasyDict()); QUAL_CAMS[-1].position = [0.14358797004960389, 0.6832940458221193, 4.533705645058711]
QUAL_CAMS[-1].focal_point = [0.0005729794502258301, -0.24407485127449036, 0.023363471031188965]
QUAL_CAMS[-1].view_angle = 30.0
QUAL_CAMS[-1].view_up = [-0.007471124657963384, 0.9795291784012184, -0.20116354281275517]
QUAL_CAMS[-1].clipping_range = [3.9025500034049414, 5.504674393476118]

QUAL_CAMS.append(EasyDict());QUAL_CAMS[-1].position = (np.array([-5.349964475006663, 1.920414677096611, 3.5135667406172577]) * 0.7).tolist()
QUAL_CAMS[-1].focal_point = (np.array([0.0005729794502258301, -0.24407485127449036, 0.023363471031188965]) * 0.7).tolist()
QUAL_CAMS[-1].view_angle = 30.0
QUAL_CAMS[-1].view_up = [0.23395229205946214, 0.9452442385240355, -0.22755143281737628]
QUAL_CAMS[-1].clipping_range = [4.639483246757514, 9.411640882069047]

QUAL_CAMS.append(EasyDict());QUAL_CAMS[-1].position = [-5.352019086504213, 1.912093122092013, 3.5155665978715565]
QUAL_CAMS[-1].focal_point = [0.0005729794502258301, -0.24407485127449036, 0.023363471031188965]
QUAL_CAMS[-1].view_angle = 30.0
QUAL_CAMS[-1].view_up = [0.23291695057437573, 0.9456621869330374, -0.226876006532553]
QUAL_CAMS[-1].clipping_range = [4.640831173609477, 9.409944119875998]

QUAL_CAMS.append(EasyDict());QUAL_CAMS[-1].position = [0.3016991172947491, 5.944678091667598, -2.6418653921878774]
QUAL_CAMS[-1].focal_point = [0.0005729794502258301, -0.24407485127449036, 0.023363471031188965]
QUAL_CAMS[-1].view_angle = 30.0
QUAL_CAMS[-1].view_up = [-0.015364204129237906, 0.3961206725293572, 0.9180699069386626]
QUAL_CAMS[-1].clipping_range = [5.025554519400046, 8.925656692134611]



RENDER_CAMS = []
RENDER_CAMS.append(EasyDict()); RENDER_CAMS[-1].position = [-19.27443412557687, 38.41618852722742, 62.92332194272892]
RENDER_CAMS[-1].focal_point = [23.017113793085723, 18.842020665076497, 23.695796429963483]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [0.20211350941774725, 0.945834333888464, -0.25406208325645113]
RENDER_CAMS[-1].clipping_range = [0.15819844912584044, 158.19844912584043]

RENDER_CAMS.append(EasyDict()); RENDER_CAMS[-1].position = [-144.1860060581795, 34.53448330277345, 86.16948584800988]
RENDER_CAMS[-1].focal_point = [21.312243456177722, 15.002081047684001, 22.222465595347277]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [0.09548726097493275, 0.9938295520535453, -0.0564358437214747]
RENDER_CAMS[-1].clipping_range = [112.11309459927534, 263.93543508539744]

RENDER_CAMS.append(EasyDict()); RENDER_CAMS[-1].position = [-51.63941001424021, 67.21827286979888, 84.85400337209043]
RENDER_CAMS[-1].focal_point = [23.449717789866693, 18.90443358379182, 24.251999898628164]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [0.2965712485568456, 0.8912608449592379, -0.3430737541287285]
RENDER_CAMS[-1].clipping_range = [25.199195102223293, 207.96568091950252]

RENDER_CAMS.append(EasyDict()); RENDER_CAMS[-1].position = [20.427899300153182, 11.237827978724097, 113.05472985145948]
RENDER_CAMS[-1].focal_point = [23.449717789866693, 18.90443358379182, 24.251999898628164]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [0.00545284037263625, 0.9962632813394683, 0.08619594414232047]
RENDER_CAMS[-1].clipping_range = [36.38982858490981, 156.46646352808165]

RENDER_CAMS.append(EasyDict()); RENDER_CAMS[-1].position = [-90.036793284838, 90.59368492899732, 107.5826573645741]
RENDER_CAMS[-1].focal_point = [23.449717789866693, 18.90443358379182, 24.251999898628164]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [0.3038803751031019, 0.8865459876957194, -0.3488451365976307]
RENDER_CAMS[-1].clipping_range = [75.04408985727662, 258.464042510806]

RENDER_CAMS.append(EasyDict()); RENDER_CAMS[-1].position = [-86.48827422812629, 89.64072568243479, 112.97939318412429]
RENDER_CAMS[-1].focal_point = [23.449717789866693, 18.90443358379182, 24.251999898628164]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [0.2965712485568456, 0.8912608449592379, -0.3430737541287285]
RENDER_CAMS[-1].clipping_range = [74.78077593982162, 258.7993218792624]

RENDER_CAMS.append(EasyDict()); RENDER_CAMS[-1].position = [27.108517842894685, 24.1244045894474, 97.90399417111082]
RENDER_CAMS[-1].focal_point = [24.00978994369507, 20.092411518096924, 24.373619973659515]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [-0.008629691549773182, 0.9984826123964274, -0.05438750928073159]
RENDER_CAMS[-1].clipping_range = [21.407154508523647, 138.84036852640696]


RENDER_CAMS.append(EasyDict()); RENDER_CAMS[-1].position = [19.658280556162836, 39.258393231752464, 255.78973339401603]
RENDER_CAMS[-1].focal_point = [24.451400098591378, 25.776652345817254, 24.911669731532935]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [-0.0013994651195821134, 0.9982967878778062, -0.0583229355419777]
RENDER_CAMS[-1].clipping_range = [179.11924775277237, 298.68985351120807]

RENDER_CAMS.append(EasyDict())
RENDER_CAMS[-1].position = [-206.51287214542666, 12.97026336431632, 25.86715957797304]
RENDER_CAMS[-1].focal_point = [24.451400098591378, 25.776652345817254, 24.911669731532935]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [-0.05549689644640207, 0.9972946303127063, -0.048202850893492126]
RENDER_CAMS[-1].clipping_range = [179.51250722414733, 297.04896304028614]

RENDER_CAMS.append(EasyDict())
RENDER_CAMS[-1].position = [38.87211053759958, 22.64955041102462, -205.93822945996436]
RENDER_CAMS[-1].focal_point = [24.451400098591378, 25.776652345817254, 24.911669731532935]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [-0.021256912847085984, 0.9996634663260102, -0.01486935604783261]
RENDER_CAMS[-1].clipping_range = [178.363686255711, 297.5664896734477]

RENDER_CAMS.append(EasyDict())
RENDER_CAMS[-1].position = [255.64174032500543, 29.325106202973558, 31.828798761657023]
RENDER_CAMS[-1].focal_point = [24.451400098591378, 25.776652345817254, 24.911669731532935]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [-0.014131192450592338, 0.9990907695991301, -0.04022366842720664]
RENDER_CAMS[-1].clipping_range = [180.13693595224146, 296.27146972182925]

RENDER_CAMS.append(EasyDict())
RENDER_CAMS[-1].position = [97.8208622599962, -112.19165848347292, -145.6494009788123]
RENDER_CAMS[-1].focal_point = [24.451400098591378, 25.776652345817254, 24.911669731532935]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [0.6685425098853874, 0.6920572599727359, -0.2722272238318125]
RENDER_CAMS[-1].clipping_range = [150.73165412928108, 330.88137847290943]

RENDER_CAMS.append(EasyDict())
RENDER_CAMS[-1].position = [38.68827819668907, 230.2832671439704, -82.2506716884857]
RENDER_CAMS[-1].focal_point = [24.451400098591378, 25.776652345817254, 24.911669731532935]
RENDER_CAMS[-1].view_angle = 30.0
RENDER_CAMS[-1].view_up = [0.5469972779447123, 0.3583578132980697, 0.7565538021643279]
RENDER_CAMS[-1].clipping_range = [164.0466819115382, 318.6776408316894]




VIDEO_CAMS = []

VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-11.969750729651192, 26.279381457184805, 177.46761678573006]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.005563053267064506, 0.9999005720422458, -0.012957564120532006]
VIDEO_CAMS[-1].clipping_range = [98.83139919783581, 231.20034420029512]

VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-5.077489065428195, 28.09838171375504, 178.9167234835192]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0031820488224001556, 0.9997091518391671, -0.02390577951678792]
VIDEO_CAMS[-1].clipping_range = [99.90909305753756, 229.84374966835918]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-0.4435730350428556, 28.132622920571905, 179.73183602831432]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0031820488224001456, 0.9997091518391669, -0.023905779516787972]
VIDEO_CAMS[-1].clipping_range = [101.0344863853643, 228.42711133358216]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [1.1062328421212686, 28.14331548852088, 179.97269410469835]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.003182048822400145, 0.9997091518391669, -0.023905779516787976]
VIDEO_CAMS[-1].clipping_range = [101.42054101182381, 227.9411480977626]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [3.4352563335068407, 28.158674543552518, 180.3049786643814]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.003182048822400139, 0.9997091518391668, -0.023905779516788024]
VIDEO_CAMS[-1].clipping_range = [102.00973245031648, 227.1994774679012]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [5.768994038259084, 28.17321480046771, 180.60239463393856]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.003182048822400135, 0.9997091518391671, -0.023905779516788087]
VIDEO_CAMS[-1].clipping_range = [102.61093787868045, 226.44268370003607]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [10.445963360887458, 29.080616958944884, 181.06884333794173]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999771655, 0.9995619473128375, -0.029479126862585053]
VIDEO_CAMS[-1].clipping_range = [103.5931670513701, 225.2062595404442]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [13.572896538543471, 29.096168825026087, 181.31764678706762]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999771694, 0.9995619473128378, -0.02947912686258501]
VIDEO_CAMS[-1].clipping_range = [104.44439522242, 224.13473865176073]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [19.838559290423444, 29.121781278544965, 181.6280077200859]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999771755, 0.9995619473128378, -0.029479126862584894]
VIDEO_CAMS[-1].clipping_range = [106.2076725465296, 221.91513578899932]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [21.406523818616137, 29.12703659386416, 181.666541093304]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999771755, 0.9995619473128377, -0.029479126862584866]
VIDEO_CAMS[-1].clipping_range = [106.6609078731097, 221.34460589297748]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [25.327440076522336, 29.13816033098807, 181.6944767268328]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999771768, 0.9995619473128378, -0.029479126862584814]
VIDEO_CAMS[-1].clipping_range = [107.32313102424163, 220.5110033836381]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [26.895799056033358, 29.141802843610208, 181.67828894285344]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999771763, 0.9995619473128377, -0.029479126862584772]
VIDEO_CAMS[-1].clipping_range = [106.86048772826045, 221.0933759748104]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [30.03164613794576, 29.147702446559734, 181.5990148114613]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999771733, 0.9995619473128375, -0.029479126862584682]
VIDEO_CAMS[-1].clipping_range = [105.94981803267677, 222.23972149613536]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [31.59882232990498, 29.149958950076698, 181.5359363491363]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.002625750099977171, 0.9995619473128375, -0.029479126862584627]
VIDEO_CAMS[-1].clipping_range = [105.50188221382659, 222.80358040378357]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [35.51333284413558, 29.153576335746177, 181.3099213337938]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.00262575009997717, 0.9995619473128379, -0.0294791268625846]
VIDEO_CAMS[-1].clipping_range = [104.40398002599112, 224.18561305731504]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [38.64012860334584, 29.15438745379088, 181.05891588239157]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999771668, 0.9995619473128379, -0.029479126862584564]
VIDEO_CAMS[-1].clipping_range = [103.54858174925275, 225.2623832498927]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [41.7612986493411, 29.15334680303339, 180.74562276821666]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999771603, 0.9995619473128379, -0.029479126862584505]
VIDEO_CAMS[-1].clipping_range = [102.71392912886189, 226.3130389353595]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [45.652958539585086, 29.149442640591225, 180.26660659433188]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999771425, 0.9995619473128378, -0.02947912686258436]
VIDEO_CAMS[-1].clipping_range = [101.70029202430327, 227.58899921019335]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [51.078650948091486, 29.139122060614234, 179.433387499419]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.002625750099977128, 0.9995619473128378, -0.029479126862584266]
VIDEO_CAMS[-1].clipping_range = [100.3377238503249, 229.30419181110568]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [52.62318923241693, 29.13513446867349, 179.16060409042075]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999771295, 0.9995619473128375, -0.029479126862584266]
VIDEO_CAMS[-1].clipping_range = [99.96073524850893, 229.7787427897735]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [55.70372053926888, 29.125776611199427, 178.56891547943525]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.002625750099977116, 0.9995619473128378, -0.02947912686258421]
VIDEO_CAMS[-1].clipping_range = [99.22342522741997, 230.7068641982296]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [61.0640433358853, 29.104974584809163, 177.3861198597297]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.002625750099977086, 0.9995619473128377, -0.029479126862584075]
VIDEO_CAMS[-1].clipping_range = [97.98736398561051, 232.2628106357836]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [66.37981847795606, 29.078562106229274, 176.0170560532764]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0026257500999770627, 0.9995619473128378, -0.029479126862584002]
VIDEO_CAMS[-1].clipping_range = [96.82160745901228, 233.73025792177776]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.6415635392273, 29.927087408073223, 174.43505397705013]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.004198403158178429, 0.9993835992138334, -0.03485390987142756]
VIDEO_CAMS[-1].clipping_range = [95.4750625903341, 235.425280482601]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [76.10873540865545, 29.894336700675318, 172.9578736976855]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.004198403158178386, 0.9993835992138334, -0.0348539098714274]
VIDEO_CAMS[-1].clipping_range = [94.59568502617171, 236.5322356575893]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [78.32524588708775, 29.876148597312973, 172.16936254082276]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.004198403158178365, 0.9993835992138334, -0.034853909871427335]
VIDEO_CAMS[-1].clipping_range = [94.17645764355873, 237.0599565588181]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [82.72163580730181, 29.836167242508957, 170.49338048295436]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.004198403158178323, 0.9993835992138334, -0.034853909871427224]
VIDEO_CAMS[-1].clipping_range = [93.37939476412446, 238.0632945050407]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [86.3410130732743, 30.679622709373685, 168.9638434591842]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.006324430883239545, 0.9991783108644487, -0.04003379412533878]
VIDEO_CAMS[-1].clipping_range = [92.5071376373115, 239.16128651643064]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [94.87096405326147, 31.454935425009097, 164.97012847295088]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.00865531419451313, 0.9989439187668533, -0.04512352706645229]
VIDEO_CAMS[-1].clipping_range = [90.92767321377691, 241.14950680836753]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [97.6619844500328, 31.41446763269851, 163.53889696933123]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008655314194513162, 0.9989439187668532, -0.04512352706645234]
VIDEO_CAMS[-1].clipping_range = [90.5353543523658, 241.64335542536998]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [101.79359402781171, 31.348613240443754, 161.28851384785727]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.00865531419451319, 0.9989439187668534, -0.045123527066452375]
VIDEO_CAMS[-1].clipping_range = [89.99552769397947, 242.32288596771062]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [109.18588893374313, 31.212004406246628, 156.84632565398658]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008655314194513243, 0.9989439187668532, -0.045123527066452444]
VIDEO_CAMS[-1].clipping_range = [89.15889415319698, 243.37603522382142]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [111.81198251937828, 31.157299041188033, 155.13153822783687]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008655314194513266, 0.9989439187668532, -0.045123527066452486]
VIDEO_CAMS[-1].clipping_range = [88.9041717447256, 243.6966782555907]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [114.40333864344805, 31.09994504002132, 153.36477846172335]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008655314194513293, 0.9989439187668532, -0.04512352706645254]
VIDEO_CAMS[-1].clipping_range = [88.67602133642089, 243.98387261378343]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [118.22298869479366, 31.008998047051016, 150.61873388318608]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008655314194513302, 0.9989439187668532, -0.04512352706645258]
VIDEO_CAMS[-1].clipping_range = [88.38381388539155, 244.3517015910842]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [123.19086879943816, 29.998411665607474, 146.8199925871744]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.005182971439474573, 0.9991566899109503, -0.040731386095375736]
VIDEO_CAMS[-1].clipping_range = [88.33666606999664, 244.41105097679457]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [127.9956120026594, 29.859000867342928, 142.78879958667906]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.005182971439474569, 0.9991566899109505, -0.04073138609537574]
VIDEO_CAMS[-1].clipping_range = [88.14809571439777, 244.64842220331417]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [131.4852826008139, 30.629057959542582, 139.603799490335]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008932869644299818, 0.9989520202602772, -0.044889476025329667]
VIDEO_CAMS[-1].clipping_range = [87.82851020504714, 245.05071451533567]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [133.7628614466791, 30.552559066676366, 137.44819316388782]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008932869644299783, 0.9989520202602772, -0.044889476025329625]
VIDEO_CAMS[-1].clipping_range = [87.81498797654534, 245.067736215033]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [135.9969695759198, 30.473652136197394, 135.24764989515973]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008932869644299714, 0.9989520202602773, -0.04488947602532956]
VIDEO_CAMS[-1].clipping_range = [87.8284710848289, 245.05076375963097]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [139.264684633272, 30.350845602080103, 131.86449817572395]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008932869644299665, 0.9989520202602772, -0.04488947602532953]
VIDEO_CAMS[-1].clipping_range = [87.8993160818604, 244.96158450457406]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [143.46134355286404, 30.178976323878036, 127.20466722485946]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008932869644299603, 0.9989520202602773, -0.04488947602532947]
VIDEO_CAMS[-1].clipping_range = [88.08815463172024, 244.72387567673576]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [147.46870409872733, 29.998070320379764, 122.381406491686]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008932869644299533, 0.9989520202602773, -0.04488947602532941]
VIDEO_CAMS[-1].clipping_range = [88.3845691054798, 244.35075092459144]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [152.20201886889856, 29.759669802903844, 116.1342236947028]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.008932869644299464, 0.9989520202602773, -0.04488947602532935]
VIDEO_CAMS[-1].clipping_range = [88.90563103630438, 243.6948413081519]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [157.43798217514865, 31.216049215087203, 108.25373591607317]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0182846610898743, 0.9985288152083032, -0.051047785138344334]
VIDEO_CAMS[-1].clipping_range = [89.25894283955694, 243.25009454074092]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [160.66304449785991, 31.880881611403726, 102.84187004068234]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.023136166823483587, 0.998281522884593, -0.05383975159613587]
VIDEO_CAMS[-1].clipping_range = [89.71152257225734, 242.68038990234146]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [162.5803902768464, 31.741008627053283, 99.42445740087535]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.023136166823483545, 0.998281522884593, -0.05383975159613588]
VIDEO_CAMS[-1].clipping_range = [90.19918196952594, 242.06652719371914]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [164.02997126649507, 32.505211101168236, 96.60828257111415]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.028055256383685407, 0.9980078744898213, -0.05650827413357014]
VIDEO_CAMS[-1].clipping_range = [90.37831159511998, 241.84103990119996]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [166.13391030448724, 32.326276283716446, 92.40350177384246]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0280552563836854, 0.9980078744898214, -0.05650827413357016]
VIDEO_CAMS[-1].clipping_range = [91.05623552909216, 240.98767334109914]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [169.05154594266457, 32.04479037242124, 85.98355394396748]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.028055256383685362, 0.9980078744898215, -0.056508274133570145]
VIDEO_CAMS[-1].clipping_range = [92.18140161840061, 239.5713210527987]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [171.9741298298561, 30.835607698106628, 78.75441740339008]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.022893899302093866, 0.9982598626733498, -0.05434257953143685]
VIDEO_CAMS[-1].clipping_range = [93.8215838319868, 237.50666957288263]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [174.0305714752984, 30.560480355564156, 72.83404052931425]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.022893899302094022, 0.9982598626733498, -0.05434257953143692]
VIDEO_CAMS[-1].clipping_range = [95.0523928550503, 235.95733459661466]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [175.8490711712162, 30.2757092157926, 66.83675165177004]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.022893899302094203, 0.9982598626733499, -0.05434257953143697]
VIDEO_CAMS[-1].clipping_range = [96.37969589144855, 234.2865335382341]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [178.10188537254362, 30.710820231920604, 57.669070288533455]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.028341930132922542, 0.9980488926157648, -0.055634008886527864]
VIDEO_CAMS[-1].clipping_range = [98.3065433827179, 231.86102953540518]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [180.2792569315598, 30.08489650136804, 45.33105236609814]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.028341930132922178, 0.9980488926157648, -0.05563400888652778]
VIDEO_CAMS[-1].clipping_range = [101.51408315078538, 227.82339781731093]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [181.07769390192735, 31.433078467695577, 37.4483163799078]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.03947009041878324, 0.9976013361071846, -0.05686550940151373]
VIDEO_CAMS[-1].clipping_range = [103.21979899838625, 225.67625299407496]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [181.55627971634433, 31.79766785898079, 28.017065145905324]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.04506554645176998, 0.9973554343270212, -0.057019594363360585]
VIDEO_CAMS[-1].clipping_range = [105.77358534121583, 222.46156214543814]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [181.44741413301765, 31.166258272507413, 17.058837499261287]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.04506554645177006, 0.9973554343270211, -0.05701959436336055]
VIDEO_CAMS[-1].clipping_range = [104.82683060973207, 223.65333129235603]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [180.66159740380306, 30.550800792172165, 6.914664978310148]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.045065546451770105, 0.9973554343270213, -0.05701959436336056]
VIDEO_CAMS[-1].clipping_range = [102.22700566281408, 226.92597525819983]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [176.93166351082752, 32.75510872586568, -13.298701578295583]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.06649377681562577, 0.9965071990251252, -0.05051712517447145]
VIDEO_CAMS[-1].clipping_range = [96.66994066263038, 233.92117517049104]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [175.3247690825301, 32.34118769344899, -19.348665983318583]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.06649377681562561, 0.9965071990251252, -0.05051712517447149]
VIDEO_CAMS[-1].clipping_range = [95.47181601709941, 235.4293672494134]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [173.4169443115211, 32.792467042475714, -25.37263166334655]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.07184255719893534, 0.9962178416270274, -0.04887390918582916]
VIDEO_CAMS[-1].clipping_range = [94.11362875390766, 237.13904518624506]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [171.33309514370998, 32.352698536150285, -31.273452448074124]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.07184255719893529, 0.9962178416270273, -0.04887390918582918]
VIDEO_CAMS[-1].clipping_range = [93.10936945853487, 238.4032007314253]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [166.83683835082303, 35.05947322876008, -41.538039607349596]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.0924785628604865, 0.9949011669084697, -0.040241564276562544]
VIDEO_CAMS[-1].clipping_range = [90.58917663742, 241.57560425750256]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [166.13413750559263, 25.16454946468267, -43.84932456736475]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.03701693559399568, 0.9970652258082845, -0.06701255078792721]
VIDEO_CAMS[-1].clipping_range = [93.05323738381712, 238.47385944859562]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [165.55259882935957, 16.262685066512535, -44.55513527558558]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.013243556962766224, 0.9956966436671231, -0.09172132788506793]
VIDEO_CAMS[-1].clipping_range = [90.76441399834705, 241.35501652427544]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [164.10207345783311, 6.454863142907135, -45.62798887635154]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.06817542777289165, 0.9905310168048223, -0.11916549750489183]
VIDEO_CAMS[-1].clipping_range = [87.96022495858942, 244.88491277783976]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [160.8790641363699, -15.176003115334186, -43.019377543332816]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.19140543147640982, 0.964802944874536, -0.18033091349724895]
VIDEO_CAMS[-1].clipping_range = [83.26209495637181, 250.79889049419972]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [157.82153619117804, -25.39450786874104, -42.34464656140966]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.24902829186613196, 0.9456403891144471, -0.2091630089803487]
VIDEO_CAMS[-1].clipping_range = [81.33331766646047, 253.22682371592234]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [153.66273106159503, -34.673517128115805, -42.9372632685957]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.3002929465913752, 0.9242345934525336, -0.2358273997929411]
VIDEO_CAMS[-1].clipping_range = [79.63452927841772, 255.36524829484046]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [150.50868604216922, -43.51381200488787, -40.52486539969928]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.350310135432318, 0.8993040870046514, -0.2617918411832603]
VIDEO_CAMS[-1].clipping_range = [78.66100290662061, 256.5907174211477]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [145.42307512601926, -54.43717789086274, -37.78027416809144]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.41185682000790247, 0.8626277441055603, -0.29367896573014973]
VIDEO_CAMS[-1].clipping_range = [77.74120675149095, 257.7485512747158]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [139.97572095276337, -64.82504261257259, -33.96961034982952]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.4708433877360929, 0.8206427127498426, -0.32380865065000763]
VIDEO_CAMS[-1].clipping_range = [77.40523078440059, 258.1714757961036]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [134.8107476801869, -72.65262562339785, -31.403550714398488]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.5151487278847706, 0.7840964518551058, -0.34614237294346994]
VIDEO_CAMS[-1].clipping_range = [77.37652024048346, 258.2076164556576]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [125.11359851742739, -76.51589910921119, -41.870118444152816]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.5337589277828544, 0.7664217256361495, -0.35735017207915915]
VIDEO_CAMS[-1].clipping_range = [75.99892348237262, 259.94172694765143]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [121.04213705208093, -78.49802619506916, -44.80601038219076]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.5442492816976733, 0.7555182016606807, -0.3646710385141353]
VIDEO_CAMS[-1].clipping_range = [75.74718757268587, 260.25861059275945]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [107.2820808474618, -90.32094823053146, -44.02428503651238]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6047418061483988, 0.6817266396688608, -0.41174766139247915]
VIDEO_CAMS[-1].clipping_range = [76.55567905252116, 259.2408863932684]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [93.95791380460263, -96.46959654008816, -48.124070559284746]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6368640610795238, 0.6313359430703831, -0.442514513539078]
VIDEO_CAMS[-1].clipping_range = [77.47029199276002, 258.08957713934973]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [78.01349496530986, -99.68975877690522, -55.78393596038481]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.658225451352688, 0.587741488849791, -0.47042448647597135]
VIDEO_CAMS[-1].clipping_range = [78.9762606220782, 256.19387290998213]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [66.37286661926295, -101.6526347500662, -59.55044014436289]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6715400736575127, 0.5521316506994796, -0.49415035137894064]
VIDEO_CAMS[-1].clipping_range = [80.73405420054227, 253.98117295317186]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [43.51997993450318, -103.76966277024512, -64.47996784590345]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6851189587156761, 0.48104212984875927, -0.547001354403184]
VIDEO_CAMS[-1].clipping_range = [85.43429113494228, 248.06454304328642]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [27.834875367387472, -103.06323041733879, -67.41736263370638]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6887007812862755, 0.4356310360180868, -0.5795833281880027]
VIDEO_CAMS[-1].clipping_range = [89.43507132348789, 243.02838506725305]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [26.970295643611045, -85.09657199245613, -88.25987650177217]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.688431487214894, 0.5275891165601272, -0.49770655159279703]
VIDEO_CAMS[-1].clipping_range = [88.83698993168852, 243.78124631672924]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [28.242761575509398, -69.35843606446531, -101.6304341429179]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6895391717997871, 0.5906798072370757, -0.4190860244341667]
VIDEO_CAMS[-1].clipping_range = [89.16263713913185, 243.37132357570144]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [36.48368519569415, -52.28746005940427, -112.22007156418293]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6834507775870744, 0.6604268202474851, -0.31101679972859925]
VIDEO_CAMS[-1].clipping_range = [88.6393995455, 244.02997190336384]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [39.17163156930662, -39.35466131098026, -118.4672205707096]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6773887374295195, 0.6950544475683798, -0.2409228368566495]
VIDEO_CAMS[-1].clipping_range = [89.82817340252755, 242.53355054062305]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [39.343042361595806, -30.622939822248753, -122.03820764757563]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6733040329373503, 0.71191418186472, -0.19959929080581962]
VIDEO_CAMS[-1].clipping_range = [91.31181962056766, 240.66594562796433]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [49.370986883484974, -2.529342361921247, -128.4154974640666]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6510059686335812, 0.7585486811259049, -0.02819799931915297]
VIDEO_CAMS[-1].clipping_range = [94.78752034940337, 236.29075450950543]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [53.15450421501366, 13.885857261808535, -129.7678192750661]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6352405109169006, 0.76957052496149, 0.06504383445440562]
VIDEO_CAMS[-1].clipping_range = [98.14159833588374, 232.06866136571998]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [40.57963121643223, 40.35275436260507, -131.13558828501237]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6255899191123726, 0.7669796789165743, 0.1427565243132657]
VIDEO_CAMS[-1].clipping_range = [99.9164142826457, 229.834533754344]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [28.940947076461494, 61.92387123858015, -128.18278407319428]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6242730218576469, 0.7543548402938214, 0.20305656627158158]
VIDEO_CAMS[-1].clipping_range = [97.84092264233506, 232.44715011563943]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [19.440222602620082, 74.62005283244933, -124.47878929069367]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6243500005970033, 0.7465330053472791, 0.22994683881645633]
VIDEO_CAMS[-1].clipping_range = [94.98295023824238, 236.04474854390702]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [14.532797979558717, 89.77440178570848, -118.23124415194124]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6289281832602793, 0.7227122928838, 0.2865942811982915]
VIDEO_CAMS[-1].clipping_range = [90.87509327419919, 241.2156941192478]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [2.8311271061498258, 102.83753472205942, -110.13280173471759]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6323093473390847, 0.7095199475437151, 0.3110728745891935]
VIDEO_CAMS[-1].clipping_range = [85.91911408064522, 247.45425084279202]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-2.12349092752583, 111.8179809653838, -103.55881708309005]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6377909466618192, 0.6916039385600179, 0.3389789086720462]
VIDEO_CAMS[-1].clipping_range = [83.73019253945165, 250.20965207931977]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-10.948330988287912, 125.22652387852166, -90.95511647066097]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6485010393746655, 0.6609844230680293, 0.37755263790815086]
VIDEO_CAMS[-1].clipping_range = [80.86666976444712, 253.8142372810215]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-17.055910535426793, 133.50365419611308, -80.96852788024151]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6568574534508187, 0.6384913075654985, 0.40108245537470877]
VIDEO_CAMS[-1].clipping_range = [79.55882595782143, 255.4605431783064]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-15.936400629501323, 144.20760545391008, -69.1337941257373]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6776344336279313, 0.581267006913356, 0.45048889113466595]
VIDEO_CAMS[-1].clipping_range = [80.22798235400134, 254.61821314190922]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-14.11320290201272, 151.5166787120062, -59.809457312797775]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6933895014832437, 0.5327961909149209, 0.48511773640980205]
VIDEO_CAMS[-1].clipping_range = [81.36940331580882, 253.18139931812635]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-12.881023251039204, 156.451614329691, -52.46577991713741]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7043545064488431, 0.49583251029881414, 0.5079713092055331]
VIDEO_CAMS[-1].clipping_range = [82.45208299748892, 251.8185286635492]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-12.594217845578381, 163.52326144901042, -38.973242238476985]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7208598114198679, 0.4376929347499586, 0.537388153153455]
VIDEO_CAMS[-1].clipping_range = [84.44671163029503, 249.30770216848447]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-11.306395301450639, 166.3584975178277, -33.174686295290854]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7286708305132605, 0.4076631222187118, 0.5503177259929092]
VIDEO_CAMS[-1].clipping_range = [85.71084918729872, 247.71641343466794]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-12.252203067533054, 168.4519704380446, -27.06943356724638]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7339493891922028, 0.38727346846197325, 0.5579763030181708]
VIDEO_CAMS[-1].clipping_range = [86.62255950540771, 246.5687579839831]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-19.970165711715005, 198.681884230034, -37.899014616368135]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7339493891922028, 0.38727346846197325, 0.5579763030181708]
VIDEO_CAMS[-1].clipping_range = [119.31796654026105, 280.0898066106257]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-12.25220306753306, 168.4519704380446, -27.069433567246378]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7339493891922028, 0.38727346846197325, 0.5579763030181708]
VIDEO_CAMS[-1].clipping_range = [86.62255950540771, 246.5687579839831]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-5.873721543415748, 143.4685706099542, -18.119366584501137]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7339493891922028, 0.38727346846197325, 0.5579763030181708]
VIDEO_CAMS[-1].clipping_range = [59.60156195594215, 218.86541201155123]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [-0.6022492094345004, 122.82113273549935, -10.722617011984404]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7339493891922028, 0.38727346846197325, 0.5579763030181708]
VIDEO_CAMS[-1].clipping_range = [37.27015902249954, 195.97008476160755]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [3.7543394963351275, 105.75713449214821, -4.60960083635074]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7339493891922028, 0.38727346846197325, 0.5579763030181708]
VIDEO_CAMS[-1].clipping_range = [18.81445411882794, 177.04832670380284]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [7.3548260300290345, 91.65465660508116, 0.4424786476440232]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7339493891922028, 0.38727346846197325, 0.5579763030181708]
VIDEO_CAMS[-1].clipping_range = [3.561805438107612, 161.41051012710477]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [10.621928374337834, 93.65758722826945, 4.323367859246148]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7495255714235877, 0.31954845905424323, 0.5797414942008235]
VIDEO_CAMS[-1].clipping_range = [6.841403970111017, 157.28217127149247]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [19.579174350264097, 94.6284737163257, 3.517820013889727]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7617594659103715, 0.234182386999644, 0.6040539096099743]
VIDEO_CAMS[-1].clipping_range = [11.418644090164968, 151.5203690098166]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [33.42579724125671, 88.51375858005434, -10.21365846303898]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7607657981144633, 0.2242106994810864, 0.6090689309572982]
VIDEO_CAMS[-1].clipping_range = [4.010641479254545, 160.8455180250583]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [46.447255503942806, 85.48090948342616, -9.881912137704939]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.7283609525483327, 0.11837349879783132, 0.6748911301723006]
VIDEO_CAMS[-1].clipping_range = [0.16859407583807176, 168.59407583807175]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [57.010838492804666, 81.1093915154425, -8.97971625584884]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.68788887115661, 0.033740004817746, 0.725031387605932]
VIDEO_CAMS[-1].clipping_range = [0.1728386071649676, 172.83860716496758]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [63.77494957384674, 74.38986635143472, -12.255255396117366]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6739663460891034, 0.013605731636897045, 0.7386367499697859]
VIDEO_CAMS[-1].clipping_range = [0.17550263805960303, 175.50263805960304]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [73.93885222190548, 59.578931985632465, -16.824782510667873]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6458073201981798, -0.010822838267130763, 0.7634237167853046]
VIDEO_CAMS[-1].clipping_range = [0.17544044924026064, 175.44044924026065]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.6495627831974486, -0.008119164947198304, 0.7602646051509597]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.5763232099272058, 0.12445227772086259, 0.8076900322953674]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.5358862499934194, 0.18717295969010117, 0.8232813675948452]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.4627720389186498, 0.28773345215278084, 0.838481663787181]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.35002819449927003, 0.41897020359459225, 0.8378211214546261]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.27377545973244777, 0.49508352167703423, 0.824584322075161]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.19077796177331208, 0.5682772393739534, 0.8004153599920101]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [0.08199589945664335, 0.6506037517232559, 0.7549777683587247]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.026568124963837952, 0.7184094775117218, 0.6951129097903741]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.13018146994106497, 0.7699625248173882, 0.6246683081930896]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.22721435478929922, 0.8061011204757749, 0.5464198207837843]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]


VIDEO_CAMS.append(EasyDict()); VIDEO_CAMS[-1].position = [71.41730210775927, 64.60646082963285, -15.157380025367386]
VIDEO_CAMS[-1].focal_point = [24.5, 24.5, 24.5]
VIDEO_CAMS[-1].view_angle = 30.0
VIDEO_CAMS[-1].view_up = [-0.3406244299464927, 0.8318122062759381, 0.438250443484069]
VIDEO_CAMS[-1].clipping_range = [0.17611331335966626, 176.11331335966625]



CIRCLE_VIDEO_CAMS_OLD = []

CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [30.324166193290644, 23.204439259732126, -73.28654503880925]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.10283072609905028, 0.9919109412931835, 0.07442127594184697]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [43.928650185859055, 165.4947465670241]

CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [32.99336828811157, 19.80556416760746, -73.25345575117547]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.09318979860951118, 0.9947377523028489, 0.04257306165399029]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [44.54279379080252, 165.31030747348606]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [38.34956017717283, 19.40844540156144, -72.63211944026955]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.09303416448706041, 0.9946960501188983, 0.04386812186607577]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [42.52806322336266, 167.87433900957905]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [42.22269744501904, 19.14960828714258, -71.99703513183874]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.0927307829411453, 0.9946674705436812, 0.045140059120001866]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [41.14589887687123, 169.63696980214416]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [47.14858126979484, 19.429818929628162, -70.94281234365177]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.09119796262586083, 0.994431245413095, 0.05281505239089408]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [39.174133857555184, 172.04667919085878]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [52.708858459331765, 19.889348538517737, -69.42895483863005]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.08856788194545194, 0.9940814308428748, 0.06291135939621346]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [37.02728988804768, 174.65479324810133]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [59.090857477711396, 20.235279803002488, -67.25996413043143]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.08469710723170319, 0.9937574206954104, 0.0726125942201947]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [34.8720509540012, 177.31085653968927]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [65.45063565543973, 20.680864888421215, -64.59695899156426]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.08013078647864211, 0.9932873703323477, 0.08342216130359972]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [32.91267126133526, 179.72311025880526]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [71.64758517727111, 18.482620251944535, -61.582903126518474]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.08654472807310938, 0.9937439200194048, 0.07059059050056511]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [32.648237370752504, 180.5029084331468]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [77.17350952398782, 18.952949532898618, -58.33742592419051]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.08058901873787105, 0.993474541402328, 0.08070777932951317]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [31.352648612346194, 182.11048911437115]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [83.60787503430717, 17.811917466703335, -53.96038816232834]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.0822700364677667, 0.9934979874427603, 0.07869809430207127]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [31.058996121231772, 182.78737775819368]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [89.23013267075737, 17.77808570488257, -49.45934699500157]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.07733904710165795, 0.9933938604695034, 0.08477800289523389]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [30.603583427577092, 183.48102358311024]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [97.10962511293664, 15.728972034482004, -41.89369404535446]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.08311801753827518, 0.9933653791791197, 0.07947715778027784]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [31.05125665039934, 183.50553336699434]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [107.80345898355782, 15.175165066136907, -28.206770462154637]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.07209595456827562, 0.9934480912971655, 0.08867391517750889]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [32.424601340902, 182.3101273925187]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [113.38882672098282, 13.366678748375742, -18.37749939538604]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.07669059661577096, 0.9933221609154097, 0.08619534224688802]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [33.85559456942193, 181.1779515992966]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [117.85407086270806, 12.904801428594347, -7.981310312368555]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06968654188727108, 0.9935259351968375, 0.08972180320882393]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [36.613053436985595, 178.1679418917962]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [120.56836771218043, 12.322153794558806, 0.7183800887859811]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06598363670781209, 0.9936639395913963, 0.09098425601453143]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [39.328232276047416, 175.18536921449785]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [121.92423288319598, 11.553273056314145, 6.696983502517062]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06725857656799175, 0.9936032723224073, 0.09071284974165442]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [41.292537222438234, 173.0860800159239]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [123.13427099884018, 9.926419860001339, 14.98068522043932]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.07502138778505787, 0.9931299173685679, 0.08980400103728806]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [44.113131087073754, 170.16455460936143]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [123.70263284050806, 11.055329779241509, 20.915966765897114]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05814896802458775, 0.9941299303164753, 0.09123803574516076]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [47.16355514412064, 166.35905159455973]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [123.84941112063406, 11.041847204298179, 26.693750988768784]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.052937777816612595, 0.9944151929489944, 0.09130178372876059]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [49.57740380203831, 163.56309786002743]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [123.58979364401344, 10.589679616266721, 33.202117000974354]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.051639883537126584, 0.9944968033796151, 0.0911560776690137]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [46.60333008019397, 167.66576677723424]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [122.51579091493954, 10.457036685750191, 42.42550198736404]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.04520432944212365, 0.9949032265113406, 0.09013400289020984]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [42.91932532040303, 172.72808483597694]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [119.69189664604187, 8.717393731242893, 54.34549937085616]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.052672352995540656, 0.9943248717167997, 0.09243198967501984]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [38.23048578326271, 179.4857636900437]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [115.16459081857786, 6.410381976157113, 66.22281591813963]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06599677075560012, 0.9930356404798438, 0.09759427742761867]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [34.147027829971805, 185.6060672847395]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [110.31095552205961, 6.05487664280725, 75.79623745534356]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06371573157821167, 0.9933052564754865, 0.0963585647860184]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [32.00766006317231, 188.82950176251776]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [106.3969752535206, 6.476670426301147, 82.15785632894897]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05710053567915014, 0.9940970155661264, 0.09225319760135559]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [31.166552995722974, 190.1293260170833]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [99.48553733497612, 5.524497606740297, 91.05841794656504]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06130946632023601, 0.9935602011843185, 0.09528523475389253]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [30.071791652846542, 192.14614640702845]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [93.94286123584399, 6.464389135707251, 97.15811635042455]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.052916700595874264, 0.9947879092355942, 0.08715984417564045]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [30.3193148912263, 191.99128632426078]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [85.29735812911315, 7.426095249460238, 104.88753509176124]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.04514949592273406, 0.9958973035308964, 0.07842247023533783]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [31.311505626653364, 191.00437844758966]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [73.38515365276213, 5.517527856856321, 112.66831425457502]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05441600677513085, 0.9940778021002723, 0.09406498593068816]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [32.714988904251356, 190.0708609843885]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [67.41036492972114, 4.1185564799949725, 115.64695834843836]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06100792878608135, 0.992417456700783, 0.10670344071669963]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [33.662756840258645, 189.33790951655757]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [63.305149497720066, 4.816175595013978, 117.62360237621353]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05834693405108019, 0.9931945339520288, 0.1007980804115402]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [34.98747614262419, 187.6797060123632]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [55.77818130836571, 4.963697319747055, 120.53396688839025]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.058479636487732356, 0.9931920705937739, 0.10074543674983469]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [37.298821300973096, 184.96469748229788]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [50.10890847092857, 5.639644740454834, 122.3436333100251]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05697733402611513, 0.9937704073735836, 0.09578079575686431]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [39.43003099021368, 182.30895759761393]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [41.90899290480095, 5.926625590749245, 124.24397512962894]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05697733402611514, 0.9937704073735838, 0.0957807957568643]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [42.56715276098957, 178.49551144049582]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [36.96101357130457, 6.7836280123833586, 125.11251361911555]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05606368045007172, 0.9944208543165192, 0.08935338960999924]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [44.849730284423785, 175.56709505106969]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [30.90175634873808, 7.197990333320451, 125.76439811750686]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05596988544987351, 0.9945417305066111, 0.08805747102684476]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [47.583950071894975, 172.15685970135942]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [24.739085996542986, 7.520803100394734, 126.0355072625189]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.0559698854498735, 0.9945417305066111, 0.08805747102684472]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [50.50259961299536, 168.51665036989374]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [18.621379853319954, 7.483601716509388, 125.8949738882692]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05577695950102895, 0.9942000136038666, 0.09195250806199953]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [48.49876673653381, 171.120935580695]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [13.194124398510462, 8.21659300989191, 125.52063265419643]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05609933033934708, 0.994533688526089, 0.088065927135125]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [46.445935868033196, 173.62689204007955]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [6.326200806958914, 8.68785856096965, 124.57357452179534]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05609933033934709, 0.9945336885260888, 0.08806592713512483]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [43.88405740898829, 176.81664328754914]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-2.2201951770426307, 9.335651307881857, 122.70219675026743]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05609933033934712, 0.994533688526089, 0.0880659271351248]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [41.030877571015935, 180.3242331746353]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-8.4665541660761, 9.593463674694304, 120.80524528301896]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05532982961944265, 0.9943534202665846, 0.0905532194813]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [39.09966435853876, 182.71118565175232]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-13.753520329780478, 10.587250607627006, 118.89165787858062]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.057271460492431966, 0.9946719539902794, 0.08571863133661958]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [37.903554731185, 184.0245945821271]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-18.187007424201703, 11.006702593892701, 116.98653126531437]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.057271460492431994, 0.9946719539902794, 0.0857186313366196]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [36.87520058583911, 185.2210966318167]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-26.346716561988185, 12.492364677150908, 112.80861398154926]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06034014218148132, 0.9949691984650426, 0.07997100316583092]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [35.56029768252341, 186.53709266902243]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-34.885504835699095, 15.54694958744112, 107.33331095791317]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.07177198476577398, 0.9954599784297602, 0.06251570640569419]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [35.26021085599899, 186.23892716884387]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-43.46051623071129, 17.125648944488095, 100.35315412992867]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.07480634950158277, 0.9954950469220258, 0.058254799183936834]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [34.36306055468708, 186.9001139568278]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-48.86846387612496, 17.450079360959744, 95.03679398909084]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.072098414477248, 0.9955254101972691, 0.061081718062165655]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [34.07645943381673, 187.05028873689884]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-53.381282962443294, 17.567222903473596, 89.91330669598372]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06824306209337577, 0.9955750518787866, 0.06460031387440139]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [34.229846861581464, 186.68081382342208]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-59.999906762081686, 18.215082358312074, 80.85952721004983]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06503825291018404, 0.9956411871547162, 0.06684797752765952]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [34.8853380918184, 185.4452878685451]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-65.1561071533578, 19.67896152208508, 71.86693416381115]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06950773381305984, 0.9955162651925329, 0.06415637674684857]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [35.74150025864787, 183.79090160826502]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-70.18228937573285, 20.652597456618842, 60.16885265571591]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06828087315312267, 0.9955726524315796, 0.06459733811695881]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [38.140662785054005, 180.16892851563819]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-73.07059007685379, 21.211351907828025, 50.571075603318704]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.0657823423435212, 0.9956922854933511, 0.06534183992532619]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [40.81066212429283, 176.34172539408925]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-75.79378150099593, 21.891212090248334, 32.973114651826634]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.05931212932118129, 0.9960488466714372, 0.06609664408944875]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [47.08215554593574, 167.62145265628953]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-75.86409643345435, 23.574937847304778, 23.450172083115746]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.06972633355786287, 0.9953836887006134, 0.06595112339709062]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [47.68912679211209, 166.15096501079066]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-74.4898093322579, 25.50776429377713, 10.565130267655405]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.08140829007955488, 0.9944463775122346, 0.06670151841687795]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [41.007879107314906, 173.6492483237837]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-72.1764523661247, 26.020750430008558, 0.09364991123304969]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.08140829007955482, 0.9944463775122346, 0.06670151841687795]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [36.70941595364196, 178.49685862295874]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-67.64677757761598, 27.283394364868123, -12.552251956581664]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.08874460562821924, 0.9936419790863342, 0.06928212157032021]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [31.967386802670298, 183.64509113464516]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-60.46276574504226, 27.59915690103167, -26.283026746720267]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.0887446056282194, 0.9936419790863344, 0.06928212157032025]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [28.441860608798663, 187.35486667367297]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-51.07082631030054, 27.644480311470705, -38.96334288825436]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.08874460562821938, 0.9936419790863342, 0.06928212157032028]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [26.59345604602131, 189.01782233486335]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-41.05274647227981, 28.484956904998512, -48.80440341479634]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.09599766152264225, 0.9924231890905326, 0.07668547931362937]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [25.977491059537606, 189.09440644919056]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-30.441119623054874, 28.080964184254057, -56.86016082434978]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.09599766152264214, 0.9924231890905328, 0.0766854793136292]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [27.18374654881599, 187.17208072202752]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-24.129587073684586, 29.198086920458035, -60.49950416458405]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.10329494671595643, 0.9906683676182082, 0.08891759884085101]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [27.717496882365282, 186.05959696181395]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [-13.994288689291627, 29.236406878377668, -65.37223490266749]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.1060972872455917, 0.9898315643642303, 0.09474618634745656]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [30.007530428699035, 182.8316530337047]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [3.9806263410548244, 26.28791843812657, -71.21882434399956]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.10233628431610511, 0.9915562120413788, 0.0796464894049062]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [37.077741378696274, 173.98901615208482]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [10.381789573102282, 25.719402514278716, -72.36583992767954]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.10233628431610507, 0.9915562120413788, 0.07964648940490614]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [39.786038818689164, 170.55121283904197]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [21.639836855074407, 24.504923900497033, -73.36932919710775]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.10264375795218846, 0.9916299798962485, 0.07832140144569877]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [45.15908452037581, 163.8191110659036]


CIRCLE_VIDEO_CAMS_OLD.append(EasyDict()); CIRCLE_VIDEO_CAMS_OLD[-1].position = [30.324166193290644, 23.204439259732126, -73.28654503880925]
CIRCLE_VIDEO_CAMS_OLD[-1].focal_point = [23.91439093779708, 16.393071745902247, 26.35420570007013]
CIRCLE_VIDEO_CAMS_OLD[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS_OLD[-1].view_up = [0.10283072609905028, 0.9919109412931835, 0.07442127594184697]
CIRCLE_VIDEO_CAMS_OLD[-1].clipping_range = [43.928650185859055, 165.4947465670241]



CIRCLE_VIDEO_CAMS = []

CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [25.138323923977495, 16.942743488795763, -49.26072254551744]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.009907699976227128, 0.9990424886335852, -0.042613887243406416]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [23.84445325003143, 137.35364071353837]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [27.415974256342558, 17.307601191624485, -49.20526335667313]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.00972969708095112, 0.9992529491286423, -0.037401559491879094]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [22.629389731230695, 138.79933986890586]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [30.608304262755293, 16.708243006959698, -48.96265758375236]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.010237482384226722, 0.9989251586479835, -0.045207536700573904]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [20.391713770777315, 141.6481020947446]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [34.09724255336393, 16.306270903692884, -48.544861484926535]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.010891183216828384, 0.9986707478249162, -0.050379753542078014]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [18.18830432229189, 144.42191138725775]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [37.475987486696766, 17.353729175654912, -48.04432427196358]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.00861162005804862, 0.9993066269082521, -0.036222995696480936]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [16.941707120647887, 145.7967579174793]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [40.50688067953988, 16.677853406557908, -47.38554358770847]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.01051782838694551, 0.9989249612053005, -0.045147504549105766]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [15.026831695812096, 148.25034094936157]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [43.22115929133639, 17.256402124576677, -46.73243288384621]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.008606867743010286, 0.9992572096530304, -0.037562624829659774]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [14.032967594463624, 149.37983013137878]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [45.60961161005817, 17.74230482904742, -46.061561117379085]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.006749408990176842, 0.9994868579072528, -0.03131239928479206]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [13.20591015056921, 150.3167012950026]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [46.88992684851277, 17.55410471310661, -45.64954788282418]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.007535159310962901, 0.9994002231540169, -0.03379963511429765]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [12.53761896919685, 151.16250231232448]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [49.09313270717898, 18.04368233137892, -44.907086009367774]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.0054567339096430675, 0.9996036441207401, -0.02761844882639525]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [11.876722991363753, 151.89222596681148]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [51.076444920619686, 18.149728138918555, -44.15964276506219]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.004987793057509447, 0.99963898372719, -0.02640121461763576]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [11.139301841133403, 152.77340020716355]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [52.847498432445335, 18.35236638832296, -43.43894371692723]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.00397264534147733, 0.9997041182734685, -0.023997791482928806]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [10.576590385161051, 153.4248972240947]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [54.55187295353583, 18.555910202095024, -42.69354381355432]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.0029046085245977744, 0.9997621005522218, -0.021617251183443025]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [10.072772654593429, 154.00304517827885]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [56.70034866457672, 18.57172815653837, -41.67330900022186]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.002904608524597784, 0.9997621005522218, -0.021617251183443043]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [9.350178178859572, 154.87420599566784]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [58.08332877437827, 18.58289967397173, -40.970820259958316]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.002904608524597815, 0.9997621005522217, -0.021617251183443112]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [8.914566028814761, 155.39749924539132]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [59.31578713914158, 19.074347251683964, -40.32283837964718]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.0001649751325182528, 0.9998741932315438, -0.01586097389113084]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [8.817244930729508, 155.43252794156905]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [61.60744998980165, 19.287706311768602, -39.02155891677026]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.0014406928825357783, 0.9999066897399467, -0.013584410819768421]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [8.294332994957585, 156.02239050531477]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [63.15219123990414, 19.495060813818213, -38.08192686194242]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.0027891936253638915, 0.9999316955498668, -0.011350093994714107]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [8.017937169035967, 156.31539793074066]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [65.35899940042825, 19.998435048238733, -36.64616378648596]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.006384727370636562, 0.9999621776987859, -0.005905796161789455]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.801352184986094, 156.48122318263117]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [68.22154299912876, 20.028801081772066, -34.60514057884308]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.006423091524726553, 0.9999617652144871, -0.005933970373371811]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.275285391290204, 157.0883269263991]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [70.98535218846396, 19.674793103452764, -32.42686613288126]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.003162129722020815, 0.9999449065525033, -0.010009235501609747]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.678235076343519, 157.8376829525523]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [73.37223635416771, 19.702975265127574, -30.365472767157637]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.00316212972202072, 0.9999449065525032, -0.010009235501609866]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.468714366179746, 158.0536987351747]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [75.22477674741873, 19.24537614566374, -28.630859987457875]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.0012772041464585793, 0.9998898115065579, -0.014789644855390168]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.115498838385271, 158.5255545185648]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [77.90642035130945, 19.282310526586446, -25.902240206271845]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.001277204146458627, 0.999889811506558, -0.014789644855390222]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.144247236198858, 158.4334348095817]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [79.29568478771452, 19.68790273017715, -24.37545349763387]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.0026275040583625016, 0.9999324057717802, -0.01132608095871448]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.453077847146567, 157.96304160441224]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [81.2968320186084, 20.00851051712053, -22.009067567264918]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.005634026255948264, 0.9999452405328078, -0.008818938934636157]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.852034939963843, 157.37821905944136]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [82.73089724508009, 20.225108993453908, -20.17949272074056]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.007709859307667689, 0.9999440895551494, -0.007237114976257671]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.0496683863057825, 157.07149364188396]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [83.46880981232914, 20.622733988019494, -19.18680919591844]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.011896565539467068, 0.9999207449130535, -0.00412014816352478]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.928963302190084, 157.1547490199291]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [84.75148062119185, 20.934021688640424, -17.374512736532214]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.015097561166897356, 0.9998842825107197, -0.0018668783663504558]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.03754096854167, 156.94967687982347]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [86.11375233008812, 21.43935122870124, -15.30575781019713]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.02053758539311425, 0.9997875787192734, 0.0017335007566268781]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.122776211702757, 156.74557444459722]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [87.50013068713429, 21.0791764272972, -13.058582377350515]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.016087247078882417, 0.9998700993433984, -0.0009924315823494631]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.905772302189533, 155.77904319806206]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [88.55016582417551, 21.097898636739014, -11.217042108048657]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.01608724707888244, 0.9998700993433983, -0.0009924315823494493]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [8.384906194235114, 155.15044326818264]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [89.81701896140359, 21.6977795777403, -8.79262786967695]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.023027392811667494, 0.9997313802012363, 0.002628045134434197]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [8.690768322932215, 154.65589960979247]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [91.37663766453318, 22.109650817511636, -5.472795182031777]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.027736076203113115, 0.9996033874713821, 0.004876251900072583]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [9.505934989192536, 153.53777928616606]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [92.19850590166932, 22.41160461753037, -3.5206710128256646]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03135504919242932, 0.9994880456970849, 0.006364542305791886]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [10.004851632238317, 152.8490092576765]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [93.3028283706553, 22.139540960421563, -0.6754046833247991]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.02771266723653359, 0.9996037488432231, 0.004934913698555133]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [11.285979783552126, 151.24545705240757]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [93.7592497174737, 22.338134000701928, 0.643048135250004]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.030180602480818493, 0.9995277478549512, 0.00578035482510264]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [11.689898709652084, 150.69819296142526]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [94.58396404285857, 22.636528410087923, 3.2552561296928353]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03390421597331466, 0.9994007084377681, 0.006980552508405405]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [12.607124923781857, 149.48074037802166]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [95.47922794533545, 22.836303349271258, 6.547698060256643]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.036414648506222225, 0.9993071953678356, 0.007687825455810385]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [13.977944712529872, 147.7019386960643]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [96.3281890802888, 22.549452822132924, 10.331392582660296]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03259048305896922, 0.999445224357829, 0.0068632297289500065]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [16.026810498816634, 145.13417365923203]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [97.02331690429882, 22.543041678337275, 14.565862324926274]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03259048305896933, 0.9994452243578287, 0.006863229728950026]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [18.26500005195522, 142.290762537214]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [97.32781974256783, 22.53498417212241, 17.185172971468603]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03259048305896948, 0.9994452243578287, 0.00686322972895004]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [19.73274999924394, 140.42953250076022]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [97.56876712038988, 22.521246375927063, 20.329868266254678]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03259048305896964, 0.9994452243578288, 0.006863229728950052]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [21.577615831893436, 138.09328211431313]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [97.57867705839811, 22.71190806061416, 20.64100725060988]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03519612960177209, 0.9993560431812314, 0.0069807892253270516]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [21.63589095631762, 137.99280660913985]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [97.68253423936942, 22.215273738469712, 23.376583983325368]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.028677429624081336, 0.9995655491918999, 0.006805725446370661]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [23.644708125329597, 135.51908048557513]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [97.67980537625124, 22.297689465506767, 25.39373697147191]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.02998166599304103, 0.9995274019097729, 0.006787675285485274]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [23.571479364037327, 135.5923557848696]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [97.61725598000334, 22.861546788898796, 27.10333425435933]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03780365749776186, 0.9992635504964774, 0.006575722695359593]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [22.132945577048094, 137.32129932805566]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [97.40174441009034, 23.312916084514868, 30.40947915882487]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04430317910604703, 0.9989996976327344, 0.00606897444408803]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [19.854867768577694, 140.11838019661445]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [97.25901964356393, 23.29720337705494, 31.954028273922834]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.044303179106047016, 0.9989996976327344, 0.006068974444088038]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [18.974715351437098, 141.22460430162982]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [96.88376433589072, 23.262171463719532, 34.98120657109321]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04430317910604696, 0.9989996976327341, 0.006068974444088044]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [17.313869901211866, 143.3139329172801]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [96.3939103006883, 23.415341945687587, 37.88968995074684]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04686419722363235, 0.9988856864174237, 0.0055796495388935816]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [15.67145365630385, 145.3568147756205]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [96.07436560996665, 23.584147462405323, 39.463133818645645]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.049414624277783856, 0.9987656391813229, 0.005039136653181767]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [14.758284645791697, 146.48221521655108]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [95.6116710983979, 23.55081862074215, 41.53170768749478]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04941462427778387, 0.9987656391813227, 0.005039136653181765]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [13.762111373282202, 147.73965563441004]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [94.43286563983426, 23.18144166549499, 45.91984134421597]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.045673994292397685, 0.9989372769174313, 0.006180859985938286]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [11.982520050943272, 150.03047433581048]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [93.8018194974787, 23.140404304165326, 47.88904176325708]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04567399429239767, 0.998937276917431, 0.0061808599859382845]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [11.16282731841827, 151.0697233847663]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [92.25525314211896, 23.14018006614602, 52.058708424407044]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04688083514623507, 0.9988842982191762, 0.005687360299414778]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [9.501910710929437, 153.16765184670393]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [91.13843754608584, 23.07287861654977, 54.67312225084378]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.046880835146235145, 0.9988842982191761, 0.0056873602994147315]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [8.601386071378506, 154.3170175894836]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [89.75977150941088, 22.991702129956355, 57.5660113574791]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04688083514623513, 0.9988842982191763, 0.005687360299414742]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.701192019739963, 155.47070698386327]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [86.97315336051416, 22.446958686363864, 62.62639267400776]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04239082018954209, 0.9990662717092182, 0.008342847038054057]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.646256177113422, 156.89463912781667]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [83.65997604592187, 22.264670007898903, 67.62111046995443]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04239082018954202, 0.9990662717092184, 0.008342847038054109]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [5.719232010426609, 158.1167582996702]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [80.52322970005554, 22.001591934286825, 71.65785856057246]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04138627880812806, 0.9991010935293709, 0.009175011434929287]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [5.344544114164023, 158.65556926405517]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [78.26036810854438, 22.269350244407335, 74.22694348991487]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04521369827787863, 0.9989614803522722, 0.005628699710778437]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [4.975108279909776, 159.10855616765917]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [76.19878377977156, 21.971275750901274, 76.39283694174259]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04338545272216721, 0.9990303318926674, 0.007489889887000674]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [5.1255965786428135, 158.98206092518478]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [73.90262384207512, 21.566086675864657, 78.61157101099272]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04073190551103831, 0.9991163302582318, 0.010366797227882063]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [5.4671370895981894, 158.63287298549818]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [72.73091682381374, 21.988827261195652, 79.65364542491704]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04509457940460722, 0.998967483859653, 0.0055177078062946736]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [5.2621239345845865, 158.84742301362826]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [71.99501181398313, 21.855764019016824, 80.29799221987318]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.044249976903802064, 0.998999266513007, 0.00651191603807386]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [5.416264953967193, 158.67995943949205]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [70.09571682417585, 21.568749957193596, 81.88132133806126]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04261835863726607, 0.9990548562291088, 0.008548084703589861]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [5.825638579201701, 158.22589768800572]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [67.37949757753711, 21.62769085941608, 83.9550247337553]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04415705385717414, 0.9990038397937933, 0.0064407043024894695]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.206786551162281, 157.77101711258456]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [63.363810021223884, 21.432592900859397, 86.68495568704364]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.0441570538571741, 0.9990038397937931, 0.006440704302489544]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.179689750897275, 156.62259482117088]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [58.36979935098697, 20.904413718551407, 89.58077264120273]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04225290770281911, 0.9990582854961454, 0.009860728784410383]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [8.905231397902995, 154.5858561166812]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [54.67794295179828, 20.34521219586303, 91.40081938529056]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04004849206400772, 0.9990911832924244, 0.014591975559949255]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [10.501902109554685, 152.70008916848164]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [50.857941217262336, 20.553533237490598, 93.01919881329837]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.041917879993279485, 0.9990737909874802, 0.009718616093558249]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [11.77664743985759, 151.11854912311588]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [44.943278224364555, 20.381902129807013, 95.04867880331169]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.0422762244826863, 0.9990701095307843, 0.008463869424566187]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [14.396160886772975, 147.92551706617962]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [39.54712767556012, 19.949320628901283, 96.42800578217143]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04172978860950084, 0.9990682031212926, 0.011015909158929407]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [17.093273296189714, 144.66477743693804]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [34.5532026631349, 20.01965664121308, 97.32635742329812]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.042292824988000806, 0.9990797322222668, 0.007141821701671796]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [19.763532244178624, 141.365683300573]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [30.016124978479233, 20.305298652888116, 97.83896261580713]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.042805390140385646, 0.999083225905455, 0.0006374082511807501]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [22.24308998665847, 138.27216119578154]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [23.664445732632423, 20.129281596260892, 98.08362252202043]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04279932385913229, 0.9990834661455915, -0.000667496600551644]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [25.843761920602052, 133.85813459436562]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [17.280391761321386, 19.56677079412103, 97.77477201909795]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04317340754655987, 0.9990623747231927, 0.0032292868057064886]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [21.752167331666072, 139.18194652257148]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [14.200644766176657, 19.434817713258884, 97.4259404051353]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04318245772327231, 0.9990619878512353, 0.003227967737722593]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [19.96454869536422, 141.4977443474465]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [9.31144915800272, 19.226163827421686, 96.59900627497376]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04318245772327232, 0.9990619878512355, 0.0032279677377225723]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [17.30088491939132, 144.95576347369467]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [5.753978667034741, 18.689959665514365, 95.7743065972104]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04446512967043961, 0.9989765663699972, 0.008286862313416816]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [15.283489493483877, 147.62446755982347]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [1.315725851877681, 17.92553018260654, 94.47058849841125]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04691874754038306, 0.9987750181458377, 0.015719232074874138]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [12.886524757212594, 150.81692928003838]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-2.5246958983445893, 17.670384819570383, 93.10031624142736]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.0473987507469547, 0.9987325306429529, 0.01693194209626598]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [11.20022285696627, 153.0372456020457]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-5.3910981370097435, 17.554366607494163, 91.91954942468493]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.047400427107941595, 0.998732463507958, 0.016931209209928057]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [10.08494254558467, 154.5045626083531]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-7.897728091033251, 17.839901622633082, 90.78255766345619]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04511528897855666, 0.9989067797670939, 0.01224156998575237]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [9.398701150712512, 155.3719565821493]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-11.60380801270124, 16.541218305425932, 88.81249646921808]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.05274457159950463, 0.9982720705659759, 0.025901414917039848]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.5779714151303, 157.90280975234077]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-12.867223494721296, 16.58929589264785, 88.10056654160756]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.052097277228920005, 0.9983347745916715, 0.02476997267518964]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.2652798694971, 158.3117061604263]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-14.199680174555256, 16.53942984209345, 87.30789902589011]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.05209727722891999, 0.9983347745916715, 0.024769972675189624]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.904256463551349, 158.7960303596606]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-16.113015989181985, 16.27711543221373, 86.09003085326106]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.05353791125517099, 0.9982022909319574, 0.026942131256181515]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.32679812038662, 159.59190356264304]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-14.592185744743848, 15.660551890279384, 87.01189123889202]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.058757303580610584, 0.9976788131057341, 0.03441751234283881]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.340527306294291, 159.6314743017884]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-16.532912306070873, 15.300676378182708, 85.7465431408478]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.06091094709998594, 0.9974318415753618, 0.037677286725096654]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [5.719950953655697, 160.495404562292]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-19.794308178615548, 15.284255221102, 83.46162916902968]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.060135533719871466, 0.9975178532708253, 0.03663127066033421]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [5.116357185786391, 161.3153759061634]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-23.182705891310963, 15.56286340574975, 80.81532661622019]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.05680575402011655, 0.9978520332273144, 0.032625543586596514]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [4.832640467914196, 161.69612965958612]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-27.57613432370037, 14.770782948419981, 76.75574944716261]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.0631275513473604, 0.9972357067446929, 0.039190017278629564]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [4.2010107633950895, 162.6801277928347]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-30.868931185464525, 14.986780869099231, 73.31825282084893]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.060229349521878216, 0.997514325335927, 0.03657316237053454]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [4.444605182979174, 162.4071896379282]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-33.82560277771912, 15.033980443997306, 69.79168604820165]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.059212351013289526, 0.9976046720534044, 0.035760533898397645]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [4.858790471750339, 161.9376073883087]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-36.31332132147975, 15.486864604560123, 66.4834317807393]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.053947724629174736, 0.9980330839542552, 0.03193127526569194]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [5.661910863373457, 160.91570667125688]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-39.530778311260434, 15.761178676590882, 61.48379415841815]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.05073525440995745, 0.9982702288132006, 0.02970663604646705]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.980088003169925, 159.28712212214555]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-41.77691275151186, 16.346312472009053, 57.41009840058052]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04378646538923071, 0.9986992675863219, 0.02612505255997879]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [8.497890689838094, 157.34703718591186]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-43.34696178061979, 16.460573178676466, 54.092492154841324]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04259370871692249, 0.9987644499644828, 0.02559979423515198]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [9.688315806203093, 155.8695790206143]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-45.25547684850955, 16.213481255889327, 49.28565807328857]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04628281413721537, 0.9985662629450649, 0.02689460212510572]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [11.40570891025683, 153.78847000287917]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-47.00159012139568, 16.568629144076276, 43.82359470561873]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04251669087627295, 0.9987614373237244, 0.025844193003075187]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [13.988333431728762, 150.53691884393248]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-48.27692456128517, 16.9423728766385, 38.446433828862524]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.0386876823725204, 0.9989373521473252, 0.025048547214035887]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [16.823223824828705, 146.9578867006594]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-49.40544432943752, 17.104060590963552, 30.2564793861783]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.038684774440947985, 0.9989366874504633, 0.025079527347002]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [21.417618147265202, 141.20133584486194]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-49.638802890933626, 17.21815626035748, 25.352010029547497]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03868477444094804, 0.9989366874504632, 0.025079527347001997]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [24.456646479157655, 137.38290531638432]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-49.58596734111511, 17.321406716511444, 21.320963673600748]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03868477444094804, 0.9989366874504632, 0.025079527347001994]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [23.239801690437766, 138.91551757406265]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-49.28273293090412, 17.441873817088553, 16.990401532903626]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03868477444094792, 0.9989366874504629, 0.02507952734700199]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [20.75934402397094, 142.0338206072572]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-48.95302322949681, 17.911697392199866, 14.126460812659955]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03351707089074896, 0.9991052207774651, 0.025794646229738344]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [19.42448313179606, 143.65676638013935]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-48.44132527521866, 18.010524684947555, 10.96346928315159]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03351707089074895, 0.9991052207774649, 0.025794646229738316]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [17.801235346397288, 145.69128467586756]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-47.78101921747981, 18.114982863395696, 7.775474931098355]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03351707089074896, 0.9991052207774652, 0.02579464622973833]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [16.258709761455993, 147.62128858128526]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-47.086606676793856, 18.113312546929553, 5.016127169695708]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03478755123310557, 0.9990693807002856, 0.025498996575398114]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [14.947769111638209, 149.2719438815882]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-45.55036935971788, 17.61853737878611, 0.11284990953950763]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.04353360629887531, 0.9987901058161593, 0.022872464807265958]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [12.525133710007346, 152.38425425970794]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-43.912912631239976, 18.46247949505461, -4.212580997166555]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03498989941312795, 0.9990470554977898, 0.026089995022139184]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [11.30180868707981, 153.80220837228498]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-42.16111979114636, 18.62457580312101, -8.069782165522772]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.03499819075735283, 0.9990463124241282, 0.02610732225761922]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [10.054291089142513, 155.33922472956317]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-39.94753773704195, 20.25317796568078, -12.270014737624294]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.017412247381710187, 0.9992457508802488, 0.03470940778644659]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [9.522339903848675, 155.77168530929805]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-38.221769364365386, 21.152126065693484, -15.11603243875864]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.008547319640280649, 0.9991542337915073, 0.04022139260981839]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [8.181538971096856, 157.32302759188514]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-35.47970817479596, 20.952574201387904, -19.141623493020553]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.012777803845574153, 0.9992274514635053, 0.037164875495983796]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.494372058819408, 158.18929128330834]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-33.88844887051949, 21.339318822509753, -21.228517044912618]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.009689407665947686, 0.9991698274825012, 0.039569827238341175]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.904719903570822, 158.86339663810995]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-31.713238973632645, 21.2722536666, -23.863955654901773]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.011670742625252459, 0.9992144544542252, 0.03787172792889523]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.649375811499702, 159.1713344306398]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-29.457340882921127, 21.489014106701713, -26.347424635059486]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.0107454187347109, 0.9991895123635449, 0.03879245234471769]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.343986581518557, 159.50211010300117]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-27.500282052414043, 21.875553453321555, -28.322114920550327]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.007999024684361665, 0.9991030926708895, 0.04158155624268005]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.05513949340537, 159.79164936851902]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-22.56130992458092, 21.71438931324651, -32.73855734260822]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.011511153518439447, 0.9992219605381931, 0.03772223380030725]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.4733635638835025, 159.22969971419263]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-19.870301063641755, 22.112476065731133, -34.818893500794076]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.009056946970355178, 0.9991274811868714, 0.040770676334271824]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [6.570469320316711, 159.02120114323537]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-17.167185817411674, 22.1187037431307, -36.72992680832519]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.009803277205105732, 0.9991635188668784, 0.039700860464192175]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [7.054097587506625, 158.3782675173727]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-13.180329612559119, 21.968846494589073, -39.23974315808158]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.01184891151692748, 0.9992683156687323, 0.03636532134381312]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [8.073104327204433, 157.0654839908068]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-8.38865890586202, 21.542092203475008, -41.81940007758984]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.015390420796369191, 0.9994496674219213, 0.02938532351228447]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [9.745278552507873, 154.95606845262665]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-5.1951683335782635, 21.634419377787466, -43.28704428868359]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.015390420796369257, 0.9994496674219212, 0.029385323512284293]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [10.75762145038594, 153.6268568008908]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [-2.9471240368766334, 21.600068771980446, -44.214435126105116]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.015879551645140176, 0.9994768438097434, 0.02817585007880867]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [11.606075141897445, 152.53321122129225]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [0.4049210190388848, 21.591615257075777, -45.43918265222227]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.016297437783311972, 0.9995041932617149, 0.02693995497300076]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [12.94156012038372, 150.80752188892075]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [3.71959917049412, 21.192780878416627, -46.480848982626675]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.01827773971385808, 0.9996181429804349, 0.0207241997505105]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [14.640198788096821, 148.67730586267632]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [7.159939742837203, 21.947079811434843, -47.361309922800594]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.01590219030173378, 0.9994370472186285, 0.029541648405534422]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [15.719915953421946, 147.167836974348]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [9.42151987036476, 22.766751479653102, -47.8294867624175]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.013721414158970578, 0.9991156221167483, 0.039745395151401124]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [16.275078606545083, 146.3259926757436]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [11.500056525296696, 23.195163973279165, -48.201425034402035]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.01279341935566165, 0.9989105619040093, 0.044877808968430286]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [17.047592089208166, 145.26580577058726]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [13.282258304975779, 23.61506997071934, -48.467084944836095]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.011990549010130285, 0.9986757383820047, 0.05002995404348914]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [17.707199602590176, 144.3529443220078]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [14.863829729136555, 24.221082642885793, -48.65078483970902]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.010915128029723463, 0.9982699839014689, 0.05777455513854057]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [18.16361681816001, 143.6734112880074]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [16.7660496452224, 24.15808641209985, -48.86744757857662]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.011068843486705983, 0.9983423292811213, 0.056480742465127416]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [19.264524233297042, 142.26830551951804]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [18.791269487822497, 22.845189555911297, -49.10042872703701]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.012689810228319186, 0.9991855527404662, 0.038304045623132825]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [21.2992703671848, 139.85480860791998]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [22.66720632325544, 22.708719773286617, -49.27821935753306]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [-0.012773968916933708, 0.9992810053208991, 0.035697312545374706]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [23.714930165219123, 136.7747604934275]


CIRCLE_VIDEO_CAMS.append(EasyDict()); CIRCLE_VIDEO_CAMS[-1].position = [25.138323923977495, 16.942743488795763, -49.26072254551744]
CIRCLE_VIDEO_CAMS[-1].focal_point = [24.00916862487793, 20.094802856445312, 24.373815536499023]
CIRCLE_VIDEO_CAMS[-1].view_angle = 30.0
CIRCLE_VIDEO_CAMS[-1].view_up = [0.009907699976227128, 0.9990424886335852, -0.042613887243406416]
CIRCLE_VIDEO_CAMS[-1].clipping_range = [23.84445325003143, 137.35364071353837]


# dummy
CAMS = []
constants/mayavi_cams.py
jellyheadandrew/CHORUS
## result-save directories
DEFAULT_PROMPT_DIR = "prompts/demo"
DEFAULT_IMAGE_DIR = "results/images"
DEFAULT_SEGMENTATION_DIR = "results/segmentations"
DEFAULT_HUMAN_DIR = "results/humans"
DEFAULT_KEYPOINT_DIR = "results/human_kps"
DEFAULT_FILTERING_DIR = "results/filtered"
DEFAULT_PERSPECTIE_CAMERA_DIR = "results/filtered_cameras"
DEFAULT_AGGREGATION_DIR = "results/aggregations"
DEFAULT_QUANT_EVAL_DIR = "results/quant_results"

## result-save name for 'total aggregation'
DEFAULT_TOTAL_AGGREGATION_NAME = "total-aggregation"

## global seed for reproducibility
DEFAULT_SEED = 42
constants/metadata.py
jellyheadandrew/CHORUS
DEFAULT_FOV = 46.4
DEFAULT_LR = 0.01
DEFAULT_EPOCHS = 2400
DEFAULT_STOP_THRES = 0.7
constants/perspective.py
jellyheadandrew/CHORUS
DEFAULT_EFT_MIN_BBOX_OVERLAP = 0.5
DEFAULT_EFT_REPROJECTION_ERROR_THRESH = 1.0
DEFAULT_EFT_FOV = 46.4
DEFAULT_EFT_LR = 0.01
DEFAULT_EFT_JOINT_LOSS_THRES = 0.5
DEFAULT_EFT_EPOCHS = 3000
constants/prepare_eft.py
jellyheadandrew/CHORUS
DEFAULT_TEMPERATURE = 1.0
DEFAULT_AVAILABLE_M = list(range(3,20+1))
DEFAULT_SWITCHING_PART = ["category is person", "word is 'A person'"]
DEFAULT_TEMPLATE = \
f"Generate at most [M] simple subject-verb-object prompt where subject's [SWITCHING-PART] and object's category is [CATEGORY]. "\
"You should use diverse and general word but no pronoun for subject. "\
"Generated prompt must align with common sense. "\
"Verb must depict physical interaction between subject and object. "\
"Simple verb preferred."
constants/prompt_generation.py
jellyheadandrew/CHORUS
DETECTRON2_CODE_DIR="./imports/detectron2_repo"

COCO_SEG_CONFIG_PTH="./imports/detectron2_repo/projects/PointRend/configs/InstanceSegmentation/pointrend_rcnn_R_50_FPN_3x_coco.yaml"
COCO_SEG_WEIGHTS_PTH="detectron2://PointRend/InstanceSegmentation/pointrend_rcnn_R_50_FPN_3x_coco/164955410/model_final_edd263.pkl"

LVIS_SEG_CONFIG_PTH="LVISv0.5-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_1x.yaml"
LVIS_SEG_WEIGHTS_PTH="LVISv0.5-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_1x.yaml"

DEFAULT_SEGMENTATION_THRESHOLD=0.8
constants/segmentation.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.
import copy
import logging
import re
from typing import Dict, List
import torch
from tabulate import tabulate


def convert_basic_c2_names(original_keys):
    """
    Apply some basic name conversion to names in C2 weights.
    It only deals with typical backbone models.

    Args:
        original_keys (list[str]):
    Returns:
        list[str]: The same number of strings matching those in original_keys.
    """
    layer_keys = copy.deepcopy(original_keys)
    layer_keys = [
        {"pred_b": "linear_b", "pred_w": "linear_w"}.get(k, k) for k in layer_keys
    ]  # some hard-coded mappings

    layer_keys = [k.replace("_", ".") for k in layer_keys]
    layer_keys = [re.sub("\\.b$", ".bias", k) for k in layer_keys]
    layer_keys = [re.sub("\\.w$", ".weight", k) for k in layer_keys]
    # Uniform both bn and gn names to "norm"
    layer_keys = [re.sub("bn\\.s$", "norm.weight", k) for k in layer_keys]
    layer_keys = [re.sub("bn\\.bias$", "norm.bias", k) for k in layer_keys]
    layer_keys = [re.sub("bn\\.rm", "norm.running_mean", k) for k in layer_keys]
    layer_keys = [re.sub("bn\\.running.mean$", "norm.running_mean", k) for k in layer_keys]
    layer_keys = [re.sub("bn\\.riv$", "norm.running_var", k) for k in layer_keys]
    layer_keys = [re.sub("bn\\.running.var$", "norm.running_var", k) for k in layer_keys]
    layer_keys = [re.sub("bn\\.gamma$", "norm.weight", k) for k in layer_keys]
    layer_keys = [re.sub("bn\\.beta$", "norm.bias", k) for k in layer_keys]
    layer_keys = [re.sub("gn\\.s$", "norm.weight", k) for k in layer_keys]
    layer_keys = [re.sub("gn\\.bias$", "norm.bias", k) for k in layer_keys]

    # stem
    layer_keys = [re.sub("^res\\.conv1\\.norm\\.", "conv1.norm.", k) for k in layer_keys]
    # to avoid mis-matching with "conv1" in other components (e.g. detection head)
    layer_keys = [re.sub("^conv1\\.", "stem.conv1.", k) for k in layer_keys]

    # layer1-4 is used by torchvision, however we follow the C2 naming strategy (res2-5)
    # layer_keys = [re.sub("^res2.", "layer1.", k) for k in layer_keys]
    # layer_keys = [re.sub("^res3.", "layer2.", k) for k in layer_keys]
    # layer_keys = [re.sub("^res4.", "layer3.", k) for k in layer_keys]
    # layer_keys = [re.sub("^res5.", "layer4.", k) for k in layer_keys]

    # blocks
    layer_keys = [k.replace(".branch1.", ".shortcut.") for k in layer_keys]
    layer_keys = [k.replace(".branch2a.", ".conv1.") for k in layer_keys]
    layer_keys = [k.replace(".branch2b.", ".conv2.") for k in layer_keys]
    layer_keys = [k.replace(".branch2c.", ".conv3.") for k in layer_keys]

    # DensePose substitutions
    layer_keys = [re.sub("^body.conv.fcn", "body_conv_fcn", k) for k in layer_keys]
    layer_keys = [k.replace("AnnIndex.lowres", "ann_index_lowres") for k in layer_keys]
    layer_keys = [k.replace("Index.UV.lowres", "index_uv_lowres") for k in layer_keys]
    layer_keys = [k.replace("U.lowres", "u_lowres") for k in layer_keys]
    layer_keys = [k.replace("V.lowres", "v_lowres") for k in layer_keys]
    return layer_keys


def convert_c2_detectron_names(weights):
    """
    Map Caffe2 Detectron weight names to Detectron2 names.

    Args:
        weights (dict): name -> tensor

    Returns:
        dict: detectron2 names -> tensor
        dict: detectron2 names -> C2 names
    """
    logger = logging.getLogger(__name__)
    logger.info("Renaming Caffe2 weights ......")
    original_keys = sorted(weights.keys())
    layer_keys = copy.deepcopy(original_keys)

    layer_keys = convert_basic_c2_names(layer_keys)

    # --------------------------------------------------------------------------
    # RPN hidden representation conv
    # --------------------------------------------------------------------------
    # FPN case
    # In the C2 model, the RPN hidden layer conv is defined for FPN level 2 and then
    # shared for all other levels, hence the appearance of "fpn2"
    layer_keys = [
        k.replace("conv.rpn.fpn2", "proposal_generator.rpn_head.conv") for k in layer_keys
    ]
    # Non-FPN case
    layer_keys = [k.replace("conv.rpn", "proposal_generator.rpn_head.conv") for k in layer_keys]

    # --------------------------------------------------------------------------
    # RPN box transformation conv
    # --------------------------------------------------------------------------
    # FPN case (see note above about "fpn2")
    layer_keys = [
        k.replace("rpn.bbox.pred.fpn2", "proposal_generator.rpn_head.anchor_deltas")
        for k in layer_keys
    ]
    layer_keys = [
        k.replace("rpn.cls.logits.fpn2", "proposal_generator.rpn_head.objectness_logits")
        for k in layer_keys
    ]
    # Non-FPN case
    layer_keys = [
        k.replace("rpn.bbox.pred", "proposal_generator.rpn_head.anchor_deltas") for k in layer_keys
    ]
    layer_keys = [
        k.replace("rpn.cls.logits", "proposal_generator.rpn_head.objectness_logits")
        for k in layer_keys
    ]

    # --------------------------------------------------------------------------
    # Fast R-CNN box head
    # --------------------------------------------------------------------------
    layer_keys = [re.sub("^bbox\\.pred", "bbox_pred", k) for k in layer_keys]
    layer_keys = [re.sub("^cls\\.score", "cls_score", k) for k in layer_keys]
    layer_keys = [re.sub("^fc6\\.", "box_head.fc1.", k) for k in layer_keys]
    layer_keys = [re.sub("^fc7\\.", "box_head.fc2.", k) for k in layer_keys]
    # 4conv1fc head tensor names: head_conv1_w, head_conv1_gn_s
    layer_keys = [re.sub("^head\\.conv", "box_head.conv", k) for k in layer_keys]

    # --------------------------------------------------------------------------
    # FPN lateral and output convolutions
    # --------------------------------------------------------------------------
    def fpn_map(name):
        """
        Look for keys with the following patterns:
        1) Starts with "fpn.inner."
           Example: "fpn.inner.res2.2.sum.lateral.weight"
           Meaning: These are lateral pathway convolutions
        2) Starts with "fpn.res"
           Example: "fpn.res2.2.sum.weight"
           Meaning: These are FPN output convolutions
        """
        splits = name.split(".")
        norm = ".norm" if "norm" in splits else ""
        if name.startswith("fpn.inner."):
            # splits example: ['fpn', 'inner', 'res2', '2', 'sum', 'lateral', 'weight']
            stage = int(splits[2][len("res") :])
            return "fpn_lateral{}{}.{}".format(stage, norm, splits[-1])
        elif name.startswith("fpn.res"):
            # splits example: ['fpn', 'res2', '2', 'sum', 'weight']
            stage = int(splits[1][len("res") :])
            return "fpn_output{}{}.{}".format(stage, norm, splits[-1])
        return name

    layer_keys = [fpn_map(k) for k in layer_keys]

    # --------------------------------------------------------------------------
    # Mask R-CNN mask head
    # --------------------------------------------------------------------------
    # roi_heads.StandardROIHeads case
    layer_keys = [k.replace(".[mask].fcn", "mask_head.mask_fcn") for k in layer_keys]
    layer_keys = [re.sub("^\\.mask\\.fcn", "mask_head.mask_fcn", k) for k in layer_keys]
    layer_keys = [k.replace("mask.fcn.logits", "mask_head.predictor") for k in layer_keys]
    # roi_heads.Res5ROIHeads case
    layer_keys = [k.replace("conv5.mask", "mask_head.deconv") for k in layer_keys]

    # --------------------------------------------------------------------------
    # Keypoint R-CNN head
    # --------------------------------------------------------------------------
    # interestingly, the keypoint head convs have blob names that are simply "conv_fcnX"
    layer_keys = [k.replace("conv.fcn", "roi_heads.keypoint_head.conv_fcn") for k in layer_keys]
    layer_keys = [
        k.replace("kps.score.lowres", "roi_heads.keypoint_head.score_lowres") for k in layer_keys
    ]
    layer_keys = [k.replace("kps.score.", "roi_heads.keypoint_head.score.") for k in layer_keys]

    # --------------------------------------------------------------------------
    # Done with replacements
    # --------------------------------------------------------------------------
    assert len(set(layer_keys)) == len(layer_keys)
    assert len(original_keys) == len(layer_keys)

    new_weights = {}
    new_keys_to_original_keys = {}
    for orig, renamed in zip(original_keys, layer_keys):
        new_keys_to_original_keys[renamed] = orig
        if renamed.startswith("bbox_pred.") or renamed.startswith("mask_head.predictor."):
            # remove the meaningless prediction weight for background class
            new_start_idx = 4 if renamed.startswith("bbox_pred.") else 1
            new_weights[renamed] = weights[orig][new_start_idx:]
            logger.info(
                "Remove prediction weight for background class in {}. The shape changes from "
                "{} to {}.".format(
                    renamed, tuple(weights[orig].shape), tuple(new_weights[renamed].shape)
                )
            )
        elif renamed.startswith("cls_score."):
            # move weights of bg class from original index 0 to last index
            logger.info(
                "Move classification weights for background class in {} from index 0 to "
                "index {}.".format(renamed, weights[orig].shape[0] - 1)
            )
            new_weights[renamed] = torch.cat([weights[orig][1:], weights[orig][:1]])
        else:
            new_weights[renamed] = weights[orig]

    return new_weights, new_keys_to_original_keys


# Note the current matching is not symmetric.
# it assumes model_state_dict will have longer names.
def align_and_update_state_dicts(model_state_dict, ckpt_state_dict, c2_conversion=True):
    """
    Match names between the two state-dict, and returns a new chkpt_state_dict with names
    converted to match model_state_dict with heuristics. The returned dict can be later
    loaded with fvcore checkpointer.
    If `c2_conversion==True`, `ckpt_state_dict` is assumed to be a Caffe2
    model and will be renamed at first.

    Strategy: suppose that the models that we will create will have prefixes appended
    to each of its keys, for example due to an extra level of nesting that the original
    pre-trained weights from ImageNet won't contain. For example, model.state_dict()
    might return backbone[0].body.res2.conv1.weight, while the pre-trained model contains
    res2.conv1.weight. We thus want to match both parameters together.
    For that, we look for each model weight, look among all loaded keys if there is one
    that is a suffix of the current weight name, and use it if that's the case.
    If multiple matches exist, take the one with longest size
    of the corresponding name. For example, for the same model as before, the pretrained
    weight file can contain both res2.conv1.weight, as well as conv1.weight. In this case,
    we want to match backbone[0].body.conv1.weight to conv1.weight, and
    backbone[0].body.res2.conv1.weight to res2.conv1.weight.
    """
    model_keys = sorted(model_state_dict.keys())
    if c2_conversion:
        ckpt_state_dict, original_keys = convert_c2_detectron_names(ckpt_state_dict)
        # original_keys: the name in the original dict (before renaming)
    else:
        original_keys = {x: x for x in ckpt_state_dict.keys()}
    ckpt_keys = sorted(ckpt_state_dict.keys())

    def match(a, b):
        # Matched ckpt_key should be a complete (starts with '.') suffix.
        # For example, roi_heads.mesh_head.whatever_conv1 does not match conv1,
        # but matches whatever_conv1 or mesh_head.whatever_conv1.
        return a == b or a.endswith("." + b)

    # get a matrix of string matches, where each (i, j) entry correspond to the size of the
    # ckpt_key string, if it matches
    match_matrix = [len(j) if match(i, j) else 0 for i in model_keys for j in ckpt_keys]
    match_matrix = torch.as_tensor(match_matrix).view(len(model_keys), len(ckpt_keys))
    # use the matched one with longest size in case of multiple matches
    max_match_size, idxs = match_matrix.max(1)
    # remove indices that correspond to no-match
    idxs[max_match_size == 0] = -1

    logger = logging.getLogger(__name__)
    # matched_pairs (matched checkpoint key --> matched model key)
    matched_keys = {}
    result_state_dict = {}
    for idx_model, idx_ckpt in enumerate(idxs.tolist()):
        if idx_ckpt == -1:
            continue
        key_model = model_keys[idx_model]
        key_ckpt = ckpt_keys[idx_ckpt]
        value_ckpt = ckpt_state_dict[key_ckpt]
        shape_in_model = model_state_dict[key_model].shape

        if shape_in_model != value_ckpt.shape:
            logger.warning(
                "Shape of {} in checkpoint is {}, while shape of {} in model is {}.".format(
                    key_ckpt, value_ckpt.shape, key_model, shape_in_model
                )
            )
            logger.warning(
                "{} will not be loaded. Please double check and see if this is desired.".format(
                    key_ckpt
                )
            )
            continue

        assert key_model not in result_state_dict
        result_state_dict[key_model] = value_ckpt
        if key_ckpt in matched_keys:  # already added to matched_keys
            logger.error(
                "Ambiguity found for {} in checkpoint!"
                "It matches at least two keys in the model ({} and {}).".format(
                    key_ckpt, key_model, matched_keys[key_ckpt]
                )
            )
            raise ValueError("Cannot match one checkpoint key to multiple keys in the model.")

        matched_keys[key_ckpt] = key_model

    # logging:
    matched_model_keys = sorted(matched_keys.values())
    if len(matched_model_keys) == 0:
        logger.warning("No weights in checkpoint matched with model.")
        return ckpt_state_dict
    common_prefix = _longest_common_prefix(matched_model_keys)
    rev_matched_keys = {v: k for k, v in matched_keys.items()}
    original_keys = {k: original_keys[rev_matched_keys[k]] for k in matched_model_keys}

    model_key_groups = _group_keys_by_module(matched_model_keys, original_keys)
    table = []
    memo = set()
    for key_model in matched_model_keys:
        if key_model in memo:
            continue
        if key_model in model_key_groups:
            group = model_key_groups[key_model]
            memo |= set(group)
            shapes = [tuple(model_state_dict[k].shape) for k in group]
            table.append(
                (
                    _longest_common_prefix([k[len(common_prefix) :] for k in group]) + "*",
                    _group_str([original_keys[k] for k in group]),
                    " ".join([str(x).replace(" ", "") for x in shapes]),
                )
            )
        else:
            key_checkpoint = original_keys[key_model]
            shape = str(tuple(model_state_dict[key_model].shape))
            table.append((key_model[len(common_prefix) :], key_checkpoint, shape))
    submodule_str = common_prefix[:-1] if common_prefix else "model"
    logger.info(
        f"Following weights matched with submodule {submodule_str} - Total num: {len(table)}"
    )

    unmatched_ckpt_keys = [k for k in ckpt_keys if k not in set(matched_keys.keys())]
    for k in unmatched_ckpt_keys:
        result_state_dict[k] = ckpt_state_dict[k]
    return result_state_dict


def _group_keys_by_module(keys: List[str], original_names: Dict[str, str]):
    """
    Params in the same submodule are grouped together.

    Args:
        keys: names of all parameters
        original_names: mapping from parameter name to their name in the checkpoint

    Returns:
        dict[name -> all other names in the same group]
    """

    def _submodule_name(key):
        pos = key.rfind(".")
        if pos < 0:
            return None
        prefix = key[: pos + 1]
        return prefix

    all_submodules = [_submodule_name(k) for k in keys]
    all_submodules = [x for x in all_submodules if x]
    all_submodules = sorted(all_submodules, key=len)

    ret = {}
    for prefix in all_submodules:
        group = [k for k in keys if k.startswith(prefix)]
        if len(group) <= 1:
            continue
        original_name_lcp = _longest_common_prefix_str([original_names[k] for k in group])
        if len(original_name_lcp) == 0:
            # don't group weights if original names don't share prefix
            continue

        for k in group:
            if k in ret:
                continue
            ret[k] = group
    return ret


def _longest_common_prefix(names: List[str]) -> str:
    """
    ["abc.zfg", "abc.zef"] -> "abc."
    """
    names = [n.split(".") for n in names]
    m1, m2 = min(names), max(names)
    ret = [a for a, b in zip(m1, m2) if a == b]
    ret = ".".join(ret) + "." if len(ret) else ""
    return ret


def _longest_common_prefix_str(names: List[str]) -> str:
    m1, m2 = min(names), max(names)
    lcp = []
    for a, b in zip(m1, m2):
        if a == b:
            lcp.append(a)
        else:
            break
    lcp = "".join(lcp)
    return lcp


def _group_str(names: List[str]) -> str:
    """
    Turn "common1", "common2", "common3" into "common{1,2,3}"
    """
    lcp = _longest_common_prefix_str(names)
    rest = [x[len(lcp) :] for x in names]
    rest = "{" + ",".join(rest) + "}"
    ret = lcp + rest

    # add some simplification for BN specifically
    ret = ret.replace("bn_{beta,running_mean,running_var,gamma}", "bn_*")
    ret = ret.replace("bn_beta,bn_running_mean,bn_running_var,bn_gamma", "bn_*")
    return ret

imports/detectron2_repo/detectron2/checkpoint/c2_model_loading.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.
import itertools
import logging
import numpy as np
import operator
import pickle
from collections import OrderedDict, defaultdict
from typing import Any, Callable, Dict, List, Optional, Union
import torch
import torch.utils.data as torchdata
from tabulate import tabulate
from termcolor import colored

from detectron2.config import configurable
from detectron2.structures import BoxMode
from detectron2.utils.comm import get_world_size
from detectron2.utils.env import seed_all_rng
from detectron2.utils.file_io import PathManager
from detectron2.utils.logger import _log_api_usage, log_first_n

from .catalog import DatasetCatalog, MetadataCatalog
from .common import AspectRatioGroupedDataset, DatasetFromList, MapDataset, ToIterableDataset
from .dataset_mapper import DatasetMapper
from .detection_utils import check_metadata_consistency
from .samplers import (
    InferenceSampler,
    RandomSubsetTrainingSampler,
    RepeatFactorTrainingSampler,
    TrainingSampler,
)

"""
This file contains the default logic to build a dataloader for training or testing.
"""

__all__ = [
    "build_batch_data_loader",
    "build_detection_train_loader",
    "build_detection_test_loader",
    "get_detection_dataset_dicts",
    "load_proposals_into_dataset",
    "print_instances_class_histogram",
]


def filter_images_with_only_crowd_annotations(dataset_dicts):
    """
    Filter out images with none annotations or only crowd annotations
    (i.e., images without non-crowd annotations).
    A common training-time preprocessing on COCO dataset.

    Args:
        dataset_dicts (list[dict]): annotations in Detectron2 Dataset format.

    Returns:
        list[dict]: the same format, but filtered.
    """
    num_before = len(dataset_dicts)

    def valid(anns):
        for ann in anns:
            if ann.get("iscrowd", 0) == 0:
                return True
        return False

    dataset_dicts = [x for x in dataset_dicts if valid(x["annotations"])]
    num_after = len(dataset_dicts)
    logger = logging.getLogger(__name__)
    logger.info(
        "Removed {} images with no usable annotations. {} images left.".format(
            num_before - num_after, num_after
        )
    )
    return dataset_dicts


def filter_images_with_few_keypoints(dataset_dicts, min_keypoints_per_image):
    """
    Filter out images with too few number of keypoints.

    Args:
        dataset_dicts (list[dict]): annotations in Detectron2 Dataset format.

    Returns:
        list[dict]: the same format as dataset_dicts, but filtered.
    """
    num_before = len(dataset_dicts)

    def visible_keypoints_in_image(dic):
        # Each keypoints field has the format [x1, y1, v1, ...], where v is visibility
        annotations = dic["annotations"]
        return sum(
            (np.array(ann["keypoints"][2::3]) > 0).sum()
            for ann in annotations
            if "keypoints" in ann
        )

    dataset_dicts = [
        x for x in dataset_dicts if visible_keypoints_in_image(x) >= min_keypoints_per_image
    ]
    num_after = len(dataset_dicts)
    logger = logging.getLogger(__name__)
    logger.info(
        "Removed {} images with fewer than {} keypoints.".format(
            num_before - num_after, min_keypoints_per_image
        )
    )
    return dataset_dicts


def load_proposals_into_dataset(dataset_dicts, proposal_file):
    """
    Load precomputed object proposals into the dataset.

    The proposal file should be a pickled dict with the following keys:

    - "ids": list[int] or list[str], the image ids
    - "boxes": list[np.ndarray], each is an Nx4 array of boxes corresponding to the image id
    - "objectness_logits": list[np.ndarray], each is an N sized array of objectness scores
      corresponding to the boxes.
    - "bbox_mode": the BoxMode of the boxes array. Defaults to ``BoxMode.XYXY_ABS``.

    Args:
        dataset_dicts (list[dict]): annotations in Detectron2 Dataset format.
        proposal_file (str): file path of pre-computed proposals, in pkl format.

    Returns:
        list[dict]: the same format as dataset_dicts, but added proposal field.
    """
    logger = logging.getLogger(__name__)
    logger.info("Loading proposals from: {}".format(proposal_file))

    with PathManager.open(proposal_file, "rb") as f:
        proposals = pickle.load(f, encoding="latin1")

    # Rename the key names in D1 proposal files
    rename_keys = {"indexes": "ids", "scores": "objectness_logits"}
    for key in rename_keys:
        if key in proposals:
            proposals[rename_keys[key]] = proposals.pop(key)

    # Fetch the indexes of all proposals that are in the dataset
    # Convert image_id to str since they could be int.
    img_ids = set({str(record["image_id"]) for record in dataset_dicts})
    id_to_index = {str(id): i for i, id in enumerate(proposals["ids"]) if str(id) in img_ids}

    # Assuming default bbox_mode of precomputed proposals are 'XYXY_ABS'
    bbox_mode = BoxMode(proposals["bbox_mode"]) if "bbox_mode" in proposals else BoxMode.XYXY_ABS

    for record in dataset_dicts:
        # Get the index of the proposal
        i = id_to_index[str(record["image_id"])]

        boxes = proposals["boxes"][i]
        objectness_logits = proposals["objectness_logits"][i]
        # Sort the proposals in descending order of the scores
        inds = objectness_logits.argsort()[::-1]
        record["proposal_boxes"] = boxes[inds]
        record["proposal_objectness_logits"] = objectness_logits[inds]
        record["proposal_bbox_mode"] = bbox_mode

    return dataset_dicts


def print_instances_class_histogram(dataset_dicts, class_names):
    """
    Args:
        dataset_dicts (list[dict]): list of dataset dicts.
        class_names (list[str]): list of class names (zero-indexed).
    """
    num_classes = len(class_names)
    hist_bins = np.arange(num_classes + 1)
    histogram = np.zeros((num_classes,), dtype=int)
    for entry in dataset_dicts:
        annos = entry["annotations"]
        classes = np.asarray(
            [x["category_id"] for x in annos if not x.get("iscrowd", 0)], dtype=int
        )
        if len(classes):
            assert classes.min() >= 0, f"Got an invalid category_id={classes.min()}"
            assert (
                classes.max() < num_classes
            ), f"Got an invalid category_id={classes.max()} for a dataset of {num_classes} classes"
        histogram += np.histogram(classes, bins=hist_bins)[0]

    N_COLS = min(6, len(class_names) * 2)

    def short_name(x):
        # make long class names shorter. useful for lvis
        if len(x) > 13:
            return x[:11] + ".."
        return x

    data = list(
        itertools.chain(*[[short_name(class_names[i]), int(v)] for i, v in enumerate(histogram)])
    )
    total_num_instances = sum(data[1::2])
    data.extend([None] * (N_COLS - (len(data) % N_COLS)))
    if num_classes > 1:
        data.extend(["total", total_num_instances])
    data = itertools.zip_longest(*[data[i::N_COLS] for i in range(N_COLS)])
    table = tabulate(
        data,
        headers=["category", "#instances"] * (N_COLS // 2),
        tablefmt="pipe",
        numalign="left",
        stralign="center",
    )
    log_first_n(
        logging.INFO,
        "Distribution of instances among all {} categories:\n".format(num_classes)
        + colored(table, "cyan"),
        key="message",
    )


def get_detection_dataset_dicts(
    names,
    filter_empty=True,
    min_keypoints=0,
    proposal_files=None,
    check_consistency=True,
):
    """
    Load and prepare dataset dicts for instance detection/segmentation and semantic segmentation.

    Args:
        names (str or list[str]): a dataset name or a list of dataset names
        filter_empty (bool): whether to filter out images without instance annotations
        min_keypoints (int): filter out images with fewer keypoints than
            `min_keypoints`. Set to 0 to do nothing.
        proposal_files (list[str]): if given, a list of object proposal files
            that match each dataset in `names`.
        check_consistency (bool): whether to check if datasets have consistent metadata.

    Returns:
        list[dict]: a list of dicts following the standard dataset dict format.
    """
    if isinstance(names, str):
        names = [names]
    assert len(names), names

    available_datasets = DatasetCatalog.keys()
    names_set = set(names)
    if not names_set.issubset(available_datasets):
        logger = logging.getLogger(__name__)
        logger.warning(
            "The following dataset names are not registered in the DatasetCatalog: "
            f"{names_set - available_datasets}. "
            f"Available datasets are {available_datasets}"
        )

    dataset_dicts = [DatasetCatalog.get(dataset_name) for dataset_name in names]

    if isinstance(dataset_dicts[0], torchdata.Dataset):
        if len(dataset_dicts) > 1:
            # ConcatDataset does not work for iterable style dataset.
            # We could support concat for iterable as well, but it's often
            # not a good idea to concat iterables anyway.
            return torchdata.ConcatDataset(dataset_dicts)
        return dataset_dicts[0]

    for dataset_name, dicts in zip(names, dataset_dicts):
        assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)

    if proposal_files is not None:
        assert len(names) == len(proposal_files)
        # load precomputed proposals from proposal files
        dataset_dicts = [
            load_proposals_into_dataset(dataset_i_dicts, proposal_file)
            for dataset_i_dicts, proposal_file in zip(dataset_dicts, proposal_files)
        ]

    dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts))

    has_instances = "annotations" in dataset_dicts[0]
    if filter_empty and has_instances:
        dataset_dicts = filter_images_with_only_crowd_annotations(dataset_dicts)
    if min_keypoints > 0 and has_instances:
        dataset_dicts = filter_images_with_few_keypoints(dataset_dicts, min_keypoints)

    if check_consistency and has_instances:
        try:
            class_names = MetadataCatalog.get(names[0]).thing_classes
            check_metadata_consistency("thing_classes", names)
            print_instances_class_histogram(dataset_dicts, class_names)
        except AttributeError:  # class names are not available for this dataset
            pass

    assert len(dataset_dicts), "No valid data found in {}.".format(",".join(names))
    return dataset_dicts


def build_batch_data_loader(
    dataset,
    sampler,
    total_batch_size,
    *,
    aspect_ratio_grouping=False,
    num_workers=0,
    collate_fn=None,
    drop_last: bool = True,
    prefetch_factor=None,
    persistent_workers=False,
    pin_memory=False,
):
    """
    Build a batched dataloader. The main differences from `torch.utils.data.DataLoader` are:
    1. support aspect ratio grouping options
    2. use no "batch collation", because this is common for detection training

    Args:
        dataset (torch.utils.data.Dataset): a pytorch map-style or iterable dataset.
        sampler (torch.utils.data.sampler.Sampler or None): a sampler that produces indices.
            Must be provided iff. ``dataset`` is a map-style dataset.
        total_batch_size, aspect_ratio_grouping, num_workers, collate_fn: see
            :func:`build_detection_train_loader`.
        drop_last (bool): if ``True``, the dataloader will drop incomplete batches.

    Returns:
        iterable[list]. Length of each list is the batch size of the current
            GPU. Each element in the list comes from the dataset.
    """
    world_size = get_world_size()
    assert (
        total_batch_size > 0 and total_batch_size % world_size == 0
    ), "Total batch size ({}) must be divisible by the number of gpus ({}).".format(
        total_batch_size, world_size
    )
    batch_size = total_batch_size // world_size

    if isinstance(dataset, torchdata.IterableDataset):
        assert sampler is None, "sampler must be None if dataset is IterableDataset"
    else:
        dataset = ToIterableDataset(dataset, sampler, shard_chunk_size=batch_size)

    if aspect_ratio_grouping:
        assert drop_last, "Aspect ratio grouping will drop incomplete batches."
        data_loader = torchdata.DataLoader(
            dataset,
            num_workers=num_workers,
            collate_fn=operator.itemgetter(0),  # don't batch, but yield individual elements
            worker_init_fn=worker_init_reset_seed,
            prefetch_factor=prefetch_factor,
            persistent_workers=persistent_workers,
            pin_memory=pin_memory,
        )  # yield individual mapped dict
        data_loader = AspectRatioGroupedDataset(data_loader, batch_size)
        if collate_fn is None:
            return data_loader
        return MapDataset(data_loader, collate_fn)
    else:
        return torchdata.DataLoader(
            dataset,
            batch_size=batch_size,
            drop_last=drop_last,
            num_workers=num_workers,
            collate_fn=trivial_batch_collator if collate_fn is None else collate_fn,
            worker_init_fn=worker_init_reset_seed,
            prefetch_factor=prefetch_factor,
            persistent_workers=persistent_workers,
            pin_memory=pin_memory,
        )


def _get_train_datasets_repeat_factors(cfg) -> Dict[str, float]:
    repeat_factors = cfg.DATASETS.TRAIN_REPEAT_FACTOR
    assert all(len(tup) == 2 for tup in repeat_factors)
    name_to_weight = defaultdict(lambda: 1, dict(repeat_factors))
    # The sampling weights map should only contain datasets in train config
    unrecognized = set(name_to_weight.keys()) - set(cfg.DATASETS.TRAIN)
    assert not unrecognized, f"unrecognized datasets: {unrecognized}"
    logger = logging.getLogger(__name__)
    logger.info(f"Found repeat factors: {list(name_to_weight.items())}")

    # pyre-fixme[7]: Expected `Dict[str, float]` but got `DefaultDict[typing.Any, int]`.
    return name_to_weight


def _build_weighted_sampler(cfg, enable_category_balance=False):
    dataset_repeat_factors = _get_train_datasets_repeat_factors(cfg)
    # OrderedDict to guarantee order of values() consistent with repeat factors
    dataset_name_to_dicts = OrderedDict(
        {
            name: get_detection_dataset_dicts(
                [name],
                filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
                min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
                if cfg.MODEL.KEYPOINT_ON
                else 0,
                proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN
                if cfg.MODEL.LOAD_PROPOSALS
                else None,
            )
            for name in cfg.DATASETS.TRAIN
        }
    )
    # Repeat factor for every sample in the dataset
    repeat_factors = [
        [dataset_repeat_factors[dsname]] * len(dataset_name_to_dicts[dsname])
        for dsname in cfg.DATASETS.TRAIN
    ]

    repeat_factors = list(itertools.chain.from_iterable(repeat_factors))

    repeat_factors = torch.tensor(repeat_factors)
    logger = logging.getLogger(__name__)
    if enable_category_balance:
        """
        1. Calculate repeat factors using category frequency for each dataset and then merge them.
        2. Element wise dot producting the dataset frequency repeat factors with
            the category frequency repeat factors gives the final repeat factors.
        """
        category_repeat_factors = [
            RepeatFactorTrainingSampler.repeat_factors_from_category_frequency(
                dataset_dict, cfg.DATALOADER.REPEAT_THRESHOLD
            )
            for dataset_dict in dataset_name_to_dicts.values()
        ]
        # flatten the category repeat factors from all datasets
        category_repeat_factors = list(itertools.chain.from_iterable(category_repeat_factors))
        category_repeat_factors = torch.tensor(category_repeat_factors)
        repeat_factors = torch.mul(category_repeat_factors, repeat_factors)
        repeat_factors = repeat_factors / torch.min(repeat_factors)
        logger.info(
            "Using WeightedCategoryTrainingSampler with repeat_factors={}".format(
                cfg.DATASETS.TRAIN_REPEAT_FACTOR
            )
        )
    else:
        logger.info(
            "Using WeightedTrainingSampler with repeat_factors={}".format(
                cfg.DATASETS.TRAIN_REPEAT_FACTOR
            )
        )

    sampler = RepeatFactorTrainingSampler(repeat_factors)
    return sampler


def _train_loader_from_config(cfg, mapper=None, *, dataset=None, sampler=None):
    if dataset is None:
        dataset = get_detection_dataset_dicts(
            cfg.DATASETS.TRAIN,
            filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS,
            min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
            if cfg.MODEL.KEYPOINT_ON
            else 0,
            proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
        )
        _log_api_usage("dataset." + cfg.DATASETS.TRAIN[0])

    if mapper is None:
        mapper = DatasetMapper(cfg, True)

    if sampler is None:
        sampler_name = cfg.DATALOADER.SAMPLER_TRAIN
        logger = logging.getLogger(__name__)
        if isinstance(dataset, torchdata.IterableDataset):
            logger.info("Not using any sampler since the dataset is IterableDataset.")
            sampler = None
        else:
            logger.info("Using training sampler {}".format(sampler_name))
            if sampler_name == "TrainingSampler":
                sampler = TrainingSampler(len(dataset))
            elif sampler_name == "RepeatFactorTrainingSampler":
                repeat_factors = RepeatFactorTrainingSampler.repeat_factors_from_category_frequency(
                    dataset, cfg.DATALOADER.REPEAT_THRESHOLD
                )
                sampler = RepeatFactorTrainingSampler(repeat_factors)
            elif sampler_name == "RandomSubsetTrainingSampler":
                sampler = RandomSubsetTrainingSampler(
                    len(dataset), cfg.DATALOADER.RANDOM_SUBSET_RATIO
                )
            elif sampler_name == "WeightedTrainingSampler":
                sampler = _build_weighted_sampler(cfg)
            elif sampler_name == "WeightedCategoryTrainingSampler":
                sampler = _build_weighted_sampler(cfg, enable_category_balance=True)
            else:
                raise ValueError("Unknown training sampler: {}".format(sampler_name))

    return {
        "dataset": dataset,
        "sampler": sampler,
        "mapper": mapper,
        "total_batch_size": cfg.SOLVER.IMS_PER_BATCH,
        "aspect_ratio_grouping": cfg.DATALOADER.ASPECT_RATIO_GROUPING,
        "num_workers": cfg.DATALOADER.NUM_WORKERS,
    }


@configurable(from_config=_train_loader_from_config)
def build_detection_train_loader(
    dataset,
    *,
    mapper,
    sampler=None,
    total_batch_size,
    aspect_ratio_grouping=True,
    num_workers=0,
    collate_fn=None,
    prefetch_factor=None,
    persistent_workers=False,
    pin_memory=False,
):
    """
    Build a dataloader for object detection with some default features.

    Args:
        dataset (list or torch.utils.data.Dataset): a list of dataset dicts,
            or a pytorch dataset (either map-style or iterable). It can be obtained
            by using :func:`DatasetCatalog.get` or :func:`get_detection_dataset_dicts`.
        mapper (callable): a callable which takes a sample (dict) from dataset and
            returns the format to be consumed by the model.
            When using cfg, the default choice is ``DatasetMapper(cfg, is_train=True)``.
        sampler (torch.utils.data.sampler.Sampler or None): a sampler that produces
            indices to be applied on ``dataset``.
            If ``dataset`` is map-style, the default sampler is a :class:`TrainingSampler`,
            which coordinates an infinite random shuffle sequence across all workers.
            Sampler must be None if ``dataset`` is iterable.
        total_batch_size (int): total batch size across all workers.
        aspect_ratio_grouping (bool): whether to group images with similar
            aspect ratio for efficiency. When enabled, it requires each
            element in dataset be a dict with keys "width" and "height".
        num_workers (int): number of parallel data loading workers
        collate_fn: a function that determines how to do batching, same as the argument of
            `torch.utils.data.DataLoader`. Defaults to do no collation and return a list of
            data. No collation is OK for small batch size and simple data structures.
            If your batch size is large and each sample contains too many small tensors,
            it's more efficient to collate them in data loader.

    Returns:
        torch.utils.data.DataLoader:
            a dataloader. Each output from it is a ``list[mapped_element]`` of length
            ``total_batch_size / num_workers``, where ``mapped_element`` is produced
            by the ``mapper``.
    """
    if isinstance(dataset, list):
        dataset = DatasetFromList(dataset, copy=False)
    if mapper is not None:
        dataset = MapDataset(dataset, mapper)

    if isinstance(dataset, torchdata.IterableDataset):
        assert sampler is None, "sampler must be None if dataset is IterableDataset"
    else:
        if sampler is None:
            sampler = TrainingSampler(len(dataset))
        assert isinstance(sampler, torchdata.Sampler), f"Expect a Sampler but got {type(sampler)}"
    return build_batch_data_loader(
        dataset,
        sampler,
        total_batch_size,
        aspect_ratio_grouping=aspect_ratio_grouping,
        num_workers=num_workers,
        collate_fn=collate_fn,
        prefetch_factor=prefetch_factor,
        persistent_workers=persistent_workers,
        pin_memory=pin_memory,
    )


def _test_loader_from_config(cfg, dataset_name, mapper=None):
    """
    Uses the given `dataset_name` argument (instead of the names in cfg), because the
    standard practice is to evaluate each test set individually (not combining them).
    """
    if isinstance(dataset_name, str):
        dataset_name = [dataset_name]

    dataset = get_detection_dataset_dicts(
        dataset_name,
        filter_empty=False,
        proposal_files=[
            cfg.DATASETS.PROPOSAL_FILES_TEST[list(cfg.DATASETS.TEST).index(x)] for x in dataset_name
        ]
        if cfg.MODEL.LOAD_PROPOSALS
        else None,
    )
    if mapper is None:
        mapper = DatasetMapper(cfg, False)
    return {
        "dataset": dataset,
        "mapper": mapper,
        "num_workers": cfg.DATALOADER.NUM_WORKERS,
        "sampler": InferenceSampler(len(dataset))
        if not isinstance(dataset, torchdata.IterableDataset)
        else None,
    }


@configurable(from_config=_test_loader_from_config)
def build_detection_test_loader(
    dataset: Union[List[Any], torchdata.Dataset],
    *,
    mapper: Callable[[Dict[str, Any]], Any],
    sampler: Optional[torchdata.Sampler] = None,
    batch_size: int = 1,
    num_workers: int = 0,
    collate_fn: Optional[Callable[[List[Any]], Any]] = None,
) -> torchdata.DataLoader:
    """
    Similar to `build_detection_train_loader`, with default batch size = 1,
    and sampler = :class:`InferenceSampler`. This sampler coordinates all workers
    to produce the exact set of all samples.

    Args:
        dataset: a list of dataset dicts,
            or a pytorch dataset (either map-style or iterable). They can be obtained
            by using :func:`DatasetCatalog.get` or :func:`get_detection_dataset_dicts`.
        mapper: a callable which takes a sample (dict) from dataset
           and returns the format to be consumed by the model.
           When using cfg, the default choice is ``DatasetMapper(cfg, is_train=False)``.
        sampler: a sampler that produces
            indices to be applied on ``dataset``. Default to :class:`InferenceSampler`,
            which splits the dataset across all workers. Sampler must be None
            if `dataset` is iterable.
        batch_size: the batch size of the data loader to be created.
            Default to 1 image per worker since this is the standard when reporting
            inference time in papers.
        num_workers: number of parallel data loading workers
        collate_fn: same as the argument of `torch.utils.data.DataLoader`.
            Defaults to do no collation and return a list of data.

    Returns:
        DataLoader: a torch DataLoader, that loads the given detection
        dataset, with test-time transformation and batching.

    Examples:
    ::
        data_loader = build_detection_test_loader(
            DatasetRegistry.get("my_test"),
            mapper=DatasetMapper(...))

        # or, instantiate with a CfgNode:
        data_loader = build_detection_test_loader(cfg, "my_test")
    """
    if isinstance(dataset, list):
        dataset = DatasetFromList(dataset, copy=False)
    if mapper is not None:
        dataset = MapDataset(dataset, mapper)
    if isinstance(dataset, torchdata.IterableDataset):
        assert sampler is None, "sampler must be None if dataset is IterableDataset"
    else:
        if sampler is None:
            sampler = InferenceSampler(len(dataset))
    return torchdata.DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        drop_last=False,
        num_workers=num_workers,
        collate_fn=trivial_batch_collator if collate_fn is None else collate_fn,
    )


def trivial_batch_collator(batch):
    """
    A batch collator that does nothing.
    """
    return batch


def worker_init_reset_seed(worker_id):
    initial_seed = torch.initial_seed() % 2**31
    seed_all_rng(initial_seed + worker_id)

imports/detectron2_repo/detectron2/data/build.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import os
import os.path as osp
import sys
import numpy as np
import cv2

import torch
import torchvision.transforms as transforms
# from PIL import Image

# Code from https://github.com/Daniil-Osokin/lightweight-human-pose-estimation.pytorch/blob/master/demo.py

# 2D body pose estimator
from detectors.body_pose_estimator.pose2d_models.with_mobilenet import PoseEstimationWithMobileNet
from detectors.body_pose_estimator.modules.load_state import load_state
from detectors.body_pose_estimator.val import normalize, pad_width
from detectors.body_pose_estimator.modules.pose import Pose, track_poses
from detectors.body_pose_estimator.modules.keypoints import extract_keypoints, group_keypoints

from constants.frankmocap import FRANKMOCAP_DIR

class BodyPoseEstimator(object):
    """
    Hand Detector for third-view input.
    It combines a body pose estimator (https://github.com/jhugestar/lightweight-human-pose-estimation.pytorch.git)
    """
    def __init__(self):
        print("Loading Body Pose Estimator")
        self.__load_body_estimator()
    

    def __load_body_estimator(self, pose2d_checkpoint=f"{FRANKMOCAP_DIR}/extra_data/body_module/body_pose_estimator/checkpoint_iter_370000.pth"):
        net = PoseEstimationWithMobileNet()
        checkpoint = torch.load(pose2d_checkpoint, map_location='cpu')
        load_state(net, checkpoint)
        net = net.eval()
        net = net.cuda()
        self.model = net
    

    #Code from https://github.com/Daniil-Osokin/lightweight-human-pose-estimation.pytorch/demo.py
    def __infer_fast(self, img, input_height_size, stride, upsample_ratio, 
        cpu=False, pad_value=(0, 0, 0), img_mean=(128, 128, 128), img_scale=1/256):
        height, width, _ = img.shape
        scale = input_height_size / height

        scaled_img = cv2.resize(img, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
        scaled_img = normalize(scaled_img, img_mean, img_scale)
        min_dims = [input_height_size, max(scaled_img.shape[1], input_height_size)]
        padded_img, pad = pad_width(scaled_img, stride, pad_value, min_dims)

        tensor_img = torch.from_numpy(padded_img).permute(2, 0, 1).unsqueeze(0).float()
        if not cpu:
            tensor_img = tensor_img.cuda()

        stages_output = self.model(tensor_img)

        stage2_heatmaps = stages_output[-2]
        heatmaps = np.transpose(stage2_heatmaps.squeeze().cpu().data.numpy(), (1, 2, 0))
        heatmaps = cv2.resize(heatmaps, (0, 0), fx=upsample_ratio, fy=upsample_ratio, interpolation=cv2.INTER_CUBIC)

        stage2_pafs = stages_output[-1]
        pafs = np.transpose(stage2_pafs.squeeze().cpu().data.numpy(), (1, 2, 0))
        pafs = cv2.resize(pafs, (0, 0), fx=upsample_ratio, fy=upsample_ratio, interpolation=cv2.INTER_CUBIC)

        return heatmaps, pafs, scale, pad
    
    def detect_body_pose(self, img):
        """
        Output:
            current_bbox: BBOX_XYWH
        """
        stride = 8
        upsample_ratio = 4
        orig_img = img.copy()

        # forward
        heatmaps, pafs, scale, pad = self.__infer_fast(img, 
            input_height_size=256, stride=stride, upsample_ratio=upsample_ratio)

        total_keypoints_num = 0
        all_keypoints_by_type = []
        num_keypoints = Pose.num_kpts
        for kpt_idx in range(num_keypoints):  # 19th for bg
            total_keypoints_num += extract_keypoints(heatmaps[:, :, kpt_idx], all_keypoints_by_type, total_keypoints_num)
        
        pose_entries, all_keypoints = group_keypoints(all_keypoints_by_type, pafs, demo=True)
        for kpt_id in range(all_keypoints.shape[0]):
            all_keypoints[kpt_id, 0] = (all_keypoints[kpt_id, 0] * stride / upsample_ratio - pad[1]) / scale
            all_keypoints[kpt_id, 1] = (all_keypoints[kpt_id, 1] * stride / upsample_ratio - pad[0]) / scale
        
        '''
        # print(len(pose_entries))
        if len(pose_entries)>1:
            pose_entries = pose_entries[:1]
            print("We only support one person currently")
            # assert len(pose_entries) == 1, "We only support one person currently"
        '''

        current_poses, current_bbox = list(), list()
        for n in range(len(pose_entries)):
            if len(pose_entries[n]) == 0:
                continue
            pose_keypoints = np.ones((num_keypoints, 2), dtype=np.int32) * -1
            for kpt_id in range(num_keypoints):
                if pose_entries[n][kpt_id] != -1.0:  # keypoint was found
                    pose_keypoints[kpt_id, 0] = int(all_keypoints[int(pose_entries[n][kpt_id]), 0])
                    pose_keypoints[kpt_id, 1] = int(all_keypoints[int(pose_entries[n][kpt_id]), 1])
            pose = Pose(pose_keypoints, pose_entries[n][18]) 
            current_poses.append(pose.keypoints)
            current_bbox.append(np.array(pose.bbox))

        # enlarge the bbox
        for i, bbox in enumerate(current_bbox):
            x, y, w, h = bbox
            margin = 0.05
            x_margin = int(w * margin)
            y_margin = int(h * margin)
            x0 = max(x-x_margin, 0)
            y0 = max(y-y_margin, 0)
            x1 = min(x+w+x_margin, orig_img.shape[1])
            y1 = min(y+h+y_margin, orig_img.shape[0])
            current_bbox[i] = np.array((x0, y0, x1-x0, y1-y0)).astype(np.int32)

        return current_poses, current_bbox
imports/frankmocap/bodymocap/body_bbox_detector.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import os
import cv2
import sys
import torch
import numpy as np
import pickle 
from torchvision.transforms import Normalize

from bodymocap.models import hmr, SMPL, SMPLX
from bodymocap import constants
from bodymocap.utils.imutils import crop, crop_bboxInfo, process_image_bbox, process_image_keypoints, bbox_from_keypoints
from mocap_utils.coordconv import convert_smpl_to_bbox, convert_bbox_to_oriIm
import mocap_utils.geometry_utils as gu

from constants.frankmocap import FRANKMOCAP_DIR

class BodyMocap(object):
    def __init__(self, regressor_checkpoint, smpl_dir, smplModelPath, device=torch.device('cuda'), use_smplx=False):

        self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

        # Load parametric model (SMPLX or SMPL)
        if use_smplx:
            smplModelPath = smpl_dir + '/SMPLX_NEUTRAL.pkl'
            self.smpl = SMPLX(smpl_dir,
                    batch_size=1,
                    num_betas = 10,
                    use_pca = False,
                    create_transl=False).to(self.device)
            self.use_smplx = True
        else:
            smplModelPath = smpl_dir + '/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl' if smpl_dir is not None else smplModelPath
            self.smpl = SMPL(smplModelPath, batch_size=1, create_transl=False).to(self.device)
            self.use_smplx = False
            
        #Load pre-trained neural network 
        SMPL_MEAN_PARAMS = os.path.join(
            FRANKMOCAP_DIR,
            './extra_data/body_module/data_from_spin/smpl_mean_params.npz')
        self.model_regressor = hmr(SMPL_MEAN_PARAMS).to(self.device)
        checkpoint = torch.load(regressor_checkpoint)
        self.model_regressor.load_state_dict(checkpoint['model'], strict=False)
        self.model_regressor.eval()
        

    def regress(self, img_original, body_bbox_list):
        """
            args: 
                img_original: original raw image (BGR order by using cv2.imread)
                body_bbox: bounding box around the target: (minX, minY, width, height)
            outputs:
                pred_vertices_img:
                pred_joints_vis_img:
                pred_rotmat
                pred_betas
                pred_camera
                bbox: [bbr[0], bbr[1],bbr[0]+bbr[2], bbr[1]+bbr[3]])
                bboxTopLeft:  bbox top left (redundant)
                boxScale_o2n: bbox scaling factor (redundant) 
        """
        pred_output_list = list()

        for body_bbox in body_bbox_list:
            img, norm_img, boxScale_o2n, bboxTopLeft, bbox = process_image_bbox(
                img_original, body_bbox, input_res=constants.IMG_RES)
            bboxTopLeft = np.array(bboxTopLeft)

            # bboxTopLeft = bbox['bboxXYWH'][:2]
            if img is None:
                pred_output_list.append(None)
                continue

            with torch.no_grad():
                # model forward
                pred_rotmat, pred_betas, pred_camera = self.model_regressor(norm_img.to(self.device))
                # -> pred_rotmat: (B,24,3,3) where 23 are joint rotation matrices and 1 is global rotation matrix
                # -> pred_betas: (B,10) where 10 are number of betas
                # -> pred_camera: (B,3) where each are tx,ty,scale


                #Convert rot_mat to aa since hands are always in aa
                # pred_aa = rotmat3x3_to_angle_axis(pred_rotmat)
                pred_aa = gu.rotation_matrix_to_angle_axis(pred_rotmat).cuda()
                # --> 
                pred_aa = pred_aa.reshape(pred_aa.shape[0], 72)
                smpl_output = self.smpl(
                    betas=pred_betas, 
                    body_pose=pred_aa[:,3:],
                    global_orient=pred_aa[:,:3], 
                    pose2rot=True)
                pred_vertices = smpl_output.vertices
                pred_joints_3d = smpl_output.joints

                pred_vertices = pred_vertices[0].cpu().numpy()

                pred_camera = pred_camera.cpu().numpy().ravel()
                camScale = pred_camera[0] # *1.15
                camTrans = pred_camera[1:]

                pred_output = dict()
                # Convert mesh to original image space (X,Y are aligned to image)
                # 1. SMPL -> 2D bbox
                # 2. 2D bbox -> original 2D image
                pred_vertices_bbox = convert_smpl_to_bbox(pred_vertices, camScale, camTrans)
                pred_vertices_img = convert_bbox_to_oriIm(
                    pred_vertices_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0])

                # Convert joint to original image space (X,Y are aligned to image)
                pred_joints_3d = pred_joints_3d[0].cpu().numpy() # (1,49,3) -> (49,3)
                pred_joints_vis = pred_joints_3d[:,:3]  # (49,3)
                pred_joints_vis_bbox = convert_smpl_to_bbox(pred_joints_vis, camScale, camTrans) 
                pred_joints_vis_img = convert_bbox_to_oriIm(
                    pred_joints_vis_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0]) 

                # Output
                pred_output['img_cropped'] = img[:, :, ::-1]
                pred_output['pred_vertices_smpl'] = smpl_output.vertices[0].cpu().numpy() # SMPL vertex in original smpl space
                pred_output['pred_vertices_img'] = pred_vertices_img # SMPL vertex in image space
                pred_output['pred_joints_img'] = pred_joints_vis_img # SMPL joints in image space
                
                pred_output['pred_joints_smpl'] = np.expand_dims(pred_joints_3d, axis=0) # shape: (1,49,3)

                pred_aa_tensor = gu.rotation_matrix_to_angle_axis(pred_rotmat.detach().cpu()[0])
                pred_output['pred_body_pose'] = pred_aa_tensor.cpu().numpy().reshape(1, 72) # (1, 72)

                pred_output['pred_rotmat'] = pred_rotmat.detach().cpu().numpy() # (1, 24, 3, 3)
                pred_output['pred_betas'] = pred_betas.detach().cpu().numpy() # (1, 10)

                pred_output['pred_camera'] = pred_camera

                # top-left coordinate in pixels, when origin is center of image (not top-left of image)
                # e.g. [-40 57]
                # given as np.ndarray
                pred_output['bbox_top_left'] = bboxTopLeft 
                
                pred_output['bbox_scale_ratio'] = boxScale_o2n
                
                # original bbox (of body)
                pred_output['body_bbox'] = body_bbox
                
                # SMPL faces
                pred_output['faces'] = self.smpl.faces # Faces. 

                if self.use_smplx:
                    assert False
                    img_center = np.array((img_original.shape[1], img_original.shape[0]) ) * 0.5
                    # right hand
                    pred_joints = smpl_output.right_hand_joints[0].cpu().numpy()     
                    pred_joints_bbox = convert_smpl_to_bbox(pred_joints, camScale, camTrans)
                    pred_joints_img = convert_bbox_to_oriIm(
                        pred_joints_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0])
                    pred_output['right_hand_joints_img_coord'] = pred_joints_img
                    # left hand 
                    pred_joints = smpl_output.left_hand_joints[0].cpu().numpy()
                    pred_joints_bbox = convert_smpl_to_bbox(pred_joints, camScale, camTrans)
                    pred_joints_img = convert_bbox_to_oriIm(
                        pred_joints_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0])
                    pred_output['left_hand_joints_img_coord'] = pred_joints_img

                # Used for predicting perspective camera when (ortho -> persp) transformation!
                # Check "test_optimize_camera_origin_fixed.ipynb" and "test_optimize_camera_origin"          
                smpl_output_zerorot = self.smpl(
                    betas=pred_betas, 
                    body_pose=pred_aa[:,3:],
                    global_orient=torch.zeros_like(pred_aa[:,:3]), 
                    pose2rot=True)
                pred_vertices_zerorot = smpl_output_zerorot.vertices
                pred_joints_3d_zerorot = smpl_output_zerorot.joints

                pred_vertices_zerorot = pred_vertices_zerorot[0].cpu().numpy()

                pred_vertices_bbox_zerorot = convert_smpl_to_bbox(pred_vertices_zerorot, camScale, camTrans)
                pred_vertices_img_zerorot = convert_bbox_to_oriIm(
                    pred_vertices_bbox_zerorot, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0])

                # Convert joint to original image space (X,Y are aligned to image)
                pred_joints_3d_zerorot = pred_joints_3d_zerorot[0].cpu().numpy() # (1,49,3) -> (49,3)
                pred_joints_vis_zerorot = pred_joints_3d_zerorot[:,:3]  # (49,3)
                pred_joints_vis_bbox_zerorot = convert_smpl_to_bbox(pred_joints_vis_zerorot, camScale, camTrans) 
                pred_joints_vis_img_zerorot = convert_bbox_to_oriIm(
                    pred_joints_vis_bbox_zerorot, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0]) 
                
                # Output
                pred_output['pred_vertices_smpl_zerorot'] = smpl_output_zerorot.vertices[0].cpu().numpy() # SMPL vertex in original smpl space
                pred_output['pred_vertices_img_zerorot'] = pred_vertices_img_zerorot # SMPL vertex in image space
                pred_output['pred_joints_img_zerorot'] = pred_joints_vis_img_zerorot # SMPL joints in image space
                pred_output['pred_joints_smpl_zerorot'] = np.expand_dims(pred_joints_3d_zerorot, axis=0) # shape: (1,49,3)
                
                smpl_output_zerorot_zerobeta = self.smpl(
                    betas=torch.zeros_like(pred_betas), 
                    body_pose=pred_aa[:,3:],
                    global_orient=torch.zeros_like(pred_aa[:,:3]), 
                    pose2rot=True)
                pred_vertices_zerorot_zerobeta = smpl_output_zerorot_zerobeta.vertices
                pred_joints_3d_zerorot_zerobeta = smpl_output_zerorot_zerobeta.joints

                pred_vertices_zerorot_zerobeta = pred_vertices_zerorot_zerobeta[0].cpu().numpy()

                pred_vertices_bbox_zerorot_zerobeta = convert_smpl_to_bbox(pred_vertices_zerorot_zerobeta, camScale, camTrans)
                pred_vertices_img_zerorot_zerobeta = convert_bbox_to_oriIm(
                    pred_vertices_bbox_zerorot_zerobeta, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0])

                # Convert joint to original image space (X,Y are aligned to image)
                pred_joints_3d_zerorot_zerobeta = pred_joints_3d_zerorot_zerobeta[0].cpu().numpy() # (1,49,3) -> (49,3)
                pred_joints_vis_zerorot_zerobeta = pred_joints_3d_zerorot_zerobeta[:,:3]  # (49,3)
                pred_joints_vis_bbox_zerorot_zerobeta = convert_smpl_to_bbox(pred_joints_vis_zerorot_zerobeta, camScale, camTrans) 
                pred_joints_vis_img_zerorot_zerobeta = convert_bbox_to_oriIm(
                    pred_joints_vis_bbox_zerorot_zerobeta, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0]) 
                
                # Output
                pred_output['pred_vertices_smpl_zerorot_zerobeta'] = smpl_output_zerorot_zerobeta.vertices[0].cpu().numpy() # SMPL vertex in original smpl space
                pred_output['pred_vertices_img_zerorot_zerobeta'] = pred_vertices_img_zerorot_zerobeta # SMPL vertex in image space
                pred_output['pred_joints_img_zerorot_zerobeta'] = pred_joints_vis_img_zerorot_zerobeta # SMPL joints in image space
                pred_output['pred_joints_smpl_zerorot_zerobeta'] = np.expand_dims(pred_joints_3d_zerorot_zerobeta, axis=0) # shape: (1,49,3)
                    
                pred_output_list.append(pred_output)

        return pred_output_list
    

    def get_hand_bboxes(self, pred_body_list, img_shape):
        """
            args: 
                pred_body_list: output of body regresion
                img_shape: img_height, img_width
            outputs:
                hand_bbox_list: 
        """
        hand_bbox_list = list()
        for pred_body in pred_body_list:
            hand_bbox = dict(
                left_hand = None,
                right_hand = None
            )
            if pred_body is None:
                hand_bbox_list.append(hand_bbox)
            else:
                for hand_type in hand_bbox:
                    key = f'{hand_type}_joints_img_coord'
                    pred_joints_vis_img = pred_body[key]

                    if pred_joints_vis_img is not None:
                        # get initial bbox
                        x0, x1 = np.min(pred_joints_vis_img[:, 0]), np.max(pred_joints_vis_img[:, 0])
                        y0, y1 = np.min(pred_joints_vis_img[:, 1]), np.max(pred_joints_vis_img[:, 1])
                        width, height = x1-x0, y1-y0
                        # extend the obtained bbox
                        margin = int(max(height, width) * 0.2)
                        img_height, img_width = img_shape
                        x0 = max(x0 - margin, 0)
                        y0 = max(y0 - margin, 0)
                        x1 = min(x1 + margin, img_width)
                        y1 = min(y1 + margin, img_height)
                        # result bbox in (x0, y0, w, h) format
                        hand_bbox[hand_type] = np.array([x0, y0, x1-x0, y1-y0]) # in (x, y, w, h ) format

                hand_bbox_list.append(hand_bbox)

        return hand_bbox_list
imports/frankmocap/bodymocap/body_mocap_api.py
jellyheadandrew/CHORUS
# Original code from SPIN: https://github.com/nkolot/SPIN

FOCAL_LENGTH = 5000.
IMG_RES = 224

# Mean and standard deviation for normalizing input image
IMG_NORM_MEAN = [0.485, 0.456, 0.406]
IMG_NORM_STD = [0.229, 0.224, 0.225]

"""
We create a superset of joints containing the OpenPose joints together with the ones that each dataset provides.
We keep a superset of 24 joints such that we include all joints from every dataset.
If a dataset doesn't provide annotations for a specific joint, we simply ignore it.
The joints used here are the following:
"""
JOINT_NAMES = [
'OP Nose', 'OP Neck', 'OP RShoulder',           #0,1,2
'OP RElbow', 'OP RWrist', 'OP LShoulder',       #3,4,5
'OP LElbow', 'OP LWrist', 'OP MidHip',          #6, 7,8
'OP RHip', 'OP RKnee', 'OP RAnkle',             #9,10,11
'OP LHip', 'OP LKnee', 'OP LAnkle',             #12,13,14
'OP REye', 'OP LEye', 'OP REar',                #15,16,17
'OP LEar', 'OP LBigToe', 'OP LSmallToe',        #18,19,20
'OP LHeel', 'OP RBigToe', 'OP RSmallToe', 'OP RHeel',  #21, 22, 23, 24  ##Total 25 joints  for openpose
'Right Ankle', 'Right Knee', 'Right Hip',               #0,1,2
'Left Hip', 'Left Knee', 'Left Ankle',                  #3, 4, 5
'Right Wrist', 'Right Elbow', 'Right Shoulder',     #6
'Left Shoulder', 'Left Elbow', 'Left Wrist',            #9
'Neck (LSP)', 'Top of Head (LSP)',                      #12, 13
'Pelvis (MPII)', 'Thorax (MPII)',                       #14, 15
'Spine (H36M)', 'Jaw (H36M)',                           #16, 17
'Head (H36M)', 'Nose', 'Left Eye',                      #18, 19, 20
'Right Eye', 'Left Ear', 'Right Ear'                    #21,22,23 (Total 24 joints)
]

# Dict containing the joints in numerical order
JOINT_IDS = {JOINT_NAMES[i]: i for i in range(len(JOINT_NAMES))}

# Map joints to SMPL joints
JOINT_MAP = {
'OP Nose': 24, 'OP Neck': 12, 'OP RShoulder': 17,
'OP RElbow': 19, 'OP RWrist': 21, 'OP LShoulder': 16,
'OP LElbow': 18, 'OP LWrist': 20, 'OP MidHip': 0,
'OP RHip': 2, 'OP RKnee': 5, 'OP RAnkle': 8,
'OP LHip': 1, 'OP LKnee': 4, 'OP LAnkle': 7,
'OP REye': 25, 'OP LEye': 26, 'OP REar': 27,
'OP LEar': 28, 'OP LBigToe': 29, 'OP LSmallToe': 30,
'OP LHeel': 31, 'OP RBigToe': 32, 'OP RSmallToe': 33, 'OP RHeel': 34,
'Right Ankle': 8, 'Right Knee': 5, 'Right Hip': 45,
'Left Hip': 46, 'Left Knee': 4, 'Left Ankle': 7,
'Right Wrist': 21, 'Right Elbow': 19, 'Right Shoulder': 17,
'Left Shoulder': 16, 'Left Elbow': 18, 'Left Wrist': 20,
'Neck (LSP)': 47, 'Top of Head (LSP)': 48,
'Pelvis (MPII)': 49, 'Thorax (MPII)': 50,
'Spine (H36M)': 51, 'Jaw (H36M)': 52,
'Head (H36M)': 53, 'Nose': 24, 'Left Eye': 26,
'Right Eye': 25, 'Left Ear': 28, 'Right Ear': 27
}

# Joint selectors
# Indices to get the 14 LSP joints from the 17 H36M joints
H36M_TO_J17 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9]
H36M_TO_J14 = H36M_TO_J17[:14]
# Indices to get the 14 LSP joints from the ground truth joints
J24_TO_J17 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 18, 14, 16, 17]
J24_TO_J14 = J24_TO_J17[:14]

# Permutation of SMPL pose parameters when flipping the shape
SMPL_JOINTS_FLIP_PERM = [0, 2, 1, 3, 5, 4, 6, 8, 7, 9, 11, 10, 12, 14, 13, 15, 17, 16, 19, 18, 21, 20, 23, 22]
SMPL_POSE_FLIP_PERM = []
for i in SMPL_JOINTS_FLIP_PERM:
    SMPL_POSE_FLIP_PERM.append(3*i)
    SMPL_POSE_FLIP_PERM.append(3*i+1)
    SMPL_POSE_FLIP_PERM.append(3*i+2)
# Permutation indices for the 24 ground truth joints
J24_FLIP_PERM = [5, 4, 3, 2, 1, 0, 11, 10, 9, 8, 7, 6, 12, 13, 14, 15, 16, 17, 18, 19, 21, 20, 23, 22]
# Permutation indices for the full set of 49 joints
J49_FLIP_PERM = [0, 1, 5, 6, 7, 2, 3, 4, 8, 12, 13, 14, 9, 10, 11, 16, 15, 18, 17, 22, 23, 24, 19, 20, 21]\
              + [25+i for i in J24_FLIP_PERM]
imports/frankmocap/bodymocap/constants.py
jellyheadandrew/CHORUS
from .hmr import hmr
from .smpl import SMPL, SMPLX

imports/frankmocap/bodymocap/models/__init__.py
jellyheadandrew/CHORUS
# Original code from SPIN: https://github.com/nkolot/SPIN

import torch
import torch.nn as nn
import torchvision.models.resnet as resnet
import numpy as np
import math

import sys
from bodymocap.utils.geometry import rot6d_to_rotmat

class Bottleneck(nn.Module):
    """ Redefinition of Bottleneck residual block
        Adapted from the official PyTorch implementation
    """
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

class HMR(nn.Module):
    """ SMPL Iterative Regressor with ResNet50 backbone
    """

    def __init__(self, block, layers, smpl_mean_params):
        self.inplanes = 64
        super(HMR, self).__init__()
        npose = 24 * 6
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AvgPool2d(7, stride=1)
        self.fc1 = nn.Linear(512 * block.expansion + npose + 13, 1024)
        self.drop1 = nn.Dropout()
        self.fc2 = nn.Linear(1024, 1024)
        self.drop2 = nn.Dropout()
        self.decpose = nn.Linear(1024, npose)
        self.decshape = nn.Linear(1024, 10)
        self.deccam = nn.Linear(1024, 3)
        nn.init.xavier_uniform_(self.decpose.weight, gain=0.01)
        nn.init.xavier_uniform_(self.decshape.weight, gain=0.01)
        nn.init.xavier_uniform_(self.deccam.weight, gain=0.01)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

        mean_params = np.load(smpl_mean_params)
        init_pose = torch.from_numpy(mean_params['pose'][:]).unsqueeze(0)
        init_shape = torch.from_numpy(mean_params['shape'][:].astype('float32')).unsqueeze(0)
        init_cam = torch.from_numpy(mean_params['cam']).unsqueeze(0)
        self.register_buffer('init_pose', init_pose)
        self.register_buffer('init_shape', init_shape)
        self.register_buffer('init_cam', init_cam)


    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)


    def forward(self, x, init_pose=None, init_shape=None, init_cam=None, n_iter=3):

        batch_size = x.shape[0]

        if init_pose is None:
            init_pose = self.init_pose.expand(batch_size, -1)
        if init_shape is None:
            init_shape = self.init_shape.expand(batch_size, -1)
        if init_cam is None:
            init_cam = self.init_cam.expand(batch_size, -1)

        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x1 = self.layer1(x)
        x2 = self.layer2(x1)
        x3 = self.layer3(x2)
        x4 = self.layer4(x3)

        xf = self.avgpool(x4)
        xf = xf.view(xf.size(0), -1)

        pred_pose = init_pose
        pred_shape = init_shape
        pred_cam = init_cam
        for i in range(n_iter):
            xc = torch.cat([xf, pred_pose, pred_shape, pred_cam],1)
            xc = self.fc1(xc)
            xc = self.drop1(xc)
            xc = self.fc2(xc)
            xc = self.drop2(xc)
            pred_pose = self.decpose(xc) + pred_pose
            pred_shape = self.decshape(xc) + pred_shape
            pred_cam = self.deccam(xc) + pred_cam
        
        pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3)

        return pred_rotmat, pred_shape, pred_cam

def hmr(smpl_mean_params, pretrained=True, **kwargs):
    """ Constructs an HMR model with ResNet50 backbone.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = HMR(Bottleneck, [3, 4, 6, 3],  smpl_mean_params, **kwargs)
    if pretrained:
        resnet_imagenet = resnet.resnet50(pretrained=True)
        model.load_state_dict(resnet_imagenet.state_dict(),strict=False)
    return model


imports/frankmocap/bodymocap/models/hmr.py
jellyheadandrew/CHORUS
# Original code from SPIN: https://github.com/nkolot/SPIN

import os
import torch
import numpy as np
import smplx
from smplx import SMPL as _SMPL
from smplx import SMPLX as _SMPLX
# from bodymocap.models.body_models import SMPLX as _SMPLX        #Use our custom SMPLX 
# from smplx.body_models import ModelOutput
# from bodymocap.models.body_models import ModelOutput
from smplx.lbs import vertices2joints

from bodymocap import constants

from constants.frankmocap import FRANKMOCAP_DIR

from collections import namedtuple
ModelOutput = namedtuple('ModelOutput',
                         ['vertices', 'joints', 'full_pose', 'betas',
                          'global_orient',
                          'body_pose', 'expression',
                          'left_hand_pose', 'right_hand_pose',
                          'right_hand_joints', 'left_hand_joints',
                          'jaw_pose'])
ModelOutput.__new__.__defaults__ = (None,) * len(ModelOutput._fields)


class SMPL(_SMPL):
    """ Extension of the official SMPL implementation to support more joints """

    def __init__(self, *args, **kwargs):
        super(SMPL, self).__init__(*args, **kwargs)
        joints = [constants.JOINT_MAP[i] for i in constants.JOINT_NAMES]
        JOINT_REGRESSOR_TRAIN_EXTRA = os.path.join(
            FRANKMOCAP_DIR,
            'extra_data/body_module/data_from_spin//J_regressor_extra.npy')
        J_regressor_extra = np.load(JOINT_REGRESSOR_TRAIN_EXTRA)
        self.register_buffer('J_regressor_extra', torch.tensor(J_regressor_extra, dtype=torch.float32))
        self.joint_map = torch.tensor(joints, dtype=torch.long)
        self.smpl_joints = None
        self.joints = None

    def forward(self, *args, **kwargs):
        kwargs['get_skin'] = True
        smpl_output = super(SMPL, self).forward(*args, **kwargs)
        extra_joints = vertices2joints(self.J_regressor_extra, smpl_output.vertices)        #Additional 9 joints #Check doc/J_regressor_extra.png
        self.smpl_joints = smpl_output.joints
        self.joints = torch.cat([self.smpl_joints, extra_joints], dim=1)               #[N, 24 + 21, 3]  + [N, 9, 3]
        joints = self.joints[:, self.joint_map, :]
        output = ModelOutput(vertices=smpl_output.vertices,
                             global_orient=smpl_output.global_orient,
                             body_pose=smpl_output.body_pose,
                             joints=joints,
                             betas=smpl_output.betas,
                             full_pose=smpl_output.full_pose)
        return output



class SMPLX(_SMPLX):
    """ Extension of the official SMPL implementation to support more joints """

    def __init__(self, *args, **kwargs):
        kwargs['ext'] = 'pkl'       #We have pkl file
        super(SMPLX, self).__init__(*args, **kwargs)
        joints = [constants.JOINT_MAP[i] for i in constants.JOINT_NAMES]
        JOINT_REGRESSOR_TRAIN_EXTRA_SMPLX = 'extra_data/body_module/J_regressor_extra_smplx.npy'
        J_regressor_extra = np.load(JOINT_REGRESSOR_TRAIN_EXTRA_SMPLX)           #(9, 10475)
        self.register_buffer('J_regressor_extra', torch.tensor(J_regressor_extra, dtype=torch.float32))
        self.joint_map = torch.tensor(joints, dtype=torch.long)

    def forward(self, *args, **kwargs):
        kwargs['get_skin'] = True

        #if pose parameter is for SMPL with 21 joints (ignoring root)
        if(kwargs['body_pose'].shape[1]==69):
            kwargs['body_pose'] = kwargs['body_pose'][:,:-2*3]        #Ignore the last two joints (which are on the palm. Not used)

        if(kwargs['body_pose'].shape[1]==23):
            kwargs['body_pose'] = kwargs['body_pose'][:,:-2]        #Ignore the last two joints (which are on the palm. Not used)

        smpl_output = super(SMPLX, self).forward(*args, **kwargs)
        extra_joints = vertices2joints(self.J_regressor_extra, smpl_output.vertices)

        #SMPL-X Joint order: https://docs.google.com/spreadsheets/d/1_1dLdaX-sbMkCKr_JzJW_RZCpwBwd7rcKkWT_VgAQ_0/edit#gid=0
        smplx_to_smpl = list(range(0,22)) + [28,43] + list(range(55,76)) # 28 left middle finger , 43: right middle finger 1
        smpl_joints = smpl_output.joints[:,smplx_to_smpl,:] # Convert SMPL-X to SMPL     127 ->45
        joints = torch.cat([smpl_joints, extra_joints], dim=1) # [N, 127, 3]->[N, 45, 3]  + [N, 9, 3]  # SMPL-X has more joints. should convert 45
        joints = joints[:, self.joint_map, :]     

        # Hand joints
        smplx_hand_to_panoptic = [0, 13,14,15,16, 1,2,3,17, 4,5,6,18, 10,11,12,19, 7,8,9,20] #Wrist Thumb to Pinky

        smplx_lhand =  [20] + list(range(25,40)) + list(range(66, 71))         #20 for left wrist. 20 finger joints
        lhand_joints = smpl_output.joints[:,smplx_lhand, :]      #(N,21,3)
        lhand_joints = lhand_joints[:, smplx_hand_to_panoptic, :]     #Convert SMPL-X hand order to paonptic hand order

        smplx_rhand = [21] + list(range(40,55)) + list(range(71, 76))     #21 for right wrist. 20 finger joints
        rhand_joints = smpl_output.joints[:, smplx_rhand, :]      #(N,21,3)
        rhand_joints = rhand_joints[:,smplx_hand_to_panoptic,:] #Convert SMPL-X hand order to paonptic hand order

        output = ModelOutput(vertices=smpl_output.vertices,
                             global_orient=smpl_output.global_orient,
                             body_pose=smpl_output.body_pose,
                             joints=joints,
                             right_hand_joints=rhand_joints,     #N,21,3
                             left_hand_joints=lhand_joints,         #N,21,3
                             betas=smpl_output.betas,
                             full_pose=smpl_output.full_pose)
        return output


"""
0	pelvis',
1	left_hip',
2	right_hip',
3	spine1',
4	left_knee',
5	right_knee',
6	spine2',
7	left_ankle',
8	right_ankle',
9	spine3',
10	left_foot',
11	right_foot',
12	neck',
13	left_collar',
14	right_collar',
15	head',
16	left_shoulder',
17	right_shoulder',
18	left_elbow',
19	right_elbow',
20	left_wrist',
21	right_wrist',
22	jaw',
23	left_eye_smplhf',
24	right_eye_smplhf',
25	left_index1',
26	left_index2',
27	left_index3',
28	left_middle1',
29	left_middle2',
30	left_middle3',
31	left_pinky1',
32	left_pinky2',
33	left_pinky3',
34	left_ring1',
35	left_ring2',
36	left_ring3',
37	left_thumb1',
38	left_thumb2',
39	left_thumb3',
40	right_index1',
41	right_index2',
42	right_index3',
43	right_middle1',
44	right_middle2',
45	right_middle3',
46	right_pinky1',
47	right_pinky2',
48	right_pinky3',
49	right_ring1',
50	right_ring2',
51	right_ring3',
52	right_thumb1',
53	right_thumb2',
54	right_thumb3',
55	nose',
56	right_eye',
57	left_eye',
58	right_ear',
59	left_ear',
60	left_big_toe',
61	left_small_toe',
62	left_heel',
63	right_big_toe',
64	right_small_toe',
65	right_heel',
66	left_thumb',
67	left_index',
68	left_middle',
69	left_ring',
70	left_pinky',
71	right_thumb',
72	right_index',
73	right_middle',
74	right_ring',
75	right_pinky',
76	right_eye_brow1',
77	right_eye_brow2',
78	right_eye_brow3',
79	right_eye_brow4',
80	right_eye_brow5',
81	left_eye_brow5',
82	left_eye_brow4',
83	left_eye_brow3',
84	left_eye_brow2',
85	left_eye_brow1',
86	nose1',
87	nose2',
88	nose3',
89	nose4',
90	right_nose_2',
91	right_nose_1',
92	nose_middle',
93	left_nose_1',
94	left_nose_2',
95	right_eye1',
96	right_eye2',
97	right_eye3',
98	right_eye4',
99	right_eye5',
100	right_eye6',
101	left_eye4',
102	left_eye3',
103	left_eye2',
104	left_eye1',
105	left_eye6',
106	left_eye5',
107	right_mouth_1',
108	right_mouth_2',
109	right_mouth_3',
110	mouth_top',
111	left_mouth_3',
112	left_mouth_2',
113	left_mouth_1',
114	left_mouth_5', # 59 in OpenPose output
115	left_mouth_4', # 58 in OpenPose output
116	mouth_bottom',
117	right_mouth_4',
118	right_mouth_5',
119	right_lip_1',
120	right_lip_2',
121	lip_top',
122	left_lip_2',
123	left_lip_1',
124	left_lip_3',
125	lip_bottom',
126	right_lip_3',
127	right_contour_1',
128	right_contour_2',
129	right_contour_3',
130	right_contour_4',
131	right_contour_5',
132	right_contour_6',
133	right_contour_7',
134	right_contour_8',
135	contour_middle',
136	left_contour_8',
137	left_contour_7',
138	left_contour_6',
139	left_contour_5',
140	left_contour_4',
141	left_contour_3',
142	left_contour_2',
143	left_contour_1'
"""


#SMPL Joints:
"""
0	pelvis',
1	left_hip',
2	right_hip',
3	spine1',
4	left_knee',
5	right_knee',
6	spine2',
7	left_ankle',
8	right_ankle',
9	spine3',
10	left_foot',
11	right_foot',
12	neck',
13	left_collar',
14	right_collar',
15	head',
16	left_shoulder',
17	right_shoulder',
18	left_elbow',
19	right_elbow',
20	left_wrist',
21	right_wrist',
22	
23	
24	nose',
25	right_eye',
26	left_eye',
27	right_ear',
28	left_ear',
29	left_big_toe',
30	left_small_toe',
31	left_heel',
32	right_big_toe',
33	right_small_toe',
34	right_heel',
35	left_thumb',
36	left_index',
37	left_middle',
38	left_ring',
39	left_pinky',
40	right_thumb',
41	right_index',
42	right_middle',
43	right_ring',
44	right_pinky',
"""
imports/frankmocap/bodymocap/models/smpl.py
jellyheadandrew/CHORUS

# Original code from SPIN: https://github.com/nkolot/SPIN

import torch
from torch.nn import functional as F
import numpy as np

import torchgeometry

"""
Useful geometric operations, e.g. Perspective projection and a differentiable Rodrigues formula
Parts of the code are taken from https://github.com/MandyMo/pytorch_HMR
"""
def batch_rodrigues(theta):
    """Convert axis-angle representation to rotation matrix.
    Args:
        theta: size = [B, 3]
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    l1norm = torch.norm(theta + 1e-8, p = 2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)
    return quat_to_rotmat(quat)

def quat_to_rotmat(quat):
    """Convert quaternion coefficients to rotation matrix.
    Args:
        quat: size = [B, 4] 4 <===>(w, x, y, z)
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """ 
    norm_quat = quat
    norm_quat = norm_quat/norm_quat.norm(p=2, dim=1, keepdim=True)
    w, x, y, z = norm_quat[:,0], norm_quat[:,1], norm_quat[:,2], norm_quat[:,3]

    B = quat.size(0)

    w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
    wx, wy, wz = w*x, w*y, w*z
    xy, xz, yz = x*y, x*z, y*z

    rotMat = torch.stack([w2 + x2 - y2 - z2, 2*xy - 2*wz, 2*wy + 2*xz,
                          2*wz + 2*xy, w2 - x2 + y2 - z2, 2*yz - 2*wx,
                          2*xz - 2*wy, 2*wx + 2*yz, w2 - x2 - y2 + z2], dim=1).view(B, 3, 3)
    return rotMat    






def cross_product(u, v):
    batch = u.shape[0]
    i = u[:, 1] * v[:, 2] - u[:, 2] * v[:, 1]
    j = u[:, 2] * v[:, 0] - u[:, 0] * v[:, 2]
    k = u[:, 0] * v[:, 1] - u[:, 1] * v[:, 0]
        
    out = torch.cat((i.view(batch, 1), j.view(batch, 1), k.view(batch, 1)), 1)
        
    return out

def normalize_vector(v):
    batch = v.shape[0]
    v_mag = torch.sqrt(v.pow(2).sum(1))  # batch
    v_mag = torch.max(v_mag, v.new([1e-8]))
    v_mag = v_mag.view(batch, 1).expand(batch, v.shape[1])
    v = v/v_mag
    return v

#Code from
def rot6d_to_rotmat(x):
    """Convert 6D rotation representation to 3x3 rotation matrix.
    Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
    Input:
        (B,6) Batch of 6-D rotation representations
    Output:
        (B,3,3) Batch of corresponding rotation matrices
    """
    x = x.view(-1,3,2)
    a1 = x[:, :, 0]
    a2 = x[:, :, 1]
    b1 = F.normalize(a1)
    b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1)
    b3 = torch.cross(b1, b2)
    return torch.stack((b1, b2, b3), dim=-1)

     
def perspective_projection(points, rotation, translation,
                           focal_length, camera_center):
    """
    This function computes the perspective projection of a set of points.
    Input:
        points (bs, N, 3): 3D points
        rotation (bs, 3, 3): Camera rotation
        translation (bs, 3): Camera translation
        focal_length (bs,) or scalar: Focal length
        camera_center (bs, 2): Camera center
    """
    batch_size = points.shape[0]
    K = torch.zeros([batch_size, 3, 3], device=points.device)
    K[:,0,0] = focal_length
    K[:,1,1] = focal_length
    K[:,2,2] = 1.
    K[:,:-1, -1] = camera_center

    # Transform points
    points = torch.einsum('bij,bkj->bki', rotation, points)
    points = points + translation.unsqueeze(1)

    # Apply perspective distortion
    projected_points = points / points[:,:,-1].unsqueeze(-1)

    # Apply camera intrinsics
    projected_points = torch.einsum('bij,bkj->bki', K, projected_points)

    return projected_points[:, :, :-1]


def estimate_translation_np(S, joints_2d, joints_conf, focal_length=5000, img_size=224):
    """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
    Input:
        S: (25, 3) 3D joint locations
        joints: (25, 3) 2D joint locations and confidence
    Returns:
        (3,) camera translation vector
    """

    num_joints = S.shape[0]
    # focal length
    f = np.array([focal_length,focal_length])
    # optical center
    center = np.array([img_size/2., img_size/2.])

    # transformations
    Z = np.reshape(np.tile(S[:,2],(2,1)).T,-1)
    XY = np.reshape(S[:,0:2],-1)
    O = np.tile(center,num_joints)
    F = np.tile(f,num_joints)
    weight2 = np.reshape(np.tile(np.sqrt(joints_conf),(2,1)).T,-1)

    # least squares
    Q = np.array([F*np.tile(np.array([1,0]),num_joints), F*np.tile(np.array([0,1]),num_joints), O-np.reshape(joints_2d,-1)]).T
    c = (np.reshape(joints_2d,-1)-O)*Z - F*XY

    # weighted least squares
    W = np.diagflat(weight2)
    Q = np.dot(W,Q)
    c = np.dot(W,c)

    # square matrix
    A = np.dot(Q.T,Q)
    b = np.dot(Q.T,c)

    # solution
    trans = np.linalg.solve(A, b)

    return trans


def estimate_translation(S, joints_2d, focal_length=5000., img_size=224.):
    """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
    Input:
        S: (B, 49, 3) 3D joint locations
        joints: (B, 49, 3) 2D joint locations and confidence
    Returns:
        (B, 3) camera translation vectors
    """

    device = S.device
    # Use only joints 25:49 (GT joints)
    S = S[:, 25:, :].cpu().numpy()
    joints_2d = joints_2d[:, 25:, :].cpu().numpy()
    joints_conf = joints_2d[:, :, -1]
    joints_2d = joints_2d[:, :, :-1]
    trans = np.zeros((S.shape[0], 3), dtype=np.float32)
    # Find the translation for each example in the batch
    for i in range(S.shape[0]):
        S_i = S[i]
        joints_i = joints_2d[i]
        conf_i = joints_conf[i]
        trans[i] = estimate_translation_np(S_i, joints_i, conf_i, focal_length=focal_length, img_size=img_size)
    return torch.from_numpy(trans).to(device)




def weakProjection_gpu(skel3D, scale, trans2D ):
    # if len(skel3D.shape)==1:
    #     skel3D = np.reshape(skel3D, (-1,3))

    skel3D = skel3D.view((skel3D.shape[0],-1,3))
    trans2D = trans2D.view((trans2D.shape[0],1,2))
    scale = scale.view((scale.shape[0],1,1))
    skel3D_proj = scale* skel3D[:,:,:2] + trans2D

    return skel3D_proj#skel3D_proj.view((skel3D.shape[0],-1))       #(N, 19*2) o



#(57) (1) (2)
def weakProjection(skel3D, scale, trans2D ):

    skel3D_proj = scale* skel3D[:,:2] + trans2D

    return skel3D_proj#skel3D_proj.view((skel3D.shape[0],-1))       #(N, 19*2) o



def rotmat_to_angleaxis(init_pred_rotmat):
    """
        init_pred_rotmat: torch.tensor with (24,3,3) dimension
    """
    device = init_pred_rotmat.device
    ones = torch.tensor([0,0,1], dtype=torch.float32,).view(1, 3, 1).expand(init_pred_rotmat.shape[1], -1, -1).to(device)

    pred_rotmat_hom = torch.cat([ init_pred_rotmat.view(-1, 3, 3),ones ], dim=-1)       #24,3,4
    pred_aa = torchgeometry.rotation_matrix_to_angle_axis(pred_rotmat_hom).contiguous().view(1, -1)  #[1,72]
    # tgm.rotation_matrix_to_angle_axis returns NaN for 0 rotation, so manually hack it
    pred_aa[torch.isnan(pred_aa)] = 0.0       #[1,72]
    pred_aa = pred_aa.view(1,24,3) 

    return pred_aa
imports/frankmocap/bodymocap/utils/geometry.py
jellyheadandrew/CHORUS

# Original code from SPIN: https://github.com/nkolot/SPIN

"""
This file contains functions that are used to perform data augmentation.
"""
import torch
import numpy as np
import scipy.misc
import cv2

from bodymocap import constants
from torchvision.transforms import Normalize

def get_transform(center, scale, res, rot=0):
    """Generate transformation matrix."""
    h = 200 * scale
    t = np.zeros((3, 3))
    t[0, 0] = float(res[1]) / h
    t[1, 1] = float(res[0]) / h
    t[0, 2] = res[1] * (-float(center[0]) / h + .5)
    t[1, 2] = res[0] * (-float(center[1]) / h + .5)
    t[2, 2] = 1
    if not rot == 0:
        rot = -rot # To match direction of rotation from cropping
        rot_mat = np.zeros((3,3))
        rot_rad = rot * np.pi / 180
        sn,cs = np.sin(rot_rad), np.cos(rot_rad)
        rot_mat[0,:2] = [cs, -sn]
        rot_mat[1,:2] = [sn, cs]
        rot_mat[2,2] = 1
        # Need to rotate around center
        t_mat = np.eye(3)
        t_mat[0,2] = -res[1]/2
        t_mat[1,2] = -res[0]/2
        t_inv = t_mat.copy()
        t_inv[:2,2] *= -1
        t = np.dot(t_inv,np.dot(rot_mat,np.dot(t_mat,t)))
    return t

def transform(pt, center, scale, res, invert=0, rot=0):
    """Transform pixel location to different reference."""
    t = get_transform(center, scale, res, rot=rot)
    if invert:
        t = np.linalg.inv(t)
    new_pt = np.array([pt[0]-1, pt[1]-1, 1.]).T
    new_pt = np.dot(t, new_pt)
    return new_pt[:2].astype(int)+1

def crop(img, center, scale, res, rot=0):
    """Crop image according to the supplied bounding box."""
    # Upper left point
    ul = np.array(transform([1, 1], center, scale, res, invert=1))-1
    # Bottom right point
    br = np.array(transform([res[0]+1,
                             res[1]+1], center, scale, res, invert=1))-1

    # Padding so that when rotated proper amount of context is included
    pad = int(np.linalg.norm(br - ul) / 2 - float(br[1] - ul[1]) / 2)
    if not rot == 0:
        ul -= pad
        br += pad

    new_shape = [br[1] - ul[1], br[0] - ul[0]]
    if new_shape[0]>15000 or new_shape[1]>15000:
        print("Image Size Too Big!  scale{}, new_shape{} br{}, ul{}".format(scale, new_shape, br, ul))
        return  None


    if len(img.shape) > 2:
        new_shape += [img.shape[2]]


    new_img = np.zeros(new_shape, dtype=np.uint8)

    # #Compute bbox for Han's format
    # bboxScale_o2n = 224/new_img.shape[0]
    # bboxTopLeft = ul *bboxScale_o2n


    # Range to fill new array
    new_x = max(0, -ul[0]), min(br[0], len(img[0])) - ul[0]
    new_y = max(0, -ul[1]), min(br[1], len(img)) - ul[1]
    # Range to sample from original image
    old_x = max(0, ul[0]), min(len(img[0]), br[0])
    old_y = max(0, ul[1]), min(len(img), br[1])
    # print("{} vs {}  || {} vs {}".format(new_y[1] - new_y[0] , old_y[1] - old_y[0], new_x[1] - new_x[0], old_x[1] -old_x[0] )   )
    if new_y[1] - new_y[0] != old_y[1] - old_y[0] or new_x[1] - new_x[0] != old_x[1] -old_x[0] or  new_y[1] - new_y[0] <0 or  new_x[1] - new_x[0] <0:
        print("Warning: maybe person is out of image boundary!")
        return None
    new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[0]:old_y[1],
                                                        old_x[0]:old_x[1]]

    if not rot == 0:
        # Remove padding
        new_img = scipy.misc.imrotate(new_img, rot)
        new_img = new_img[pad:-pad, pad:-pad]

    new_img = cv2.resize(new_img, tuple(res))
    # new_img = scipy.misc.imresize(new_img, res)     #Need this to get the same number with the old model (trained with this resize)

    return new_img#, bboxScale_o2n, bboxTopLeft


def crop_bboxInfo(img, center, scale, res =(224,224)):
    """Crop image according to the supplied bounding box."""
    # Upper left point
    ul = np.array(transform([1, 1], center, scale, res, invert=1))-1
    # Bottom right point
    br = np.array(transform([res[0]+1,
                             res[1]+1], center, scale, res, invert=1))-1


    # Padding so that when rotated proper amount of context is included
    pad = int(np.linalg.norm(br - ul) / 2 - float(br[1] - ul[1]) / 2)

    new_shape = [br[1] - ul[1], br[0] - ul[0]]
    if len(img.shape) > 2:
        new_shape += [img.shape[2]]
    # new_img = np.zeros(new_shape)
    if new_shape[0] <1  or new_shape[1] <1:
        return None, None, None
    new_img = np.zeros(new_shape, dtype=np.uint8)

    if new_img.shape[0] ==0:
        return None, None, None

    #Compute bbox for Han's format
    bboxScale_o2n = res[0]/new_img.shape[0]             #224/ 531

    # Range to fill new array
    new_x = max(0, -ul[0]), min(br[0], len(img[0])) - ul[0]
    new_y = max(0, -ul[1]), min(br[1], len(img)) - ul[1]
    # Range to sample from original image
    old_x = max(0, ul[0]), min(len(img[0]), br[0])
    old_y = max(0, ul[1]), min(len(img), br[1])

    if new_y[0] <0 or new_y[1]<0 or new_x[0] <0 or new_x[1]<0 :
        return None, None, None

    new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[0]:old_y[1],
                                                        old_x[0]:old_x[1]]

    bboxTopLeft_inOriginal = (ul[0], ul[1] )

    if new_img.shape[0] <20 or new_img.shape[1]<20:
        return None, None, None
    # print(bboxTopLeft_inOriginal)
    # from renderer import viewer2D
    # viewer2D.ImShow(new_img.astype(np.uint8),name='cropped')

    new_img = cv2.resize(new_img, res)

    # viewer2D.ImShow(new_img.astype(np.uint8),name='original')

    return new_img, bboxScale_o2n, np.array(bboxTopLeft_inOriginal)


def uncrop(img, center, scale, orig_shape, rot=0, is_rgb=True):
    """'Undo' the image cropping/resizing.
    This function is used when evaluating mask/part segmentation.
    """
    res = img.shape[:2]
    # Upper left point
    ul = np.array(transform([1, 1], center, scale, res, invert=1))-1
    # Bottom right point
    br = np.array(transform([res[0]+1,res[1]+1], center, scale, res, invert=1))-1
    # size of cropped image
    crop_shape = [br[1] - ul[1], br[0] - ul[0]]

    new_shape = [br[1] - ul[1], br[0] - ul[0]]
    if len(img.shape) > 2:
        new_shape += [img.shape[2]]
    new_img = np.zeros(orig_shape, dtype=np.uint8)
    # Range to fill new array
    new_x = max(0, -ul[0]), min(br[0], orig_shape[1]) - ul[0]
    new_y = max(0, -ul[1]), min(br[1], orig_shape[0]) - ul[1]
    # Range to sample from original image
    old_x = max(0, ul[0]), min(orig_shape[1], br[0])
    old_y = max(0, ul[1]), min(orig_shape[0], br[1])
    img = cv2.resize(img, crop_shape, interp='nearest')
    new_img[old_y[0]:old_y[1], old_x[0]:old_x[1]] = img[new_y[0]:new_y[1], new_x[0]:new_x[1]]
    return new_img

def rot_aa(aa, rot):
    """Rotate axis angle parameters."""
    # pose parameters
    R = np.array([[np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0],
                  [np.sin(np.deg2rad(-rot)), np.cos(np.deg2rad(-rot)), 0],
                  [0, 0, 1]])
    # find the rotation of the body in camera frame
    per_rdg, _ = cv2.Rodrigues(aa)
    # apply the global rotation to the global orientation
    resrot, _ = cv2.Rodrigues(np.dot(R,per_rdg))
    aa = (resrot.T)[0]
    return aa

def flip_img(img):
    """Flip rgb images or masks.
    channels come last, e.g. (256,256,3).
    """
    img = np.fliplr(img)
    return img

def flip_kp(kp):
    """Flip keypoints."""
    if len(kp) == 24:
        flipped_parts = constants.J24_FLIP_PERM
    elif len(kp) == 49:
        flipped_parts = constants.J49_FLIP_PERM
    kp = kp[flipped_parts]
    kp[:,0] = - kp[:,0]
    return kp

def flip_pose(pose):
    """Flip pose.
    The flipping is based on SMPL parameters.
    """
    flipped_parts = constants.SMPL_POSE_FLIP_PERM
    pose = pose[flipped_parts]
    # we also negate the second and the third dimension of the axis-angle
    pose[1::3] = -pose[1::3]
    pose[2::3] = -pose[2::3]
    return pose


""" Extract bbox information """
def bbox_from_openpose(openpose_file, rescale=1.2, detection_thresh=0.2):
    """Get center and scale for bounding box from openpose detections."""
    with open(openpose_file, 'r') as f:
        data = json.load(f)
        if 'people' not in data or len(data['people'])==0:
            return None, None
        # keypoints = json.load(f)['people'][0]['pose_keypoints_2d']
        keypoints = data['people'][0]['pose_keypoints_2d']
    keypoints = np.reshape(np.array(keypoints), (-1,3))
    valid = keypoints[:,-1] > detection_thresh

    valid_keypoints = keypoints[valid][:,:-1]           #(25,2)

    # min_pt = np.min(valid_keypoints, axis=0)
    # max_pt = np.max(valid_keypoints, axis=0)
    # bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]

    center = valid_keypoints.mean(axis=0)
    bbox_size = (valid_keypoints.max(axis=0) - valid_keypoints.min(axis=0)).max()
    # adjust bounding box tightness
    scale = bbox_size / 200.0
    scale *= rescale
    return center, scale#, bbox


# keypoints: (Nx3)
def bbox_from_keypoint2d(keypoints, rescale=1.2, detection_thresh=0.2):
    """
        output:
            center: bbox center
            scale: scale_n2o: 224x224 -> original bbox size (max length if not a square bbox)
    """
    # """Get center and scale for bounding box from openpose detections."""

    if len(keypoints.shape)==2 and keypoints.shape[1]==2:       #(X,2)
        valid_keypoints = keypoints
    else:
        keypoints = np.reshape(np.array(keypoints), (-1,3))
        valid = keypoints[:,-1] > detection_thresh

        valid_keypoints = keypoints[valid][:,:-1]           #(25,2)

    # min_pt = np.min(valid_keypoints, axis=0)
    # max_pt = np.max(valid_keypoints, axis=0)
    # bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]

    center = valid_keypoints.mean(axis=0)
    bbox_size = (valid_keypoints.max(axis=0) - valid_keypoints.min(axis=0)).max()


    # adjust bounding box tightness
    scale = bbox_size / 200.0
    scale *= rescale
    return center, scale#, bbox




# def bbox_from_bboxXYXY(bboxXYXY, rescale=1.2):
#     """
#     bboxXYXY
#     """
    # pass

def bbox_from_keypoints(keypoints, rescale=1.2, detection_thresh=0.2, imageHeight= None):
    """Get center and scale for bounding box from openpose detections."""
    # with open(openpose_file, 'r') as f:
    #     data = json.load(f)
    #     if 'people' not in data or len(data['people'])==0:
    #         return None, None
    #     # keypoints = json.load(f)['people'][0]['pose_keypoints_2d']
    #     keypoints = data['people'][0]['pose_keypoints_2d']
    keypoints = np.reshape(np.array(keypoints), (-1,3))
    valid = keypoints[:,-1] > detection_thresh

    # if g_debugUpperBodyOnly:    #Intentionally remove lower bodies
    #     valid[ [ 9,10,11,12,13,14, 22,23,24, 19,20,21] ] = False

    valid_keypoints = keypoints[valid][:,:-1]           #(25,2)

    if len(valid_keypoints)<2:
        return None, None, None


    if False:            #Should have all limbs and nose
        if np.sum(valid[ [ 2,3,4, 5,6,7, 9,10, 12,13,1,0] ]) <12:
            return None, None, None

    min_pt = np.min(valid_keypoints, axis=0)
    max_pt = np.max(valid_keypoints, axis=0)

    
    bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]



    # print(valid_keypoints)
    # print(valid)
    print(bbox)

    if imageHeight is not None:

        if valid[10]==False and valid[13] == False:  # No knees ub ioeb
            max_pt[1] = min(max_pt[1] + (max_pt[1]- min_pt[1]), imageHeight )
            bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]
            valid_keypoints = np.vstack( (valid_keypoints, np.array(max_pt)) )


        elif valid[11]==False and valid[14] == False: #No foot
            max_pt[1] = min(max_pt[1] + (max_pt[1]- min_pt[1])*0.2, imageHeight )
            bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]

            valid_keypoints = np.vstack( (valid_keypoints, np.array(max_pt)) )


    center = valid_keypoints.mean(axis=0)
    bbox_size = (valid_keypoints.max(axis=0) - valid_keypoints.min(axis=0)).max()
    # adjust bounding box tightness
    scale = bbox_size / 200.0
    scale *= rescale
    return center, scale, bbox


def process_image_keypoints(img, keypoints, input_res=224):
    """Read image, do preprocessing and possibly crop it according to the bounding box.
    If there are bounding box annotations, use them to crop the image.
    If no bounding box is specified but openpose detections are available, use them to get the bounding box.
    """
    normalize_img = Normalize(mean=constants.IMG_NORM_MEAN, std=constants.IMG_NORM_STD)
    img = img[:,:,::-1].copy() # PyTorch does not support negative stride at the moment

    center, scale, bbox = bbox_from_keypoints(keypoints, imageHeight = img.shape[0])
    if center is None:
        return None, None, None, None, None

    img, boxScale_o2n, bboxTopLeft = crop_bboxInfo(img, center, scale, (input_res, input_res))

    # viewer2D.ImShow(img, name='cropped', waitTime=1)        #224,224,3


    if img is None:
        return None, None, None, None, None


    # unCropped = uncrop(img, center, scale, (input_res, input_res))

    # if True:
    #     viewer2D.ImShow(img)
    img = img.astype(np.float32) / 255.
    img = torch.from_numpy(img).permute(2,0,1)
    norm_img = normalize_img(img.clone())[None]
    # return img, norm_img, img_original, boxScale_o2n, bboxTopLeft, bbox
    bboxInfo ={"center": center, "scale": scale, "bboxXYWH":bbox}
    return img, norm_img, boxScale_o2n, bboxTopLeft, bboxInfo


#bbr: (minX, minY, width, height)
def bbox_from_bbr(bbox_XYWH, rescale=1.2, detection_thresh=0.2, imageHeight= None):
    """Get center and scale for bounding box from openpose detections."""

    # bbox= bbr
    # if imageHeight is not None:

    #     if valid[10]==False and valid[13] == False:  # No knees ub ioeb
    #         max_pt[1] = min(max_pt[1] + (max_pt[1]- min_pt[1]), imageHeight )
    #         bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]
    #         valid_keypoints = np.vstack( (valid_keypoints, np.array(max_pt)) )


    #     elif valid[11]==False and valid[14] == False: #No foot
    #         max_pt[1] = min(max_pt[1] + (max_pt[1]- min_pt[1])*0.2, imageHeight )
    #         bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]

    #         valid_keypoints = np.vstack( (valid_keypoints, np.array(max_pt)) )


    center = bbox_XYWH[:2] + 0.5 * bbox_XYWH[2:]
    bbox_size = max(bbox_XYWH[2:])
    # adjust bounding box tightness
    scale = bbox_size / 200.0
    scale *= rescale
    return center, scale#, bbox_XYWH

def bbox_from_json(bbox_file):
    """Get center and scale of bounding box from bounding box annotations.
    The expected format is [top_left(x), top_left(y), width, height].
    """
    with open(bbox_file, 'r') as f:
        bbox = np.array(json.load(f)['bbox']).astype(np.float32)
    ul_corner = bbox[:2]
    center = ul_corner + 0.5 * bbox[2:]
    width = max(bbox[2], bbox[3])
    scale = width / 200.0
    # make sure the bounding box is rectangular
    return center, scale



def process_image_bbox(img_original, bbox_XYWH, input_res=224):
    """Read image, do preprocessing and possibly crop it according to the bounding box.
    If there are bounding box annotations, use them to crop the image.
    If no bounding box is specified but openpose detections are available, use them to get the bounding box.
    """
    normalize_img = Normalize(mean=constants.IMG_NORM_MEAN, std=constants.IMG_NORM_STD)
    # -> normalized_img: Normalized image

    img_original = img_original[:,:,::-1].copy() # PyTorch does not support negative stride at the moment
    # -> img_original: image converted from RGB to BGR
     
    img = img_original.copy()
    # -> img: copy of img_original


    center, scale = bbox_from_bbr(bbox_XYWH, imageHeight = img.shape[0])
    # -> center: center of bbox in pixel scale
    # -> scale:

    if center is None:
        return None, None,  None, None, None

    img, boxScale_o2n, bboxTopLeft = crop_bboxInfo(img, center, scale, (input_res, input_res))

    # viewer2D.ImShow(img, name='cropped', waitTime=1)        #224,224,3

    if img is None:
        return None, None,  None, None, None


    # unCropped = uncrop(img, center, scale, (input_res, input_res))

    # if True:
    #     viewer2D.ImShow(img)
    norm_img = (img.copy()).astype(np.float32) / 255.
    norm_img = torch.from_numpy(norm_img).permute(2,0,1)
    norm_img = normalize_img(norm_img.clone())[None]

    bboxInfo ={"center": center, "scale": scale, "bboxXYWH":bbox_XYWH}
    return img, norm_img, boxScale_o2n, bboxTopLeft, bboxInfo



g_de_normalize_img = Normalize(mean=[ -constants.IMG_NORM_MEAN[0]/constants.IMG_NORM_STD[0]    , -constants.IMG_NORM_MEAN[1]/constants.IMG_NORM_STD[1], -constants.IMG_NORM_MEAN[2]/constants.IMG_NORM_STD[2]], std=[1/constants.IMG_NORM_STD[0], 1/constants.IMG_NORM_STD[1], 1/constants.IMG_NORM_STD[2]])
def deNormalizeBatchImg(normTensorImg):
    """
    Normalized Batch Img to original image
    Input: 
        normImg: normTensorImg in cpu
    Input: 
        deNormImg: numpy form
    """
    deNormImg = g_de_normalize_img(normTensorImg).numpy()
    deNormImg = np.transpose(deNormImg , (1,2,0) )*255.0
    deNormImg = deNormImg[:,:,[2,1,0]] 
    deNormImg = np.ascontiguousarray(deNormImg, dtype=np.uint8)

    return deNormImg

imports/frankmocap/bodymocap/utils/imutils.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import os
import sys
import os.path as osp
import torch
from torchvision.transforms import Normalize
import numpy as np
import cv2
import argparse
import json
import pickle
from datetime import datetime

from demo.demo_options import DemoOptions
from bodymocap.body_mocap_api import BodyMocap
from bodymocap.body_bbox_detector import BodyPoseEstimator
import mocap_utils.demo_utils as demo_utils
import mocap_utils.general_utils as gnu
from mocap_utils.timer import Timer

import renderer.image_utils as imu
from renderer.viewer2D import ImShow

def run_body_mocap(args, body_bbox_detector, body_mocap, visualizer):
    #Setup input data to handle different types of inputs
    input_type, input_data = demo_utils.setup_input(args)

    cur_frame = args.start_frame # --> Default: 0 (Check demo_options.py)
    # video_frame = 0
    # timer = Timer()
    while True:
        # timer.tic()
        # load data
        load_bbox = False

        if input_type =='image_dir':
            if cur_frame < len(input_data):
                image_path = input_data[cur_frame]
                img_original_bgr  = cv2.imread(image_path)
            else:
                img_original_bgr = None

        elif input_type == 'bbox_dir':
            if cur_frame < len(input_data):
                print("Use pre-computed bounding boxes")
                image_path = input_data[cur_frame]['image_path']
                hand_bbox_list = input_data[cur_frame]['hand_bbox_list']
                body_bbox_list = input_data[cur_frame]['body_bbox_list']
                img_original_bgr  = cv2.imread(image_path)
                load_bbox = True
            else:
                img_original_bgr = None

        elif input_type == 'video':      
            _, img_original_bgr = input_data.read()
            if video_frame < cur_frame:
                video_frame += 1
                continue
            # save the obtained video frames
            image_path = osp.join(args.out_dir, "frames", f"{cur_frame:05d}.jpg")
            if img_original_bgr is not None:
                video_frame += 1
                if args.save_frame:
                    gnu.make_subdir(image_path)
                    cv2.imwrite(image_path, img_original_bgr)

        elif input_type == 'webcam':    
            _, img_original_bgr = input_data.read()

            if video_frame < cur_frame:
                video_frame += 1
                continue
            # save the obtained video frames
            image_path = osp.join(args.out_dir, "frames", f"scene_{cur_frame:05d}.jpg")
            if img_original_bgr is not None:
                video_frame += 1
                if args.save_frame:
                    gnu.make_subdir(image_path)
                    cv2.imwrite(image_path, img_original_bgr)
        else:
            assert False, "Unknown input_type"

        cur_frame +=1
        if img_original_bgr is None or cur_frame > args.end_frame:
            break   
        print("--------------------------------------")

        if load_bbox:
            body_pose_list = None
        else:
            body_pose_list, body_bbox_list = body_bbox_detector.detect_body_pose(
                img_original_bgr)
        hand_bbox_list = [None, ] * len(body_bbox_list)

        # save the obtained body & hand bbox to json file
        if args.save_bbox_output: 
            demo_utils.save_info_to_json(args, image_path, body_bbox_list, hand_bbox_list)

        if len(body_bbox_list) < 1: 
            print(f"No body deteced: {image_path}")
            continue

        #Sort the bbox using bbox size 
        # (to make the order as consistent as possible without tracking)
        bbox_size =  [ (x[2] * x[3]) for x in body_bbox_list]
        idx_big2small = np.argsort(bbox_size)[::-1]
        body_bbox_list = [ body_bbox_list[i] for i in idx_big2small ] # If bbox is big, it comes in front
        if args.single_person and len(body_bbox_list)>0: # args.single_person: Just reconstruct one person (with biggest bbox)
            body_bbox_list = [body_bbox_list[0], ]       

        # Body Pose Regression
        pred_output_list = body_mocap.regress(img_original_bgr, body_bbox_list)
        assert len(body_bbox_list) == len(pred_output_list)

        # extract mesh for rendering (vertices in image space and faces) from pred_output_list
        pred_mesh_list = demo_utils.extract_mesh_from_output(pred_output_list)

        # visualization
        res_img = visualizer.visualize(
            img_original_bgr,
            pred_mesh_list = pred_mesh_list, 
            body_bbox_list = body_bbox_list)
        
        # show result in the screen
        if not args.no_display:
            res_img = res_img.astype(np.uint8)
            ImShow(res_img)

        # save result image
        if args.out_dir is not None:
            demo_utils.save_res_img(args.out_dir, image_path, res_img)

        # save predictions to pkl
        if args.save_pred_pkl:
            demo_type = 'body'
            demo_utils.save_pred_to_pkl(
                args, demo_type, image_path, body_bbox_list, hand_bbox_list, pred_output_list)

        timer.toc(bPrint=True,title="Time")
        print(f"Processed : {image_path}")

    #save images as a video
    if not args.no_video_out and input_type in ['video', 'webcam']:
        demo_utils.gen_video_out(args.out_dir, args.seq_name)

    if input_type =='webcam' and input_data is not None:
        input_data.release()
    cv2.destroyAllWindows()


def main():
    args = DemoOptions().parse()

    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    assert torch.cuda.is_available(), "Current version only supports GPU"

    # Set bbox detector
    body_bbox_detector = BodyPoseEstimator()

    # Set mocap regressor
    use_smplx = args.use_smplx
    checkpoint_path = args.checkpoint_body_smplx if use_smplx else args.checkpoint_body_smpl
    print("use_smplx", use_smplx)
    body_mocap = BodyMocap(checkpoint_path, args.smpl_dir, device, use_smplx)

    # Set Visualizer
    if args.renderer_type in ['pytorch3d', 'opendr']:
        from renderer.screen_free_visualizer import Visualizer
    else:
        from renderer.visualizer import Visualizer
    visualizer = Visualizer(args.renderer_type)
  
    run_body_mocap(args, body_bbox_detector, body_mocap, visualizer)


if __name__ == '__main__':
    main()
imports/frankmocap/demo/demo_bodymocap.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import os
import sys
import os.path as osp
import torch
from torchvision.transforms import Normalize
import numpy as np
import cv2
import argparse
import json
import pickle

############# input parameters  #############
from demo.demo_options import DemoOptions
from bodymocap.body_mocap_api import BodyMocap
from handmocap.hand_mocap_api import HandMocap
import mocap_utils.demo_utils as demo_utils
import mocap_utils.general_utils as gnu
from mocap_utils.timer import Timer
from datetime import datetime

from bodymocap.body_bbox_detector import BodyPoseEstimator
from handmocap.hand_bbox_detector import HandBboxDetector
from integration.copy_and_paste import integration_copy_paste

import renderer.image_utils as imu
from renderer.viewer2D import ImShow


def __filter_bbox_list(body_bbox_list, hand_bbox_list, single_person):
    # (to make the order as consistent as possible without tracking)
    bbox_size =  [ (x[2] * x[3]) for x in body_bbox_list]
    idx_big2small = np.argsort(bbox_size)[::-1]
    body_bbox_list = [ body_bbox_list[i] for i in idx_big2small ]
    hand_bbox_list = [hand_bbox_list[i] for i in idx_big2small]

    if single_person and len(body_bbox_list)>0:
        body_bbox_list = [body_bbox_list[0], ]
        hand_bbox_list = [hand_bbox_list[0], ]

    return body_bbox_list, hand_bbox_list


def run_regress(
    args, img_original_bgr, 
    body_bbox_list, hand_bbox_list, bbox_detector,
    body_mocap, hand_mocap
):
    cond1 = len(body_bbox_list) > 0 and len(hand_bbox_list) > 0
    cond2 = not args.frankmocap_fast_mode

    # use pre-computed bbox or use slow detection mode
    if cond1 or cond2:
        if not cond1 and cond2:
            # run detection only when bbox is not available
            body_pose_list, body_bbox_list, hand_bbox_list, _ = \
                bbox_detector.detect_hand_bbox(img_original_bgr.copy())
        else:
            print("Use pre-computed bounding boxes")
        assert len(body_bbox_list) == len(hand_bbox_list)

        if len(body_bbox_list) < 1: 
            return list(), list(), list()

        # sort the bbox using bbox size 
        # only keep on bbox if args.single_person is set
        body_bbox_list, hand_bbox_list = __filter_bbox_list(
            body_bbox_list, hand_bbox_list, args.single_person)

        # hand & body pose regression
        pred_hand_list = hand_mocap.regress(
            img_original_bgr, hand_bbox_list, add_margin=True)
        pred_body_list = body_mocap.regress(img_original_bgr, body_bbox_list)
        assert len(hand_bbox_list) == len(pred_hand_list)
        assert len(pred_hand_list) == len(pred_body_list)

    else:
        _, body_bbox_list = bbox_detector.detect_body_bbox(img_original_bgr.copy())

        if len(body_bbox_list) < 1: 
            return list(), list(), list()

        # sort the bbox using bbox size 
        # only keep on bbox if args.single_person is set
        hand_bbox_list = [None, ] * len(body_bbox_list)
        body_bbox_list, _ = __filter_bbox_list(
            body_bbox_list, hand_bbox_list, args.single_person)

        # body regression first 
        pred_body_list = body_mocap.regress(img_original_bgr, body_bbox_list)
        assert len(body_bbox_list) == len(pred_body_list)

        # get hand bbox from body
        hand_bbox_list = body_mocap.get_hand_bboxes(pred_body_list, img_original_bgr.shape[:2])
        assert len(pred_body_list) == len(hand_bbox_list)

        # hand regression
        pred_hand_list = hand_mocap.regress(
            img_original_bgr, hand_bbox_list, add_margin=True)
        assert len(hand_bbox_list) == len(pred_hand_list) 

    # integration by copy-and-paste
    integral_output_list = integration_copy_paste(
        pred_body_list, pred_hand_list, body_mocap.smpl, img_original_bgr.shape)
    
    return body_bbox_list, hand_bbox_list, integral_output_list


def run_frank_mocap(args, bbox_detector, body_mocap, hand_mocap, visualizer):
    #Setup input data to handle different types of inputs
    input_type, input_data = demo_utils.setup_input(args)

    cur_frame = args.start_frame
    video_frame = 0
    while True:
        # load data
        load_bbox = False

        if input_type =='image_dir':
            if cur_frame < len(input_data):
                image_path = input_data[cur_frame]
                img_original_bgr  = cv2.imread(image_path)
            else:
                img_original_bgr = None

        elif input_type == 'bbox_dir':
            if cur_frame < len(input_data):
                image_path = input_data[cur_frame]['image_path']
                hand_bbox_list = input_data[cur_frame]['hand_bbox_list']
                body_bbox_list = input_data[cur_frame]['body_bbox_list']
                img_original_bgr  = cv2.imread(image_path)
                load_bbox = True
            else:
                img_original_bgr = None

        elif input_type == 'video':      
            _, img_original_bgr = input_data.read()
            if video_frame < cur_frame:
                video_frame += 1
                continue
          # save the obtained video frames
            image_path = osp.join(args.out_dir, "frames", f"{cur_frame:05d}.jpg")
            if img_original_bgr is not None:
                video_frame += 1
                if args.save_frame:
                    gnu.make_subdir(image_path)
                    cv2.imwrite(image_path, img_original_bgr)
        
        elif input_type == 'webcam':
            _, img_original_bgr = input_data.read()

            if video_frame < cur_frame:
                video_frame += 1
                continue
            # save the obtained video frames
            image_path = osp.join(args.out_dir, "frames", f"scene_{cur_frame:05d}.jpg")
            if img_original_bgr is not None:
                video_frame += 1
                if args.save_frame:
                    gnu.make_subdir(image_path)
                    cv2.imwrite(image_path, img_original_bgr)
        else:
            assert False, "Unknown input_type"

        cur_frame +=1
        if img_original_bgr is None or cur_frame > args.end_frame:
            break   
        print("--------------------------------------")
        
        # bbox detection
        if not load_bbox:
            body_bbox_list, hand_bbox_list = list(), list()
        
        # regression (includes integration)
        body_bbox_list, hand_bbox_list, pred_output_list = run_regress(
            args, img_original_bgr, 
            body_bbox_list, hand_bbox_list, bbox_detector,
            body_mocap, hand_mocap)

        # save the obtained body & hand bbox to json file
        if args.save_bbox_output: 
            demo_utils.save_info_to_json(args, image_path, body_bbox_list, hand_bbox_list)

        if len(body_bbox_list) < 1: 
            print(f"No body deteced: {image_path}")
            continue

        pred_mesh_list = demo_utils.extract_mesh_from_output(pred_output_list)

        # visualization
        res_img = visualizer.visualize(
            img_original_bgr,
            pred_mesh_list = pred_mesh_list,
            body_bbox_list = body_bbox_list,
            hand_bbox_list = hand_bbox_list)
        
       # show result in the screen
        if not args.no_display:
            res_img = res_img.astype(np.uint8)
            ImShow(res_img)

        # save result image
        if args.out_dir is not None:
            demo_utils.save_res_img(args.out_dir, image_path, res_img)

        # save predictions to pkl
        if args.save_pred_pkl:
            demo_type = 'frank'
            demo_utils.save_pred_to_pkl(
                args, demo_type, image_path, body_bbox_list, hand_bbox_list, pred_output_list)

        print(f"Processed : {image_path}")

    # save images as a video
    if not args.no_video_out and input_type in ['video', 'webcam']:
        demo_utils.gen_video_out(args.out_dir, args.seq_name)

    if input_type =='webcam' and input_data is not None:
        input_data.release()
    cv2.destroyAllWindows()

def main():
    args = DemoOptions().parse()
    args.use_smplx = True

    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    assert torch.cuda.is_available(), "Current version only supports GPU"

    hand_bbox_detector =  HandBboxDetector('third_view', device)
    
    #Set Mocap regressor
    body_mocap = BodyMocap(args.checkpoint_body_smplx, args.smpl_dir, device = device, use_smplx= True)
    hand_mocap = HandMocap(args.checkpoint_hand, args.smpl_dir, device = device)

    # Set Visualizer
    if args.renderer_type in ['pytorch3d', 'opendr']:
        from renderer.screen_free_visualizer import Visualizer
    else:
        from renderer.visualizer import Visualizer
    visualizer = Visualizer(args.renderer_type)

    run_frank_mocap(args, hand_bbox_detector, body_mocap, hand_mocap, visualizer)


if __name__ == '__main__':
    main()
imports/frankmocap/demo/demo_frankmocap.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import os, sys, shutil
import os.path as osp
import numpy as np
import cv2
import json
import torch
from torchvision.transforms import Normalize

from demo.demo_options import DemoOptions
import mocap_utils.general_utils as gnu
import mocap_utils.demo_utils as demo_utils

from handmocap.hand_mocap_api import HandMocap
from handmocap.hand_bbox_detector import HandBboxDetector

import renderer.image_utils as imu
from renderer.viewer2D import ImShow
import time


def run_hand_mocap(args, bbox_detector, hand_mocap, visualizer):
    #Set up input data (images or webcam)
    input_type, input_data = demo_utils.setup_input(args)
 
    assert args.out_dir is not None, "Please specify output dir to store the results"
    cur_frame = args.start_frame
    video_frame = 0

    while True:
        # load data
        load_bbox = False

        if input_type =='image_dir':
            if cur_frame < len(input_data):
                image_path = input_data[cur_frame]
                img_original_bgr  = cv2.imread(image_path)
            else:
                img_original_bgr = None

        elif input_type == 'bbox_dir':
            if cur_frame < len(input_data):
                print("Use pre-computed bounding boxes")
                image_path = input_data[cur_frame]['image_path']
                hand_bbox_list = input_data[cur_frame]['hand_bbox_list']
                body_bbox_list = input_data[cur_frame]['body_bbox_list']
                img_original_bgr  = cv2.imread(image_path)
                load_bbox = True
            else:
                img_original_bgr = None

        elif input_type == 'video':      
            _, img_original_bgr = input_data.read()
            if video_frame < cur_frame:
                video_frame += 1
                continue
            # save the obtained video frames
            image_path = osp.join(args.out_dir, "frames", f"{cur_frame:05d}.jpg")
            if img_original_bgr is not None:
                video_frame += 1
                if args.save_frame:
                    gnu.make_subdir(image_path)
                    cv2.imwrite(image_path, img_original_bgr)
        
        elif input_type == 'webcam':
            _, img_original_bgr = input_data.read()

            if video_frame < cur_frame:
                video_frame += 1
                continue
            # save the obtained video frames
            image_path = osp.join(args.out_dir, "frames", f"scene_{cur_frame:05d}.jpg")
            if img_original_bgr is not None:
                video_frame += 1
                if args.save_frame:
                    gnu.make_subdir(image_path)
                    cv2.imwrite(image_path, img_original_bgr)
        else:
            assert False, "Unknown input_type"

        cur_frame +=1
        if img_original_bgr is None or cur_frame > args.end_frame:
            break   
        print("--------------------------------------")

        # bbox detection
        if load_bbox:
            body_pose_list = None
            raw_hand_bboxes = None
        elif args.crop_type == 'hand_crop':
            # hand already cropped, thererore, no need for detection
            img_h, img_w = img_original_bgr.shape[:2]
            body_pose_list = None
            raw_hand_bboxes = None
            hand_bbox_list = [ dict(right_hand = np.array([0, 0, img_w, img_h])) ]
        else:            
            # Input images has other body part or hand not cropped.
            # Use hand detection model & body detector for hand detection
            assert args.crop_type == 'no_crop'
            detect_output = bbox_detector.detect_hand_bbox(img_original_bgr.copy())
            body_pose_list, body_bbox_list, hand_bbox_list, raw_hand_bboxes = detect_output # hand_bbox_list: list of detected dictionary of left_hand, right_hand bbox (x,y,w,h) form

        # save the obtained body & hand bbox to json file
        if args.save_bbox_output:
            demo_utils.save_info_to_json(args, image_path, body_bbox_list, hand_bbox_list)

        if len(hand_bbox_list) < 1:
            print(f"No hand deteced: {image_path}")
            continue
    
        # Hand Pose Regression
        pred_output_list = hand_mocap.regress(
                img_original_bgr, hand_bbox_list, add_margin=True)
        assert len(hand_bbox_list) == len(body_bbox_list)
        assert len(body_bbox_list) == len(pred_output_list)

        # extract mesh for rendering (vertices in image space and faces) from pred_output_list
        pred_mesh_list = demo_utils.extract_mesh_from_output(pred_output_list)

        # visualize
        res_img = visualizer.visualize(
            img_original_bgr, 
            pred_mesh_list = pred_mesh_list, 
            hand_bbox_list = hand_bbox_list)

        # show result in the screen
        if not args.no_display:
            res_img = res_img.astype(np.uint8)
            ImShow(res_img)

        # save the image (we can make an option here)
        if args.out_dir is not None:
            demo_utils.save_res_img(args.out_dir, image_path, res_img)

        # save predictions to pkl
        if args.save_pred_pkl or True:
            demo_type = 'hand'
            demo_utils.save_pred_to_pkl(
                args, demo_type, image_path, body_bbox_list, hand_bbox_list, pred_output_list)

        print(f"Processed : {image_path}")
        
    #save images as a video
    if not args.no_video_out and input_type in ['video', 'webcam']:
        demo_utils.gen_video_out(args.out_dir, args.seq_name)

    # When everything done, release the capture
    if input_type =='webcam' and input_data is not None:
        input_data.release()
    cv2.destroyAllWindows()

  
def main():
    args = DemoOptions().parse()
    args.use_smplx = True

    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    assert torch.cuda.is_available(), "Current version only supports GPU"

    #Set Bbox detector
    bbox_detector =  HandBboxDetector(args.view_type, device)

    # Set Mocap regressor
    hand_mocap = HandMocap(args.checkpoint_hand, args.smpl_dir, device = device)

    # Set Visualizer
    if args.renderer_type in ['pytorch3d', 'opendr']:
        from renderer.screen_free_visualizer import Visualizer
    else:
        from renderer.visualizer import Visualizer
    visualizer = Visualizer(args.renderer_type)

    # run
    run_hand_mocap(args, bbox_detector, hand_mocap, visualizer)
   

if __name__ == '__main__':
    main()
imports/frankmocap/demo/demo_handmocap.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import argparse

class DemoOptions():

    def __init__(self):
        parser = argparse.ArgumentParser()
        
        # parser.add_argument('--checkpoint', required=False, default=default_checkpoint, help='Path to pretrained checkpoint')
        default_checkpoint_body_smpl ='./extra_data/body_module/pretrained_weights/2020_05_31-00_50_43-best-51.749683916568756.pt'
        parser.add_argument('--checkpoint_body_smpl', required=False, default=default_checkpoint_body_smpl, help='Path to pretrained checkpoint')
        default_checkpoint_body_smplx ='./extra_data/body_module/pretrained_weights/smplx-03-28-46060-w_spin_mlc3d_46582-2089_2020_03_28-21_56_16.pt'
        parser.add_argument('--checkpoint_body_smplx', required=False, default=default_checkpoint_body_smplx, help='Path to pretrained checkpoint')
        default_checkpoint_hand = "./extra_data/hand_module/pretrained_weights/pose_shape_best.pth"
        parser.add_argument('--checkpoint_hand', required=False, default=default_checkpoint_hand, help='Path to pretrained checkpoint')

        # input options
        parser.add_argument('--input_path', type=str, default=None, help="""Path of video, image, or a folder where image files exists""")
        parser.add_argument('--start_frame', type=int, default=0, help='given a sequence of frames, set the starting frame')
        parser.add_argument('--end_frame', type=int, default=float('inf'), help='given a sequence of frames, set the last frame')
        parser.add_argument('--pkl_dir', type=str, help='Path of storing pkl files that store the predicted results')
        parser.add_argument('--openpose_dir', type=str, help='Directory of storing the prediction of openpose prediction')

        # output options
        parser.add_argument('--out_dir', type=str, default=None, help='Folder of output images.')
        # parser.add_argument('--pklout', action='store_true', help='Export mocap output as pkl file')
        parser.add_argument('--save_bbox_output', action='store_true', help='Save the bboxes in json files (bbox_xywh format)')
        parser.add_argument('--save_pred_pkl', action='store_true', help='Save the predictions (bboxes, params, meshes in pkl format')
        parser.add_argument("--save_mesh", action='store_true', help="Save the predicted vertices and faces")
        parser.add_argument("--save_frame", action='store_true', help='Save the extracted frames from video input or webcam')

        # Other options
        parser.add_argument('--single_person', action='store_true', help='Reconstruct only one person in the scene with the biggest bbox')
        parser.add_argument('--no_display', action='store_true', help='Do not visualize output on the screen')
        parser.add_argument('--no_video_out', action='store_true', help='Do not merge rendered frames to video (ffmpeg)')
        parser.add_argument('--smpl_dir', type=str, default='./extra_data/smpl/', help='Folder where smpl files are located.')
        parser.add_argument('--skip', action='store_true', help='Skip there exist already processed outputs')
        parser.add_argument('--video_url', type=str, default=None, help='URL of YouTube video, or image.')
        parser.add_argument('--download', '-d', action='store_true', help='Download YouTube video first (in webvideo folder), and process it')

        # Body mocap specific options
        parser.add_argument('--use_smplx', action='store_true', help='Use SMPLX model for body mocap')

        # Hand mocap specific options
        parser.add_argument('--view_type', type=str, default='third_view', choices=['third_view', 'ego_centric'],
            help = "The view type of input. It could be ego-centric (such as epic kitchen) or third view")
        parser.add_argument('--crop_type', type=str, default='no_crop', choices=['hand_crop', 'no_crop'],
            help = """ 'hand_crop' means the hand are central cropped in input. (left hand should be flipped to right). 
                        'no_crop' means hand detection is required to obtain hand bbox""")
        
        # Whole motion capture (FrankMocap) specific options
        parser.add_argument('--frankmocap_fast_mode', action='store_true', help="Use fast hand detection mode for whole body motion capture (frankmocap)")

        # renderer
        parser.add_argument("--renderer_type", type=str, default="opengl", 
            choices=['pytorch3d', 'opendr', 'opengl_gui', 'opengl'], help="type of renderer to use")

        self.parser = parser
    

    def parse(self):
        self.opt = self.parser.parse_args()
        return self.opt
imports/frankmocap/demo/demo_options.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import os
import sys
import os.path as osp
import torch
import numpy as np
import cv2
import argparse
import json
import pickle
import smplx
from datetime import datetime

from demo.demo_options import DemoOptions
from bodymocap.body_mocap_api import BodyMocap
import mocap_utils.demo_utils as demo_utils
import mocap_utils.general_utils as gnu
from bodymocap.models import SMPL, SMPLX
from handmocap.hand_modules.h3dw_model import extract_hand_output
from mocap_utils.coordconv import convert_smpl_to_bbox, convert_bbox_to_oriIm


def __get_data_type(pkl_files):
    for pkl_file in pkl_files:
        saved_data = gnu.load_pkl(pkl_file)
        return saved_data['demo_type'], saved_data['smpl_type']


def __get_smpl_model(demo_type, smpl_type):
    smplx_model_path = './extra_data/smpl/SMPLX_NEUTRAL.pkl'
    smpl_model_path = './extra_data/smpl//basicModel_neutral_lbs_10_207_0_v1.0.0.pkl'

    if demo_type == 'hand':
        # use original smpl-x
        smpl = smplx.create(
            smplx_model_path, 
            model_type = "smplx", 
            batch_size = 1,
            gender = 'neutral',
            num_betas = 10,
            use_pca = False,
            ext='pkl'
        )
    else:
        if smpl_type == 'smplx':
            # use modified smpl-x from body module
            smpl = SMPLX(
                smplx_model_path,
                batch_size=1,
                num_betas = 10,
                use_pca = False,
                create_transl=False)
        else:
            # use modified smpl from body module
            assert smpl_type == 'smpl'
            smpl = SMPL(
                smpl_model_path, 
                batch_size=1, 
                create_transl=False)
    return smpl
    

def __calc_hand_mesh(hand_type, pose_params, betas, smplx_model):
    hand_rotation = pose_params[:, :3]
    hand_pose = pose_params[:, 3:]
    body_pose = torch.zeros((1, 63)).float()

    assert hand_type in ['left_hand', 'right_hand']
    if hand_type == 'right_hand':
        body_pose[:, 60:] = hand_rotation # set right hand rotation
        right_hand_pose = hand_pose
        left_hand_pose = torch.zeros((1, 45), dtype=torch.float32)
    else:
        body_pose[:, 57:60] = hand_rotation # set right hand rotation
        left_hand_pose = hand_pose
        right_hand_pose = torch.zeros((1, 45), dtype=torch.float32)

    output = smplx_model(
        global_orient = torch.zeros((1,3)),
        body_pose = body_pose,
        betas = betas,
        left_hand_pose = left_hand_pose,
        right_hand_pose = right_hand_pose,
        return_verts = True)
    
    hand_info_file = "extra_data/hand_module/SMPLX_HAND_INFO.pkl"
    hand_info = gnu.load_pkl(hand_info_file)
    hand_output = extract_hand_output(
        output, 
        hand_type = hand_type.split("_")[0],
        hand_info = hand_info,
        top_finger_joints_type = 'ave',
        use_cuda = False)

    pred_verts = hand_output['hand_vertices_shift'].detach().numpy()
    faces = hand_info[f'{hand_type}_faces_local']
    return pred_verts[0], faces


def _calc_body_mesh(smpl_type, smpl_model, body_pose, betas, 
    left_hand_pose, right_hand_pose):
    if smpl_type == 'smpl':
        smpl_output = smpl_model(
            global_orient = body_pose[:, :3],
            body_pose = body_pose[:, 3:], 
            betas = betas, 
        )
    else:
        smpl_output = smpl_model(
            global_orient = body_pose[:, :3],
            body_pose = body_pose[:, 3:], 
            betas = betas, 
            left_hand_pose = left_hand_pose,
            right_hand_pose = right_hand_pose,
        )

    vertices = smpl_output.vertices.detach().cpu().numpy()[0]
    faces = smpl_model.faces
    return vertices, faces


def __calc_mesh(demo_type, smpl_type, smpl_model, img_shape, pred_output_list):
    for pred_output in pred_output_list:
        if pred_output is not None:
            # hand
            if demo_type == 'hand':
                assert 'left_hand' in pred_output and 'right_hand' in pred_output
                for hand_type in pred_output:
                    hand_pred = pred_output[hand_type]
                    if hand_pred is not None:
                        pose_params = torch.from_numpy(hand_pred['pred_hand_pose'])
                        betas = torch.from_numpy(hand_pred['pred_hand_betas'])
                        pred_verts, hand_faces = __calc_hand_mesh(hand_type, pose_params, betas, smpl_model)
                        hand_pred['pred_vertices_smpl'] = pred_verts

                        cam_scale = hand_pred['pred_camera'][0]
                        cam_trans = hand_pred['pred_camera'][1:]
                        vert_bboxcoord = convert_smpl_to_bbox(
                            pred_verts, cam_scale, cam_trans, bAppTransFirst=True) # SMPL space -> bbox space

                        bbox_scale_ratio = hand_pred['bbox_scale_ratio']
                        bbox_top_left = hand_pred['bbox_top_left']
                        vert_imgcoord = convert_bbox_to_oriIm(
                                vert_bboxcoord, bbox_scale_ratio, bbox_top_left,
                                img_shape[1], img_shape[0]) 
                        pred_output[hand_type]['pred_vertices_img'] = vert_imgcoord
            # body
            else:
                pose_params = torch.from_numpy(pred_output['pred_body_pose'])
                betas = torch.from_numpy(pred_output['pred_betas'])
                if 'pred_right_hand_pose' in pred_output:
                    pred_right_hand_pose = torch.from_numpy(pred_output['pred_right_hand_pose'])
                else:
                    pred_right_hand_pose = torch.zeros((1, 45), dtype=torch.float32)
                if 'pred_left_hand_pose' in pred_output:
                    pred_left_hand_pose = torch.from_numpy(pred_output['pred_left_hand_pose'])
                else:
                    pred_left_hand_pose = torch.zeros((1, 45), dtype=torch.float32)
                pred_verts, faces = _calc_body_mesh(
                    smpl_type, smpl_model, pose_params, betas, pred_left_hand_pose, pred_right_hand_pose)

                pred_output['pred_vertices_smpl'] = pred_verts
                pred_output['faces'] = faces

                cam_scale = pred_output['pred_camera'][0]
                cam_trans = pred_output['pred_camera'][1:]
                vert_bboxcoord = convert_smpl_to_bbox(
                    pred_verts, cam_scale, cam_trans, bAppTransFirst=False) # SMPL space -> bbox space

                bbox_scale_ratio = pred_output['bbox_scale_ratio']
                bbox_top_left = pred_output['bbox_top_left']
                vert_imgcoord = convert_bbox_to_oriIm(
                        vert_bboxcoord, bbox_scale_ratio, bbox_top_left,
                        img_shape[1], img_shape[0]) 
                pred_output['pred_vertices_img'] = vert_imgcoord


def visualize_prediction(args, demo_type, smpl_type, smpl_model, pkl_files, visualizer):
    for pkl_file in pkl_files:
        # load data
        saved_data = gnu.load_pkl(pkl_file)
        
        image_path = saved_data['image_path']
        img_original_bgr = cv2.imread(image_path)
        if img_original_bgr is None:
            print(f"{image_path} does not exists, skip")
        
        print("--------------------------------------")

        demo_type = saved_data['demo_type']
        assert saved_data['smpl_type'] == smpl_type

        hand_bbox_list = saved_data['hand_bbox_list']        
        body_bbox_list = saved_data['body_bbox_list']
        pred_output_list = saved_data['pred_output_list']

        if not saved_data['save_mesh']:
            __calc_mesh(
                demo_type, smpl_type, smpl_model, img_original_bgr.shape[:2], pred_output_list)
        else:
            pass

        pred_mesh_list = demo_utils.extract_mesh_from_output(pred_output_list)

        # visualization
        res_img = visualizer.visualize(
            img_original_bgr,
            pred_mesh_list = pred_mesh_list,
            body_bbox_list = body_bbox_list,
            hand_bbox_list = hand_bbox_list)

        # save result image
        demo_utils.save_res_img(args.out_dir, image_path, res_img)

        # save predictions to pkl
        if args.save_pred_pkl:
            args.use_smplx = smpl_type == 'smplx'
            demo_utils.save_pred_to_pkl(
                args, demo_type, image_path, body_bbox_list, hand_bbox_list, pred_output_list)


def main():
    args = DemoOptions().parse()

    # load pkl files
    pkl_files = gnu.get_all_files(args.pkl_dir, ".pkl", "full")

    # get smpl type
    demo_type, smpl_type = __get_data_type(pkl_files)

    # get smpl model
    smpl_model = __get_smpl_model(demo_type, smpl_type)

   # Set Visualizer
    assert args.renderer_type in ['pytorch3d', 'opendr'], \
        f"{args.renderer_type} not implemented yet."
    from renderer.screen_free_visualizer import Visualizer
    visualizer = Visualizer(args.renderer_type)

    # load smpl model
    visualize_prediction(args, demo_type, smpl_type, smpl_model, pkl_files, visualizer)


if __name__ == '__main__':
    main()
imports/frankmocap/demo/demo_visualize_prediction.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import sys
import numpy as np
import torch
import pdb
import mocap_utils.geometry_utils as gu
from mocap_utils.coordconv import convert_smpl_to_bbox, convert_bbox_to_oriIm


def get_kinematic_map(smplx_model, dst_idx):
    cur = dst_idx
    kine_map = dict()
    while cur>=0:
        parent = int(smplx_model.parents[cur])
        if cur != dst_idx: # skip the dst_idx itself
            kine_map[parent] = cur
        cur = parent
    return kine_map


def __transfer_rot(body_pose_rotmat, part_rotmat, kinematic_map, transfer_type):

    rotmat= body_pose_rotmat[0] 
    parent_id = 0
    while parent_id in kinematic_map:
        child_id = kinematic_map[parent_id]
        local_rotmat = body_pose_rotmat[child_id]
        rotmat = torch.matmul(rotmat, local_rotmat)
        parent_id = child_id

    if transfer_type == 'g2l':
        part_rot_new = torch.matmul(rotmat.T, part_rotmat)
    else:
        assert transfer_type == 'l2g'
        part_rot_new = torch.matmul(rotmat, part_rotmat)

    return part_rot_new


def transfer_rotation(
    smplx_model, body_pose, part_rot, part_idx, 
    transfer_type="g2l", result_format="rotmat"):

    assert transfer_type in ["g2l", "l2g"]
    assert result_format in ['rotmat', 'aa']

    assert type(body_pose) == type(part_rot)
    return_np = False

    if isinstance(body_pose, np.ndarray):
        body_pose = torch.from_numpy(body_pose)
        return_np = True
    
    if isinstance(part_rot, np.ndarray):
        part_rot = torch.from_numpy(part_rot)
        return_np = True

    if body_pose.dim() == 2:
        # aa
        assert body_pose.size(0) == 1 and body_pose.size(1) in [66, 72]
        body_pose_rotmat = gu.angle_axis_to_rotation_matrix(body_pose.view(22, 3)).clone()
    else:
        # rotmat
        assert body_pose.dim() == 4
        assert body_pose.size(0) == 1 and body_pose.size(1) in [22, 24]
        assert body_pose.size(2) == 3 and body_pose.size(3) == 3
        body_pose_rotmat = body_pose[0].clone()

    if part_rot.dim() == 2:
        # aa
        assert part_rot.size(0) == 1 and part_rot.size(1) == 3
        part_rotmat = gu.angle_axis_to_rotation_matrix(part_rot)[0,:3,:3].clone()
    else:
        # rotmat
        assert part_rot.dim() == 3
        assert part_rot.size(0) == 1 and part_rot.size(1) == 3 and part_rot.size(2) == 3
        part_rotmat = part_rot[0,:3,:3].clone()

    kinematic_map = get_kinematic_map(smplx_model, part_idx)
    part_rot_trans = __transfer_rot(
        body_pose_rotmat, part_rotmat, kinematic_map, transfer_type)

    if result_format == 'rotmat':    
        return_value = part_rot_trans
    else:
        part_rot_aa = gu.rotation_matrix_to_angle_axis(part_rot_trans)
        return_value = part_rot_aa
    if return_np:
        return_value = return_value.numpy()
    return return_value


def integration_copy_paste(pred_body_list, pred_hand_list, smplx_model, image_shape):
    integral_output_list = list()
    for i in range(len(pred_body_list)):
        body_info = pred_body_list[i]
        hand_info = pred_hand_list[i]
        if body_info is None:
            integral_output_list.append(None)
            continue
    
        # copy and paste 
        pred_betas = torch.from_numpy(body_info['pred_betas']).cuda()
        pred_rotmat = torch.from_numpy(body_info['pred_rotmat']).cuda()

        # integrate right hand pose
        hand_output = dict()
        if hand_info is not None and hand_info['right_hand'] is not None:
            right_hand_pose = torch.from_numpy(hand_info['right_hand']['pred_hand_pose'][:, 3:]).cuda()
            right_hand_global_orient = torch.from_numpy(hand_info['right_hand']['pred_hand_pose'][:,: 3]).cuda()
            right_hand_local_orient = transfer_rotation(
                smplx_model, pred_rotmat, right_hand_global_orient, 21)
            pred_rotmat[0, 21] = right_hand_local_orient
        else:
            right_hand_pose = torch.from_numpy(np.zeros( (1,45) , dtype= np.float32)).cuda()
            right_hand_global_orient = None
            right_hand_local_orient = None

        # integrate left hand pose
        if hand_info is not None and hand_info['left_hand'] is not None:
            left_hand_pose = torch.from_numpy(hand_info['left_hand']['pred_hand_pose'][:, 3:]).cuda()
            left_hand_global_orient = torch.from_numpy(hand_info['left_hand']['pred_hand_pose'][:, :3]).cuda()
            left_hand_local_orient = transfer_rotation(
                smplx_model, pred_rotmat, left_hand_global_orient, 20)
            pred_rotmat[0, 20] = left_hand_local_orient
        else:
            left_hand_pose = torch.from_numpy(np.zeros((1,45), dtype= np.float32)).cuda()
            left_hand_global_orient = None
            left_hand_local_orient = None

        pred_aa = gu.rotation_matrix_to_angle_axis(pred_rotmat).cuda()
        pred_aa = pred_aa.reshape(pred_aa.shape[0], 72)
        smplx_output = smplx_model(
            betas = pred_betas, 
            body_pose = pred_aa[:,3:], 
            global_orient = pred_aa[:,:3],
            right_hand_pose = right_hand_pose, 
            left_hand_pose= left_hand_pose,
            pose2rot = True)

        pred_vertices = smplx_output.vertices
        pred_vertices = pred_vertices[0].detach().cpu().numpy()
        pred_body_joints = smplx_output.joints
        pred_body_joints = pred_body_joints[0].detach().cpu().numpy()   
        pred_lhand_joints = smplx_output.left_hand_joints
        pred_lhand_joints = pred_lhand_joints[0].detach().cpu().numpy()
        pred_rhand_joints = smplx_output.right_hand_joints
        pred_rhand_joints = pred_rhand_joints[0].detach().cpu().numpy()

        camScale = body_info['pred_camera'][0]
        camTrans = body_info['pred_camera'][1:]
        bbox_top_left = body_info['bbox_top_left']
        bbox_scale_ratio = body_info['bbox_scale_ratio']

        integral_output = dict()
        integral_output['pred_vertices_smpl'] = pred_vertices
        integral_output['faces'] = smplx_model.faces
        integral_output['bbox_scale_ratio'] = bbox_scale_ratio
        integral_output['bbox_top_left'] = bbox_top_left
        integral_output['pred_camera'] = body_info['pred_camera']

        pred_aa_tensor = gu.rotation_matrix_to_angle_axis(pred_rotmat.detach().cpu()[0])
        integral_output['pred_body_pose'] = pred_aa_tensor.cpu().numpy().reshape(1, 72)
        integral_output['pred_betas'] = pred_betas.detach().cpu().numpy()

        # convert mesh to original image space (X,Y are aligned to image)
        pred_vertices_bbox = convert_smpl_to_bbox(
            pred_vertices, camScale, camTrans)
        pred_vertices_img = convert_bbox_to_oriIm(
            pred_vertices_bbox, bbox_scale_ratio, bbox_top_left, image_shape[1], image_shape[0])
        integral_output['pred_vertices_img'] = pred_vertices_img

        # convert joints to original image space (X, Y are aligned to image)
        pred_body_joints_bbox = convert_smpl_to_bbox(
            pred_body_joints, camScale, camTrans)
        pred_body_joints_img = convert_bbox_to_oriIm(
            pred_body_joints_bbox, bbox_scale_ratio, bbox_top_left, image_shape[1], image_shape[0])
        integral_output['pred_body_joints_img'] = pred_body_joints_img

        # convert left /right joints to original image space (X, Y are aligned to image)
        pred_lhand_joints_bbox = convert_smpl_to_bbox(
            pred_lhand_joints, camScale, camTrans)
        pred_lhand_joints_img = convert_bbox_to_oriIm(
            pred_lhand_joints_bbox, bbox_scale_ratio, bbox_top_left, image_shape[1], image_shape[0])
        integral_output['pred_lhand_joints_img'] = pred_lhand_joints_img

        pred_rhand_joints_bbox = convert_smpl_to_bbox(
            pred_rhand_joints, camScale, camTrans)
        pred_rhand_joints_img = convert_bbox_to_oriIm(
            pred_rhand_joints_bbox, bbox_scale_ratio, bbox_top_left, image_shape[1], image_shape[0])
        integral_output['pred_rhand_joints_img'] = pred_rhand_joints_img

        # keep hand info
        r_hand_local_orient_body = body_info['pred_rotmat'][:, 21] # rot-mat
        r_hand_global_orient_body = transfer_rotation(
            smplx_model, pred_rotmat,
            torch.from_numpy(r_hand_local_orient_body).cuda(),
            21, 'l2g', 'aa').numpy().reshape(1, 3) # aa
        r_hand_local_orient_body = gu.rotation_matrix_to_angle_axis(r_hand_local_orient_body) # rot-mat -> aa

        l_hand_local_orient_body = body_info['pred_rotmat'][:, 20]
        l_hand_global_orient_body = transfer_rotation(
            smplx_model, pred_rotmat,
            torch.from_numpy(l_hand_local_orient_body).cuda(),
            20, 'l2g', 'aa').numpy().reshape(1, 3)
        l_hand_local_orient_body = gu.rotation_matrix_to_angle_axis(l_hand_local_orient_body) # rot-mat -> aa

        r_hand_local_orient_hand = None
        r_hand_global_orient_hand = None
        if right_hand_local_orient is not None:
            r_hand_local_orient_hand = gu.rotation_matrix_to_angle_axis(
                right_hand_local_orient).detach().cpu().numpy().reshape(1, 3)
            r_hand_global_orient_hand = right_hand_global_orient.detach().cpu().numpy().reshape(1, 3)

        l_hand_local_orient_hand = None
        l_hand_global_orient_hand = None
        if left_hand_local_orient is not None:
            l_hand_local_orient_hand = gu.rotation_matrix_to_angle_axis(
                left_hand_local_orient).detach().cpu().numpy().reshape(1, 3)
            l_hand_global_orient_hand = left_hand_global_orient.detach().cpu().numpy().reshape(1, 3)

        # poses and rotations related to hands
        integral_output['left_hand_local_orient_body'] = l_hand_local_orient_body
        integral_output['left_hand_global_orient_body'] = l_hand_global_orient_body
        integral_output['right_hand_local_orient_body'] = r_hand_local_orient_body
        integral_output['right_hand_global_orient_body'] = r_hand_global_orient_body

        integral_output['left_hand_local_orient_hand'] = l_hand_local_orient_hand
        integral_output['left_hand_global_orient_hand'] = l_hand_global_orient_hand
        integral_output['right_hand_local_orient_hand'] = r_hand_local_orient_hand
        integral_output['right_hand_global_orient_hand'] = r_hand_global_orient_hand

        integral_output['pred_left_hand_pose'] = left_hand_pose.detach().cpu().numpy()
        integral_output['pred_right_hand_pose'] = right_hand_pose.detach().cpu().numpy()

        # predicted hand betas, cameras, top-left corner and center
        left_hand_betas = None
        left_hand_camera = None
        left_hand_bbox_scale = None
        left_hand_bbox_top_left = None
        if hand_info is not None and hand_info['left_hand'] is not None:
            left_hand_betas = hand_info['left_hand']['pred_hand_betas']
            left_hand_camera = hand_info['left_hand']['pred_camera']
            left_hand_bbox_scale = hand_info['left_hand']['bbox_scale_ratio']
            left_hand_bbox_top_left = hand_info['left_hand']['bbox_top_left']

        right_hand_betas = None
        right_hand_camera = None
        right_hand_bbox_scale = None
        right_hand_bbox_top_left = None
        if hand_info is not None and hand_info['right_hand'] is not None:
            right_hand_betas = hand_info['right_hand']['pred_hand_betas']
            right_hand_camera = hand_info['right_hand']['pred_camera']
            right_hand_bbox_scale = hand_info['right_hand']['bbox_scale_ratio']
            right_hand_bbox_top_left = hand_info['right_hand']['bbox_top_left']

        integral_output['pred_left_hand_betas'] = left_hand_betas
        integral_output['left_hand_camera'] = left_hand_camera
        integral_output['left_hand_bbox_scale_ratio'] = left_hand_bbox_scale
        integral_output['left_hand_bbox_top_left'] = left_hand_bbox_top_left

        integral_output['pred_right_hand_betas'] = right_hand_betas
        integral_output['right_hand_camera'] = right_hand_camera
        integral_output['right_hand_bbox_scale_ratio'] = right_hand_bbox_scale
        integral_output['right_hand_bbox_top_left'] = right_hand_bbox_top_left

        integral_output_list.append(integral_output)

    return integral_output_list
imports/frankmocap/integration/copy_and_paste.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

"""
This code is used to visually compare the results
"""
import os, sys, shutil
import os.path as osp
import ry_utils
import cv2
import numpy as np

def check_keywords(subdir, keywords):
    if len(keywords) == 0:
        return True
    else:
        for keyword in keywords:
            if subdir.find(keyword)>=0:
                return True
    return False

def main():
    dir_list = [
        'samples/output/body/third_view_thresh_0.3_distance_2.0',
        'samples/output/body/third_view_thresh_0.5_distance_1.5',
        'samples/output/body/third_view_thresh_0.7_distance_1.0',
    ]
    dir1 = dir_list[0]

    keywords = ['cj_dance', 'body_capture']

    res_dir = "samples/output/body/third_view_compare"
    res_dir = osp.join(res_dir, '_&&_'.join(['_'.join(item.split('/')[-1:]) for item in dir_list]))

    for subdir in os.listdir(dir1):
        if osp.isdir(osp.join(dir1, subdir)):
            if check_keywords(subdir, keywords):
                dir_path1 = osp.join(dir1, subdir)
                for img_name in ry_utils.get_all_files(dir_path1, ('.jpg','.png'), 'name_only'):
                    img_list = list()
                    #print(img_name)
                    for dir in dir_list:
                        dir_path = dir_path1.replace(dir1, dir)
                        img_path = osp.join(dir_path, img_name)
                        img = cv2.imread(img_path)
                        img_list.append(img)
                        if img_path.find(dir1)>=0:
                            res_img_path = img_path.replace(dir1, res_dir)
                        #print(img_path, osp.exists(img_path))
                    if any([img is None for img in img_list]):
                        continue
                    res_img = np.concatenate(img_list, axis=0)
                    h, w = res_img.shape[:2]
                    res_img = cv2.resize(res_img, (int(w*0.7), int(h*0.7)))
                    res_img_path = res_img_path.replace('.png', '.jpg')
                    ry_utils.make_subdir(res_img_path)
                    cv2.imwrite(res_img_path, res_img)
                    print(res_img_path)



if __name__ == '__main__':
     main()
imports/frankmocap/mocap_utils/compare_results.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import os, sys, shutil
import os.path as osp
import cv2
import numpy as np
import general_utils as gnu


def main():
    in_dir = "./sample_data/images/single_person"
    out_dir = "./sample_data/images/multi_person"
    gnu.renew_dir(out_dir)

    all_imgs = gnu.get_all_files(in_dir, (".jpg", ".png", ".jpeg"), "full")
    num_img = len(all_imgs)

    for i in range(num_img):
        for j in range(num_img):
            img1 = cv2.imread(all_imgs[i])
            img2 = cv2.imread(all_imgs[j])
            img2 = cv2.resize(img2, img1.shape[:2][::-1])
            res_img = np.concatenate((img1, img2), axis=1)
            res_img_path = osp.join(out_dir, f"{i:02d}_{j:02d}.jpg")
            cv2.imwrite(res_img_path, res_img)


if __name__ == '__main__':
    main()
imports/frankmocap/mocap_utils/compose_image.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

# Original code from SPIN: https://github.com/nkolot/SPIN

"""
This file contains functions that are used to perform data augmentation.
"""
import sys
import torch
import numpy as np
import scipy.misc
import cv2
from torchvision.transforms import Normalize

# For converting coordinate between SMPL 3D coord <-> 2D bbox <-> original 2D image 
# data3D: (N,3), where N is number of 3D points in "smpl"'s 3D coordinate (vertex or skeleton)

def convert_smpl_to_bbox(data3D, scale, trans, bAppTransFirst=False):
    data3D = data3D.copy()
    resnet_input_size_half = 224 *0.5
    if bAppTransFirst:      # Hand model
        data3D[:,0:2] += trans
        data3D *= scale   # apply scaling
    else:
        data3D *= scale # apply scaling
        data3D[:,0:2] += trans
    
    data3D*= resnet_input_size_half # 112 is originated from hrm's input size (224,24)
    # data3D[:,:2]*= resnet_input_size_half # 112 is originated from hrm's input size (224,24)
    return data3D


def convert_bbox_to_oriIm(data3D, boxScale_o2n, bboxTopLeft, imgSizeW, imgSizeH):
    data3D = data3D.copy()
    resnet_input_size_half = 224 *0.5
    imgSize = np.array([imgSizeW,imgSizeH])

    data3D /= boxScale_o2n

    if not isinstance(bboxTopLeft, np.ndarray):
        assert isinstance(bboxTopLeft, tuple)
        assert len(bboxTopLeft) == 2
        bboxTopLeft = np.array(bboxTopLeft)

    data3D[:,:2] += (bboxTopLeft + resnet_input_size_half/boxScale_o2n)

    return data3D


def convert_smpl_to_bbox_perspective(data3D, scale_ori, trans_ori, focalLeng, scaleFactor=1.0):
    data3D = data3D.copy()
    resnet_input_size_half = 224 *0.5

    scale = scale_ori* resnet_input_size_half
    trans = trans_ori *resnet_input_size_half

    if False:   #Weak perspective
        data3D *= scale           #apply scaling
        data3D[:,0:2] += trans
    else:
        # delta = (trans - imageShape*0.5)/scale            
        # Current projection already consider camera center during the rendering. 
        # Thus no need to consider principle axis
        delta = (trans )/scale
        data3D[:,0:2] +=delta

        newZ = focalLeng/scale
        deltaZ =  newZ - np.mean(data3D[:,2])
        data3D[:,2] +=deltaZ
        # data3D[:,2] +=16.471718554146534        #debug

    if False:   #Scaling to be a certain dist from camera
        texture_plan_depth = 500
        ratio = texture_plan_depth /np.mean(data3D[:,2])
        data3D *=ratio  
    else:
        data3D *=scaleFactor

    return data3D

## extract bbox information
def bbox_from_openpose(openpose_file, rescale=1.2, detection_thresh=0.2):
    with open(openpose_file, 'r') as f:
        data = json.load(f)
        if 'people' not in data or len(data['people'])==0:
            return None, None
        # keypoints = json.load(f)['people'][0]['pose_keypoints_2d']
        keypoints = data['people'][0]['pose_keypoints_2d']
    keypoints = np.reshape(np.array(keypoints), (-1,3))
    valid = keypoints[:,-1] > detection_thresh

    valid_keypoints = keypoints[valid][:,:-1]           #(25,2)

    # min_pt = np.min(valid_keypoints, axis=0)
    # max_pt = np.max(valid_keypoints, axis=0)
    # bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]

    center = valid_keypoints.mean(axis=0)
    bbox_size = (valid_keypoints.max(axis=0) - valid_keypoints.min(axis=0)).max()
    # adjust bounding box tightness
    scale = bbox_size / 200.0
    scale *= rescale
    return center, scale#, bbox


# keypoints: (Nx3)
def bbox_from_keypoint2d(keypoints, rescale=1.2, detection_thresh=0.2):
    """
        output:
            center: bbox center
            scale: scale_n2o: 224x224 -> original bbox size (max length if not a square bbox)
    """
    # """Get center and scale for bounding box from openpose detections."""

    if len(keypoints.shape)==2 and keypoints.shape[1]==2:       #(X,2)
        valid_keypoints = keypoints
    else:
        keypoints = np.reshape(np.array(keypoints), (-1,3))
        valid = keypoints[:,-1] > detection_thresh

        valid_keypoints = keypoints[valid][:,:-1]           #(25,2)

    # min_pt = np.min(valid_keypoints, axis=0)
    # max_pt = np.max(valid_keypoints, axis=0)
    # bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]

    center = valid_keypoints.mean(axis=0)
    bbox_size = (valid_keypoints.max(axis=0) - valid_keypoints.min(axis=0)).max()


    # adjust bounding box tightness
    scale = bbox_size / 200.0
    scale *= rescale
    return center, scale#, bbox


def bbox_from_keypoints(keypoints, rescale=1.2, detection_thresh=0.2, imageHeight= None):
    """Get center and scale for bounding box from openpose detections."""
  
    keypoints = np.reshape(np.array(keypoints), (-1,3))
    valid = keypoints[:,-1] > detection_thresh

    valid_keypoints = keypoints[valid][:,:-1]           #(25,2)

    if len(valid_keypoints)<2:
        return None, None, None


    if False:            #Should have all limbs and nose
        if np.sum(valid[ [ 2,3,4, 5,6,7, 9,10, 12,13,1,0] ]) <12:
            return None, None, None

    min_pt = np.min(valid_keypoints, axis=0)
    max_pt = np.max(valid_keypoints, axis=0)

    bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]

    if imageHeight is not None:

        if valid[10]==False and valid[13] == False:  # No knees ub ioeb
            max_pt[1] = min(max_pt[1] + (max_pt[1]- min_pt[1]), imageHeight )
            bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]
            valid_keypoints = np.vstack( (valid_keypoints, np.array(max_pt)) )


        elif valid[11]==False and valid[14] == False: #No foot
            max_pt[1] = min(max_pt[1] + (max_pt[1]- min_pt[1])*0.2, imageHeight )
            bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]

            valid_keypoints = np.vstack( (valid_keypoints, np.array(max_pt)) )


    center = valid_keypoints.mean(axis=0)
    bbox_size = (valid_keypoints.max(axis=0) - valid_keypoints.min(axis=0)).max()
    # adjust bounding box tightness
    scale = bbox_size / 200.0
    scale *= rescale
    return center, scale, bbox


def bbox_from_bbr(bbox_XYWH, rescale=1.2, detection_thresh=0.2, imageHeight= None):
    #bbr: (minX, minY, width, height)
    """Get center and scale for bounding box from openpose detections."""

    center = bbox_XYWH[:2] + 0.5 * bbox_XYWH[2:]
    bbox_size = max(bbox_XYWH[2:])
    # adjust bounding box tightness
    scale = bbox_size / 200.0
    scale *= rescale
    return center, scale#, bbox_XYWH


def bbox_from_json(bbox_file):
    """Get center and scale of bounding box from bounding box annotations.
    The expected format is [top_left(x), top_left(y), width, height].
    """
    with open(bbox_file, 'r') as f:
        bbox = np.array(json.load(f)['bbox']).astype(np.float32)
    ul_corner = bbox[:2]
    center = ul_corner + 0.5 * bbox[2:]
    width = max(bbox[2], bbox[3])
    scale = width / 200.0
    # make sure the bounding box is rectangular
    return center, scale
imports/frankmocap/mocap_utils/coordconv.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import os, sys, shutil
import os.path as osp
import cv2
from collections import OrderedDict
import mocap_utils.general_utils as gnu
import numpy as np
import json
import subprocess as sp


def setup_render_out(out_dir):
    if out_dir is not None:
        gnu.build_dir(out_dir)
        outputFileName = 'scene_%08d.jpg' # Hardcoded in glViewer.py

        overlaidImageFolder= osp.join(out_dir, 'overlaid')
        gnu.build_dir(overlaidImageFolder)

        sideImageFolder= osp.join(out_dir, 'side')
        gnu.build_dir(sideImageFolder)

        mergedImageFolder= osp.join(out_dir, 'merged')
        gnu.build_dir(mergedImageFolder)

        res_subdirs = \
            [outputFileName, overlaidImageFolder, sideImageFolder, mergedImageFolder]
        return res_subdirs
    
    else:
        return None


def __get_input_type(args):
    input_type =None
    image_exts = ('jpg', 'png', 'jpeg', 'bmp')
    video_exts = ('mp4', 'avi', 'mov')
    extension = osp.splitext(args.input_path)[1][1:]

    if extension.lower() in video_exts:
        input_type ='video'
    elif osp.isdir(args.input_path):
        file_list = os.listdir(args.input_path)
        assert len(file_list) >0, f"{args.input_path} is a blank folder"
        extension = osp.splitext(file_list[0])[1][1:]
        if extension == 'json':
            input_type ='bbox_dir'
        else:
            assert extension.lower() in image_exts
            input_type ='image_dir'
    elif args.input_path =='webcam':
        input_type ='webcam'
    else:
        assert False, "Unknown input path. It should be an image," + \
            "or an image folder, or a video file, or \'webcam\' "
    return input_type


def __video_setup(args):
    video_path = args.input_path
    video_dir, video_name, video_basename, ext = gnu.analyze_path(video_path)
    args.seq_name = video_basename

    if args.save_frame:
        frame_dir = osp.join(args.out_dir, "frames")
        gnu.build_dir(frame_dir)    

    render_out_dir = osp.join(args.out_dir, "rendered")
    gnu.build_dir(render_out_dir)

    mocap_out_dir = osp.join(args.out_dir, "mocap")
    gnu.build_dir(mocap_out_dir)


def __img_seq_setup(args):
    seq_dir_path = args.input_path
    args.seq_name = os.path.basename(args.input_path)

    render_out_dir = osp.join(args.out_dir, 'rendered')
    gnu.build_dir(render_out_dir)

    mocap_out_dir = osp.join(args.out_dir, "mocap")
    gnu.build_dir(mocap_out_dir)


def setup_input(args):
    """
    Input type can be 
        an image file
        a video file
        a folder with image files
        a folder with bbox (json) files
        "webcam"
    
    """
    image_exts = ('jpg', 'png', 'jpeg', 'bmp')
    video_exts = ('mp4', 'avi', 'mov')

    # get type of input 
    input_type = __get_input_type(args)

    if input_type =='video':
        cap = cv2.VideoCapture(args.input_path)
        assert cap.isOpened(), f"Failed in opening video: {args.input_path}"
        __video_setup(args)
        return input_type, cap

    elif input_type =='webcam':
        cap = cv2.VideoCapture(0)       #webcam input
        return input_type, cap

    elif input_type =='image_dir':
        image_list = gnu.get_all_files(args.input_path, image_exts, "relative") 
        image_list = [ osp.join(args.input_path, image_name) for image_name in image_list ]
        __img_seq_setup(args)
        return input_type, image_list

    elif input_type =='bbox_dir':
        __img_seq_setup(args)
        json_files = gnu.get_all_files(args.input_path, '.json', "relative") 
        input_data = list()
        for json_file in json_files:
            json_path = osp.join(args.input_path, json_file)
            image_path, body_bbox_list, hand_bbox_list = load_info_from_json(json_path)
            input_data.append(dict(
                image_path = image_path,
                hand_bbox_list = hand_bbox_list,
                body_bbox_list = body_bbox_list
            ))
        return input_type, input_data

    else:
        assert False, "Unknown input type"


def extract_mesh_from_output(pred_output_list):
    pred_mesh_list = list()
    for pred_output in pred_output_list:
        if pred_output is not None:
            if 'left_hand' in pred_output: # hand mocap
                for hand_type in pred_output:
                    if pred_output[hand_type] is not None:
                        vertices = pred_output[hand_type]['pred_vertices_img']
                        faces = pred_output[hand_type]['faces'].astype(np.int32)
                        pred_mesh_list.append(dict(
                            vertices = vertices,
                            faces = faces
                        ))
            else: # body mocap (includes frank/whole/total mocap)
                vertices = pred_output['pred_vertices_img']
                faces = pred_output['faces'].astype(np.int32)
                pred_mesh_list.append(dict(
                    vertices = vertices,
                    faces = faces
                ))
    return pred_mesh_list
                
    
def load_info_from_json(json_path):
    data = gnu.load_json(json_path)
    # image path
    assert ('image_path' in data), "Path of input image should be specified"
    image_path = data['image_path']
    assert osp.exists(image_path), f"{image_path} does not exists"
    # body bboxes
    body_bbox_list = list()
    if 'body_bbox_list' in data:
        body_bbox_list = data['body_bbox_list']
        assert isinstance(body_bbox_list, list)
        for b_id, body_bbox in enumerate(body_bbox_list):
            if isinstance(body_bbox, list) and len(body_bbox) == 4:
                body_bbox_list[b_id] = np.array(body_bbox)
    # hand bboxes
    hand_bbox_list = list()
    if 'hand_bbox_list' in data:
        hand_bbox_list = data['hand_bbox_list']
        assert isinstance(hand_bbox_list, list)
        for hand_bbox in hand_bbox_list:
            for hand_type in ['left_hand', 'right_hand']:
                if hand_type in hand_bbox:
                    bbox = hand_bbox[hand_type]
                    if isinstance(bbox, list) and len(bbox) == 4:
                        hand_bbox[hand_type] = np.array(bbox)
                    else:
                        hand_bbox[hand_type] = None
    return image_path, body_bbox_list, hand_bbox_list


def save_info_to_json(args, image_path, body_bbox_list, hand_bbox_list):
    saved_data = dict()

    # image_path
    saved_data['image_path'] = image_path 

    # body_bbox_list
    saved_body_bbox_list = list()
    for body_bbox in body_bbox_list:
        if body_bbox is not None:
            saved_body_bbox_list.append(body_bbox.tolist())
    saved_data['body_bbox_list'] = saved_body_bbox_list

    # hand_bbox_list
    saved_hand_bbox_list = list()
    for hand_bbox in hand_bbox_list:
        if hand_bbox is not None:
            saved_hand_bbox = dict(
                left_hand = None,
                right_hand = None)
            for hand_type in saved_hand_bbox:
                bbox = hand_bbox[hand_type]
                if bbox is not None:
                    saved_hand_bbox[hand_type] = bbox.tolist()
            saved_hand_bbox_list.append(saved_hand_bbox)
    saved_data['hand_bbox_list'] = saved_hand_bbox_list

    # write data to json
    img_name = osp.basename(image_path)
    record = img_name.split('.')
    json_name = f"{'.'.join(record[:-1])}_bbox.json"
    json_path = osp.join(args.out_dir, 'bbox', json_name)
    gnu.make_subdir(json_path)
    gnu.save_json(json_path, saved_data)
    print(f"Bbox saved: {json_path}")


def save_pred_to_pkl(
    args, demo_type, image_path, 
    body_bbox_list, hand_bbox_list, pred_output_list):

    smpl_type = 'smplx' if args.use_smplx else 'smpl'
    assert demo_type in ['hand', 'body', 'frank']
    if demo_type in ['hand', 'frank']:
        assert smpl_type == 'smplx'

    assert len(hand_bbox_list) == len(body_bbox_list)
    assert len(body_bbox_list) == len(pred_output_list)

    saved_data = dict()
    # demo type / smpl type / image / bbox
    saved_data = OrderedDict()
    saved_data['demo_type'] = demo_type
    saved_data['smpl_type'] = smpl_type
    saved_data['image_path'] = osp.abspath(image_path)
    saved_data['body_bbox_list'] = body_bbox_list
    saved_data['hand_bbox_list'] = hand_bbox_list
    saved_data['save_mesh'] = args.save_mesh

    saved_data['pred_output_list'] = list()
    num_subject = len(hand_bbox_list)
    for s_id in range(num_subject):
        # predict params
        pred_output = pred_output_list[s_id]
        if pred_output is None:
            saved_pred_output = None
        else:
            saved_pred_output = dict()
            if demo_type == 'hand':
                for hand_type in ['left_hand', 'right_hand']:
                    pred_hand = pred_output[hand_type]
                    saved_pred_output[hand_type] = dict()
                    saved_data_hand = saved_pred_output[hand_type]
                    if pred_hand is None:
                        saved_data_hand = None
                    else:
                        for pred_key in pred_hand:
                            if pred_key.find("vertices")<0 or pred_key == 'faces' :
                                saved_data_hand[pred_key] = pred_hand[pred_key]
                            else:
                                if args.save_mesh:
                                    if pred_key != 'faces':
                                        saved_data_hand[pred_key] = \
                                            pred_hand[pred_key].astype(np.float16)
                                    else:
                                        saved_data_hand[pred_key] = pred_hand[pred_key]
            else:
                for pred_key in pred_output:
                    if pred_key.find("vertices")<0 or pred_key == 'faces' :
                        saved_pred_output[pred_key] = pred_output[pred_key]
                    else:
                        if args.save_mesh:
                            if pred_key != 'faces':
                                saved_pred_output[pred_key] = \
                                    pred_output[pred_key].astype(np.float16)
                            else:
                                saved_pred_output[pred_key] = pred_output[pred_key]

        saved_data['pred_output_list'].append(saved_pred_output)

    # write data to pkl
    img_name = osp.basename(image_path)
    record = img_name.split('.')
    pkl_name = f"{'.'.join(record[:-1])}_prediction_result.pkl"
    pkl_path = osp.join(args.out_dir, 'mocap', pkl_name)
    gnu.make_subdir(pkl_path)
    gnu.save_pkl(pkl_path, saved_data)
    print(f"Prediction saved: {pkl_path}")
 

def save_res_img(out_dir, image_path, res_img):
    out_dir = osp.join(out_dir, "rendered")
    img_name = osp.basename(image_path)
    img_name = img_name[:-4] + '.jpg'           #Always save as jpg
    res_img_path = osp.join(out_dir, img_name)
    gnu.make_subdir(res_img_path)
    cv2.imwrite(res_img_path, res_img)
    print(f"Visualization saved: {res_img_path}")


def gen_video_out(out_dir, seq_name):
    outVideo_fileName = osp.join(out_dir, seq_name+'.mp4')
    print(f">> Generating video in {outVideo_fileName}")

    in_dir = osp.abspath(osp.join(out_dir, "rendered"))
    out_path = osp.abspath(osp.join(out_dir, seq_name+'.mp4'))
    ffmpeg_cmd = f'ffmpeg -y -f image2 -framerate 25 -pattern_type glob -i "{in_dir}/*.jpg"  -pix_fmt yuv420p -c:v libx264 -x264opts keyint=25:min-keyint=25:scenecut=-1 -vf "scale=trunc(iw/2)*2:trunc(ih/2)*2" {out_path}'
    os.system(ffmpeg_cmd)
imports/frankmocap/mocap_utils/demo_utils.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import os, sys, shutil
import os.path as osp
import subprocess as sp
import general_utils as gnu


def extract_frame(video_dir, frame_dir):
    for file in os.listdir(video_dir):
        if file.endswith((".mov", ".mp4")):
            file_path = osp.join(video_dir, file)
            file_name = file[:-4]
            # if file_name != 'legao_02_01': continue
            res_dir = osp.join(frame_dir, file_name)
            gnu.build_dir(res_dir)
            command = f"ffmpeg -i {file_path} {res_dir}/{file_name}_%05d.png"
            command = command.split()
            sp.run(command)


def main():
    root_dir = "./sample_data/"

    video_dir = osp.join(root_dir, "videos")
    frame_dir = osp.join(root_dir, "frames")
    gnu.renew_dir(frame_dir)

    extract_frame(video_dir, frame_dir)

if __name__ == '__main__':
    main()
imports/frankmocap/mocap_utils/extract_frame.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

# file to store some often use functions
import os, sys, shutil
import os.path as osp
import multiprocessing as mp
import numpy as np
import cv2
import pickle
import json


def save_mesh_to_obj(obj_path, verts, faces=None):
    assert isinstance(verts, np.ndarray)
    assert isinstance(faces, np.ndarray)

    with open(obj_path, 'w') as out_f:
        # write verts
        for v in verts:
            out_f.write(f"v {v[0]:.4f} {v[1]:.4f} {v[2]:.4f}\n")
        # write faces 
        if faces is not None:
            faces = faces.copy() + 1
            for f in faces:
                out_f.write(f"f {f[0]} {f[1]} {f[2]}\n")


def renew_dir(target_dir):
    if osp.exists(target_dir):
        shutil.rmtree(target_dir)
    os.makedirs(target_dir)


def build_dir(target_dir):
    if not osp.exists(target_dir):
        os.makedirs(target_dir)


def get_subdir(in_path):
    subdir_path = '/'.join(in_path.split('/')[:-1])
    return subdir_path

def make_subdir(in_path):
    subdir_path = get_subdir(in_path)
    build_dir(subdir_path)


def update_extension(file_path, new_extension):
    assert new_extension[0] == '.'
    old_extension = '.' + file_path.split('.')[-1]
    new_file_path = file_path.replace(old_extension, new_extension)
    return new_file_path


def get_all_files(in_dir, extension, path_type='full', keywords=''):
    assert path_type in ['full', 'relative', 'name_only']
    assert isinstance(extension, str) or isinstance(extension, tuple)
    assert isinstance(keywords, str)

    all_files = list()
    for subdir, dirs, files in os.walk(in_dir):
        for file in files:
            if len(keywords)>0:
                if file.find(keywords)<0: 
                    continue
            if file.endswith(extension):
                if path_type == 'full':
                    file_path = osp.join(subdir, file)
                elif path_type == 'relative':
                    file_path = osp.join(subdir, file).replace(in_dir, '')
                    if file_path.startswith('/'):
                        file_path = file_path[1:]
                else:
                    file_path = file
                all_files.append(file_path)
    return sorted(all_files)


def remove_swp(in_dir):
    remove_files = list()
    for subdir, dirs, files in os.walk(in_dir):
        for file in files:
            if file.endswith('.swp'):
                full_path = osp.join(subdir,file)
                os.remove(full_path)


def remove_pyc(in_dir):
    remove_files = list()
    for subdir, dirs, files in os.walk(in_dir):
        for file in files:
            if file.endswith('.pyc'):
                full_path = osp.join(subdir,file)
                os.remove(full_path)


def md5sum(file_path):
    import hashlib
    hash_md5 = hashlib.md5()
    with open(file_path, 'rb') as in_f:
        hash_md5.update(in_f.read())
    return hash_md5.hexdigest()


# save data to pkl
def save_pkl(res_file, data_list, protocol=-1):
    assert res_file.endswith(".pkl")
    res_file_dir = '/'.join(res_file.split('/')[:-1])
    if len(res_file_dir)>0:
        if not osp.exists(res_file_dir):
            os.makedirs(res_file_dir)
    with open(res_file, 'wb') as out_f:
        if protocol==2:
            pickle.dump(data_list, out_f, protocol=2)
        else:
            pickle.dump(data_list, out_f)


def load_pkl(pkl_file, res_list=None):
    assert pkl_file.endswith(".pkl")
    with open(pkl_file, 'rb') as in_f:
        try:
            data = pickle.load(in_f)
        except UnicodeDecodeError:
            in_f.seek(0)
            data = pickle.load(in_f, encoding='latin1')
    return data


def load_json(in_file):
    assert in_file.endswith(".json")
    with open(in_file, 'r') as in_f:
        all_data = json.load(in_f)
        return all_data


def save_json(out_file, data):
    assert out_file.endswith(".json")
    with open(out_file, "w") as out_f:
        json.dump(data, out_f)


def load_npz(npz_file):
    res_data = dict()
    assert npz_file.endswith(".npz")
    raw_data = np.load(npz_file, mmap_mode='r')
    for key in raw_data.files:
        res_data[key] = raw_data[key]
    return res_data


def update_npz_file(npz_file, new_key, new_data):
    # load original data
    assert npz_file.endswith(".npz")
    raw_data = np.load(npz_file, mmap_mode='r')
    all_data = dict()
    for key in raw_data.files:
        all_data[key] = raw_data[key]
    # add new data && save
    all_data[new_key] = new_data
    np.savez(npz_file, **all_data)


def analyze_path(input_path):
    # assume input_path is the path of a file not a directory
    record = input_path.split('/')
    input_dir = '/'.join(record[:-1])
    file_name = record[-1]
    assert file_name.find(".")>0
    ext = file_name.split('.')[-1]
    file_basename = '.'.join(file_name.split('.')[:-1])
    return input_dir, file_name, file_basename, ext
imports/frankmocap/mocap_utils/general_utils.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import os, sys, shutil
import os.path as osp
import numpy as np
import torch
from torch.nn import functional as F
import cv2
import numpy.matlib as npm
import mocap_utils.geometry_utils_torch as gut


def flip_hand_pose(pose):
    pose = pose.copy()
    if len(pose.shape) == 1:
        pose = pose.reshape(-1, 3)
        pose[:, 1] *= -1
        pose[:, 2] *= -1
        return pose.reshape(-1,)
    else:
        assert len(pose.shape) == 2
        pose[:, 1] *= -1
        pose[:, 2] *= -1
        return pose


def flip_hand_joints_3d(joints_3d):
    assert joints_3d.shape[1] == 3
    assert len(joints_3d.shape) == 2
    rot_mat = np.diag([-1, 1, 1])
    return np.matmul(rot_mat, joints_3d.T).T


def __quaternion_to_angle_axis_torch(quat):
    quat = quat.clone()
    if quat.dim() == 1:
        assert quat.size(0) == 4
        quat = quat.view(1, 4)
        angle_axis = gut.quaternion_to_angle_axis(quat)[0]
    elif quat.dim() == 2:
        assert quat.size(1) == 4
        angle_axis = gut.quaternion_to_angle_axis(quat)
    else:
        assert quat.dim() == 3
        dim0 = quat.size(0)
        dim1 = quat.size(1)
        assert quat.size(2) == 4
        quat = quat.view(dim0*dim1, 4)
        angle_axis = gut.quaternion_to_angle_axis(quat)
        angle_axis = angle_axis.view(dim0, dim1, 3)
    return angle_axis


def quaternion_to_angle_axis(quaternion):
    quat = quaternion
    if isinstance(quat, torch.Tensor):
        return __quaternion_to_angle_axis_torch(quaternion)
    else:
        assert isinstance(quat, np.ndarray)
        quat_torch = torch.from_numpy(quat)
        angle_axis_torch = __quaternion_to_angle_axis_torch(quat_torch)
        return angle_axis_torch.numpy()


def __angle_axis_to_quaternion_torch(aa):
    aa = aa.clone()
    if aa.dim() == 1:
        assert aa.size(0) == 3 
        aa = aa.view(1, 3)
        quat = gut.angle_axis_to_quaternion(aa)[0]
    elif aa.dim() == 2:
        assert aa.size(1) == 3
        quat = gut.angle_axis_to_quaternion(aa)
    else:
        assert aa.dim() == 3
        dim0 = aa.size(0)
        dim1 = aa.size(1)
        assert aa.size(2) == 3
        aa = aa.view(dim0*dim1, 3)
        quat = gut.angle_axis_to_quaternion(aa)
        quat = quat.view(dim0, dim1, 4)
    return quat


def angle_axis_to_quaternion(angle_axis):
    aa = angle_axis
    if isinstance(aa, torch.Tensor):
        return __angle_axis_to_quaternion_torch(aa)
    else:
        assert isinstance(aa, np.ndarray)
        aa_torch = torch.from_numpy(aa)
        quat_torch = __angle_axis_to_quaternion_torch(aa_torch)
        return quat_torch.numpy()


def __angle_axis_to_rotation_matrix_torch(aa):
    aa = aa.clone()
    if aa.dim() == 1:
        assert aa.size(0) ==3 
        aa = aa.view(1, 3)
        rotmat = gut.angle_axis_to_rotation_matrix(aa)[0][:3, :3]
    elif aa.dim() == 2:
        assert aa.size(1) == 3
        rotmat = gut.angle_axis_to_rotation_matrix(aa)[:, :3, :3]
    else:
        assert aa.dim() == 3
        dim0 = aa.size(0)
        dim1 = aa.size(1)
        assert aa.size(2) == 3
        aa = aa.view(dim0*dim1, 3)
        rotmat = gut.angle_axis_to_rotation_matrix(aa)
        rotmat = rotmat.view(dim0, dim1, 4, 4)
        rotmat = rotmat[:, :, :3, :3]
    return rotmat


def angle_axis_to_rotation_matrix(angle_axis):
    aa = angle_axis
    if isinstance(aa, torch.Tensor):
        return __angle_axis_to_rotation_matrix_torch(aa)
    else:
        assert isinstance(aa, np.ndarray)
        aa_torch = torch.from_numpy(aa)
        rotmat_torch = __angle_axis_to_rotation_matrix_torch(aa_torch)
        return rotmat_torch.numpy()


def __rotation_matrix_to_angle_axis_torch(rotmat):
    rotmat = rotmat.clone()
    if rotmat.dim() == 2:
        assert rotmat.size(0) == 3
        assert rotmat.size(1) == 3
        rotmat0 = torch.zeros((1, 3, 4))
        rotmat0[0, :, :3] = rotmat
        rotmat0[:, 2, 3] = 1.0
        aa = gut.rotation_matrix_to_angle_axis(rotmat0)[0]
    elif rotmat.dim() == 3:
        dim0 = rotmat.size(0)
        assert rotmat.size(1) == 3
        assert rotmat.size(2) == 3
        rotmat0 = torch.zeros((dim0, 3, 4))
        rotmat0[:, :, :3] = rotmat
        rotmat0[:, 2, 3] = 1.0
        aa = gut.rotation_matrix_to_angle_axis(rotmat0)
    else:
        assert rotmat.dim() == 4
        dim0 = rotmat.size(0)
        dim1 = rotmat.size(1)
        assert rotmat.size(2) == 3
        assert rotmat.size(3) == 3
        rotmat0 = torch.zeros((dim0*dim1, 3, 4))
        rotmat0[:, :, :3] = rotmat.view(dim0*dim1, 3, 3)
        rotmat0[:, 2, 3] = 1.0
        aa = gut.rotation_matrix_to_angle_axis(rotmat0)
        aa = aa.view(dim0, dim1, 3)
    return aa


def rotation_matrix_to_angle_axis(rotmat):
    if isinstance(rotmat, torch.Tensor):
        return __rotation_matrix_to_angle_axis_torch(rotmat)
    else:
        assert isinstance(rotmat, np.ndarray)
        rotmat_torch = torch.from_numpy(rotmat)
        aa_torch = __rotation_matrix_to_angle_axis_torch(rotmat_torch)
        return aa_torch.numpy()
    

def rot6d_to_rotmat(x):
    """Convert 6D rotation representation to 3x3 rotation matrix.
    Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
    Input:
        (B,6) Batch of 6-D rotation representations
    Output:
        (B,3,3) Batch of corresponding rotation matrices
    """
    assert isinstance(x, torch.Tensor), "Current version only supports torch.tensor"

    x = x.view(-1,3,2)
    a1 = x[:, :, 0]
    a2 = x[:, :, 1]
    b1 = F.normalize(a1)
    b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1)
    b3 = torch.cross(b1, b2)
    return torch.stack((b1, b2, b3), dim=-1)


def angle_axis_to_rot6d(aa):
    assert aa.dim() == 2
    assert aa.size(1) == 3
    bs = aa.size(0)

    rotmat = angle_axis_to_rotation_matrix(aa)
    rot6d = rotmat[:, :3, :2]
    return rot6d
imports/frankmocap/mocap_utils/geometry_utils.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import torch
import torch.nn as nn
import numpy as np
import torchgeometry as tgm

__all__ = [
    # flip 3d hand
    "flip_hand_pose",
    "flip_hand_joints_3d",
    # functional api
    "pi",
    "rad2deg",
    "deg2rad",
    "convert_points_from_homogeneous",
    "convert_points_to_homogeneous",
    "angle_axis_to_rotation_matrix",
    "rotation_matrix_to_angle_axis",
    "rotation_matrix_to_quaternion",
    "quaternion_to_angle_axis",
    "angle_axis_to_quaternion",
    "rtvec_to_pose",
    # layer api
    "RadToDeg",
    "DegToRad",
    "ConvertPointsFromHomogeneous",
    "ConvertPointsToHomogeneous",
]


"""Constant with number pi
"""
pi = torch.Tensor([3.14159265358979323846])


def flip_hand_pose(pose):
    if len(pose.shape) == 1:
        pose = pose.reshape(-1, 3)
        pose[:, 1] *= -1
        pose[:, 2] *= -1
        return pose.reshape(-1,)
    else:
        assert len(pose.shape) == 2
        pose[:, 1] *= -1
        pose[:, 2] *= -1
        return pose


def flip_hand_joints_3d(joints_3d):
    assert joints_3d.shape[1] == 3
    assert len(joints_3d.shape) == 2
    rot_mat = np.diag([-1, 1, 1])
    return np.matmul(rot_mat, joints_3d.T).T


def rad2deg(tensor):
    r"""Function that converts angles from radians to degrees.

    See :class:`~torchgeometry.RadToDeg` for details.

    Args:
        tensor (Tensor): Tensor of arbitrary shape.

    Returns:
        Tensor: Tensor with same shape as input.

    Example:
        >>> input = tgm.pi * torch.rand(1, 3, 3)
        >>> output = tgm.rad2deg(input)
    """
    if not torch.is_tensor(tensor):
        raise TypeError("Input type is not a torch.Tensor. Got {}"
                        .format(type(tensor)))

    return 180. * tensor / pi.to(tensor.device).type(tensor.dtype)


def deg2rad(tensor):
    r"""Function that converts angles from degrees to radians.

    See :class:`~torchgeometry.DegToRad` for details.

    Args:
        tensor (Tensor): Tensor of arbitrary shape.

    Returns:
        Tensor: Tensor with same shape as input.

    Examples::

        >>> input = 360. * torch.rand(1, 3, 3)
        >>> output = tgm.deg2rad(input)
    """
    if not torch.is_tensor(tensor):
        raise TypeError("Input type is not a torch.Tensor. Got {}"
                        .format(type(tensor)))

    return tensor * pi.to(tensor.device).type(tensor.dtype) / 180.


def convert_points_from_homogeneous(points):
    r"""Function that converts points from homogeneous to Euclidean space.

    See :class:`~torchgeometry.ConvertPointsFromHomogeneous` for details.

    Examples::

        >>> input = torch.rand(2, 4, 3)  # BxNx3
        >>> output = tgm.convert_points_from_homogeneous(input)  # BxNx2
    """
    if not torch.is_tensor(points):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(points)))
    if len(points.shape) < 2:
        raise ValueError("Input must be at least a 2D tensor. Got {}".format(
            points.shape))

    return points[..., :-1] / points[..., -1:]


def convert_points_to_homogeneous(points):
    r"""Function that converts points from Euclidean to homogeneous space.

    See :class:`~torchgeometry.ConvertPointsToHomogeneous` for details.

    Examples::

        >>> input = torch.rand(2, 4, 3)  # BxNx3
        >>> output = tgm.convert_points_to_homogeneous(input)  # BxNx4
    """
    if not torch.is_tensor(points):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(points)))
    if len(points.shape) < 2:
        raise ValueError("Input must be at least a 2D tensor. Got {}".format(
            points.shape))

    return nn.functional.pad(points, (0, 1), "constant", 1.0)


def angle_axis_to_rotation_matrix(angle_axis):
    """Convert 3d vector of axis-angle rotation to 4x4 rotation matrix

    Args:
        angle_axis (Tensor): tensor of 3d vector of axis-angle rotations.

    Returns:
        Tensor: tensor of 4x4 rotation matrices.

    Shape:
        - Input: :math:`(N, 3)`
        - Output: :math:`(N, 4, 4)`

    Example:
        >>> input = torch.rand(1, 3)  # Nx3
        >>> output = tgm.angle_axis_to_rotation_matrix(input)  # Nx4x4
    """
    def _compute_rotation_matrix(angle_axis, theta2, eps=1e-6):
        # We want to be careful to only evaluate the square root if the
        # norm of the angle_axis vector is greater than zero. Otherwise
        # we get a division by zero.
        k_one = 1.0
        theta = torch.sqrt(theta2)
        wxyz = angle_axis / (theta + eps)
        wx, wy, wz = torch.chunk(wxyz, 3, dim=1)
        cos_theta = torch.cos(theta)
        sin_theta = torch.sin(theta)

        r00 = cos_theta + wx * wx * (k_one - cos_theta)
        r10 = wz * sin_theta + wx * wy * (k_one - cos_theta)
        r20 = -wy * sin_theta + wx * wz * (k_one - cos_theta)
        r01 = wx * wy * (k_one - cos_theta) - wz * sin_theta
        r11 = cos_theta + wy * wy * (k_one - cos_theta)
        r21 = wx * sin_theta + wy * wz * (k_one - cos_theta)
        r02 = wy * sin_theta + wx * wz * (k_one - cos_theta)
        r12 = -wx * sin_theta + wy * wz * (k_one - cos_theta)
        r22 = cos_theta + wz * wz * (k_one - cos_theta)
        rotation_matrix = torch.cat(
            [r00, r01, r02, r10, r11, r12, r20, r21, r22], dim=1)
        return rotation_matrix.view(-1, 3, 3)

    def _compute_rotation_matrix_taylor(angle_axis):
        rx, ry, rz = torch.chunk(angle_axis, 3, dim=1)
        k_one = torch.ones_like(rx)
        rotation_matrix = torch.cat(
            [k_one, -rz, ry, rz, k_one, -rx, -ry, rx, k_one], dim=1)
        return rotation_matrix.view(-1, 3, 3)

    # stolen from ceres/rotation.h
    _angle_axis = torch.unsqueeze(angle_axis, dim=1)
    theta2 = torch.matmul(_angle_axis, _angle_axis.transpose(1, 2))
    theta2 = torch.squeeze(theta2, dim=1)

    # compute rotation matrices
    rotation_matrix_normal = _compute_rotation_matrix(angle_axis, theta2)
    rotation_matrix_taylor = _compute_rotation_matrix_taylor(angle_axis)

    # create mask to handle both cases
    eps = 1e-6
    mask = (theta2 > eps).view(-1, 1, 1).to(theta2.device)
    mask_pos = (mask).type_as(theta2)
    mask_neg = (mask == False).type_as(theta2)  # noqa

    # create output pose matrix
    batch_size = angle_axis.shape[0]
    rotation_matrix = torch.eye(4).to(angle_axis.device).type_as(angle_axis)
    rotation_matrix = rotation_matrix.view(1, 4, 4).repeat(batch_size, 1, 1)
    # fill output matrix with masked values
    rotation_matrix[..., :3, :3] = \
        mask_pos * rotation_matrix_normal + mask_neg * rotation_matrix_taylor
    return rotation_matrix  # Nx4x4


def rtvec_to_pose(rtvec):
    """
    Convert axis-angle rotation and translation vector to 4x4 pose matrix

    Args:
        rtvec (Tensor): Rodrigues vector transformations

    Returns:
        Tensor: transformation matrices

    Shape:
        - Input: :math:`(N, 6)`
        - Output: :math:`(N, 4, 4)`

    Example:
        >>> input = torch.rand(3, 6)  # Nx6
        >>> output = tgm.rtvec_to_pose(input)  # Nx4x4
    """
    assert rtvec.shape[-1] == 6, 'rtvec=[rx, ry, rz, tx, ty, tz]'
    pose = angle_axis_to_rotation_matrix(rtvec[..., :3])
    pose[..., :3, 3] = rtvec[..., 3:]
    return pose


def rotation_matrix_to_angle_axis(rotation_matrix):
    """Convert 3x4 rotation matrix to Rodrigues vector

    Args:
        rotation_matrix (Tensor): rotation matrix.

    Returns:
        Tensor: Rodrigues vector transformation.

    Shape:
        - Input: :math:`(N, 3, 4)`
        - Output: :math:`(N, 3)`

    Example:
        >>> input = torch.rand(2, 3, 4)  # Nx4x4
        >>> output = tgm.rotation_matrix_to_angle_axis(input)  # Nx3
    """
    # todo add check that matrix is a valid rotation matrix
    quaternion = rotation_matrix_to_quaternion(rotation_matrix)
    return quaternion_to_angle_axis(quaternion)


def rotmat3x3_to_angle_axis(init_pred_rotmat):
    """
        init_pred_rotmat: torch.tensor with (1, N,3,3) dimension
        output: (1, N,3)
    """
    assert init_pred_rotmat.shape[0] == 1, "Sould be fixed to handle general batch size.. not confirmed yet"

    device = init_pred_rotmat.device
    jointNum = init_pred_rotmat.shape[1]
    ones = torch.tensor([0,0,1], dtype=torch.float32,).view(1, 3, 1).expand(jointNum, -1, -1).to(device)

    pred_rotmat_hom = torch.cat([ init_pred_rotmat.view(-1, 3, 3),ones ], dim=-1)       #24,3,4
    pred_aa = rotation_matrix_to_angle_axis(pred_rotmat_hom).contiguous().view(1, -1)  #[1,72]
    # tgm.rotation_matrix_to_angle_axis returns NaN for 0 rotation, so manually hack it
    pred_aa[torch.isnan(pred_aa)] = 0.0       #[1,72]
    pred_aa = pred_aa.view(1,jointNum,3) 

    return pred_aa


def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
    """Convert 3x4 rotation matrix to 4d quaternion vector

    This algorithm is based on algorithm described in
    https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201

    Args:
        rotation_matrix (Tensor): the rotation matrix to convert.

    Return:
        Tensor: the rotation in quaternion

    Shape:
        - Input: :math:`(N, 3, 4)`
        - Output: :math:`(N, 4)`

    Example:
        >>> input = torch.rand(4, 3, 4)  # Nx3x4
        >>> output = tgm.rotation_matrix_to_quaternion(input)  # Nx4
    """
    if not torch.is_tensor(rotation_matrix):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(rotation_matrix)))

    if len(rotation_matrix.shape) > 3:
        raise ValueError(
            "Input size must be a three dimensional tensor. Got {}".format(
                rotation_matrix.shape))
    if not rotation_matrix.shape[-2:] == (3, 4):
        raise ValueError(
            "Input size must be a N x 3 x 4  tensor. Got {}".format(
                rotation_matrix.shape))

    rmat_t = torch.transpose(rotation_matrix, 1, 2)

    mask_d2 = rmat_t[:, 2, 2] < eps

    mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1]
    mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1]

    t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
    q0 = torch.stack([rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
                      t0, rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2]], -1)
    t0_rep = t0.repeat(4, 1).t()

    t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
    q1 = torch.stack([rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
                      rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
                      t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1]], -1)
    t1_rep = t1.repeat(4, 1).t()

    t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
    q2 = torch.stack([rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
                      rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2], -1)
    t2_rep = t2.repeat(4, 1).t()

    t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
    q3 = torch.stack([t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
                      rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
                      rmat_t[:, 0, 1] - rmat_t[:, 1, 0]], -1)
    t3_rep = t3.repeat(4, 1).t()

    mask_d2 = mask_d2.int()
    mask_d0_d1 = mask_d0_d1.int()
    mask_d0_nd1 = mask_d0_nd1.int()
    mask_c0 = mask_d2 * mask_d0_d1
    mask_c1 = mask_d2 * (1 - mask_d0_d1)
    mask_c2 = (1 - mask_d2) * mask_d0_nd1
    mask_c3 = (1 - mask_d2) * (1 - mask_d0_nd1)
    mask_c0 = mask_c0.view(-1, 1).type_as(q0)
    mask_c1 = mask_c1.view(-1, 1).type_as(q1)
    mask_c2 = mask_c2.view(-1, 1).type_as(q2)
    mask_c3 = mask_c3.view(-1, 1).type_as(q3)

    q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3
    q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 +  # noqa
                    t2_rep * mask_c2 + t3_rep * mask_c3)  # noqa
    q *= 0.5
    return q


def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor:
    """Convert quaternion vector to angle axis of rotation.

    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h

    Args:
        quaternion (torch.Tensor): tensor with quaternions.

    Return:
        torch.Tensor: tensor with angle axis of rotation.

    Shape:
        - Input: :math:`(*, 4)` where `*` means, any number of dimensions
        - Output: :math:`(*, 3)`

    Example:
        >>> quaternion = torch.rand(2, 4)  # Nx4
        >>> angle_axis = tgm.quaternion_to_angle_axis(quaternion)  # Nx3
    """
    if not torch.is_tensor(quaternion):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(quaternion)))

    if not quaternion.shape[-1] == 4:
        raise ValueError("Input must be a tensor of shape Nx4 or 4. Got {}"
                         .format(quaternion.shape))
    # unpack input and compute conversion
    q1: torch.Tensor = quaternion[..., 1]
    q2: torch.Tensor = quaternion[..., 2]
    q3: torch.Tensor = quaternion[..., 3]
    sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3

    sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta)
    cos_theta: torch.Tensor = quaternion[..., 0]
    two_theta: torch.Tensor = 2.0 * torch.where(
        cos_theta < 0.0,
        torch.atan2(-sin_theta, -cos_theta),
        torch.atan2(sin_theta, cos_theta))

    k_pos: torch.Tensor = two_theta / sin_theta
    k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta)
    k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)

    angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3]
    angle_axis[..., 0] += q1 * k
    angle_axis[..., 1] += q2 * k
    angle_axis[..., 2] += q3 * k
    return angle_axis

# based on:
# https://github.com/facebookresearch/QuaterNet/blob/master/common/quaternion.py#L138


def angle_axis_to_quaternion(angle_axis: torch.Tensor) -> torch.Tensor:
    """Convert an angle axis to a quaternion.

    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h

    Args:
        angle_axis (torch.Tensor): tensor with angle axis.

    Return:
        torch.Tensor: tensor with quaternion.

    Shape:
        - Input: :math:`(*, 3)` where `*` means, any number of dimensions
        - Output: :math:`(*, 4)`

    Example:
        >>> angle_axis = torch.rand(2, 4)  # Nx4
        >>> quaternion = tgm.angle_axis_to_quaternion(angle_axis)  # Nx3
    """
    if not torch.is_tensor(angle_axis):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(angle_axis)))

    if not angle_axis.shape[-1] == 3:
        raise ValueError("Input must be a tensor of shape Nx3 or 3. Got {}"
                         .format(angle_axis.shape))
    # unpack input and compute conversion
    a0: torch.Tensor = angle_axis[..., 0:1]
    a1: torch.Tensor = angle_axis[..., 1:2]
    a2: torch.Tensor = angle_axis[..., 2:3]
    theta_squared: torch.Tensor = a0 * a0 + a1 * a1 + a2 * a2

    theta: torch.Tensor = torch.sqrt(theta_squared)
    half_theta: torch.Tensor = theta * 0.5

    mask: torch.Tensor = theta_squared > 0.0
    ones: torch.Tensor = torch.ones_like(half_theta)

    k_neg: torch.Tensor = 0.5 * ones
    k_pos: torch.Tensor = torch.sin(half_theta) / theta
    k: torch.Tensor = torch.where(mask, k_pos, k_neg)
    w: torch.Tensor = torch.where(mask, torch.cos(half_theta), ones)

    quaternion: torch.Tensor = torch.zeros_like(angle_axis)
    quaternion[..., 0:1] += a0 * k
    quaternion[..., 1:2] += a1 * k
    quaternion[..., 2:3] += a2 * k
    return torch.cat([w, quaternion], dim=-1)

# TODO: add below funtionalities
#  - pose_to_rtvec


# layer api


class RadToDeg(nn.Module):
    r"""Creates an object that converts angles from radians to degrees.

    Args:
        tensor (Tensor): Tensor of arbitrary shape.

    Returns:
        Tensor: Tensor with same shape as input.

    Examples::

        >>> input = tgm.pi * torch.rand(1, 3, 3)
        >>> output = tgm.RadToDeg()(input)
    """

    def __init__(self):
        super(RadToDeg, self).__init__()

    def forward(self, input):
        return rad2deg(input)


class DegToRad(nn.Module):
    r"""Function that converts angles from degrees to radians.

    Args:
        tensor (Tensor): Tensor of arbitrary shape.

    Returns:
        Tensor: Tensor with same shape as input.

    Examples::

        >>> input = 360. * torch.rand(1, 3, 3)
        >>> output = tgm.DegToRad()(input)
    """

    def __init__(self):
        super(DegToRad, self).__init__()

    def forward(self, input):
        return deg2rad(input)


class ConvertPointsFromHomogeneous(nn.Module):
    r"""Creates a transformation that converts points from homogeneous to
    Euclidean space.

    Args:
        points (Tensor): tensor of N-dimensional points.

    Returns:
        Tensor: tensor of N-1-dimensional points.

    Shape:
        - Input: :math:`(B, D, N)` or :math:`(D, N)`
        - Output: :math:`(B, D, N + 1)` or :math:`(D, N + 1)`

    Examples::

        >>> input = torch.rand(2, 4, 3)  # BxNx3
        >>> transform = tgm.ConvertPointsFromHomogeneous()
        >>> output = transform(input)  # BxNx2
    """

    def __init__(self):
        super(ConvertPointsFromHomogeneous, self).__init__()

    def forward(self, input):
        return convert_points_from_homogeneous(input)


class ConvertPointsToHomogeneous(nn.Module):
    r"""Creates a transformation to convert points from Euclidean to
    homogeneous space.

    Args:
        points (Tensor): tensor of N-dimensional points.

    Returns:
        Tensor: tensor of N+1-dimensional points.

    Shape:
        - Input: :math:`(B, D, N)` or :math:`(D, N)`
        - Output: :math:`(B, D, N + 1)` or :math:`(D, N + 1)`

    Examples::

        >>> input = torch.rand(2, 4, 3)  # BxNx3
        >>> transform = tgm.ConvertPointsToHomogeneous()
        >>> output = transform(input)  # BxNx4
    """

    def __init__(self):
        super(ConvertPointsToHomogeneous, self).__init__()

    def forward(self, input):
        return convert_points_to_homogeneous(input)
imports/frankmocap/mocap_utils/geometry_utils_torch.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
##############################################################################
#
# Based on:
# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

"""Timing related functions."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import time


class Timer(object):
    """A simple timer."""

    def __init__(self):
        self.reset()

    def tic(self):
        # using time.time instead of time.clock because time time.clock
        # does not normalize for multithreading
        self.start_time = time.time()

    def toc(self, average=True, bPrint=False,title="Time"):
        self.diff = time.time() - self.start_time
        self.total_time += self.diff
        self.calls += 1
        self.average_time = self.total_time / self.calls
        if average:
            if bPrint:
                # print("Avg Time: {}".format(self.average_time))
                print("{}: {:0.2f} sec/frame, FPS {:0.2f}".format(title, self.diff, 1.0/self.diff))

            return self.average_time
        else:
            if bPrint:
                print("{}: {}, FPS {:0.2f}".format(title, self.diff , 1.0/self.diff))
            return self.diff

    def reset(self):
        self.total_time = 0.
        self.calls = 0
        self.start_time = 0.
        self.diff = 0.
        self.average_time = 0.
imports/frankmocap/mocap_utils/timer.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import numpy as np
from OpenGL.GLUT import *
from OpenGL.GLU import *
from renderer.shaders.framework import *

from renderer.glRenderer import glRenderer

_glut_window = None

'''
#Usage:
    render.set_smpl_mesh(v)     #v for vertex locations in(6890,3)
    render.setBackgroundTexture(rawImg) #Optional BG texture
    render.setWindowSize(rawImg.shape[1], rawImg.shape[0])      #Optional: window size
    render.show_once()
'''

class denseposeRenderer(glRenderer):

    def __init__(self, width=1600, height=1200, name='GL Renderer',
                 program_files=['renderer/shaders/simple140.fs', 'renderer/shaders/simple140.vs'], color_size=1, ms_rate=1):
        glRenderer.__init__(self, width, height, name, program_files, color_size, ms_rate)

        self.densepose_info = self.loadDensepose_info()

        #Densepose Specific 
        self.dp_faces = self.densepose_info['All_Faces']-1 #0~7828
        self.dp_vertexIndices = self.densepose_info['All_vertices']-1    #(1,7829)       #Vertex orders used in denpose info. There are repeated vetices

        #DP color information
        dp_color_seg = self.densepose_info['All_FaceIndices']     #(13774,1)
        dp_color_seg = np.repeat(dp_color_seg,3,axis=1) /100.0#24.0   #(13774,3)
        self.dp_color_seg = np.repeat( dp_color_seg.flatten()[:,None],3,axis=1)        #(41332,3)

        dp_color_U = self.densepose_info['All_U_norm']     #(7289,1)
        dp_color_U =  np.repeat(dp_color_U,3,axis=1)    #(13774,3)
        self.dp_color_U = dp_color_U[self.dp_faces.reshape([-1])]  #(41332,3)

        dp_color_V = self.densepose_info['All_V_norm']     #(7829,3)
        dp_color_V =  np.repeat(dp_color_V,3,axis=1)    #(13774,3)
        self.dp_color_V = dp_color_V[self.dp_faces.reshape([-1])]  #(41332,3)

    #make sure you have: /yourpath/renderer/densepose_uv_data/UV_Processed.mat
    def loadDensepose_info(self, dp_data_path= 'extra_data/densepose_uv_data/UV_Processed.mat'):
        
        #Load densepose data
        import scipy.io as sio
        densepose_info = None
        densepose_info = sio.loadmat(dp_data_path)      #All_FaceIndices (13774), All_Faces(13774), All_U(7829), All_U_norm(7829), All_V(7829), All_V_norm (7829), All_vertices (7829)
        assert densepose_info is not None
        # All_FaceIndices - part labels for each face
        # All_Faces - vertex indices for each face
        # All_vertices - SMPL vertex IDs for all vertices (note that one SMPL vertex can be shared across parts and thus appear in faces with different part labels)
        # All_U - U coordinates for all vertices
        # All_V - V coordinates for all vertices
        # All_U_norm - normalization factor for U coordinates to map them to [0, 1] interval
        # All_V_norm - normalization factor for V coordinates to map them to [0, 1] interval
        # vertexColor = densepose_info['All_U_norm']*255
        # vertexColor = np.zeros((v.shape[1], 3))
        # vertexColor[:,0] = densepose_info['All_U_norm'][:v.shape[1]].flatten()       #(6890,3)
        # vertexColor[:,1] = densepose_info['All_V_norm'][:v.shape[1]].flatten()       #(6890,3)

        # # faces = smplWrapper.f
        # v =v[0]  #(6890,3)
        # dp_vertex = v[densepose_info['All_vertices']-1]  #(1,7829,3)        #Considering repeatation
        # faces =densepose_info['All_Faces']-1 #0~7828
        # # vertexColor = densepose_info['All_FaceIndices']     #(13774,1)
        # # vertexColor = np.repeat(vertexColor,3,axis=1) /24.0   #(13774,3)

        # # vertexColor = densepose_info['All_U_norm']     #(13774,1)
        # vertexColor = densepose_info['All_V_norm']     #(13774,1)
        # vertexColor =  np.repeat(vertexColor,3,axis=1) 

        # # vertexColor[vertexColor!=2]*=0
        # vertexColor[vertexColor==2]=24
        return densepose_info


    #vertice: (6890,3)
    #colormode: ['seg', 'u', 'v']
    def set_mesh(self, vertices, _):
        
        if vertices.dtype != np.dtype('float64'):
            vertices = vertices.astype(np.float64)      #Should be DOUBLE

        #Change the vertex and 
        dp_vertex = vertices[self.dp_vertexIndices][0]  #(7829,3)        #Considering repeatation

        # if colormode=='seg': #segment   
        #     self.color_data = self.dp_color_seg
        # elif colormode=='v':
        #     self.color_data[:,1] = self.dp_color_V
        # elif colormode=='u':
        #     self.color_data = self.dp_color_U       #(41322,3)
        # else:
        #     assert False

        self.color_data = self.dp_color_U       #(41322,3)       
        self.color_data[:,1] = self.dp_color_V[:,1]
        self.color_data[:,2] = self.dp_color_seg[:,2]

        self.vertex_data = dp_vertex[self.dp_faces.reshape([-1])]       #(41322,3)
        self.vertex_dim = self.vertex_data.shape[1]
        self.n_vertices = self.vertex_data.shape[0]

        glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer)
        glBufferData(GL_ARRAY_BUFFER, self.vertex_data, GL_STATIC_DRAW)

        glBindBuffer(GL_ARRAY_BUFFER, self.color_buffer)
        glBufferData(GL_ARRAY_BUFFER, self.color_data, GL_STATIC_DRAW)

        glBindBuffer(GL_ARRAY_BUFFER, 0)


    def add_mesh(self, vertices, _, color=None):
        """
        Concatenate the new mesh data to self.vertex_data  (as if a single giant mesh)

        Args:
            input_vertices (np.ndarray): (verNum, 3).
            input_faces (np.ndarray): (faceNum, 3).
        """

        if vertices.dtype != np.dtype('float64'):
            vertices = vertices.astype(np.float64)      #Should be DOUBLE
        
        dp_vertex = vertices[self.dp_vertexIndices][0]  #(7829,3)        #Considering repeatation

        color_data = self.dp_color_U       #(41322,3)       
        color_data[:,1] = self.dp_color_V[:,1]
        color_data[:,2] = self.dp_color_seg[:,2]

        if self.vertex_data is None:
            self.vertex_data = dp_vertex[self.dp_faces.reshape([-1])]       #(41322,3)
            self.color_data = color_data
        
        else:       #Add the data
            input_vertices = dp_vertex[self.dp_faces.reshape([-1])]       #(41322,3)
            self.vertex_data = np.concatenate( (self.vertex_data, input_vertices), axis=0)    #(6870,3)
            self.color_data = np.concatenate( (self.color_data, color_data), axis=0)    #(6870,3)

        self.vertex_dim = self.vertex_data.shape[1]
        self.n_vertices = self.vertex_data.shape[0]

        glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer)
        glBufferData(GL_ARRAY_BUFFER, self.vertex_data, GL_STATIC_DRAW)

        glBindBuffer(GL_ARRAY_BUFFER, self.color_buffer)
        glBufferData(GL_ARRAY_BUFFER, self.color_data, GL_STATIC_DRAW)


        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, self.index_buffer)
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, self.meshindex_data, GL_STATIC_DRAW)

        glBindBuffer(GL_ARRAY_BUFFER, 0)
imports/frankmocap/renderer/denseposeRenderer.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import numpy as np
from OpenGL.GLUT import *
from OpenGL.GLU import *
from OpenGL.GL import *
from renderer.shaders.framework import createProgram, loadShader

from renderer.render_utils import ComputeNormal

import cv2

_glut_window = None

class glRenderer:

    def __init__(self, width=640, height=480, name='GL Renderer',
                 program_files=['renderer/shaders/simple140.fs', 'renderer/shaders/simple140.vs'],
                 color_size=1, ms_rate=1,
                 # frankmocap_path='/home/hansookwan/MAIN/Ego4d_Codes/github_codes/frankmocap/'
                 ):  ## ADDED
        self.width = width
        self.height = height
        self.name = name
        self.display_mode = GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH | GLUT_MULTISAMPLE
        self.use_inverse_depth = False

        global _glut_window
        if _glut_window is None:
            glutInit()
            glutInitDisplayMode(self.display_mode)
            glutInitWindowSize(self.width, self.height)
            glutInitWindowPosition(0, 0)
            _glut_window = glutCreateWindow("GL_Renderer")


            glEnable(GL_DEPTH_CLAMP)
            glEnable(GL_DEPTH_TEST)

            glClampColor(GL_CLAMP_READ_COLOR, GL_FALSE)
            glClampColor(GL_CLAMP_FRAGMENT_COLOR, GL_FALSE)
            glClampColor(GL_CLAMP_VERTEX_COLOR, GL_FALSE)


        self.program = None
        self.initShaderProgram(program_files)

        # Init Uniform variables
        self.model_mat_unif = glGetUniformLocation(self.program, 'ModelMat')
        self.persp_mat_unif = glGetUniformLocation(self.program, 'PerspMat')

        self.model_view_matrix = None
        self.projection_matrix = None

        self.vertex_buffer = glGenBuffers(1)
        self.color_buffer = glGenBuffers(1)
        self.normal_buffer = glGenBuffers(1)

        self.index_buffer = glGenBuffers(1)     #for Mesh face indices. Without this vertices should be repeated and ordered (3x times bigger)

        #Create Background Texture
        glPixelStorei(GL_UNPACK_ALIGNMENT, 1)       #So that texture doesnt have to be power of 2
        self.backgroundTextureID = glGenTextures(1)
        K = np.array([[2000, 0, 960],[0, 2000, 540],[0,0,1]])       #MTC default camera. for 1920 x 1080 input image
        self.camView_K = K
        
        # Inner storage for buffer data
        self.vertex_data = None
        self.vertex_dim = None
        self.n_vertices = None

        #Variables for view change
        self.m_xTrans = 0.
        self.m_yTrans = 0.
        self.m_zTrans = 0.
        self.m_zoom = 378
        self.m_xRotate = 59.
        self.m_yRotate = 0.
        self.m_zRotate = 0.
        self.m_xrot = 0.0
        self.m_yrot = 0.0

        #Camera view
        self.m_viewmode = "cam"

        #To compute sideview
        self.m_meshCenter = None

        #Option
        self.bOffscreenMode = False
        self.bAntiAliasing = True #apply anti aliasing

        #Textures
        self.data_texture = None


        #Visualization option
        self.bShowBackground = False
        self.bShowFloor = False

        self.nearPlane = 1  # original
        self.farPlane = 10000.0

        self.counter=1


        self.bOrthoCam = True
        
        glutMouseFunc(self.mouseButton)
        glutMotionFunc(self.mouseMotion)
        glutDisplayFunc(self.display)
        glutReshapeFunc(self.reshape)
    
    def initShaderProgram(self, program_files):
         # Init shader programs
        shader_list = []
        for program_file in program_files:
            _, ext = os.path.splitext(program_file)
            if ext == '.vs':
                shader_list.append(loadShader(GL_VERTEX_SHADER, program_file))
            elif ext == '.fs':
                shader_list.append(loadShader(GL_FRAGMENT_SHADER, program_file))
            elif ext == '.gs':
                shader_list.append(loadShader(GL_GEOMETRY_SHADER, program_file))

        if self.program is not None:
            glDeleteProgram(self.program)

        self.program = createProgram(shader_list)
        for shader in shader_list:
            glDeleteShader(shader)

    def reshape(self,width, height):

        if self.bOffscreenMode ==False:     #Offscreen mode doesn't allow reshape
            self.width = width
            self.height = height
            glViewport(0, 0, self.width, self.height)
            print("reshape: {}, {}".format(self.width, self.height))

        # glViewport(0, 0, 500,500)

    def setWindowSize(self,new_width, new_height):

        if new_width != self.width or  new_height!=self.height:
            self.width = new_width
            self.height =new_height
            glutReshapeWindow(self.width,self.height)

        #Neet to refresh opengl to apply resizing (the number of required iterations is a bit random)
        # glutPostRedisplay()
        glutMainLoopEvent()
        iterNum =0
        while True:
            curWidth =glutGet(GLUT_WINDOW_WIDTH)
            curHeight =  glutGet(GLUT_WINDOW_HEIGHT)
            # print("curWindSize: {}, {}".format(curWidth, curHeight))

            if curWidth == new_width and curHeight == new_height:
                break
            glutPostRedisplay()
            glutMainLoopEvent()
            iterNum+=1  
            if iterNum>20:
                print("Wraning: Cannot resize the gl window")
                break

        # print("{} refreshing is done to resize window".format(iterNum))

       
    def setViewportSize(self,new_width, new_height):
        if new_width != self.width or  new_height!=self.height:
            self.width = new_width
            self.height =new_height
        glViewport(0, 0, self.width, self.height)
        # print("{}, {}".format(self.width, self.height))

    def setCameraViewOrth(self):

        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()

        # texHeight,texWidth =  self.data_texture.shape[:2]
        # # texHeight,texWidth =   1024, 1024
        # texHeight*=0.5
        # texWidth*=0.5
        texHeight,texWidth =  self.height, self.width
        # texHeight,texWidth =   1024, 1024
        texHeight*=0.5
        texWidth*=0.5
        # texHeight *=BACKGROUND_IMAGE_PLANE_DEPTH
        # texWidth *=BACKGROUND_IMAGE_PLANE_DEPTH

        glOrtho(-texWidth, texWidth, -texHeight, texHeight, -1500, 1500)

        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()
        gluLookAt(0, 0, 0, 0, 0, 1, 0, -1, 0)
        # glMultMatrixd(camMatrix.T)
    

    # Show the world in a camera cooridnate (defined by K)
    def setCameraView(self):
     
        invR = np.eye(3)
        invT = np.zeros((3, 1))
        # invT[2] = 400
        camMatrix = np.hstack((invR, invT))
        # denotes camera matrix, [R|t]
        camMatrix = np.vstack((camMatrix, [0, 0, 0, 1]))

        if self.camView_K is None:
            # print(
            #     "## Warning: no K is set, so I use a default cam parameter defined for MTC"
            # )
            # setCamView_K_DefaultForMTC()
            assert False
        K = self.camView_K.copy()

        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        # Kscale = 1920.0/g_Width
        Kscale = 1.0  # 1920.0/g_Width        :: why do we need this?
        K = K / Kscale
        ProjM = np.zeros((4, 4))
        ProjM[0, 0] = 2 * K[0, 0] / self.width
        ProjM[0, 2] = (self.width - 2 * K[0, 2]) / self.width
        ProjM[1, 1] = 2 * K[1, 1] / self.height
        ProjM[1, 2] = (-self.height + 2 * K[1, 2]) / self.height

        ProjM[2, 2] = (-self.farPlane - self.nearPlane) / (self.farPlane - self.nearPlane)
        ProjM[2, 3] = -2 * self.farPlane * self.nearPlane / (self.farPlane - self.nearPlane)
        ProjM[3, 2] = -1

        glLoadMatrixd(ProjM.T)
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()
        gluLookAt(0, 0, 0, 0, 0, 1, 0, -1, 0)
        glMultMatrixd(camMatrix.T)

    def setWorldCenterBySceneCenter(self):
        self.m_meshCenter  =np.mean(self.vertex_data,axis=0)

        meshMin = np.min(self.vertex_data,axis=0)
        meshMax = np.max(self.vertex_data,axis=0)

        self.meshWidth = max(meshMax-meshMin)



        self.m_xTrans = -self.m_meshCenter[0]# +self.counter
        self.m_yTrans = -self.m_meshCenter[1]
        self.m_zTrans = -self.m_meshCenter[2]

        self.m_zoom = min(self.m_zoom, 1000)
        self.m_zoom = 120 * max(self.meshWidth,100)/100

    def setZoom(self, z):
        self.m_zoom = z

    def setSideView(self):
        
        if self.m_meshCenter is None:
            self.setWorldCenterBySceneCenter()

        # self.m_zoom =400

        self.m_yRotate = 0 
        self.m_xRotate = 90 #+ self.counter


        glMatrixMode(GL_MODELVIEW)

        #Zoom out first
        glTranslatef(0,0,self.m_zoom)
        print(f"Zoom: {self.m_zoom}")
        
        #Rotate (to this end, around the center point)
        glRotatef( -self.m_yRotate, 1.0, 0.0, 0.0)
        glRotatef( -self.m_xRotate, 0.0, 1.0, 0.0)
        # glRotatef( self.m_zRotate, 0.0, 0.0, 1.0)

        # print(f"{self.m_xTrans}, {self.m_yTrans}, {self.m_zTrans}, {self.m_zoom}")
        # glTranslatef(0,0,self.m_zoom)
        glTranslatef( self.m_xTrans,  self.m_yTrans, self.m_zTrans)


    def setViewAngle(self,azimuthAng, elevationAng):
        """
        azimuthAng: 90 for sideview
        elevationAng: 90 for topview
        """
        self.m_xRotate = azimuthAng
        self.m_yRotate = -elevationAng

    def setFree3DView(self):
        if self.m_meshCenter is None:
            self.setWorldCenterBySceneCenter()

        glMatrixMode(GL_MODELVIEW)
        # glLoadIdentity()
        # gluLookAt(0, 0, 0, 0, 0, 1, 0, -1, 0)

        #Zoom out first
        glTranslatef(0,0,self.m_zoom)

        #Rotate (to this end, around the center point)
        glRotatef( -self.m_yRotate, 1.0, 0.0, 0.0)
        glRotatef( -self.m_xRotate, 0.0, 1.0, 0.0)
        # glRotatef( self.m_zRotate, 0.0, 0.0, 1.0)

        # print(f"{self.m_xTrans}, {self.m_yTrans}, {self.m_zTrans}, {self.m_zoom}")
        # glTranslatef(0,0,self.m_zoom)
        glTranslatef( self.m_xTrans,  self.m_yTrans, self.m_zTrans)
        

    def showBackground(self, bShow):
        self.bShowBackground = bShow

    def setBackgroundTexture(self,img):
        self.data_texture = img
    
    
    # 3x3 intrinsic camera matrix
    def setCamView_K(self, K):
        self.camView_K = K

    def setOrthoCamera(self, bOrtho = True):
        self.bOrthoCam = bOrtho

    def drawBackgroundOrth(self):

        if self.data_texture is None:
            return

        glDisable(GL_CULL_FACE)
        glDisable(GL_DEPTH_TEST)
        glDisable(GL_LIGHTING)
        glEnable(GL_TEXTURE_2D)

        # glUseProgram(0)

        glBindTexture(GL_TEXTURE_2D, self.backgroundTextureID)
        # glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, 1920, 1080, 0, GL_RGB, GL_UNSIGNED_BYTE, self.data_texture)
        # glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, 480, 640, 0, GL_RGB, GL_UNSIGNED_BYTE, self.data_texture.data)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, self.data_texture.shape[1], self.data_texture.shape[0], 0, GL_BGR, GL_UNSIGNED_BYTE, self.data_texture.data)
        texHeight,texWidth =   self.data_texture.shape[:2]
        texHeight*=0.5
        texWidth*=0.5

        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)

        glPolygonMode(GL_FRONT_AND_BACK, GL_FILL)
        glBegin(GL_QUADS)
        glColor3f(1.0, 1.0, 1.0)
        # d = BACKGROUND_IMAGE_PLANE_DEPTH
        d = 10

        glTexCoord2f(0, 0)
        # #Eigen::Map<Eigen::Matrix<double, 3, 3, Eigen::RowMajor>> K(m_options.m_pK);
        P = np.array([-texWidth, -texHeight, d])
        # P = np.matmul(K_inv,P)
        # P = P / P[2]
        glVertex3f(P[0] , P[1] , P[2] );  # K^{-1} [0, 0, 1]^T

        glTexCoord2f(1, 0)
        # P = [1920, 0, 1]
        P = [texWidth, -texHeight, d]
        glVertex3f(P[0] , P[1] , P[2] );  # K^{-1} [0, 0, 1]^T

        glTexCoord2f(1, 1)
        # P = [1920, 1080, 1]
        P = [texWidth, texHeight, d]
        # P = np.matmul(K_inv,P)
        # P = P / P[2]
        # glVertex3f(P[0] * d, P[1] * d, P[2] * d)
        glVertex3f(P[0] , P[1] , P[2] );  # K^{-1} [0, 0, 1]^T

        glTexCoord2f(0, 1)
        # P = [0, 1080, 1]
        P = [-texWidth, texHeight, d]
        # glVertex3f(P[0] * d, P[1] * d, P[2] * d)
        glVertex3f(P[0] , P[1] , P[2] );  # K^{-1} [0, 0, 1]^T
        glEnd()

        glEnable(GL_LIGHTING)
        glEnable(GL_CULL_FACE)
        glEnable(GL_DEPTH_TEST)
        glDisable(GL_TEXTURE_2D)
        
    def drawBackgroundPersp(self):

        if self.data_texture is None:
            return

        # if self.camView_K is None:
        #     self.setCamView_K_DefaultForMTC()

        K_inv = np.linalg.inv(self.camView_K)

        glDisable(GL_CULL_FACE)
        glDisable(GL_DEPTH_TEST)
        glDisable(GL_LIGHTING)
        glEnable(GL_TEXTURE_2D)

        # glUseProgram(0)

        glBindTexture(GL_TEXTURE_2D, self.backgroundTextureID)
        # glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, 1920, 1080, 0, GL_RGB, GL_UNSIGNED_BYTE, self.data_texture)
        # glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, 480, 640, 0, GL_RGB, GL_UNSIGNED_BYTE, self.data_texture.data)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, self.data_texture.shape[1], self.data_texture.shape[0], 0, GL_BGR, GL_UNSIGNED_BYTE, self.data_texture.data)
        texHeight,texWidth =   self.data_texture.shape[:2]

        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)


        glPolygonMode(GL_FRONT_AND_BACK, GL_FILL)
        glBegin(GL_QUADS)
        glColor3f(1.0, 1.0, 1.0)
        BACKGROUND_IMAGE_PLANE_DEPTH=3000
        d = BACKGROUND_IMAGE_PLANE_DEPTH

        glTexCoord2f(0, 0)
        # #Eigen::Map<Eigen::Matrix<double, 3, 3, Eigen::RowMajor>> K(m_options.m_pK);
        P = np.array([0, 0, 1])
        P = np.matmul(K_inv,P)
        P = P / P[2]
        glVertex3f(P[0] * d, P[1] * d, P[2] * d);  # K^{-1} [0, 0, 1]^T

        glTexCoord2f(1, 0)
        # P = [1920, 0, 1]
        P = [texWidth, 0, 1]

        P = np.matmul(K_inv,P)
        P = P / P[2]
        glVertex3f(P[0] * d, P[1] * d, P[2] * d);  # K^{-1} [0, 0, 1]^T

        glTexCoord2f(1, 1)
        # P = [1920, 1080, 1]
        P = [texWidth, texHeight, 1]
        P = np.matmul(K_inv,P)
        P = P / P[2]
        glVertex3f(P[0] * d, P[1] * d, P[2] * d)

        glTexCoord2f(0, 1)
        # P = [0, 1080, 1]
        P = [0, texHeight, 1]
        P = np.matmul(K_inv,P)
        P = P / P[2]
        glVertex3f(P[0] * d, P[1] * d, P[2] * d)
        glEnd()

        glEnable(GL_LIGHTING)
        glEnable(GL_CULL_FACE)
        glEnable(GL_DEPTH_TEST)
        glDisable(GL_TEXTURE_2D)


    def mouseButton(self, button, state, x, y):
        # global self.m_action, self.m_xMousePtStart, self.m_yMousePtStart
        if (button==GLUT_LEFT_BUTTON):
            if (glutGetModifiers() == GLUT_ACTIVE_SHIFT):
                self.m_action = "TRANS"
            else:
                self.m_action = "MOVE_EYE"
        #elif (button==GLUT_MIDDLE_BUTTON):
        #    action = "TRANS"
        elif (button==GLUT_RIGHT_BUTTON):
            self.m_action = "ZOOM"
        self.m_xMousePtStart = x
        self.m_yMousePtStart = y

    def mouseMotion(self, x, y):
        # global self.m_zoom, self.m_xMousePtStart, self.m_yMousePtStart, self.m_xRotate, self.m_yRotate, self.m_zRotate, self.m_xTrans, self.m_zTrans
        if (self.m_action=="MOVE_EYE"):
            self.m_xRotate += x - self.m_xMousePtStart
            self.m_yRotate -= y - self.m_yMousePtStart
        elif (self.m_action=="MOVE_EYE_2"):
            self.m_zRotate += y - self.m_yMousePtStart
        elif (self.m_action=="TRANS"):
            self.m_xTrans += x - self.m_xMousePtStart
            self.m_zTrans += y - self.m_yMousePtStart
        elif (self.m_action=="ZOOM"):
            self.m_zoom -= y - self.m_yMousePtStart
            # print(self.m_zoom)
        else:
            print("unknown action\n", self.m_action)
        self.m_xMousePtStart = x
        self.m_yMousePtStart = y

        # print ('xTrans {},  yTrans {}, zoom {} xRotate{} yRotate {} zRotate {}'.format(self.m_xTrans,  self.m_yTrans,  self.m_zoom,  self.m_xRotate,  self.m_yRotate,  self.m_zRotate))
        glutPostRedisplay()

    def set_viewpoint(self, projection, model_view):
        self.projection_matrix = projection
        self.model_view_matrix = model_view

    def get_screen_color_fbgra(self):
        glReadBuffer(GL_BACK)   #GL_BACK is Default in double buffering
        data = glReadPixels(0, 0, self.width, self.height, GL_RGBA, GL_FLOAT, outputType=None)
        rgb = data.reshape(self.height, self.width, -1)
        rgb = np.flip(rgb, 0)

        rgb = cv2.cvtColor(rgb, cv2.COLOR_RGBA2BGRA)
        return rgb
    
    def get_screen_color_ibgr(self):
        glReadBuffer(GL_BACK)   #GL_BACK is Default in double buffering
        data = glReadPixels(0, 0, self.width, self.height, GL_RGBA, GL_FLOAT, outputType=None)
        rgb = data.reshape(self.height, self.width, -1)
        rgb = np.flip(rgb, 0)

        rgb = cv2.cvtColor(rgb, cv2.COLOR_RGBA2BGR)
        rgb = (rgb*255).astype(np.uint8)
        return rgb
    
    def setCameraViewMode(self, viewmode="cam"):
        """
        viewmode
            cam: cam view
            side: side view
            free: free view point
        """
        self.m_viewmode = viewmode

    def get_z_value(self):
        glReadBuffer(GL_BACK)   #GL_BACK is Default in double buffering

        data = glReadPixels(0, 0, self.width, self.height, GL_DEPTH_COMPONENT, GL_FLOAT, outputType=None)
        glBindFramebuffer(GL_FRAMEBUFFER, 0)
        z = data.reshape(self.height, self.width)
        z = np.flip(z, 0)
        return z
    
    def drawFloor(self):
        glDisable(GL_LIGHTING)

        glPolygonMode(GL_FRONT, GL_FILL)
        glPolygonMode(GL_BACK, GL_FILL)
        # glPolygonMode(GL_FRONT, GL_FILL)
        gridNum = 10
        width = 200
        halfWidth =width/2
        # g_floorCenter = np.array([0,0.5,0])

        g_floorCenter = np.array([100,100,0])
        g_floorAxis1 = np.array([1,0,0])
        g_floorAxis2 = np.array([0,0,1])

        origin = g_floorCenter - g_floorAxis1*(width*gridNum/2 ) - g_floorAxis2*(width*gridNum/2)
        axis1 =  g_floorAxis1 * width
        axis2 =  g_floorAxis2 * width
        for y in range(gridNum+1):
            for x in range(gridNum+1):

                if (x+y) % 2 ==0:
                    glColor(1.0,1.0,1.0,1.0) #white
                else:
                    # glColor(0.95,0.95,0.95,0.3) #grey
                    glColor(0.3,0.3,0.3,0.5) #grey

                p1 = origin + axis1*x + axis2*y
                p2 = p1+ axis1
                p3 = p1+ axis2
                p4 = p1+ axis1 + axis2

                glBegin(GL_QUADS)
                glVertex3f(   p1[0], p1[1], p1[2])
                glVertex3f(   p2[0], p2[1], p2[2])
                glVertex3f(   p4[0], p4[1], p4[2])
                glVertex3f(   p3[0], p3[1], p3[2])
                glEnd()

        glEnable(GL_LIGHTING)
        
                
    def set_mesh(self, input_vertices, input_faces, color=None):
          
        #Compute normal
        if True:#bComputeNormal:
        # print("## setMeshData: Computing face normals automatically.")
            vertices_temp = input_vertices[np.newaxis,:,:]
            input_normal = ComputeNormal(vertices_temp,input_faces) #output: (N, 18540, 3)

        if input_vertices.dtype != np.dtype('float64'):
            input_vertices = input_vertices.astype(np.float64)      #Should be DOUBLE

        #Change the vertex and 
        # dp_vertex = vertices[self.dp_vertexIndices][0]  #(7829,3)        #Considering repeatation

        # if colormode=='normal': #segment   
        #     self.color_data = self.dp_color_seg
        # else:
        #     assert False
        
        input_faces_ = input_faces.flatten()

        if False:#Without index buffer
            self.vertex_data = input_vertices[input_faces_,:]   
            self.normal_data = input_normal[0][input_faces_,:]  

        else: #With index buffer
            self.vertex_data = input_vertices   #(6870,3)
            self.normal_data = input_normal[0]  #(18540,3)

            if color is None:
                self.color_data = np.ones(input_vertices.shape)   #(6870,3)
            else:
                self.color_data = np.tile(color, (input_vertices.shape[0], 1) )   #(6870,3)

            self.meshindex_data = input_faces_
            self.meshindex_data = self.meshindex_data.astype(np.int32)      #Should be DOUBLE
            

        self.vertex_dim = self.vertex_data.shape[1]
        self.n_vertices = self.vertex_data.shape[0]

        glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer)
        glBufferData(GL_ARRAY_BUFFER, self.vertex_data, GL_STATIC_DRAW)

        glBindBuffer(GL_ARRAY_BUFFER, self.normal_buffer)
        glBufferData(GL_ARRAY_BUFFER, self.normal_data, GL_STATIC_DRAW)

        glBindBuffer(GL_ARRAY_BUFFER, self.color_buffer)
        glBufferData(GL_ARRAY_BUFFER, self.color_data, GL_STATIC_DRAW)

        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, self.index_buffer)
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, self.meshindex_data, GL_STATIC_DRAW)

        glBindBuffer(GL_ARRAY_BUFFER, 0)

    def clear_mesh(self):
        """
        Clear up all mesh 
        """
        self.vertex_data = None
        self.normal_data = None
        self.meshindex_data = None

    def add_mesh(self, input_vertices, input_faces, color=None):
        """
        Concatenate the new mesh data to self.vertex_data  (as if a single giant mesh)

        Args:
            input_vertices (np.ndarray): (verNum, 3).

            input_faces (np.ndarray): (faceNum, 3).
        """

          
        #Compute normal
        if True:#bComputeNormal:
            vertices_temp = input_vertices[np.newaxis,:,:]
            input_normal = ComputeNormal(vertices_temp,input_faces) #output: (N, 18540, 3)

        if input_vertices.dtype != np.dtype('float64'):
            input_vertices = input_vertices.astype(np.float64)      #Should be DOUBLE

        #Change the vertex and 
        # dp_vertex = vertices[self.dp_vertexIndices][0]  #(7829,3)        #Considering repeatation
        # if colormode=='normal': #segment   
        #     self.color_data = self.dp_color_seg
        # else:
        #     assert False
        
        input_faces_ = input_faces.flatten()

        if self.vertex_data is None:
            self.vertex_data = input_vertices   #(6870,3)
            self.normal_data = input_normal[0]  #(6870,3)

            if color is None:
                self.color_data = np.ones(input_vertices.shape)   #(6870,3)
            else:
                self.color_data = np.tile(color, (input_vertices.shape[0], 1) )   #(6870,3)

            self.meshindex_data = input_faces_
            self.meshindex_data = self.meshindex_data.astype(np.int32)      #Should be DOUBLE
        
        else:       #Add the data
            existingVerNum = self.vertex_data.shape[0]
            input_faces_ += existingVerNum
            input_faces_ = input_faces_.astype(np.int32)
            self.meshindex_data = np.concatenate( (self.meshindex_data, input_faces_), axis=0)  #(Vx3,)
            # self.meshindex_data = self.meshindex_data.astype(np.int32)      #Should be DOUBLE

            self.vertex_data = np.concatenate( (self.vertex_data, input_vertices), axis=0)    #(6870,3)
            self.normal_data =  np.concatenate( (self.normal_data, input_normal[0]), axis=0)     #(6870,3)

            if color is None:
                color_data = np.ones(input_vertices.shape)   #(6870,3)
            else:
                color_data = np.tile(color, (input_vertices.shape[0], 1) )   #(6870,3)
            
            self.color_data =  np.concatenate( (self.color_data, color_data), axis=0)     #(6870,3)
            
            
        self.vertex_dim = self.vertex_data.shape[1]
        self.n_vertices = self.vertex_data.shape[0]

        glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer)
        glBufferData(GL_ARRAY_BUFFER, self.vertex_data, GL_STATIC_DRAW)

        glBindBuffer(GL_ARRAY_BUFFER, self.normal_buffer)
        glBufferData(GL_ARRAY_BUFFER, self.normal_data, GL_STATIC_DRAW)

        glBindBuffer(GL_ARRAY_BUFFER, self.color_buffer)
        glBufferData(GL_ARRAY_BUFFER, self.color_data, GL_STATIC_DRAW)


        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, self.index_buffer)
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, self.meshindex_data, GL_STATIC_DRAW)

        glBindBuffer(GL_ARRAY_BUFFER, 0)


    def drawMesh(self):
        # self.draw_init()

        glColor3f(1,1,0)
        glUseProgram(self.program)
        
        mvMat = glGetFloatv(GL_MODELVIEW_MATRIX)
        pMat = glGetFloatv(GL_PROJECTION_MATRIX)
        # mvpMat = pMat*mvMat

        self.model_view_matrix = mvMat
        self.projection_matrix = pMat

        # glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, self.model_view_matrix.transpose())
        # glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, self.projection_matrix.transpose())
        glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, self.model_view_matrix)
        glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, self.projection_matrix)

        # Handle vertex buffer
        glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer)
        glEnableVertexAttribArray(0)
        glVertexAttribPointer(0, self.vertex_dim, GL_DOUBLE, GL_FALSE, 0, None)

        # # Handle color buffer
        glBindBuffer(GL_ARRAY_BUFFER, self.color_buffer)
        glEnableVertexAttribArray(1)
        glVertexAttribPointer(1, 3, GL_DOUBLE, GL_FALSE, 0, None)

        glDrawArrays(GL_TRIANGLES, 0, self.n_vertices)

        glDisableVertexAttribArray(0)

        glBindBuffer(GL_ARRAY_BUFFER, 0)

        glUseProgram(0)

    def display(self):
        # First we draw a scene.
        # Notice the result is stored in the texture buffer.
        # Then we return to the default frame buffer since we will display on the screen.
        # glBindFramebuffer(GL_FRAMEBUFFER, 0)

        glUseProgram(0)
        glClear(GL_COLOR_BUFFER_BIT|GL_DEPTH_BUFFER_BIT)

        #Some anti-aliasing code (seems not working, though)
        if self.bAntiAliasing:
            glEnable (GL_BLEND)
            glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
            glEnable (GL_LINE_SMOOTH)
            glHint (GL_LINE_SMOOTH_HINT, GL_NICEST)
            glEnable(GL_POLYGON_SMOOTH)
            glEnable(GL_MULTISAMPLE)
            # glHint(GL_MULTISAMPLE_FILTER_HINT_NV, GL_NICEST)
        else:
            glDisable(GL_BLEND)
            glDisable(GL_MULTISAMPLE)

        # Set up viewing transformation, looking down -Z axis
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()
        #gluLookAt(0, 0, -g_fViewDistance, 0, 0, 0, -.1, 0, 0)   #-.1,0,0
        gluLookAt(0,0,0, 0, 0, 1, 0, -1, 0)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        #gluPerspective(zoom, float(g_Width)/float(g_Height), g_nearPlane, g_farPlane)
        gluPerspective(65, float(self.width)/float(self.height), 10, 6000)         # This should be called here (not in the reshpe)
        # glMatrixMode(GL_MODELVIEW)
        # Render the scene

        if self.m_viewmode =="cam":
            if self.bOrthoCam:
                self.setCameraViewOrth()
            else:
                self.setCameraView()

        elif self.m_viewmode =="side":
            # self.setCameraViewOrth()
            self.setSideView()
            # self.setFree3DView()
            # 
        elif self.m_viewmode =="free":
            self.setFree3DView()
        else:
            assert False
            
        glClearColor(1.0, 1.0, 1.0, 1.0)
        # glClearColor(0.0, 0.0, 0.0, 1.0)
        glClear(GL_COLOR_BUFFER_BIT|GL_DEPTH_BUFFER_BIT)
    
        # self.RenderDomeFloor()
        # glColor3f(0,1,0)
        # glTranslated(500,0,0)
        # glutSolidSphere(10, 10, 200)
        # glPointSize(30)
        # glBegin(GL_POINTS)
        # # glVertex3f(0,0,0)
        # if self.m_meshCenter is not None:
        #     glVertex3fv(self.m_meshCenter)
        # glEnd()
        # glutSolidTeapot(100, 10, 10)

        if self.bShowBackground:
            if self.bOrthoCam:
                self.drawBackgroundOrth()
            else:
                self.drawBackgroundPersp()

        if self.bShowFloor:
            self.drawFloor()

        self.drawMesh()

        glutSwapBuffers()
        glutPostRedisplay()

    def offscreenMode(self, bTrue):
        self.bOffscreenMode = bTrue

    def show_once(self):
        # for i in [1,2,3,4,5]:
        glutPostRedisplay()
        glutMainLoopEvent()


    def show(self):
        glutMainLoop()
imports/frankmocap/renderer/glRenderer.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

'''
This code has opengl visualization of 3D skeletons, including floor visualization and mouse+keyboard control
See the main function for demo codes.

showSkeleton: visualize 3D human body skeletons. This can handle holden's formation, 3.6m formation, and domeDB
setSpeech: to set speechAnnotation. Should be called before showSkeleton

renderscene: main function to render scenes using openGL

Note: this visualizer is assuming CM metric (joint, mesh).

Hanbyul Joo (jhugestar@gmail.com)
'''

from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *

import json
import numpy as np
from PIL import Image, ImageOps
import cv2
import numpy as np

import sys, math

import threading

import time
#from time import time

import pickle
#import cPickle as pickle
# import _pickle as pickle


from renderer.render_utils import ComputeNormal

#-----------
# Global Variables
#-----------
#g_camView_fileName = '/ssd/camInfo/camInfo.pkl'
g_camView_fileName = './camInfo/camInfo.pkl'

g_fViewDistance = 50.
# g_Width = 1280
# g_Height = 720
g_Width = 1920
g_Height = 1080

g_nearPlane = 0.01        #original
# g_nearPlane = 1000          #cameramode
g_farPlane = 9000.

g_show_fps = False

g_action = ""
g_xMousePtStart = g_yMousePtStart = 0.

g_xTrans = 0.
g_yTrans = 0.
g_zTrans = 0.
# g_zoom = 378
g_zoom = 600.
g_xRotate = 59.
g_yRotate = -41.
g_zRotate = 0.
g_xrot = 0.0
g_yrot = 0.0
# the usual screen dimension variables and lighting

# Generic Lighting values and coordinates
g_ambientLight = (0.35, 0.35, 0.35, 1.0)
g_diffuseLight = (0.75, 0.75, 0.75, 0.7)
g_specular = (0.2, 0.2, 0.2, 1.0)
g_specref = (0.5, 0.5, 0.5, 1.0)


# To visualize in Dome View point
#g_viewMode = 'camView'#free' #'camView'
# g_viewMode = 'camView' #'camView'
g_viewMode = 'free' #'camView'
g_bOrthoCam = False     #If true, draw by ortho camera mode
from collections import deque
#g_camid = deque('00',maxlen=2)
g_camid = deque('27',maxlen=2)


g_onlyDrawHumanIdx = -1


# To save rendered scene into file
g_stopMainLoop = False
g_winID = None

g_bSaveToFile = False
g_bSaveToFile_done = False      #If yes, file has been saved
g_saveImageName_last = None     #Remember the rendered image 

# g_savedImg = None       # Keep the save image as a global variable, if ones want to obtain this
# g_haggling_render = False       #Automatically Load next seqeucne, when frames are done
g_bSaveOnlyMode = False  #If true, load camera turn on save mode and exit after farmes are done
g_saveFolderName = None
g_saveImageName = None

"Visualization Options"
g_bApplyRootOffset = False
ROOT_OFFSET_DIST = 160
#ROOT_OFFSET_DIST = 30

""" Mesh Drawing Option """

#A tutorial for VAO and VBO: https://www.khronos.org/opengl/wiki/Tutorial2:_VAOs,_VBOs,_Vertex_and_Fragment_Shaders_(C_/_SDL)
g_vao= None
g_vertex_buffer= None
g_normal_buffer= None
g_tangent_buffer = None
g_index_buffer = None

g_hagglingseq_name= None
# global vts_num

# global SMPL_vts
# global face_num

# global SMPLModel
g_saveFrameIdx = 0
# global MoshParam


g_bGlInitDone = False       #To initialize opengl only once

BACKGROUND_IMAGE_PLANE_DEPTH=500


######################################################
#  MTC Camera view
g_camView_K = None
g_camView_K_list = None #if  each mesh uses diff K.... very ugly
g_bShowBackground = True

g_backgroundTextureID = None
g_textureData = None
# g_textureImgOriginal = None  #Original Size

g_renderOutputSize = None  #(width, height) If not None, crop the output to this size
######################################################
#  Visualizeation Option

g_bShowFloor = False
g_bShowWiredMesh = False

g_bRotateView = False   #Rotate view 360 degrees
g_rotateView_counter = 0
g_rotateInterval = 2

g_bShowSkeleton = True
g_bShowMesh = True





########################################################
# 3D keypoints Visualization Setting
g_skeletons = None # list of np.array. skeNum x  (skelDim, skelFrames)
# g_skeletons_GT = None
g_trajectory = None # list of np.array. skeNum x  (trajDim:3, skelFrames)
HOLDEN_DATA_SCALING = 5


g_faces = None #(faceNum, skelDim, skelFrames)
g_faceNormals = None
g_bodyNormals = None
g_hands = None #(handNum, skelDim, skelFrames)
g_hands_left = None
g_hands_right = None

g_posOnly = None

g_meshes = None

g_cameraPoses = None
g_cameraRots = None
g_ptCloud =None
g_ptCloudColor =None
g_ptSize = 2


g_frameLimit = -1
g_frameIdx = 0
g_lastframetime = g_currenttime = time.time()
g_fps = 0

g_speech = None # list of np.array. humanNum x  speechUnit, where elmenet is a dict with 'indicator', 'word', 'root' with a size of (1, skelFrames)
g_speechGT = None # list of np.array. humanNum x  speechUnit, , where elmenet is a dict with 'indicator', 'word', 'root' with a size of (1, skelFrames)

# Original
# g_colors = [ (0, 255, 127), (209, 206, 0), (0, 0, 128), (153, 50, 204), (60, 20, 220),
#     (0, 128, 0), (180, 130, 70), (147, 20, 255), (128, 128, 240), (154, 250, 0), (128, 0, 0),
#     (30, 105, 210), (0, 165, 255), (170, 178, 32), (238, 104, 123)]

# g_colors = [ (0, 255, 127), (170, 170, 0), (0, 0, 128), (153, 50, 204), (60, 20, 220),
#     (0, 128, 0), (180, 130, 70), (147, 20, 255), (128, 128, 240), (154, 250, 0), (128, 0, 0),
#     (30, 105, 210), (0, 165, 255), (170, 178, 32), (238, 104, 123)]

# g_colors = [ (250,0,0), (255,0, 0 ), (0, 255, 127), (209, 206, 0), (0, 0, 128), (153, 50, 204), (60, 20, 220),
#     (0, 128, 0), (180, 130, 70), (147, 20, 255), (128, 128, 240), (154, 250, 0), (128, 0, 0),
#     (30, 105, 210), (0, 165, 255), (170, 178, 32), (238, 104, 123)]


#GT visualization (by red)
g_colors = [ (255,0,0), (0, 255, 127), (170, 170, 0), (0, 0, 128), (153, 50, 204), (60, 20, 220),
    (0, 128, 0), (180, 130, 70), (147, 20, 255), (128, 128, 240), (154, 250, 0), (128, 0, 0),
    (30, 105, 210), (0, 165, 255), (170, 178, 32), (238, 104, 123)]

#Custom (non brl) originalPanoptic data
g_colors = [(0, 255, 127), (255,0,0), (170, 170, 0), (0, 0, 128), (153, 50, 204), (60, 20, 220),
    (0, 128, 0), (180, 130, 70), (147, 20, 255), (128, 128, 240), (154, 250, 0), (128, 0, 0),
    (30, 105, 210), (0, 165, 255), (170, 178, 32), (238, 104, 123)]


#RGB order!!
g_colors = [ (0,0,255), (255,0,0), (0, 255, 127), (170, 170, 0), (0, 0, 128), (153, 50, 204), (60, 20, 220),
    (0, 128, 0), (180, 130, 70), (147, 20, 255), (128, 128, 240), (154, 250, 0), (128, 0, 0),
    (30, 105, 210), (0, 165, 255), (170, 178, 32), (238, 104, 123)]

g_colorSet={} #RGB
g_colorSet['eft'] = (int(0.53*255), int(0.53*255), int(0.8*255) )  #bluish color
g_colorSet['spin'] = (int(0.7*255), int(0.5*255), int(0.5*255) )        #SPIN color
g_colorSet['hand'] = (30, 178, 166)         #Hand + Body (cyan like)


# g_meshColor = (0.4, 0.4, 0.7) #Blue Default (R,G,B)
# g_meshColor = (0.53, 0.53, 0.8)   #prediction: blue
# g_meshColor = (30/255.0, 178/255.0, 166/255.0)   #hand+body

# g_meshColor = (0.4, 0.4, 0.7) #Blue Default (R,G,B)
g_meshColor = (0.53, 0.53, 0.8)   #prediction: blue
# g_meshColor = (30/255.0, 178/255.0, 166/255.0)   #hand+body

# g_colors = [(0, 255, 127), (255,0,0), (170, 170, 0)  , (0, 0, 128), (153, 50, 204), (60, 20, 220),
#      (0, 128, 0), (180, 130, 70), (147, 20, 255), (128, 128, 240), (154, 250, 0), (128, 0, 0),
#      (30, 105, 210), (0, 165, 255), (170, 178, 32), (238, 104, 123)]

import timeit



#######################################v
# Parametric Mesh Models
g_faceModel = None

########################################################3
# Opengl Setting
def init():
    #global width
    #global height

    glClearColor(1.0, 1.0, 1.0, 1.0)
    # Enable depth testing
    glEnable(GL_DEPTH_TEST)
    glDepthFunc(GL_LESS)
    #glShadeModel(GL_FLAT)
    glShadeModel(GL_SMOOTH)
    glEnable(GL_LIGHTING)
    glLightfv(GL_LIGHT0, GL_AMBIENT, g_ambientLight)
    glLightfv(GL_LIGHT0, GL_DIFFUSE, g_diffuseLight)
    glLightfv(GL_LIGHT0, GL_SPECULAR, g_specular)
    glEnable(GL_LIGHT0)
    glEnable(GL_COLOR_MATERIAL)
    glColorMaterial(GL_FRONT, GL_AMBIENT_AND_DIFFUSE)
    glMaterialfv(GL_FRONT, GL_SPECULAR, g_specref)
    glMateriali(GL_FRONT, GL_SHININESS, 128)

    # # #Mesh Rendering
    global g_vao
    global g_vertex_buffer
    global g_normal_buffer
    global g_tangent_buffer
    global g_index_buffer

    g_vao = glGenVertexArrays(1)
    # glBindVertexArray(g_vao)

    g_vertex_buffer = glGenBuffers(40)
    # #glBindBuffer(GL_ARRAY_BUFFER, g_vertex_buffer)

    g_normal_buffer = glGenBuffers(40)
    # #glBindBuffer(GL_ARRAY_BUFFER, g_normal_buffer)

    g_tangent_buffer = glGenBuffers(40)
    # #glBindBuffer(GL_ARRAY_BUFFER, g_tangent_buffer)

    g_index_buffer = glGenBuffers(40)
    # # #glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, g_index_buffer)


    #Create Background Texture
    global g_backgroundTextureID
    glPixelStorei(GL_UNPACK_ALIGNMENT, 1)       #So that texture doesnt have to be power of 2
    g_backgroundTextureID = glGenTextures(1)

    glBindTexture(GL_TEXTURE_2D, g_backgroundTextureID)
    # glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB32F, 1920, 1080, 0, GL_RGB, GL_FLOAT, 0)
    #glTexImage2D(GL_TEXTURE_2D, 0,GL_RGB, width, height, 0, GL_BGR, GL_UNSIGNED_BYTE, data);

    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)



def init_minimum():
    #global width
    #global height

    glClearColor(1.0, 1.0, 1.0, 1.0)
    # Enable depth testing
    glEnable(GL_DEPTH_TEST)
    glEnable(GL_LIGHTING)
    glLightfv(GL_LIGHT0, GL_AMBIENT, g_ambientLight)
    glLightfv(GL_LIGHT0, GL_DIFFUSE, g_diffuseLight)
    glLightfv(GL_LIGHT0, GL_SPECULAR, g_specular)
    glEnable(GL_LIGHT0)
    glEnable(GL_COLOR_MATERIAL)
    glColorMaterial(GL_FRONT, GL_AMBIENT_AND_DIFFUSE)
    glMaterialfv(GL_FRONT, GL_SPECULAR, g_specref)
    glMateriali(GL_FRONT, GL_SHININESS, 128)


"""Render text on bottom left corner"""
def RenderText(testStr):

    glMatrixMode(GL_MODELVIEW)
    glPushMatrix()
    glLoadIdentity()
    gluOrtho2D(0.0, glutGet(GLUT_WINDOW_WIDTH), 0.0, glutGet(GLUT_WINDOW_HEIGHT))
    glMatrixMode(GL_PROJECTION)
    glPushMatrix()
    glLoadIdentity()
    # glRasterPos2i(10, 10)

    glPushMatrix()
    glColor4f(1,1,1,1)
    glClearDepth(1)
    if testStr is not None:
        glRasterPos2d(5,5)
        for c in testStr:
            #glutBitmapCharacter(GLUT_BITMAP_8_BY_13, ord(c))
            glutBitmapCharacter(GLUT_BITMAP_HELVETICA_18, ord(c))
    raster_pos = glGetIntegerv(GL_CURRENT_RASTER_POSITION)
    glPopMatrix()

    glClearDepth(0.99)
    glColor4f(.3,.3,.3,.8)
    glPushMatrix()
    glBegin(GL_QUADS)
    glVertex2i(0, 25)
    glVertex2i(raster_pos[0]+5, 25)
    glVertex2i(raster_pos[0]+5, 0)
    glVertex2i(0, 0)
    glEnd()
    glPopMatrix()
    glClearDepth(1)

    glMatrixMode(GL_PROJECTION)
    glPopMatrix()
    glMatrixMode(GL_MODELVIEW)
    glPopMatrix()


def RenderDomeFloor():
    glPolygonMode(GL_FRONT, GL_FILL)
    glPolygonMode(GL_BACK, GL_FILL)
    # glPolygonMode(GL_FRONT, GL_FILL)
    gridNum = 50
    width = 500
    halfWidth =width/2
    # g_floorCenter = np.array([0,0.5,0])

    # g_floorCenter = np.array([0,500,0])
    g_floorCenter = np.array([0,1500,0])
    g_floorAxis1 = np.array([1,0,0])
    g_floorAxis2 = np.array([0,0,1])

    origin = g_floorCenter - g_floorAxis1*(width*gridNum/2 ) - g_floorAxis2*(width*gridNum/2)
    axis1 =  g_floorAxis1 * width
    axis2 =  g_floorAxis2 * width
    for y in range(gridNum+1):
        for x in range(gridNum+1):

            if (x+y) % 2 ==0:
                glColor(1.0,1.0,1.0,1.0) #white
            else:
                # glColor(0.95,0.95,0.95,0.3) #grey
                glColor(0.3,0.3,0.3,0.5) #grey

            p1 = origin + axis1*x + axis2*y
            p2 = p1+ axis1
            p3 = p1+ axis2
            p4 = p1+ axis1 + axis2


            glBegin(GL_QUADS)
            glVertex3f(   p1[0], p1[1], p1[2])
            glVertex3f(   p2[0], p2[1], p2[2])
            glVertex3f(   p4[0], p4[1], p4[2])
            glVertex3f(   p3[0], p3[1], p3[2])
            glEnd()


def setFree3DView():
    glTranslatef(0,0,g_zoom)

    glRotatef( -g_yRotate, 1.0, 0.0, 0.0)
    glRotatef( -g_xRotate, 0.0, 1.0, 0.0)

    glRotatef( g_zRotate, 0.0, 0.0, 1.0)
    glTranslatef( g_xTrans,  0.0, 0.0 )
    glTranslatef(  0.0, g_yTrans, 0.0)
    glTranslatef(  0.0, 0, g_zTrans)


# g_hdCams = None
# def load_panoptic_cameras():
#     global g_hdCams
#     with open('/media/posefs3b/Users/xiu/domedb/171204_pose3/calibration_171204_pose3.json') as f:
#         rawCalibs = json.load(f)
#     cameras = rawCalibs['cameras']
#     allPanel = map(lambda x:x['panel'],cameras)
#     hdCamIndices = [i for i,x in enumerate(allPanel) if x==0]
#     g_hdCams = [cameras[i] for i in hdCamIndices]

# def setPanopticCameraView(camid):
#     if g_hdCams==None:
#         load_cameras()

#     if camid>=len(g_hdCams):
#         camid = 0
#     cam = g_hdCams[camid]
#     invR = np.array(cam['R'])
#     invT = np.array(cam['t'])
#     camMatrix = np.hstack((invR, invT))
#     # denotes camera matrix, [R|t]
#     camMatrix = np.vstack((camMatrix, [0, 0, 0, 1]))
#     #camMatrix = numpy.linalg.inv(camMatrix)
#     K = np.array(cam['K'])
#     #K = K.flatten()
#     glMatrixMode(GL_PROJECTION)
#     glLoadIdentity()
#     Kscale = 1920.0/g_Width
#     K = K/Kscale
#     ProjM = np.zeros((4,4))
#     ProjM[0,0] = 2*K[0,0]/g_Width
#     ProjM[0,2] = (g_Width - 2*K[0,2])/g_Width
#     ProjM[1,1] = 2*K[1,1]/g_Height
#     ProjM[1,2] = (-g_Height+2*K[1,2])/g_Height

#     ProjM[2,2] = (-g_farPlane-g_nearPlane)/(g_farPlane-g_nearPlane)
#     ProjM[2,3] = -2*g_farPlane*g_nearPlane/(g_farPlane-g_nearPlane)
#     ProjM[3,2] = -1

#     glLoadMatrixd(ProjM.T)
#     glMatrixMode(GL_MODELVIEW)
#     glLoadIdentity()
#     gluLookAt(0, 0, 0, 0, 0, 1, 0, -1, 0)
#     glMultMatrixd(camMatrix.T)


# def load_MTC_default_camera():
#     global g_Width,  g_Height

#     camRender_width = 1920
#     camRender_height = 1080

#     g_Width = camRender_width
#     g_Height = camRender_height


# 3x3 intrinsic camera matrix
def setCamView_K(K):
    global g_camView_K
    g_camView_K = K


# 3x3 intrinsic camera matrix
# Set a default camera matrix used for MTC
def setCamView_K_DefaultForMTC():
    global g_camView_K

    K = np.array([[2000, 0, 960],[0, 2000, 540],[0,0,1]])       #MTC default camera. for 1920 x 1080 input image
    g_camView_K = K


#Show the world in a camera cooridnate (defined by K)
def setCameraView():

    # camRender_width = 1920
    # camRender_height = 1080

    # global g_Width,  g_Height
    # if camRender_width != g_Width or  camRender_height!=g_Height:
    #     g_Width = camRender_width
    #     g_Height = camRender_height
    #     #reshape(g_Width, g_Height)
    #     glutReshapeWindow(g_Width,g_Height)

    # invR = np.array(cam['R'])
    # invT = np.array(cam['t'])
    invR = np.eye(3)
    invT = np.zeros((3,1))
    # invT[2] = 400
    camMatrix = np.hstack((invR, invT))
    # denotes camera matrix, [R|t]
    camMatrix = np.vstack((camMatrix, [0, 0, 0, 1]))
    #camMatrix = numpy.linalg.inv(camMatrix)
    # K = np.array(cam['K'])
    # K = np.array(cam['K'])
    # K = np.array([[2000, 0, 960],[0, 2000, 540],[0,0,1]])       #MTC default camera
    # global g_camView_K
    # g_camView_K = K
    #K = K.flatten()

    if g_camView_K is None:
        print("## Warning: no K is set, so I use a default cam parameter defined for MTC")
        setCamView_K_DefaultForMTC()
    K = g_camView_K.copy()

    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    # Kscale = 1920.0/g_Width
    Kscale = 1.0 #1920.0/g_Width        :: why do we need this?
    K = K/Kscale
    ProjM = np.zeros((4,4))
    ProjM[0,0] = 2*K[0,0]/g_Width
    ProjM[0,2] = (g_Width - 2*K[0,2])/g_Width
    ProjM[1,1] = 2*K[1,1]/g_Height
    ProjM[1,2] = (-g_Height+2*K[1,2])/g_Height

    ProjM[2,2] = (-g_farPlane-g_nearPlane)/(g_farPlane-g_nearPlane)
    ProjM[2,3] = -2*g_farPlane*g_nearPlane/(g_farPlane-g_nearPlane)
    ProjM[3,2] = -1

    glLoadMatrixd(ProjM.T)
    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()
    gluLookAt(0, 0, 0, 0, 0, 1, 0, -1, 0)
    glMultMatrixd(camMatrix.T)


def SetOrthoCamera(bOrtho=True):
    global g_bOrthoCam
    g_bOrthoCam = bOrtho

#Show the world in a camera cooridnate (defined by K)
def setCameraViewOrth():

    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()

    texHeight,texWidth =   g_textureData.shape[:2]
    # texHeight,texWidth =   1024, 1024
    texHeight*=0.5
    texWidth*=0.5
    # texHeight *=BACKGROUND_IMAGE_PLANE_DEPTH
    # texWidth *=BACKGROUND_IMAGE_PLANE_DEPTH

    glOrtho(-texWidth, texWidth, -texHeight, texHeight, -1500, 1500)

    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()
    gluLookAt(0, 0, 0, 0, 0, 1, 0, -1, 0)
    # glMultMatrixd(camMatrix.T)

def setRenderOutputSize(imWidth, imHeight):
    global g_renderOutputSize
    g_renderOutputSize = (imWidth, imHeight)

def setWindowSize(new_width, new_height):
    global g_Width, g_Height

    if new_height>1600:    #Max height of screen
        new_width = int(new_width *0.7)
        new_height = int(new_height *0.7)

    if new_width != g_Width or  new_height!=g_Height:
        g_Width = new_width
        g_Height =new_height
        #reshape(g_Width, g_Height)

        if g_bGlInitDone:
            glutReshapeWindow(g_Width,g_Height)
            

def reshape(width, height):
    #lightPos = (-50.0, 50.0, 100.0, 1.0)
    nRange = 250.0
    global g_Width, g_Height
    g_Width = width
    g_Height = height
    glViewport(0, 0, g_Width, g_Height)

    # # Set perspective (also zoom)
    # glMatrixMode(GL_PROJECTION)
    # glLoadIdentity()
    # #gluPerspective(zoom, float(g_Width)/float(g_Height), g_nearPlane, g_farPlane)
    # gluPerspective(65, float(g_Width)/float(g_Height), g_nearPlane, g_farPlane)
    # print("here: {}".format(float(g_Width)/float(g_Height)))


def SaveScenesToFile():
    global g_saveImageName_last
    #global g_Width, g_Height, g_bSaveToFile, g_fameIdx, g_hagglingseq_name
    # global g_bSaveToFile
    # global g_saveFrameIdx

    glReadBuffer(GL_FRONT)
    img = glReadPixels(0, 0, g_Width, g_Height, GL_RGBA, GL_UNSIGNED_BYTE, outputType=None)
    img = Image.frombytes("RGBA", (g_Width, g_Height), img)
    img = ImageOps.flip(img)

    if False:
        pix = np.array(img)
        from renderer import viewer2D
        viewer2D.ImShow(pix)


    #Crop and keep the original size
    # if g_textureImgOriginal is not None and g_viewMode=='camView':
    #     width, height = img.size
    #     img = img.crop( (0,0,g_textureImgOriginal.shape[1], g_textureImgOriginal.shape[0]))     #top, left, bottom, right where origin seems top left

    if g_renderOutputSize is not None and g_viewMode=='camView':
        width, height = g_renderOutputSize
        img = img.crop( (0,0,width, height))     #top, left, bottom, right where origin seems top left


    #img.save('/ssd/render_general/frame_{}.png'.format(g_frameIdx), 'PNG')
    if g_saveFolderName is not None:
        folderPath = g_saveFolderName
    else:# g_haggling_render == False:
        # folderPath = '/home/hjoo//temp/render_general/'
        folderPath = '/home/hjoo//temp/render_general/'
    # else:
    #     if g_meshes == None and g_hagglingFileList == None: #Initial loading
    #         LoadHagglingData_Caller() #Load New Scene
    #     folderPath = '/hjoo/home/temp/render_mesh/' + g_hagglingseq_name
    #     if os.path.exists(folderPath) == False:
    #         os.mkdir(folderPath)

    if os.path.exists(folderPath) == False:
        os.mkdir(folderPath)
    # img.save('{0}/scene_{1:08d}.png'.format(folderPath,g_saveFrameIdx), 'PNG')
    img = img.convert("RGB")

    if g_saveImageName is not None:
        g_saveImageName_last = '{0}/{1}.jpg'.format(folderPath,g_saveImageName)
    else:
        g_saveImageName_last = '{0}/scene_{1:08d}.jpg'.format(folderPath,g_saveFrameIdx)
    
    img.save(g_saveImageName_last, "JPEG")

    print(f"Render image to:{g_saveImageName_last}")

    # print(img.size)
    # g_saveFrameIdx +=1
    #Done saving
    global g_bSaveToFile_done
    g_bSaveToFile_done = True


    # global g_savedImg = img     #save rendered one into gobal variable

    # if g_frameIdx+1 >= g_frameLimit:
    #     if g_haggling_render == False:
    #         g_bSaveToFile = False
    #     else:
    #         bLoaded = LoadHagglingData_Caller() #Load New Scene
    #         if bLoaded ==False:
    #             g_bSaveToFile = False

        #cur_ind += 1
        #print(cur_ind)

    #glutPostRedisplay()


def SaveCamViewInfo():
    global g_Width, g_Height
    global g_nearPlane,g_farPlane
    global g_zoom,g_yRotate,g_xRotate, g_zRotate, g_xTrans, g_yTrans, g_zTrans

    fileName = '/ssd/camInfo/camInfo.pkl'
    if os.path.exists('/ssd/camInfo/') == False:
            os.mkdir('/ssd/camInfo/')

    if os.path.exists('/ssd/camInfo/archive/') == False:
        os.mkdir('/ssd/camInfo/archive/')

    if os.path.exists(fileName):
        #resave it
        for i in range(1000):
            newName = '/ssd/camInfo/archive/camInfo_old{}.pkl'.format(i)
            if os.path.exists(newName) == False:
                import shutil
                shutil.copy2(fileName,newName)
                break

    camInfo = dict()
    camInfo['g_Width'] = g_Width
    camInfo['g_Height']  =g_Height
    camInfo['g_nearPlane'] = g_nearPlane
    camInfo['g_farPlane'] = g_farPlane
    camInfo['g_zoom'] = g_zoom
    camInfo['g_yRotate'] = g_yRotate
    camInfo['g_xRotate'] =g_xRotate
    camInfo['g_zRotate'] = g_zRotate
    camInfo['g_xTrans'] = g_xTrans
    camInfo['g_yTrans'] = g_yTrans

    pickle.dump( camInfo, open(fileName, "wb" ) )

    print('camInfo')


def LoadCamViewInfo():
    global g_Width, g_Height
    global g_nearPlane,g_farPlane
    global g_zoom,g_yRotate,g_xRotate, g_zRotate, g_xTrans, g_yTrans
    global g_camView_fileName
    fileName = g_camView_fileName
    if not os.path.exists(fileName):
        print("No cam info: {}".format(fileName))
        return

    camInfo = pickle.load( open( fileName, "rb" ) , encoding='latin1')
    g_yTrans = camInfo['g_Width']
    g_Height = camInfo['g_Height']
    g_nearPlane = camInfo['g_nearPlane']
    g_farPlane = camInfo['g_farPlane']
    g_zoom = camInfo['g_zoom']
    g_yRotate = camInfo['g_yRotate']
    g_xRotate = camInfo['g_xRotate']
    g_zRotate = camInfo['g_zRotate']
    g_xTrans = camInfo['g_xTrans']
    g_yTrans= camInfo['g_yTrans']

    reshape(g_Width, g_Height)


def PuttingObjectCenter():
    global g_zoom, g_xTrans, g_yTrans, g_zTrans
    global g_xRotate, g_yRotate, g_zRotate
    if (g_skeletons is not None) and len(g_skeletons)>0:


        g_xRotate =0
        g_yRotate =0
        g_zRotate =0

        # if g_skeletonType == 'smplcoco':
        #     g_xTrans = -(g_skeletons[0][9,0] + g_skeletons[0][12,0] )*0.5
        #     g_yTrans = -(g_skeletons[0][10,0] + g_skeletons[0][13,0] )*0.5
        #     g_zTrans = -(g_skeletons[0][11,0] + g_skeletons[0][14,0] )*0.5
        #     g_zoom = 300
        # else: #Adam
        g_xTrans = -g_skeletons[0]['skeleton'][3*39]
        g_yTrans = -g_skeletons[0]['skeleton'][3*39+1]# 0#100
        g_zTrans = -g_skeletons[0]['skeleton'][3*39+2]
        g_zoom = 300

    elif g_meshes is not None and len(g_meshes)>0:

        g_xRotate =0
        g_yRotate =0
        g_zRotate =0

        g_xTrans = -g_meshes[0]['ver'][0,1767,0]
        g_yTrans = -g_meshes[0]['ver'][0,1767,1]
        g_zTrans = -g_meshes[0]['ver'][0,1767,2]

        # print("{} {} {}".format(g_xTrans, g_yTrans, g_zTrans))
        g_zoom = 300


def keyboard(key, x, y):
    global g_stopMainLoop,g_frameIdx,g_show_fps
    global g_ptSize
    global g_xRotate,g_yRotate,g_zRotate


    if isinstance(key, bytes):
        key = key.decode()  #Python3: b'X' -> 'X' (bytes -> str)
    if key == chr(27) or key == 'q':
        #sys.exit()
        g_stopMainLoop= True
        g_frameIdx = 0
        # glutIdleFunc(0); # Turn off Idle function if used.
        # glutSetOption(GLUT_ACTION_ON_WINDOW_CLOSE,GLUT_ACTION_CONTINUE_EXECUTION)
        # glutLeaveMainLoop()
        # glutDestroyWindow(g_winID) # Close open windows
    elif key == 'p':
        #global width, height
        glReadBuffer(GL_FRONT)
        img = glReadPixels(0, 0, g_Width, g_Height, GL_RGBA, GL_UNSIGNED_BYTE, outputType=None)
        img = Image.frombytes("RGBA", (g_Width, g_Height), img)
        img = ImageOps.flip(img)
        img.save('temp.jpg', 'JPG')

    elif key == 's':
        g_frameIdx = 0
  

    elif key == 't':
        
        g_xRotate =0
        g_yRotate =-90
        g_zRotate =0
        print('showTopView')
        
    elif key == 'r':
           global g_bSaveToFile
           g_frameIdx = 0
           g_bSaveToFile = not g_bSaveToFile

    elif key == 'w':
        global g_bShowWiredMesh
        g_bShowWiredMesh = not g_bShowWiredMesh

    elif key == 'h':        #Load Next Haggling Data
            LoadHagglingData_Caller()
    elif key =='o':
        global g_bApplyRootOffset
        g_bApplyRootOffset = not g_bApplyRootOffset

    elif key =='f':
        global g_bShowFloor
        g_bShowFloor = not g_bShowFloor

    elif key =='V':
        SaveCamViewInfo()

    elif key =='v':
        LoadCamViewInfo()

    # elif key =='c':
    #     g_xRotate =0
    #     g_yRotate =0
    #     g_zRotate =180

    elif key =='c':     #put the target human in the center
       PuttingObjectCenter()
    elif key=='+':
        g_ptSize +=1
    elif key=='-':
        # global g_ptSize
        if g_ptSize >=2:
            g_ptSize -=1
    elif key =='R':     #rotate cameras
        global g_bRotateView, g_rotateView_counter

        g_bRotateView = not g_bRotateView
        g_rotateView_counter =0

    elif key =='j':  #Toggle joint
        global g_bShowSkeleton
        g_bShowSkeleton = not g_bShowSkeleton
    elif key =='m':
        global g_bShowMesh
        g_bShowMesh = not g_bShowMesh
    elif key =='b':
        global g_bShowBackground
        g_bShowBackground = not g_bShowBackground

    # elif key =='C':
    #     g_xRotate =0
    #     g_yRotate =0
    #     g_zRotate =0
    elif key == 'C':
        print('Toggle camview / freeview')
        global g_viewMode, g_nearPlane
        if g_viewMode=='free':
            g_viewMode = 'camView'
            g_nearPlane = 500          #cameramode
        else:
            g_viewMode ='free'
            g_nearPlane = 0.01        #original

      
    elif key == 'S':
        global g_zoom, g_xTrans, g_yTrans
        g_xTrans=  -4.126092433929443
        g_yTrans= 12
        # g_zoom= 190
        g_zoom= 1028
        # g_xRotate = 90
        g_xRotate = -58
        # g_yRotate= 0
        g_yRotate= 9
        g_zRotate= 0.0
        g_viewMode = 'free'

    # elif key>='0' and key<='9':
    #     g_camid.popleft()
    #     g_camid.append(key)
    #     print('camView: CamID:{}'.format(g_camid))

    elif key =='0':
        glutReshapeWindow(1920,720)

    elif key =='z': # Toggle fps printing
        g_show_fps = not g_show_fps

    glutPostRedisplay()


def mouse(button, state, x, y):
    global g_action, g_xMousePtStart, g_yMousePtStart
    if (button==GLUT_LEFT_BUTTON):
        if (glutGetModifiers() == GLUT_ACTIVE_SHIFT):
            g_action = "TRANS"
        else:
            g_action = "MOVE_EYE"
    #elif (button==GLUT_MIDDLE_BUTTON):
    #    action = "TRANS"
    elif (button==GLUT_RIGHT_BUTTON):
        g_action = "ZOOM"
    g_xMousePtStart = x
    g_yMousePtStart = y

def motion(x, y):
    global g_zoom, g_xMousePtStart, g_yMousePtStart, g_xRotate, g_yRotate, g_zRotate, g_xTrans, g_zTrans
    if (g_action=="MOVE_EYE"):
        g_xRotate += x - g_xMousePtStart
        g_yRotate -= y - g_yMousePtStart
    elif (g_action=="MOVE_EYE_2"):
        g_zRotate += y - g_yMousePtStart
    elif (g_action=="TRANS"):
        g_xTrans += x - g_xMousePtStart
        g_zTrans += y - g_yMousePtStart
    elif (g_action=="ZOOM"):
        g_zoom -= y - g_yMousePtStart
        # print(g_zoom)
    else:
        print("unknown action\n", g_action)
    g_xMousePtStart = x
    g_yMousePtStart = y

    # print ('xTrans {},  yTrans {}, zoom {} xRotate{} yRotate {} zRotate {}'.format(g_xTrans,  g_yTrans,  g_zoom,  g_xRotate,  g_yRotate,  g_zRotate))
    glutPostRedisplay()

def setBackgroundTexture(img):
    global g_textureData#, g_textureImgOriginal
    g_textureData = img

    #In MTC, the background should be always 1920x1080
    # g_textureData = np.ones( (1080, 1920, 3), dtype=img.dtype)*0     #dtype==np.unit8
    # g_textureData[:img.shape[0],:img.shape[1] ] = img
    # g_textureImgOriginal = img  #keep the original image


    # import cv2
    # cv2.imshow('here??',img)
    # cv2.waitKey(0)

def SetCameraPoses(camRots, camPoses):
    global g_cameraPoses,g_cameraRots
    
    # g_cameraPoses = camPoses
    # g_cameraRots = camRots

    g_cameraPoses = camPoses        #for cam Vis
    g_cameraRots = []           #for cam Vis
    
    for r in camRots:
        cam_extR_4x4 = np.eye(4,dtype= r.dtype)
        # cam_extR_4x4[:3,:3] = cam_R_rot.transpose()
        cam_extR_4x4[:3,:3] = r.transpose()         #For visualizing cameras, R should be inversed
        cam_extR_4x4[3,3] =1.0
        g_cameraRots.append(cam_extR_4x4)

        

def SetPtCloud(ptCloud, ptCloudColor = None):
    global g_ptCloud, g_ptCloudColor
    g_ptCloud = ptCloud
    g_ptCloudColor = ptCloudColor

def DrawBackgroundOrth():

    if g_textureData is None:
        return

    glDisable(GL_CULL_FACE)
    glDisable(GL_DEPTH_TEST)
    glDisable(GL_LIGHTING)
    glEnable(GL_TEXTURE_2D)

    # glUseProgram(0)

    glBindTexture(GL_TEXTURE_2D, g_backgroundTextureID)
    # glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, 1920, 1080, 0, GL_RGB, GL_UNSIGNED_BYTE, g_textureData)
    # glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, 480, 640, 0, GL_RGB, GL_UNSIGNED_BYTE, g_textureData.data)
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, g_textureData.shape[1], g_textureData.shape[0], 0, GL_BGR, GL_UNSIGNED_BYTE, g_textureData.data)
    texHeight,texWidth =   g_textureData.shape[:2]
    texHeight*=0.5
    texWidth*=0.5

    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)

    glPolygonMode(GL_FRONT_AND_BACK, GL_FILL)
    glBegin(GL_QUADS)
    glColor3f(1.0, 1.0, 1.0)
    # d = BACKGROUND_IMAGE_PLANE_DEPTH
    d = 10

    glTexCoord2f(0, 0)
    # #Eigen::Map<Eigen::Matrix<double, 3, 3, Eigen::RowMajor>> K(m_options.m_pK);
    P = np.array([-texWidth, -texHeight, d])
    # P = np.matmul(K_inv,P)
    # P = P / P[2]
    glVertex3f(P[0] , P[1] , P[2] );  # K^{-1} [0, 0, 1]^T

    glTexCoord2f(1, 0)
    # P = [1920, 0, 1]
    P = [texWidth, -texHeight, d]
    glVertex3f(P[0] , P[1] , P[2] );  # K^{-1} [0, 0, 1]^T

    glTexCoord2f(1, 1)
    # P = [1920, 1080, 1]
    P = [texWidth, texHeight, d]
    # P = np.matmul(K_inv,P)
    # P = P / P[2]
    # glVertex3f(P[0] * d, P[1] * d, P[2] * d)
    glVertex3f(P[0] , P[1] , P[2] );  # K^{-1} [0, 0, 1]^T

    glTexCoord2f(0, 1)
    # P = [0, 1080, 1]
    P = [-texWidth, texHeight, d]
    # glVertex3f(P[0] * d, P[1] * d, P[2] * d)
    glVertex3f(P[0] , P[1] , P[2] );  # K^{-1} [0, 0, 1]^T
    glEnd()

    glEnable(GL_LIGHTING)
    glEnable(GL_CULL_FACE)
    glEnable(GL_DEPTH_TEST)
    glDisable(GL_TEXTURE_2D)



def DrawBackground():

    if g_camView_K is None or g_textureData is None:
        return

    K_inv = np.linalg.inv(g_camView_K)

    glDisable(GL_CULL_FACE)
    glDisable(GL_DEPTH_TEST)
    glDisable(GL_LIGHTING)
    glEnable(GL_TEXTURE_2D)

    # glUseProgram(0)

    glBindTexture(GL_TEXTURE_2D, g_backgroundTextureID)
    # glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, 1920, 1080, 0, GL_RGB, GL_UNSIGNED_BYTE, g_textureData)
    # glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, 480, 640, 0, GL_RGB, GL_UNSIGNED_BYTE, g_textureData.data)
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, g_textureData.shape[1], g_textureData.shape[0], 0, GL_BGR, GL_UNSIGNED_BYTE, g_textureData.data)
    texHeight,texWidth =   g_textureData.shape[:2]

    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)


    glPolygonMode(GL_FRONT_AND_BACK, GL_FILL)
    glBegin(GL_QUADS)
    glColor3f(1.0, 1.0, 1.0)
    d = BACKGROUND_IMAGE_PLANE_DEPTH

    glTexCoord2f(0, 0)
    # #Eigen::Map<Eigen::Matrix<double, 3, 3, Eigen::RowMajor>> K(m_options.m_pK);
    P = np.array([0, 0, 1])
    P = np.matmul(K_inv,P)
    P = P / P[2]
    glVertex3f(P[0] * d, P[1] * d, P[2] * d);  # K^{-1} [0, 0, 1]^T

    glTexCoord2f(1, 0)
    # P = [1920, 0, 1]
    P = [texWidth, 0, 1]

    P = np.matmul(K_inv,P)
    P = P / P[2]
    glVertex3f(P[0] * d, P[1] * d, P[2] * d);  # K^{-1} [0, 0, 1]^T

    glTexCoord2f(1, 1)
    # P = [1920, 1080, 1]
    P = [texWidth, texHeight, 1]
    P = np.matmul(K_inv,P)
    P = P / P[2]
    glVertex3f(P[0] * d, P[1] * d, P[2] * d)

    glTexCoord2f(0, 1)
    # P = [0, 1080, 1]
    P = [0, texHeight, 1]
    P = np.matmul(K_inv,P)
    P = P / P[2]
    glVertex3f(P[0] * d, P[1] * d, P[2] * d)
    glEnd()

    glEnable(GL_LIGHTING)
    glEnable(GL_CULL_FACE)
    glEnable(GL_DEPTH_TEST)
    glDisable(GL_TEXTURE_2D)


def specialkeys(key, x, y):
    global g_xrot
    global g_yrot
    global g_cur_ind
    if key == GLUT_KEY_UP:
        g_xrot -= 2.0
    if key == GLUT_KEY_DOWN:
        g_xrot += 2.0
    # if key == GLUT_KEY_LEFT:
    #     cur_ind -=1
    # if key == GLUT_KEY_RIGHT:
    #     cur_ind +=1
    glutPostRedisplay()

from multiprocessing import Pool

def init_gl_util():
    global g_bGlInitDone, g_lastframetime, g_currenttime, g_fps

    g_lastframetime = g_currenttime
    g_currenttime = time.time()
    refresh_fps = 0.15
    g_fps = (1-refresh_fps)*g_fps + refresh_fps*1/(g_currenttime-g_lastframetime)

    if g_bGlInitDone==False:
        glutInit()
        glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH | GLUT_MULTISAMPLE) #GLUT_MULTISAMPLE is required for anti-aliasing
        #glutInitDisplayMode(GLUT_RGB |GLUT_DOUBLE|GLUT_DEPTH)
        glutInitWindowPosition(100,100)
        glutInitWindowSize(g_Width,g_Height)

        global g_winID
        g_winID = glutCreateWindow("3D View")
        init()
        # init_minimum()
        glutReshapeFunc(reshape)
        glutDisplayFunc(renderscene)
        glutKeyboardFunc(keyboard)
        glutMouseFunc(mouse)
        glutMotionFunc(motion)
        glutSpecialFunc(specialkeys)
        # glutIdleFunc(idlefunc)
        glutIdleFunc(renderscene)

        #Ver 1: Infinite loop (termination is not possible)
        #glutMainLoop()

        #Ver 2: for better termination (by pressing 'q')
        g_bGlInitDone = True
    else:
        glutReshapeWindow(g_Width, g_Height)     #Just doing resize
        # glutReshapeWindow(int(g_Width*0.5), int(g_Height*0.5))     #Just doing resize


def init_gl(maxIter=-10):
    # Init_Haggling()
    # global width
    # global height
    # Setup for double-buffered display and depth testing
    init_gl_util()

    global g_stopMainLoop
    g_stopMainLoop=False
    while True:
        glutPostRedisplay()
        if bool(glutMainLoopEvent)==False:
            continue
        glutMainLoopEvent()
        if g_stopMainLoop:
            break
        if maxIter>0:
            maxIter -=1
            if maxIter<=0:
                g_stopMainLoop = True

    # print("Escaped glut loop")


#g_faces should be: #(faceNum, faceDim, faceFrames)
def DrawFaces():
    #global g_colors
    #global g_faces, g_faceNormals, g_frameIdx#, g_normals

    if g_faces is None:
        return

    #g_frameLimit = g_faces.shape[2]

    for humanIdx in range(len(g_faces)):

        if(g_frameIdx >= g_faces[humanIdx].shape[1]):
            continue

        if g_onlyDrawHumanIdx>=0 and humanIdx!=g_onlyDrawHumanIdx:
            continue

        face3D = g_faces[humanIdx][:, g_frameIdx]  #210x1
        drawface_70(face3D, g_colors[humanIdx % len(g_colors)])

        if g_faceNormals is not None and len(g_faceNormals)> humanIdx:

            if g_faceNormals[humanIdx].shape[1]<=g_frameIdx:
                print("Warning: g_faceNormals[humanIdx].shape[2]<=g_frameId")
                continue

            normal3D = g_faceNormals[humanIdx][:,g_frameIdx]  #3x1
            #drawfaceNormal_70(normal3D, face3D, g_colors[humanIdx % len(g_colors)])
            #drawfaceNormal_70(normal3D, face3D, [0, 255, 255])
            eyeCenterPoint = 0.5 *(face3D[(45*3):(45*3+3)] + face3D[(36*3):(36*3+3)])
            #drawNormal(normal3D, eyeCenterPoint, [0, 255, 0],normalLength=25)
            drawNormal(normal3D, eyeCenterPoint, [0, 255, 255],normalLength=25)

        #drawbody_joint_ptOnly(face3D, g_colors[humanIdx])
    #g_frameIdx +=1

    #if g_frameIdx>=g_frameLimit:
    #    g_frameIdx =0


#g_faces should be: #(faceNum, faceDim, faceFrames)
def DrawHands():
    #global g_colors
    #global g_hands_left,g_hands_right, g_frameIdx#, g_normals


    #g_frameLimit = g_faces.shape[2]

    if g_hands_left is not None:
        for humanIdx in range(len(g_hands_left)):
            if g_onlyDrawHumanIdx>=0 and humanIdx!=g_onlyDrawHumanIdx:
               continue

            if(g_frameIdx >= g_hands_left[humanIdx].shape[1]):
                continue

            hand3D = g_hands_left[humanIdx][:, g_frameIdx]  #210x1
            drawhand_21(hand3D, g_colors[humanIdx % len(g_colors)])
            #drawbody_joint_ptOnly(face3D, g_colors[humanIdx])

    if g_hands_right is not None:
        for humanIdx in range(len(g_hands_right)):
            if g_onlyDrawHumanIdx>=0 and humanIdx!=g_onlyDrawHumanIdx:
               continue

            if(g_frameIdx >= g_hands_right[humanIdx].shape[1]):
                continue

            hand3D = g_hands_right[humanIdx][:, g_frameIdx]  #210x1
            drawhand_21(hand3D, g_colors[humanIdx % len(g_colors)])
            #drawbody_joint_ptOnly(face3D, g_colors[humanIdx])

    #g_frameIdx +=1

    #if g_frameIdx>=g_frameLimit:
    #    g_frameIdx =0


# Face keypoint orders follow Openpose keypoint output
# https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/output.md
# Face outline points (0-16) are unstable
face_edges = np.array([ #[0,1],[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9],[9,10],[11,12],[12,13],[14,15],[15,16], #outline (ignored)
                [17,18],[18,19],[19,20],[20,21], #right eyebrow
                [22,23],[23,24],[24,25],[25,26], #left eyebrow
                [27,28],[28,29],[29,30],   #nose upper part
                [31,32],[32,33],[33,34],[34,35], #nose lower part
                [36,37],[37,38],[38,39],[39,40],[40,41],[41,36], #right eye
                [42,43],[43,44],[44,45],[45,46],[46,47],[47,42], #left eye
                [48,49],[49,50],[50,51],[51,52],[52,53],[53,54],[54,55],[55,56],[56,57],[57,58],[58,59],[59,48], #Lip outline
                [60,61],[61,62],[62,63],[63,64],[64,65],[65,66],[66,67],[67,60] #Lip inner line
                ])
#joints70: 3x70 =210 dim
def drawface_70(joints,  color):

    glLineWidth(2.0)
    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    #for i in range(len(joints)/3):
    #    glPushMatrix()
    #    glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
    #    glutSolidSphere(0.5, 10, 10)
    #    glPopMatrix()

    # connMat_coco19 = g_connMat_smc19
    #Visualize Bones
    for conn in face_edges:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]

        x1 = joints[3*conn[1]]
        y1 = joints[3*conn[1]+1]
        z1 = joints[3*conn[1]+2]

        glBegin(GL_LINES)
        glVertex3f(x0, y0, z0)
        glVertex3f(x1,y1,z1)
        glEnd()



def SetMeshColor(colorName='blue'):
    global g_meshColor
    if colorName=='blue':
        # g_meshColor = (0.4, 0.4, 0.)   #prediction: blue
        g_meshColor = (0.53, 0.53, 0.8)   #prediction: blue
        # glColor3f(0.53, 0.53, 0.8)
    elif colorName=='red':
        g_meshColor = (0.7, 0.5, 0.5)   #targer: red
    else:
        assert False


g_firsttime = True
""" With normal"""
def DrawMeshes():

    global g_firsttime
    global g_meshes, g_frameIdx#, g_normals

    global g_vao
    global g_vertex_buffer
    global g_normal_buffer
    global g_tangent_buffer
    global g_index_buffer

    # MESH_SCALING = 100.0  #from meter (model def) to CM (panoptic def)
    # MESH_SCALING = 1.0  #from meter (model def) to CM (panoptic def)

    if g_meshes is None:
        return

    #g_meshes[humanIdx]['ver']: frames x 6890 x 3
    for humanIdx in range(len(g_meshes)):


        if False:#humanIdx ==1:        #Second one red
            glColor3f(0.7, 0.5, 0.5)   #targer: spin color
            # glColor3f(0.53, 0.53, 0.8)   #SPIN: red
            # glColor3f(0.4, 0.4, 0.7)   #prediction: blue
        elif 'color' in g_meshes[humanIdx].keys():
            glColor3f(g_meshes[humanIdx]['color'][0]/255.0, g_meshes[humanIdx]['color'][1]/255.0, g_meshes[humanIdx]['color'][2]/255.0)
        else:
        #     glColor3f(0.4, 0.4, 0.7)   #prediction: blue
            glColor3f(g_meshColor[0],g_meshColor[1],g_meshColor[2])
            # glColor3f(0.53, 0.53, 0.8)
        # glColor3f(g_meshColor[0],g_meshColor[1],g_meshColor[2])


        if(g_frameIdx >= g_meshes[humanIdx]['ver'].shape[0]):
            continue

        # if(humanIdx==0 or humanIdx==1):       #Debug
        #     continue

        SMPL_vts = g_meshes[humanIdx]['ver'][g_frameIdx,:,: ]  #6890x3
        SMPL_inds = g_meshes[humanIdx]['f']
        SMPL_vts = SMPL_vts.flatten() #* MESH_SCALING
        SMPL_vts = SMPL_vts.astype(np.float32)
        SMPL_inds = SMPL_inds.flatten()
        vts_num = int(len(SMPL_vts) / 3)
        face_num = int(len(SMPL_inds) / 3)

        if 'normal' in g_meshes[humanIdx].keys() and g_meshes[humanIdx]['normal'] is None:
            here=0
        if 'normal' in g_meshes[humanIdx].keys() and len(g_meshes[humanIdx]['normal'])>0:
            SMPL_normals = g_meshes[humanIdx]['normal'][g_frameIdx,:,: ]
            SMPL_normals = SMPL_normals.flatten().astype(np.float32)
            tangent = np.zeros(SMPL_normals.shape)
            tangent.astype(np.float32)
            normal_num = len(SMPL_normals) / 3
        else :
            SMPL_normals = None
        
        glBindVertexArray(g_vao)

        glEnableVertexAttribArray(0)
        glBindBuffer(GL_ARRAY_BUFFER, g_vertex_buffer[humanIdx])
        glBufferData(GL_ARRAY_BUFFER, len(SMPL_vts) * sizeof(ctypes.c_float), (ctypes.c_float * len(SMPL_vts))(*SMPL_vts), GL_STATIC_DRAW)
        glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, None)

        if SMPL_normals is not None:
            #this is not needed.. .but normal should be the third attribute...
            if g_firsttime:
                glEnableVertexAttribArray(1)
                glBindBuffer(GL_ARRAY_BUFFER, g_tangent_buffer[humanIdx])
                glBufferData(GL_ARRAY_BUFFER, len(tangent) * sizeof(ctypes.c_float), (ctypes.c_float * len(tangent))(*tangent), GL_STATIC_DRAW)
                glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 0, None)

            glEnableVertexAttribArray(2)
            glBindBuffer(GL_ARRAY_BUFFER, g_normal_buffer[humanIdx])
            glVertexAttribPointer(2, 3, GL_FLOAT, GL_FALSE, 0, None)
            #glBufferData(GL_ARRAY_BUFFER, len(SMPL_normals) * sizeof(ctypes.c_float), (ctypes.c_float * len(SMPL_normals))(*SMPL_normals), GL_STATIC_DRAW)
            #glBufferData(GL_ARRAY_BUFFER, len(SMPL_normals) * sizeof(ctypes.c_float), (ctypes.c_float * len(SMPL_normals))(*SMPL_normals), GL_STATIC_DRAW)
            glBufferData(GL_ARRAY_BUFFER,
                            len(SMPL_normals) * sizeof(ctypes.c_float),
                            SMPL_normals,
                            GL_STATIC_DRAW)

        if True:#g_firsttime:
            glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, g_index_buffer[humanIdx])
            glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(ctypes.c_uint) * len(SMPL_inds), (ctypes.c_uint * len(SMPL_inds))(*SMPL_inds), GL_STATIC_DRAW)

        g_firsttime = False
        # set the dimensions of the position attribute, so it consumes 2 floats at a time (default is 4)

        #Draw by vertex array
        glPushMatrix()

        # if humanIdx ==0:
        #     #glColor3f(0.5, 0.2, 0.2)   #targer: red
        #     glColor3f(0.4, 0.4, 0.7)   #prediction: blue
        # elif humanIdx ==1:
        #     glColor3f(0.5, 0.2, 0.2)   #targer: red
        #     #glColor3f(0.0, 0.8, 0.3)   #buyer
        # elif humanIdx ==2:
        #     glColor3f(0.0, 0.8, 0.3)   #buyer
        #     #glColor3f(0.5, 0.5, 0)   #another: yellow
        # else:
        # #elif humanIdx ==2:
        #     glColor3f(0.5, 0.5, 0)   #another: yellow
        # glColor3f(0.6, 0.6, 0.4)   #another: yellow

        #glColor3f(0.8, 0.8, 0.8)
        glLineWidth(.5)

        if SMPL_normals is not None and g_bShowWiredMesh==False:
            glPolygonMode(GL_FRONT, GL_FILL)
            glPolygonMode(GL_BACK, GL_FILL)
        else:
            glPolygonMode(GL_FRONT, GL_LINE)
            glPolygonMode(GL_BACK, GL_LINE)
        # for vao_object in g_vao_object_list:
        #     glBindVertexArray(vao_object)
        #     glDrawElements(GL_TRIANGLES, face_num * 3, GL_UNSIGNED_INT, None)
        if g_bApplyRootOffset:
            #glTranslatef(40*humanIdx,0,0)
            glTranslatef(ROOT_OFFSET_DIST*humanIdx,0,0)

        glDrawElements(GL_TRIANGLES, face_num * 3, GL_UNSIGNED_INT, None)
        # glPolygonMode(GL_FRONT, GL_FILL)
        # glPolygonMode(GL_BACK, GL_FILL)
        glPopMatrix()

        glDisableVertexAttribArray(0)

        if SMPL_normals is not None:
            glDisableVertexAttribArray(1)
            glDisableVertexAttribArray(2)


    glBindBuffer(GL_ARRAY_BUFFER, 0)
    glBindVertexArray(0)    #Unbind

# #normal3D: 3 dim
# #joints70: 3x70 =210 dim
# def drawfaceNormal_70(normal3D, joints,  color):

#     glLineWidth(2.0)
#     #Visualize Joints
#     glColor3ub(color[0], color[1], color[2])

#     #faceEye
#     eyeCenterPoint = 0.5 *(joints[(45*3):(45*3+3)] + joints[(36*3):(36*3+3)])
#     normalEndPoint = eyeCenterPoint+ normal3D * 20

#     glBegin(GL_LINES)
#     glVertex3f(eyeCenterPoint[0], eyeCenterPoint[1], eyeCenterPoint[2])
#     glVertex3f(normalEndPoint[0], normalEndPoint[1], normalEndPoint[2])
#     glEnd()

#     glPushMatrix()
#     glTranslate(normalEndPoint[0], normalEndPoint[1], normalEndPoint[2])
#     glutSolidSphere(1, 10, 10)
#     glPopMatrix()


# #normal3D: 3 dim
# #joints: 3x21
# def drawbodyNormal(normal3D, rootPts,  color):

#     glLineWidth(2.0)
#     #Visualize Joints
#     glColor3ub(color[0], color[1], color[2])

#     #neckPoint = joints[(0*3):(0*3+3)]
#     normalEndPoint = rootPts+ normal3D * 40

#     glBegin(GL_LINES)
#     glVertex3f(rootPts[0], rootPts[1], rootPts[2])
#     glVertex3f(normalEndPoint[0], normalEndPoint[1], normalEndPoint[2])
#     glEnd()

#     glPushMatrix()
#     glTranslate(normalEndPoint[0], normalEndPoint[1], normalEndPoint[2])
#     glutSolidSphere(1, 10, 10)
#     glPopMatrix()


#normal3D: 3 dim
#rootPt: 3x1
def drawNormal(normal3D, rootPt,  color, normalLength=40):

    glLineWidth(2.0)
    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])

    #neckPoint = joints[(0*3):(0*3+3)]
    normalEndPoint = rootPt+ normal3D * normalLength

    glBegin(GL_LINES)
    glVertex3f(rootPt[0], rootPt[1], rootPt[2])
    glVertex3f(normalEndPoint[0], normalEndPoint[1], normalEndPoint[2])
    glEnd()

    glPushMatrix()
    glTranslate(normalEndPoint[0], normalEndPoint[1], normalEndPoint[2])
    glutSolidSphere(1, 10, 10)
    glPopMatrix()


g_connMat_hand21 = [ [0,1], [1,2], [2,3], [3,4],
                     [0,5], [5,6], [6,7], [7,8],
                     [0,9], [9,10], [10,11],[11,12],
                     [0,13], [13,14], [14, 15], [15, 16],
                     [0, 17], [17, 18], [18, 19], [19, 20]]
g_connMat_hand21 = np.array(g_connMat_hand21, dtype=int)
#joints70: 3x70 =210 dim
def drawhand_21(joints,  color, normal=None):

    # #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    # for i in range(len(joints)/3):

    #     glPushMatrix()
    #     glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
    #     glutSolidSphere(1, 10, 10)
    #     glPopMatrix()

    connMat_coco19 = g_connMat_hand21
    #Visualize Bones
    for conn in connMat_coco19:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]]
        y1 = joints[3*conn[1]+1]
        z1 = joints[3*conn[1]+2]


        if (x0==0) or (x1==0):
            continue

        x1 -= x0
        y1 -= y0
        z1 -= z0

        if abs(x1) >20 or abs(y1)>20 or abs(z1)>20:
            continue

        if g_bSimpleHead and conn[0] == 0 and conn[1]==1:
            x1 = x1*0.5
            y1 = y1*0.5
            z1 = z1*0.5

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(0.5, length, 10, 10)
        glPopMatrix()


# draw human location only
def DrawPosOnly():
    global g_colors
    global g_frameIdx#, g_normals
    global g_speech,g_speechGT
    global g_bApplyRootOffset

    global g_posOnly
    global g_bodyNormals, g_faceNormals

    #print(g_frameIdx)
    if g_posOnly is None:
        return

    for humanIdx in range(len(g_posOnly)):

        if(g_frameIdx >= g_posOnly[humanIdx].shape[1]):
            continue

        # if(humanIdx==1):
        #     continue


        skel = g_posOnly[humanIdx][:, g_frameIdx]
        # normal = g_normals[skelIdx, :, g_idx]
        color = g_colors[humanIdx % len(g_colors)]
        glColor3ub(color[0], color[1], color[2])


        glPushMatrix()
        glTranslate(skel[0], skel[1], skel[2])

        if humanIdx==0:#or humanIdx==1:
            glutSolidCube(10)
        elif False:#humanIdx==1:#or humanIdx==1:
            glutWireSphere(10,10,10)
        else:
            glutSolidSphere(10,10,10)
        glPopMatrix()

        #Draw body normal
        if g_bodyNormals is not None and len(g_bodyNormals)> humanIdx:

            if g_bodyNormals[humanIdx].shape[1]>g_frameIdx:
                normal3D = g_bodyNormals[humanIdx][:,g_frameIdx]  #3x1
                #drawbodyNormal(normal3D, skel, [255, 0, 0])
                drawNormal(normal3D, skel, [255, 0, 0])


        #Draw face normal
        if g_faceNormals is not None and len(g_faceNormals)> humanIdx:

            if g_faceNormals[humanIdx].shape[1]>g_frameIdx:
                normal3D = g_faceNormals[humanIdx][:,g_frameIdx]  #3x1
                drawNormal(normal3D, skel, [0, 255, 0], normalLength = 30)


        #Draw Speaking Info
        if g_speech is not None and g_skeletons is None:    #If g_skeletons is valid, draw speech there
            if len(g_speech[humanIdx]['indicator'])>g_frameIdx:
                draw_speaking_general(skel, g_speech[humanIdx]['indicator'][g_frameIdx], g_speech[humanIdx]['word'][g_frameIdx],  g_colors[humanIdx % len(g_colors)], offset=[0,0,23])





#g_trajectory: a list of np.array. skeNum x  (trajDim:3, frames)
def DrawTrajectory():
    global g_colors
    global g_trajectory, g_frameIdx#, g_normals


    if g_trajectory is None:
        return

    for humanIdx in range(len(g_trajectory)):

        if(g_frameIdx >= g_trajectory[humanIdx].shape[1]):
            continue

        color = g_colors[humanIdx % len(g_colors)]

        glLineWidth(2.0)

        #Visualize Joints
        glColor3ub(color[0], color[1], color[2])

        glPushMatrix()
        if g_bApplyRootOffset:
            #glTranslatef(40*humanIdx,0,0)
            glTranslatef(ROOT_OFFSET_DIST*humanIdx,0,0)

        #Visualize All Trajctory locations
        #interval=10
        for idx in range(0,g_trajectory[humanIdx].shape[1]-interval, interval):

            # root location
            x0 = g_trajectory[humanIdx][0,idx]
            y0 = g_trajectory[humanIdx][1,idx]
            z0 = g_trajectory[humanIdx][2,idx]

            x1 = g_trajectory[humanIdx][0,idx+interval]
            y1 = g_trajectory[humanIdx][1,idx+interval]
            z1 = g_trajectory[humanIdx][2,idx+interval]

            glBegin(GL_LINES)
            glVertex3f(x0, y0, z0)
            glVertex3f(x1,y1,z1)
            glEnd()

        """
        #Visualize current location and normal dirctio
        # x0, y0, z0 for the root location
        x0 = g_trajectory[humanIdx][0,g_frameIdx]
        y0 = g_trajectory[humanIdx][1,g_frameIdx]
        z0 = g_trajectory[humanIdx][2,g_frameIdx]

        x1 = g_trajectory[humanIdx][3,g_frameIdx]
        y1 = g_trajectory[humanIdx][4,g_frameIdx]
        z1 = g_trajectory[humanIdx][5,g_frameIdx]

        glBegin(GL_LINES)
        glVertex3f(x0, y0, z0)
        glVertex3f(x1,y1,z1)
        glEnd()
        """

        glPopMatrix()
        # drawface_70(face3D, g_colors[humanIdx % len(g_colors)])

        # if g_faceNormals is not None and len(g_faceNormals)> humanIdx:

        #     if g_faceNormals[humanIdx].shape[1]<=g_frameIdx:
        #         print("Warning: g_faceNormals[humanIdx].shape[2]<=g_frameId")
        #         continue

        #     normal3D = g_faceNormals[humanIdx][:,g_frameIdx]  #3x1
        #     #drawfaceNormal_70(normal3D, face3D, g_colors[humanIdx % len(g_colors)])
        #     #drawfaceNormal_70(normal3D, face3D, [0, 255, 255])
        #     eyeCenterPoint = 0.5 *(face3D[(45*3):(45*3+3)] + face3D[(36*3):(36*3+3)])
        #     drawNormal(normal3D, eyeCenterPoint, [0, 255, 255])


#g_skeletons should be: #(skeNum, skelDim, skelFrames)
def DrawSkeletons():
    # global g_colors
    global g_skeletons, g_frameIdx#, g_normals
    global g_speech,g_speechGT
    global g_bApplyRootOffset
    global g_bodyNormals
    if g_skeletons is None:
        return

    #for humanIdx in range(g_skeletons.shape[0]):
    for humanIdx in range(len(g_skeletons)):

        if g_onlyDrawHumanIdx>=0 and humanIdx!=g_onlyDrawHumanIdx:
            continue
       
        if(g_frameIdx >= g_skeletons[humanIdx]['skeleton'].shape[1]):
            continue

        #     continue

        skel = g_skeletons[humanIdx]['skeleton'][:, g_frameIdx]
        skel_type = g_skeletons[humanIdx]['type']
        skel_color = g_skeletons[humanIdx]['color']

        if skel_color is None:
            skel_color = g_colors[humanIdx % len(g_colors)]

        if g_bApplyRootOffset:
            skel = skel.copy()
            #skel[0::3] = skel[0::3]+ 70 *humanIdx
            skel[0::3] = skel[0::3]+ ROOT_OFFSET_DIST *humanIdx


        #If type is not specified
        if skel_type == None:   
            if skel.shape[0]==78:       #SMPlCOCO19 + headtop (19) + (leftFoot --toe20 pink21 heel22) + (rightFoot--toe23-pink24-heel25)
                drawbody_SMPLCOCO_TotalCap26(skel, [0,255,0])
            elif skel.shape[0]==57:
                drawbody_SMC19(skel, skel_color) #Panoptic Studio (SMC19) with 19 joints. Note SMC21 includes headtop
            
            elif skel.shape[0]==42: #LSP14joints also used in HMR
                drawbody_LSP14(skel, skel_color)
            elif skel.shape[0]==51: #simpler human36m (17joints)
                drawbody_joint17_human36m(skel, [0,255,0])#skel_color)
            elif skel.shape[0]==72: #SMPL LBS skeleton
                drawbody_joint24_smplLBS(skel, skel_color)
            elif skel.shape[0]==96: #human36m (32joints)
                drawbody_joint32_human36m(skel, skel_color)
            elif skel.shape[0]==66: #Holden's converted form (21joints)
                drawbody_joint22(skel, skel_color)
            elif skel.shape[0]==93: #CMU Mocap Raw data (31joints)
                drawbody_joint31(skel, skel_color)
            elif skel.shape[0]==186: #Adam model (62 joints: 22 for body 20x2 for fingers)
                drawbody_jointAdam(skel, skel_color)
            elif skel.shape[0]==189: #MTC skeleton from Hand2Body paper vis
                drawbody_jointMTC86(skel, skel_color)
            
            elif skel.shape[0]==54: #OpenPose 18
                drawbody_jointOpenPose18(skel, skel_color)

            elif skel.shape[0]==147: #SPIN 49 (25 openpose +  24 superset)
                drawbody_jointSpin49(skel, skel_color)
            else:
                drawbody_joint_ptOnly(skel, skel_color)

        elif skel_type == "smplcoco":
            drawbody_SMPLCOCO19(skel, skel_color)   #from HMR, SMPL->COCO19 regression. Same as MTC20's first 19.
        elif skel_type == "spin":
                if skel.shape[0]==72: #SPIN 24 (without openpose, 24 superset)
                    drawbody_jointSpin24(skel, skel_color)
                elif skel.shape[0]==147: #SPIN 49 (25 openpose +  24 superset)
                    drawbody_jointSpin49(skel, skel_color)
                else:
                    assert False
        #finger
        elif skel_type == "hand_smplx":
            if skel.shape[0]==63: #SPIN 49 (25 openpose +  24 superset)
                drawhand_joint21(skel, skel_color, type=skel_type)
            else:
                assert False

        elif skel_type == "hand_panopticdb":
            if skel.shape[0]==63: #SPIN 49 (25 openpose +  24 superset)
                drawhand_joint21(skel, skel_color, type=skel_type)
            else:
                assert False
                
        if g_bodyNormals is not None and len(g_bodyNormals)> humanIdx:

            if g_bodyNormals[humanIdx].shape[1]<=g_frameIdx:
                print("Warning: g_bodyNormals[humanIdx].shape[2]<=g_frameId")
                continue

            normal3D = g_bodyNormals[humanIdx][:,g_frameIdx]  #3x1
            #drawbodyNormal(normal3D, skel, [255, 0, 0])
            rootPt = skel[(0*3):(0*3+3)]

            if skel.shape[0]==66:
                i=12
                rootPt = skel[(3*i):(3*i+3)]
            # drawNormal(normal3D, rootPt, [0, 255, 255])
            #drawNormal(normal3D, rootPt, [255, 0, 0])

            # if g_onlyDrawHumanIdx>=0 and humanIdx!=g_onlyDrawHumanIdx:
            drawNormal(normal3D, rootPt, [0, 255, 0])

        # if g_faceNormals is not None and len(g_faceNormals)> humanIdx:

        #     if g_faceNormals[humanIdx].shape[1]<=g_frameIdx:
        #         print("Warning: g_bodyNormals[humanIdx].shape[2]<=g_frameId")
        #         continue

        #     normal3D = g_faceNormals[humanIdx][:,g_frameIdx]  #3x1
        #     #drawbodyNormal(normal3D, skel, [255, 0, 0])
        #     rootPt = skel[(0*3):(0*3+3)]

        #     if skel.shape[0]==66:
        #         i=13
        #         rootPt = skel[(3*i):(3*i+3)]
        #     # drawNormal(normal3D, rootPt, [0, 255, 255])
        #     drawNormal(normal3D, rootPt, [0, 255, 0])

        # #Draw Speeking Annotation
        # if skel.shape[0]==57 and g_speech is not None: #Coco19
        #     if(g_frameIdx < len(g_speech[humanIdx]['word'])):
        #         #draw_speaking_joint19(skel, g_speech[humanIdx]['indicator'][g_frameIdx], g_speech[humanIdx]['word'][g_frameIdx],  g_colors[humanIdx % len(g_colors)])
        #         draw_speaking_joint19(skel, g_speech[humanIdx]['indicator'][g_frameIdx], None,  skel_color)

        # if skel.shape[0]==66 and g_speech is not None: #Holden's
        #     if(len(g_speech)> humanIdx and  g_frameIdx < len(g_speech[humanIdx]['word'])):
        #         #draw_speaking_joint22(skel, g_speech[humanIdx][g_frameIdx],None,  g_colors[humanIdx % len(g_colors)])
        #         i=13
        #         facePt = skel[(3*i):(3*i+3)]
        #         draw_speaking_general(facePt, g_speech[humanIdx]['indicator'][g_frameIdx], g_speech[humanIdx]['word'][g_frameIdx],  [0, 0, 255], offset=[0,-43,0])

        # if skel.shape[0]==66 and g_speechGT is not None: #Holden's
        #     if(g_frameIdx < len(g_speechGT[humanIdx]['word'])):
        #         #draw_speaking_joint22(skel, g_speechGT[humanIdx][g_frameIdx],"GT: Speaking", [255, 0, 0], 40)
        #         i=13
        #         facePt = skel[(3*i):(3*i+3)]
        #         draw_speaking_general(facePt, g_speechGT[humanIdx]['indicator'][g_frameIdx], g_speechGT[humanIdx]['word'][g_frameIdx],  g_colors[humanIdx % len(g_colors)], offset=[0,-33,0])




#g_skeletons should be: #(skeNum, skelDim, skelFrames)
def DrawSkeletonsGT():
    
    assert False       #Deprecated

    global g_colors
    # global g_skeletons_GT, g_frameIdx#, g_normals
    global g_speech,g_speechGT
    global g_bApplyRootOffset
    global g_bodyNormals
    #print(g_frameIdx)
    if g_skeletons_GT is None:
        return

    #frameLimit = g_skeletons.shape[2]
    #frameLens = [l.shape[1] for l in g_skeletons]
    #g_frameLimit = min(frameLens)

    #for humanIdx in range(g_skeletons.shape[0]):
    for humanIdx in range(len(g_skeletons_GT)):
        # if skelIdx ==0:
        #     # if g_idx+time_offset>=g_skeletons.shape[2]:
        #     #     continue
        #     skel = g_skeletons[skelIdx, :, g_idx+time_offset]
        #     # normal = g_normals[skelIdx, :, g_idx+time_offset]
        # else:
        #skel = g_skeletons[humanIdx, :, g_frameIdx]
        if(g_frameIdx >= g_skeletons_GT[humanIdx].shape[1]):
            continue

        skel = g_skeletons_GT[humanIdx][:, g_frameIdx]
        # normal = g_normals[skelIdx, :, g_idx]

        if g_bApplyRootOffset:
            skel = skel.copy()
            #skel[0::3] = skel[0::3]+ 70 *humanIdx
            skel[0::3] = skel[0::3]+ ROOT_OFFSET_DIST *humanIdx

        if skel.shape[0]==78:       #SMPlCOCO19 + headtop (19) + (leftFoot --toe20 pink21 heel22) + (rightFoot--toe23-pink24-heel25)
            drawbody_SMPLCOCO_TotalCap26(skel, [0,255,0])
        elif skel.shape[0]==57: #Panoptic Studio (SMC19) with 19 joints. Note SMC21 includes headtop
            drawbody_SMC19(skel, g_colors[humanIdx % len(g_colors)])
        elif skel.shape[0]==96: #human36
            drawbody_joint32_human36m(skel, g_colors[humanIdx % len(g_colors)])
        elif skel.shape[0]==66: #Holden's converted form
            drawbody_joint22(skel, g_colors[humanIdx % len(g_colors)])
        elif skel.shape[0]==93: #CMU Mocap Raw data (31joints)
            drawbody_joint31(skel, g_colors[humanIdx % len(g_colors)])
        else:
            drawbody_joint_ptOnly(skel, g_colors[humanIdx % len(g_colors)])

        # if False:#g_bodyNormals is not None and len(g_bodyNormals)> humanIdx:

        #     if g_bodyNormals[humanIdx].shape[1]<=g_frameIdx:
        #         print("Warning: g_bodyNormals[humanIdx].shape[2]<=g_frameId")
        #         continue

        #     normal3D = g_bodyNormals[humanIdx][:,g_frameIdx]  #3x1
        #     #drawbodyNormal(normal3D, skel, [255, 0, 0])
        #     rootPt = skel[(0*3):(0*3+3)]
        #     drawNormal(normal3D, rootPt, [255, 0, 0])


        # #Draw Speeking Annotation
        # if skel.shape[0]==57 and g_speech is not None: #Coco19
        #     if(g_frameIdx < len(g_speech[humanIdx]['word'])):
        #         draw_speaking_joint19(skel, g_speech[humanIdx]['indicator'][g_frameIdx], g_speech[humanIdx]['word'][g_frameIdx],  g_colors[humanIdx % len(g_colors)])

        # # if skel.shape[0]==66 and g_speech is not None: #Holden's
        # #     if(g_frameIdx < len(g_speech[humanIdx])):
        # #         draw_speaking_joint22(skel, g_speech[humanIdx][g_frameIdx],None,  g_colors[humanIdx % len(g_colors)])

        # # if skel.shape[0]==66 and g_speechGT is not None: #Holden's
        # #     if(g_frameIdx < len(g_speechGT[humanIdx])):
        # #         draw_speaking_joint22(skel, g_speechGT[humanIdx][g_frameIdx],"GT: Speaking", [255, 0, 0], 40)



#LSP14 format (used in HMR)
# Right ankle 1
# Right knee 2
# Right hip 3
# Left hip 4
# Left knee 5
# Left ankle 6
# Right wrist 7
# Right elbow 8
# Right shoulder 9
# Left shoulder 10
# Left elbow 11
# Left wrist 12
# Neck 13
# Head top 14
g_connMat_lsp14 = [ [13,3], [3,2], [2,1], #Right leg
                     [13,4], [4,5], [5,6], #Left leg
                     [13,10], [10,11], [11,12], #Left Arm
                     [13,9], [9,8], [8,7], #Right shoulder
                     [13,14]    #Nect -> Headtop
                     ]
g_connMat_lsp14 = np.array(g_connMat_lsp14, dtype=int) - 1 #zero Idx

#LSP14 format (used in HMR)
def drawbody_LSP14(joints,  color):

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(int(len(joints)/3)):

        if g_bSimpleHead and (i>=15 or i==1):
            continue
        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(2, 10, 10)
        glPopMatrix()

    connMat = g_connMat_lsp14
    #Visualize Bones
    for conn in connMat:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0


        if g_bSimpleHead and conn[0] == 0 and conn[1]==1:
            x1 = x1*0.5
            y1 = y1*0.5
            z1 = z1*0.5

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(2, length, 10, 10)
        glPopMatrix()




#In Zero index
#This is regressed joint from SMPL model defined from HMR
#MTC20 is exactly same as this except it has one more joint on spine (19)
g_connMat_smplcoco19 = [ [12,2], [2,1], [1,0], #Right leg
                     [12,3], [3,4], [4,5], #Left leg
                     [12,9], [9,10], [10,11], #Left Arm
                     [12,8], [8,7], [7,6], #Right shoulder
                      [12,14],[14,16],[16,18],  #Neck(12)->Nose(14)->rightEye(16)->rightEar(18)
                      [14,15],[15,17],   #Nose(14)->leftEye(15)->leftEar(17).
                      [14,13] #Nose->headTop(13)
                     ]
g_connMat_smplcoco19 = np.array(g_connMat_smplcoco19, dtype=int)  #zero Idx
def drawbody_SMPLCOCO19(joints,  color, normal=None):


    bBoneIsLeft =  [ 0,0,0,
            1, 1, 1,
            1,1,1,
            0,0,0,
            1,0,0,
            1,1,
            1] #To draw left as different color. Torso is treated as left


    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(int(len(joints)/3)):

        # if i!=2:
        #     continue
        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(2, 10, 10)
        glPopMatrix()

    connMat = g_connMat_smplcoco19
    #Visualize Bones
    for i, conn in enumerate(connMat):
        if bBoneIsLeft[i]:          #Left as a color
            glColor3ub(color[0], color[1], color[2])
        else:       #Right as black
            glColor3ub(0,0,0)

        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(2, length, 10, 10)
        glPopMatrix()

    # Visualize Normals
    if normal is not None:
        i=1
        facePt = joints[(3*i):(3*i+3)]
        normalPt = facePt + normal*50

        glColor3ub(0, 255, 255)
        glPushMatrix()
        glTranslate(normalPt[0], normalPt[1], normalPt[2])
        glutSolidSphere(1, 10, 10)
        glPopMatrix()

        glBegin(GL_LINES)
        glVertex3f(facePt[0], facePt[1], facePt[2])
        glVertex3f(normalPt[0], normalPt[1],  normalPt[2])
        glEnd()




#Panoptic Studio's Skeleton Output (AKA SMC19). Note SM21 has headTop(19) and spine(20)
g_bSimpleHead = False
if g_bSimpleHead==False:
    g_connMat_smc19 = [ [1,2], [1,4], [4,5], [5,6], [1,3], [3,7], [7,8], [8,9], [3,13],[13,14], [14,15], [1,10], [10, 11], [11, 12], [2, 16], [16, 17], [2, 18], [18, 19] ]
else:
    g_connMat_smc19 = [ [1,2], [1,4], [4,5], [5,6], [1,3], [3,7], [7,8], [8,9], [3,13],[13,14], [14,15], [1,10], [10, 11], [11, 12]]
g_connMat_smc19 = np.array(g_connMat_smc19, dtype=int) - 1 #zero Idx
def drawbody_SMC19(joints,  color, normal=None):

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(int(len(joints)/3)):

        if g_bSimpleHead and (i>=15 or i==1):
            continue
        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(2, 10, 10)
        glPopMatrix()

    connMat = g_connMat_smc19
    #Visualize Bones
    for conn in connMat:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0


        if g_bSimpleHead and conn[0] == 0 and conn[1]==1:
            x1 = x1*0.5
            y1 = y1*0.5
            z1 = z1*0.5

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(2, length, 10, 10)
        glPopMatrix()

    # Visualize Normals
    if normal is not None:
        i=1
        facePt = joints[(3*i):(3*i+3)]
        normalPt = facePt + normal*50

        glColor3ub(0, 255, 255)
        glPushMatrix()
        glTranslate(normalPt[0], normalPt[1], normalPt[2])
        glutSolidSphere(1, 10, 10)
        glPopMatrix()

        glBegin(GL_LINES)
        glVertex3f(facePt[0], facePt[1], facePt[2])
        glVertex3f(normalPt[0], normalPt[1],  normalPt[2])
        glEnd()



#Panoptic Studio's Skeleton Output (AKA SMC19). Note SM21 has headTop(19) and spine(20)
def drawbody_SMPLCOCO_TotalCap26(joints,  color, normal=None):
    connMat = [ [12,2], [2,1], [1,0], #Right leg
                     [12,3], [3,4], [4,5], #Left leg
                     [12,9], [9,10], [10,11], #Left Arm
                     [12,8], [8,7], [7,6], #Right shoulder
                      [12,14],[14,16],[16,18],  #Neck(12)->Nose(14)->rightEye(16)->rightEar(18)
                      [14,15],[15,17],   #Nose(14)->leftEye(15)->leftEar(17).
                      [14,13], #Nose->headMidle(13)
                      [12,19],       #headTop19
                      [5,20], [5,21], [5,22],       #leftFoot
                      [0,23], [0,24], [0,25]       #rightFoot
                     ]
    connMat = np.array(connMat, dtype=int)  #zero Idx

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(int(len(joints)/3)):

        if g_bSimpleHead and (i>=15 or i==1):
            continue
        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(2, 10, 10)
        glPopMatrix()

    #Visualize Bones
    for conn in connMat:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0


        if g_bSimpleHead and conn[0] == 0 and conn[1]==1:
            x1 = x1*0.5
            y1 = y1*0.5
            z1 = z1*0.5

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(2, length, 10, 10)
        glPopMatrix()

    # Visualize Normals
    if normal is not None:
        i=1
        facePt = joints[(3*i):(3*i+3)]
        normalPt = facePt + normal*50

        glColor3ub(0, 255, 255)
        glPushMatrix()
        glTranslate(normalPt[0], normalPt[1], normalPt[2])
        glutSolidSphere(1, 10, 10)
        glPopMatrix()

        glBegin(GL_LINES)
        glVertex3f(facePt[0], facePt[1], facePt[2])
        glVertex3f(normalPt[0], normalPt[1],  normalPt[2])
        glEnd()



g_connMat_coco14 = [ [13,3], [3,2], [2,1], #Right leg
                     [13,4], [4,5], [5,6], #Left leg
                     [13,10], [10,11], [11,12], #Left Arm
                     [13,9], [9,8], [8,7], #Right shoulder
                     ]
g_connMat_coco14 = np.array(g_connMat_coco14, dtype=int) - 1 #zero Idx

def drawbody_joint14(joints,  color, normal=None):

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(int(len(joints)/3)):

        if g_bSimpleHead and (i>=15 or i==1):
            continue
        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(2, 10, 10)
        glPopMatrix()

    connMat_coco14 = g_connMat_coco14
    #Visualize Bones
    for conn in connMat_coco14:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0


        if g_bSimpleHead and conn[0] == 0 and conn[1]==1:
            x1 = x1*0.5
            y1 = y1*0.5
            z1 = z1*0.5

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(2, length, 10, 10)
        glPopMatrix()

    # Visualize Normals
    if normal is not None:
        i=1
        facePt = joints[(3*i):(3*i+3)]
        normalPt = facePt + normal*50

        glColor3ub(0, 255, 255)
        glPushMatrix()
        glTranslate(normalPt[0], normalPt[1], normalPt[2])
        glutSolidSphere(1, 10, 10)
        glPopMatrix()

        glBegin(GL_LINES)
        glVertex3f(facePt[0], facePt[1], facePt[2])
        glVertex3f(normalPt[0], normalPt[1],  normalPt[2])
        glEnd()


def RenderString(str):
    # glRasterPos3d(0,-2,0)
    for c in str:
        #glutBitmapCharacter(GLUT_BITMAP_8_BY_13, ord(c))
        glutBitmapCharacter(GLUT_BITMAP_HELVETICA_18, ord(c))


def draw_speaking_joint19(joints, bSpeak, word,  color, normal=None):

    # Visualize Speaking signal
    if bSpeak:
        i=1
        facePt = joints[(3*i):(3*i+3)]
        normalPt = facePt + np.array([0,-1,0])*20

        #glColor3ub(0, 255, 255)
        glColor3ub(color[0], color[1], color[2])

        glPushMatrix()

        glTranslate(normalPt[0], normalPt[1], normalPt[2])
        glutSolidSphere(1, 10, 10)
        #Render String
        if word is not None:
            RenderString(word)
        else:
            RenderString('speaking')


        glPopMatrix()

        glBegin(GL_LINES)
        glVertex3f(facePt[0], facePt[1], facePt[2])
        glVertex3f(normalPt[0], normalPt[1],  normalPt[2])
        glEnd()



def draw_speaking_joint22(joints, bSpeak, word,  color, offset=20, normal=None):

    # Visualize Speaking signal
    if bSpeak:
        i=13
        facePt = joints[(3*i):(3*i+3)]
        normalPt = facePt + np.array([0,-1,0])*offset

        #glColor3ub(0, 255, 255)
        glColor3ub(color[0], color[1], color[2])

        glPushMatrix()

        glTranslate(normalPt[0], normalPt[1], normalPt[2])
        glutSolidSphere(1, 10, 10)
        #Render String
        if word is not None:
            RenderString(word)
        else:
            RenderString('speaking')


        glPopMatrix()

        glBegin(GL_LINES)
        glVertex3f(facePt[0], facePt[1], facePt[2])
        glVertex3f(normalPt[0], normalPt[1],  normalPt[2])
        glEnd()

"""
    flagLength: show info on the upright direction
    offset: show info on the designated location
"""
def draw_speaking_general(facePt, bSpeak, word,  color, offsetLength=20, offset=None):

    # Visualize Speaking signal
    if bSpeak:
        #facePt = joints[(3*i):(3*i+3)]

        if offset is None:
            normalPt = facePt + np.array([0,-1,0])*offsetLength
        else:
            normalPt = facePt + offset

        #glColor3ub(0, 255, 255)
        glColor3ub(color[0], color[1], color[2])

        glPushMatrix()

        glTranslate(normalPt[0], normalPt[1], normalPt[2])
        glutSolidSphere(1, 10, 10)
        #Render String
        if word is not None:
            RenderString(word)
        else:
            RenderString('speaking')


        glPopMatrix()

        glBegin(GL_LINES)
        glVertex3f(facePt[0], facePt[1], facePt[2])
        glVertex3f(normalPt[0], normalPt[1],  normalPt[2])
        glEnd()

def drawbody_joint_ptOnly(joints,  color, normal=None):

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(int(len(joints)/3)):

        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(2, 10, 10)
        glPopMatrix()



#Human 36m DB's mocap data. 32 joints
g_connMat_joint32_human36m = [ [0,1],[1,2],[2,3],[3,4],[4,5], #RightLeg: root(0), rHip(1), rKnee(2), rAnkle(3), rFootMid(4), rFootEnd(5)
                     [0,6],[6,7],[7,8],[8,9], [9,10], #LeftLeg: root, lHip(6), lKnee(7), lAnkle(8), lFootMid(9), lFootEnd(10)
                     [11,12], [12,13], [13,14], [14,15], #root2(11), spineMid(12), neck(13), nose(14), head(15) #0,11 are the same points?
                     [16,17], [17,18], [18,19], [20,21], [20,22],   #Left Arms. neck(16==13), lshoulder(17),  lElbow(18), lWrist (19=20), lThumb(21), lMiddleFinger(22)
                     [24,25], [25,26], [26,27], [27,29], [27,30]   #Right Arm, neck(24==13), rshoulder(25),  rElbow(26), rWrist (27=28), rThumb(29), rMiddleFinger(30)
                     ]

def drawbody_joint32_human36m(joints,  color, normal=None):

    bBoneIsLeft = [0 ,0, 0, 0, 0,
                    1, 1, 1, 1, 1,
                    1, 1, 1, 1,
                    1, 1, 1, 1, 1,
                    0, 0, 0, 0, 0] #To draw left as different color. Torso is treated as left

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])

    for i in range(int(len(joints)/3)):

        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(2, 10, 10)
        glPopMatrix()

    #Visualize Bones
    for i, conn in enumerate(g_connMat_joint32_human36m):
        if bBoneIsLeft[i]:          #Left as a color
            glColor3ub(color[0], color[1], color[2])
        else:       #Right as black
            glColor3ub(0,0,0)
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(2, length, 10, 10)
        glPopMatrix()


# Simpler version of Human 36m DB's mocap data. with 17 joints
# 5 for torso + neck (root, spine, nect, nose, headtop)
g_connMat_joint17_human36m = [ [0,1],[1,2],[2,3],#,#root(0), rHip(1), rKnee(2), rAnkle(3)
                     [0,4],[4,5],[5,6],#,[8,9], [9,10], #root(0, lHip(4), lKnee(5), lAnkle(6)
                      [0,7], [7,8], [8,9], [9,10], #root(0, spineMid(7), neck(8), nose(9), head(10) #0,11 are the same points?
                     [8,11], [11,12], [12,13], #Left Arms. neck(8). lshoulder(11),  lElbow(12), lWrist (13)
                     [8,14], [14,15], [15,16] #Right Arm, neck(8), rshoulder(14),  rElbow(15), rWrist (16)
                     ]
def drawbody_joint17_human36m(joints,  color, normal=None):

    bBoneIsLeft = [0 ,0, 0,
                    1, 1, 1,
                    1, 1, 1, 1,
                    1, 1, 1,
                    0, 0, 0] #To draw left as different color. Torso is treated as left


    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(int(len(joints)/3)):

        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(2, 10, 10)
        glPopMatrix()

    #Visualize Bones
    for i, conn in enumerate(g_connMat_joint17_human36m):
        if bBoneIsLeft[i]:          #Left as a color
            glColor3ub(color[0], color[1], color[2])
        else:       #Right as black
            glColor3ub(0,0,0)

        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(2, length, 10, 10)
        glPopMatrix()


#SMPL 24 joints used for LBS
g_connMat_joint24_smpl = [ [0,3],[3,6],[6,9],[9,12],[12,15],  #root-> torso -> head
                     [9,13],[13,16],[16,18],[18,20],[20,22], #Nect-> left hand
                     [9,14], [14,17], [17,19], [19,21], [21,23],  #Nect-> right hand
                     [0,1], [1,4], [4,7], [7,10], # left Leg
                     [0,2], [2,5], [5,8], [8,11] #right leg
                     ]
def drawbody_joint24_smplLBS(joints, color, normal=None):

    bBoneIsLeft = [1,1,1,1,1,
                    1,1,1,1,1,
                    0,0,0,0,0,
                    0,0,0,0,
                    1,1,1,1] #To draw left as different color. Torso is treated as left


    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(int(len(joints)/3)):

        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(2, 10, 10)
        glPopMatrix()

    #Visualize Bones
    for i, conn in enumerate(g_connMat_joint24_smpl):
        if True:#bBoneIsLeft[i]:          #Left as a color
            glColor3ub(color[0], color[1], color[2])
        else:       #Right as black
            glColor3ub(0,0,0)

        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(2, length, 10, 10)
        glPopMatrix()


# g_connMat_coco31 = [ [0,1],[1,2],[2,3],[3,4],[4,5], #root(0), rHip(1), rKnee(2), rAnkle(3), RFootMid(4), rFootEnd(5)
#                      [0,6],[6,7],[7,8],[8,9], [9,10], #root, lHip(6), lKnee(7), lAnkle(8), lFootMid(9), lFootEnd(10)
#                      [11,12], [12,13], [13,14], [14,15], #root2(11), spineMid(12), neck(13), nose(14), head(15) #0,11 are the same points?
#                      [16,17], [17,18], [18,19], [18,19], [20,21], [20,22] ,
#                      [24,25], [25,26], [26,27], [27,29], [27,30] ]

#Note that left right of torso are flipped, compared human36m format
g_connMat_coco31 = [ [0,1],[1,2],[2,3],[3,4],[4,5], #root(0), lHip(1), lKnee(2), lAnkle(3), lFootMid(4), lFootEnd(5)
                     [0,6],[6,7],[7,8],[8,9], [9,10], #root, rHip(6), rKnee(7), lAnkle(8), lFootMid(9), lFootEnd(10)
                     [11,12], [12,13], [14,15],[15,16], #root2(11), lowerback(12), spine(13=14), neck(15), head(16) #(0,11), (13=14) are the same points
                     [17,18],[18,19],[19,20],[21,22],   # spine2(14=17), lShoulder(18), lElbow(19), lWrist(20==21), lHand(22), LThumb(23, not valid and same as 20)
                     [24,25],[25,26],[26,27],[28,29]   # spine2(14=24), rShoulder(25), rElbow(26), rWrist(27==28), rHand(29), rThumb(30, not valid and same as 27)
                     ]
# g_connMat_coco31 = [ [0,18]
#                       ]
def drawbody_joint31(joints,  color, normal=None):

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(int(len(joints)/3)):

        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(2, 10, 10)
        glPopMatrix()

    #Visualize Bones
    for conn in g_connMat_coco31:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(2, length, 10, 10)
        glPopMatrix()

# Adam model
# Body (22) + Lhand (20) +  Rhand
g_connMat_Adam = [ 
    [0,1], [1,4], [4,7], [7,10], #left leg
                    [0,2], [2,5], [5,8], [8,11], #right leg
                    [9,13], [13,16], [16,18], [18,20],   #left arm
                    [9,14], [14,17], [17,19], [19,21],   #right arm
                    [0,3], [3,6], [6,9], [9,12], [12,15],   #torso -> head

                    [20,22], [22,23], [23,24], [24,25],     #LeftHand Thumb
                    [20,26], [26,27], [27,28], [28,29],
                    [20,30], [30,31], [31,32], [32,33],
                    [20,34], [34,35], [35,36], [36,37],
                    [20,38], [38,39], [39,40], [40,41],     #Left Pinky


                    [21,42], [42,43], [43,44], [44,45],     #RightHand Thumb
                    [21,46], [46,47], [47,48], [48,49],
                    [21,50], [50,51], [51,52], [52,53],
                    [21,54], [54,55], [55,56], [56,57],
                    [21,58], [58,59], [59,60], [60,61]     #Right Pinky
                    ]
def drawbody_jointAdam(joints,  color, normal=None, ignore_root=False):

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(1,int(len(joints)/3)):
    # for i in range(22):

        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])

        if i<22:
            glutSolidSphere(2, 10, 10)
        else:
            glutSolidSphere(0.5, 10, 10)
        glPopMatrix()

    connMat_adam = g_connMat_Adam
    #Visualize Bones
    for conn in connMat_adam:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        if conn[1]<22:    #body
            glutSolidCone(2, length, 10, 10)
        else:
            glutSolidCone(0.4, length, 10, 10)
        glPopMatrix()

    #Spine to ground projection
    conn = [0,1]
    x0 = joints[3*conn[0]]
    y0 = joints[3*conn[0]+1]
    z0 = joints[3*conn[0]+2]
    # x1, y1, z1 is the vector points from the base to the target
    x1 = joints[3*conn[1]]
    y1 = joints[3*conn[1]+1]
    z1 = joints[3*conn[1]+2]

    glBegin(GL_LINES)
    glVertex3f(x0, y0, z0)
    glVertex3f(x1, y1, z1)
    glEnd()


# SMC21
g_connMat_MTC = [ 
                    [2,20], [20,0],  #Right Leg
                    [0,3], [3,4], [4,5], #Left Arm
                    [0,9], [9,10], [10,11], #Left Arm
                    [2,6], [6,7], [7,8],   #Left Leg
                    [2,12], [12,13], [13,14],  #Right Leg

                    [5,22], [22,23], [23,24], [24,25],     #LeftHand Thumb
                    [5,26], [26,27], [27,28], [28,29],
                    [5,30], [30,31], [31,32], [32,33],
                    [5,34], [34,35], [35,36], [36,37],
                    [5,38], [38,39], [39,40], [40,41],     #Left Pinky

                    [11,43], [43,44], [44,45], [45,46],    #RightHand Thumb
                    [11,47], [47,48], [48,49], [49,50],
                    [11,51], [51,52], [52,53], [53,54],
                    [11,55], [55,56], [56,57], [57,58],
                    [11,59], [59,60], [60,61], [61,62]     #Right Pinky
                    ]
def drawbody_jointMTC86(joints,  color, normal=None, ignore_root=False):

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(1,int(len(joints)/3)):
    # for i in range(22):

        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])

        if i<22:
            glutSolidSphere(2, 10, 10)
        else:
            glutSolidSphere(0.5, 10, 10)
        glPopMatrix()

    connMat_adam = g_connMat_MTC
    #Visualize Bones
    for conn in connMat_adam:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        if conn[1]<22:    #body
            glutSolidCone(2, length, 10, 10)
        else:
            glutSolidCone(0.4, length, 10, 10)
        glPopMatrix()

    #Spine to ground projection
    conn = [0,1]
    x0 = joints[3*conn[0]]
    y0 = joints[3*conn[0]+1]
    z0 = joints[3*conn[0]+2]
    # x1, y1, z1 is the vector points from the base to the target
    x1 = joints[3*conn[1]]
    y1 = joints[3*conn[1]+1]
    z1 = joints[3*conn[1]+2]

    glBegin(GL_LINES)
    glVertex3f(x0, y0, z0)
    glVertex3f(x1, y1, z1)
    glEnd()




def drawbody_jointSpin49(joints,  color, normal=None, ignore_root=False):

    # skelSize = 10
    skelSize = 2

    #Openpose25 + SPIN global 24
    link_openpose = [  [8,1], [1,0] , [0,16] , [16,18] , [0,15], [15,17],
                [1,2],[2,3],[3,4],      #Right Arm
                [1,5], [5,6], [6,7],       #Left Arm
                [8,12], [12,13], [13,14], [14,21], [14,19], [14,20],
                [8,9], [9,10], [10,11], [11,24], [11,22], [11,23]
                ]

    link_spin24 =[  [14,16], [16,12], [12,17] , [17,18] ,
                [12,9],[9,10],[10,11],      #Right Arm
                [12,8], [8,7], [7,6],       #Left Arm
                [14,3], [3,4], [4,5],
                [14,2], [2,1], [1,0]
                ]
    link_spin24 = np.array(link_spin24) + 25



    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    # for i in range(1,int(len(joints)/3)):
    for i in range(25,int(len(joints)/3)):

        # if i-25 == 13:
        #     glColor3f(1.0,0,0)
        # elif i-25 == 18:
        #     glColor3f(0,1.0,0)
        # elif i-25 == 15:
        #     glColor3f(1.0,1.0,0)
        # elif i-25 == 17:
        #     glColor3f(0.0,1.0,1.0)
        # else:
        #     glColor3f(0,0.0,1.0)

        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(skelSize, 10, 10)
        glPopMatrix()

    # #Visualize Bones
    # glColor3ub(0,0,255)
    # for conn in link_openpose:
    #     # x0, y0, z0 is the coordinate of the base point
    #     x0 = joints[3*conn[0]]
    #     y0 = joints[3*conn[0]+1]
    #     z0 = joints[3*conn[0]+2]
    #     # x1, y1, z1 is the vector points from the base to the target
    #     x1 = joints[3*conn[1]] - x0
    #     y1 = joints[3*conn[1]+1] - y0
    #     z1 = joints[3*conn[1]+2] - z0

    #     length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
    #     theta = math.degrees(math.acos(z1/length))
    #     phi = math.degrees(math.atan2(y1, x1))

    #     glPushMatrix()
    #     glTranslate(x0, y0, z0)
    #     glRotatef(phi, 0, 0, 1)
    #     glRotatef(theta, 0, 1, 0)
    #     glutSolidCone(0.4, length, 10, 10)
    #     glPopMatrix()

    #Visualize Bones
    # glColor3ub(255,0,0)
    glColor3ub(color[0], color[1], color[2])
    for conn in link_spin24:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(skelSize, length, 10, 10)

        glPopMatrix()

    #Spine to ground projection
    conn = [0,1]
    x0 = joints[3*conn[0]]
    y0 = joints[3*conn[0]+1]
    z0 = joints[3*conn[0]+2]
    # x1, y1, z1 is the vector points from the base to the target
    x1 = joints[3*conn[1]]
    y1 = joints[3*conn[1]+1]
    z1 = joints[3*conn[1]+2]

    glBegin(GL_LINES)
    glVertex3f(x0, y0, z0)
    glVertex3f(x1, y1, z1)
    glEnd()


def drawbody_jointOpenPose18(joints,  color, normal=None, ignore_root=False):

    skelSize = 2
    #Openpose 18
    link_openpose = [    [1,0] , [0,14] , [14,16] , [0,15], [15,17],
                [1,2],[2,3],[3,4],      #Right Arm
                [1,5], [5,6], [6,7],       #Left Arm
                [1,11], [11,12], [12,13],       #Left Leg
                [8,1], [8,9], [9,10]       #Right Leg
                ]

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(1,int(len(joints)/3)):
        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(skelSize, 10, 10)
        glPopMatrix()

    #Visualize Bones
    glColor3ub(color[0], color[1], color[2])
    for conn in link_openpose:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(skelSize, length, 10, 10)

        glPopMatrix()

    #Spine to ground projection
    conn = [0,1]
    x0 = joints[3*conn[0]]
    y0 = joints[3*conn[0]+1]
    z0 = joints[3*conn[0]+2]
    # x1, y1, z1 is the vector points from the base to the target
    x1 = joints[3*conn[1]]
    y1 = joints[3*conn[1]+1]
    z1 = joints[3*conn[1]+2]

    glBegin(GL_LINES)
    glVertex3f(x0, y0, z0)
    glVertex3f(x1, y1, z1)
    glEnd()


def drawbody_jointSpin24(joints,  color, normal=None, ignore_root=False):

    link_spin24 =[  #[14,16], [16,12]       #TODO: debug
                    # [12,17] , [17,18] ,
                [12,9],[9,10],[10,11],      #Right Arm
                [12,8], [8,7], [7,6],       #Left Arm
                [14,3], [3,4], [4,5],
                [14,2], [2,1], [1,0]
                ]
    link_spin24 = np.array(link_spin24)

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(1,int(len(joints)/3)):
    # for i in range(22):

        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])

        glutSolidSphere(2, 10, 10)
        glPopMatrix()

    #Visualize Bones
    # glColor3ub(255,0,0)
    glColor3ub(color[0], color[1], color[2])
    for conn in link_spin24:
        if joints[3*conn[0]]==0 or joints[3*conn[1]]==0 :
            continue
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        if length>0.001:
            theta = math.degrees(math.acos(z1/length))
            phi = math.degrees(math.atan2(y1, x1))

            glPushMatrix()
            glTranslate(x0, y0, z0)
            glRotatef(phi, 0, 0, 1)
            glRotatef(theta, 0, 1, 0)
            glutSolidCone(2, length, 10, 10)
            # glutSolidCone(20, length, 10, 10)

            glPopMatrix()

    #Spine to ground projection
    conn = [0,1]
    x0 = joints[3*conn[0]]
    y0 = joints[3*conn[0]+1]
    z0 = joints[3*conn[0]+2]
    # x1, y1, z1 is the vector points from the base to the target
    x1 = joints[3*conn[1]]
    y1 = joints[3*conn[1]+1]
    z1 = joints[3*conn[1]+2]

    glBegin(GL_LINES)
    glVertex3f(x0, y0, z0)
    glVertex3f(x1, y1, z1)
    glEnd()




# 0 : Wrist
# 1 : Index_00
# 2 : Index_01
# 3 : Index_02
# 4 : Middle_00
# 5 : Middle_01
# 6 : Middle_02
# 7 : Little_00
# 8 : Little_01
# 9 : Little_02
# 10 : Ring_00
# 11 : Ring_01
# 12 : Ring_02
# 13 : Thumb_00
# 14 : Thumb_01
# 15 : Thumb_02
# 16 : Index_03
# 17 : Middle_03
# 18 : Little_03
# 19 : Ring_03
# 20 : Thumb_03
def drawhand_joint21(joints,  color, normal=None, ignore_root=False, type = 'hand_smplx'):
    if type=="hand_panopticdb":
        link_panoptic_hand =  [ [0,1], [1,2], [2,3], [3,4],  #thumb
                    [0,5], [5,6],[6,7],[7,8],   #index
                    [0,9],[9,10],[10,11],[11,12],
                    [0,13],[13,14],[14,15],[15,16],
                    [0,17],[17,18],[18,19],[19,20]
                    ]
        link_hand = np.array(link_panoptic_hand)
    else:
        link_smplx_hand =[  [0,1], [1,2], [2,3], [3,16],
                    [0,4],[4,5],[5,6], [6,17],
                    [0,7], [7,8], [8,9],  [9,18],
                    [0,10], [10,11], [11,12],[12,19],
                    [0,13], [13,14], [14,15],[15,20]
                    ]
        link_hand = np.array(link_smplx_hand)

    

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(1,int(len(joints)/3)):
    # for i in range(22):

        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])

        glutSolidSphere(0,5, 10, 10)
        glPopMatrix()

    #Visualize Bones
    # glColor3ub(255,0,0)
    glColor3ub(color[0], color[1], color[2])
    for conn in link_hand:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        if length>0.001:
            theta = math.degrees(math.acos(z1/length))
            phi = math.degrees(math.atan2(y1, x1))

            glPushMatrix()
            glTranslate(x0, y0, z0)
            glRotatef(phi, 0, 0, 1)
            glRotatef(theta, 0, 1, 0)
            # glutSolidCone(2, length, 10, 10)
            glutSolidCone(0.5, length, 10, 10)

            glPopMatrix()

    #Spine to ground projection
    conn = [0,1]
    x0 = joints[3*conn[0]]
    y0 = joints[3*conn[0]+1]
    z0 = joints[3*conn[0]+2]
    # x1, y1, z1 is the vector points from the base to the target
    x1 = joints[3*conn[1]]
    y1 = joints[3*conn[1]+1]
    z1 = joints[3*conn[1]+2]

    glBegin(GL_LINES)
    glVertex3f(x0, y0, z0)
    glVertex3f(x1, y1, z1)
    glEnd()







# D. Holden's Data type
# root (3pts on the floor) + 21joints
parents = np.array([-1,0,1,2,3,4,1,6,7,8,1,10,11,12,12,14,15,16,12,18,19,20])

g_connMat_coco22 = [ [1,2], [2,3], [3,4], [4,5],    #right leg
                     [1,6], [6,7], [7,8], [8,9],    #left leg
                     [1,10], [10,11], [11,12], [12,13], #spine, head
                      [12,14], [14,15], [15,16], [16,17], #right arm
                       [12,18] , [18,19], [19,20], [20,21]] #left arm
#g_connMat_coco22 = np.array(g_connMat_coco22, dtype=int) - 1 #zero Idx
def drawbody_joint22(joints,  color, normal=None, ignore_root=False):

    #Visualize Joints
    glColor3ub(color[0], color[1], color[2])
    for i in range(1,int(len(joints)/3)):

        glPushMatrix()
        glTranslate(joints[3*i], joints[3*i+1], joints[3*i+2])
        glutSolidSphere(2, 10, 10)
        glPopMatrix()

    connMat_coco22 = g_connMat_coco22
    #Visualize Bones
    for conn in connMat_coco22:
        # x0, y0, z0 is the coordinate of the base point
        x0 = joints[3*conn[0]]
        y0 = joints[3*conn[0]+1]
        z0 = joints[3*conn[0]+2]
        # x1, y1, z1 is the vector points from the base to the target
        x1 = joints[3*conn[1]] - x0
        y1 = joints[3*conn[1]+1] - y0
        z1 = joints[3*conn[1]+2] - z0

        length = math.sqrt(x1*x1 + y1*y1 + z1*z1)
        theta = math.degrees(math.acos(z1/length))
        phi = math.degrees(math.atan2(y1, x1))

        glPushMatrix()
        glTranslate(x0, y0, z0)
        glRotatef(phi, 0, 0, 1)
        glRotatef(theta, 0, 1, 0)
        glutSolidCone(2, length, 10, 10)
        glPopMatrix()




    #Spine to ground projection
    conn = [0,1]
    x0 = joints[3*conn[0]]
    y0 = joints[3*conn[0]+1]
    z0 = joints[3*conn[0]+2]
    # x1, y1, z1 is the vector points from the base to the target
    x1 = joints[3*conn[1]]
    y1 = joints[3*conn[1]+1]
    z1 = joints[3*conn[1]+2]

    glBegin(GL_LINES)
    glVertex3f(x0, y0, z0)
    glVertex3f(x1, y1, z1)
    glEnd()

    # # Visualize Normals
    # if normal is not None:
    #     i=1
    #     facePt = joints19[(3*i):(3*i+3)]
    #     normalPt = facePt + normal*50

    #     glColor3ub(0, 255, 255)
    #     glPushMatrix()
    #     glTranslate(normalPt[0], normalPt[1], normalPt[2])
    #     glutSolidSphere(1, 10, 10)
    #     glPopMatrix()

    #     glBegin(GL_LINES);
    #     glVertex3f(facePt[0], facePt[1], facePt[2]);
    #     glVertex3f(normalPt[0], normalPt[1],  normalPt[2]);
    #     glEnd()


# The following is for drawbody_joint73 (Holden's format)
# sys.path.append('../motion/')
# from Quaternions import Quaternions

# skel_list is a list of skeletonElement
# each skeletonElement: (73,frameNum)
# D. Holden's Data type
# Joints73: 73 dim = 22joint*3 + 3 (X_velocity,Z_velocity,Rot_Velocity) +  4 (footStep)
# Input
#   - initRot: Quaternion for the first frame in global coordinate
#   - initTrans: 3x1 vector or np array
# def set_Holden_Data_73(skel_list, ignore_root=False, initRot = None, initTrans=None, bIsGT= False):
#     #global HOLDEN_DATA_SCALING

#     skel_list_output = []
#     #footsteps_output = []

#     for ai in range(len(skel_list)):
#         anim = np.swapaxes(skel_list[ai].copy(), 0, 1)  # frameNum x 73

#         if anim.shape[1]==73:
#             joints, root_x, root_z, root_r = anim[:,:-7], anim[:,-7], anim[:,-6], anim[:,-5]
#         elif anim.shape[1]==69:
#             joints, root_x, root_z, root_r = anim[:,:-3], anim[:,-3], anim[:,-2], anim[:,-1]
#         joints = joints.reshape((len(joints), -1, 3)) #(frameNum,66) -> (frameNum, 22, 3)

#         if initRot is None:
#             rotation = Quaternions.id(1)
#         else:
#             rotation = initRot[ai]
#         offsets = []

#         if initTrans is None:
#             translation = np.array([[0,0,0]])   #1x3
#         else:
#             translation = np.array(initTrans[ai])
#             if translation.shape[0]==3:
#                 translation = np.swapaxes(translation,0,1)


#         if not ignore_root:
#             for i in range(len(joints)):
#                 joints[i,:,:] = rotation * joints[i]
#                 joints[i,:,0] = joints[i,:,0] + translation[0,0]
#                 joints[i,:,2] = joints[i,:,2] + translation[0,2]
#                 rotation = Quaternions.from_angle_axis(-root_r[i], np.array([0,1,0])) * rotation
#                 offsets.append(rotation * np.array([0,0,1]))
#                 translation = translation + rotation * np.array([root_x[i], 0, root_z[i]])

#         #joints dim:(frameNum, 22, 3)
#         #Scaling
#         joints[:,:,:] = joints[:,:,:] * HOLDEN_DATA_SCALING#5 #m -> cm
#         joints[:,:,1] = joints[:,:,1] *-1 #Flip Y axis

#         #Reshaping
#         joints = joints.reshape(joints.shape[0], joints.shape[1]*joints.shape[2]) # frameNum x 66
#         joints =  np.swapaxes(joints, 0, 1)  # 66  x frameNum

#         skel_list_output.append(joints)
#         #footsteps_output.append(anim[:,-4:])


#     #adjust length
#     length = min( [ i.shape[1] for i in skel_list_output])
#     skel_list_output = [ f[:,:length] for f in skel_list_output ]

#     skel_list_output = np.asarray(skel_list_output)
#     #return skel_list_output #(skelNum, 66, frameNum)
#     #showSkeleton(skel_list_output)
#     setSkeleton(skel_list_output, bIsGT)


# traj_list is a list of 3 dim tran+rot infor
# each trajectory data: (3,frameNum)
# D. Holden's Data type
# 3 dim = 3 (X_velocity,Z_velocity,Rot_Velocity)
def set_Holden_Trajectory_3(traj_list, initRot = None, initTrans=None ):

    global HOLDEN_DATA_SCALING

    traj_list_output = []

    for ai in range(len(traj_list)):

        root_x, root_z, root_r = traj_list[ai][0,:], traj_list[ai][1,:], traj_list[ai][2,:]

        if initRot is None:
            rotation = Quaternions.id(1)
        else:
            rotation = initRot[ai]
        offsets = []

        if initTrans is None:
            translation = np.array([[0,0,0]])   #1x3
        else:
            translation = np.array(initTrans[ai])
            if translation.shape[0]==3:
                translation = np.swapaxes(translation,0,1)

        # joints = np.array([0,0,0])
        # joints = np.repeat(joints, )
        # joints = joints.reshape((len(joints), -1, 3)) #(frameNum,66) -> (frameNum, 22, 3)
        joints = np.zeros((len(root_x),2,3))        #(frames, 2,3) for original and directionPt
        joints[:,1,2] = 10 #Normal direction

        for i in range(len(joints)):
            joints[i,:,:] = rotation * joints[i]
            joints[i,:,0] = joints[i,:,0] + translation[0,0]
            joints[i,:,2] = joints[i,:,2] + translation[0,2]
            rotation = Quaternions.from_angle_axis(-root_r[i], np.array([0,1,0])) * rotation
            offsets.append(rotation * np.array([0,0,1]))
            translation = translation + rotation * np.array([root_x[i], 0, root_z[i]])

        #Reshaping
        joints = joints.reshape(joints.shape[0], joints.shape[1]*joints.shape[2]) # (frameNum,jointDim,3) -> (frameNum, jointDim*3)
        joints =  np.swapaxes(joints, 0, 1)  # jointDim*3  x frameNum

        joints = joints*HOLDEN_DATA_SCALING
        traj_list_output.append(joints)

    traj_list_output = np.asarray(traj_list_output)
    setTrajectory(traj_list_output)         #(trajNum, joitnDim*3, frames)

"""
    input:
     - speech_list: list of speechData
     - speechData: {'indicator': i, 'word': w}
     - i: frames of binary value
"""
def setSpeech_binary(speech_list):

    global g_speech#,g_frameLimit #nparray: (skelNum, skelDim, frameNum)

    for i, _ in enumerate(speech_list):
        if len(speech_list[i].shape) ==1:
            no = [None] * len(speech_list[i])
            speech_list[i] = {'indicator': speech_list[i], 'word':no}

    g_speech = speech_list #List of 2dim np.array

"""
    input:
     - speech_list: list of speechData
     - speechData: {'indicator': i, 'word': w}
     - i: frames of binary value
"""
def setSpeechGT_binary(speech_list):

    global g_speechGT#,g_frameLimit #nparray: (skelNum, skelDim, frameNum)

    for i, _ in enumerate(speech_list):
        if len(speech_list[i].shape) ==1:
            no = [None] * len(speech_list[i])
            speech_list[i] = {'indicator': speech_list[i], 'word':no}

    g_speechGT = speech_list #List of 2dim np.array

"""
    input:
     - speech_list: list of speechData
     - speechData: {'indicator': i, 'word': w}
     - i: frames of binary value
"""
def setSpeech(speech_list):

    global g_speech#,g_frameLimit #nparray: (skelNum, skelDim, frameNum)

    g_speech = speech_list #List of 2dim np.array

    # #g_frameLimit = g_skeletons.shape[2]
    # frameLens = [l.shape[1] for l in g_skeletons]
    # g_frameLimit = min(frameLens)


"""
    input:
     - speech_list: list of speechData
     - speechData: {'indicator': i, 'word': w}
     - i: frames of binary value
     - root_list: the root location to draw the speech data. same size as speech data
"""
def setSpeech_withRoot(speech_list, root_list):

    global g_speech#,g_frameLimit #nparray: (skelNum, skelDim, frameNum)

    g_speech = speech_list #List of 2dim np.array

    #Set root location
    for u,r in zip(g_speech,root_list):
        u['root'] = r

    # #g_frameLimit = g_skeletons.shape[2]
    # frameLens = [l.shape[1] for l in g_skeletons]
    # g_frameLimit = min(frameLens)



"""
    input:
     - speech_list: list of speechData
     - speechData: {'indicator': i, 'word': w}
     - i: frames of binary value
"""
def setSpeechGT(speech_list):

    global g_speechGT#,g_frameLimit #nparray: (skelNum, skelDim, frameNum)


    g_speechGT = speech_list #List of 2dim np.array

    # #g_frameLimit = g_skeletons.shape[2]
    # frameLens = [l.shape[1] for l in g_skeletons]
    # g_frameLimit = min(frameLens)



#Input: face_list (faceNum, dim, frames): nparray or list of arrays (dim, frames)
def showFace(face_list):

    #Add Skeleton Data
    global g_faces,g_frameLimit#nparray: (faceNum, faceDim, frameNum)

    #g_faces = np.asarray(face_list)  #no effect if face_list is already np.array
    g_faces  = face_list

    frameLens = [l.shape[1] for l in g_faces]
    g_frameLimit = max(g_frameLimit,min(frameLens))
    #init_gl()



def setFace(face_list):

    global g_faces,g_frameLimit#nparray: (faceNum, faceDim, frameNum)

    g_faces  = face_list

    # frameLens = [l.shape[1] for l in g_faces]
    # g_frameLimit = max(g_frameLimit,min(frameLens))
    setFrameLimit()




"""faceNormal_list: each element should have 3xframe"""
def setFaceNormal(faceNormal_list):

    global g_faceNormals,g_frameLimit#nparray: (faceNum, faceDim, frameNum)

    #g_faceNormals  = faceNormal_list
    g_faceNormals  = [ x.copy() for x in faceNormal_list]


    frameLens = [l.shape[1] for l in g_faceNormals if len(l)>0]
    maxFrameLength = max(frameLens)
    for i, p in enumerate(g_faceNormals):
        if len(p)==0:
            newData = np.zeros((3, maxFrameLength))
            g_faceNormals[i] = newData
        elif p.shape[0]==2:
            newData = np.zeros((3, p.shape[1]))
            newData[0,:] = p[0,:]
            newData[1,:] = 0 #some fixed number
            newData[2,:] = p[1,:]

            g_faceNormals[i] = newData



    g_frameLimit = max(g_frameLimit,min(frameLens))


"""bodyNormal_list: each element should have 3xframe"""
def setBodyNormal(bodyNormal_list):

    global g_bodyNormals,g_frameLimit#nparray: (faceNum, faceDim, frameNum)

    #g_bodyNormals  = bodyNormal_list
    g_bodyNormals  = [ x.copy() for x in bodyNormal_list]

    frameLens = [l.shape[1] for l in g_bodyNormals  if len(l)>0]
    maxFrameLength = max(frameLens)


    for i, p in enumerate(g_bodyNormals):
        if len(p)==0:
            newData = np.zeros((3, maxFrameLength))
            g_bodyNormals[i] = newData
        elif p.shape[0]==2:
            newData = np.zeros((3, p.shape[1]))
            newData[0,:] = p[0,:]
            newData[1,:] = 0 #some fixed number
            newData[2,:] = p[1,:]

            g_bodyNormals[i] = newData

        elif p.shape[0]==3:
            g_bodyNormals[i] = p





    g_frameLimit = max(g_frameLimit,min(frameLens))



"""pos_list: each element should have 3xframe"""
def setPosOnly(pos_list):
    global g_posOnly,g_frameLimit#nparray: (faceNum, faceDim, frameNum)

    g_posOnly  = [ x.copy() for x in pos_list]

    for i, p in enumerate(g_posOnly):
        if p.shape[0]==2:
            newData = np.zeros((3, p.shape[1]))
            newData[0,:] = p[0,:]
            #newData[1,:] = -100 #some fixed number
            newData[1,:] = 0 #some fixed number
            newData[2,:] = p[1,:]

            g_posOnly[i] = newData

    frameLens = [l.shape[1] for l in g_posOnly]
    g_frameLimit = max(g_frameLimit,min(frameLens))




def setHand_left(hand_list):

    global g_hands_left, g_frameLimit#nparray: (faceNum, faceDim, frameNum)

    g_hands_left  = hand_list

    frameLens = [l.shape[1] for l in g_hands_left]
    g_frameLimit = max(g_frameLimit,min(frameLens))


def setHand_right(hand_list):

    global g_hands_right, g_frameLimit#nparray: (faceNum, faceDim, frameNum)

    g_hands_right  = hand_list

    frameLens = [l.shape[1] for l in g_hands_right]
    g_frameLimit = max(g_frameLimit,min(frameLens))


#Input: skel_list (skelNum, dim, frames): nparray or list of arrays (dim, frames)
def setTrajectory(traj_list):
    #Add Skeleton Data
    global g_trajectory,g_frameLimit #nparray: (skelNum, skelDim, frameNum)

    # if len(skel_list)>1:
    #     lens = [len(l) for l in skel_list]
    #     minLeng=max(lens)

    #     for i in range(0,len(skel_list)):
    #         skel_list[i] = skel_list[i][:,:minLeng]

    #g_skeletons = np.asarray(skel_list)  #no effect if skel_list is already np.array
    g_trajectory = traj_list #List of 2dim np.array

    #g_frameLimit = g_skeletons.shape[2]
    frameLens = [l.shape[1] for l in g_trajectory]
    g_frameLimit = max(g_frameLimit,min(frameLens))


# def addSkeleton(skel_list, bIsGT = False, jointType=None):
def addSkeleton(skel_list, jointType=None, colorRGB=None): 
    setSkeleton(skel_list, jointType= jointType, colorRGB= colorRGB, bReset = False)

def resetSkeleton(): 
    global g_skeletons
    g_skeletons =[]



#Input: skel_list (skelNum, dim, frames): nparray or list of arrays (dim, frames)
# def setSkeleton(skel_list, bIsGT = False, jointType=None):
def setSkeleton(skel_list, jointType=None, colorRGB=None, bReset= True):
    global g_skeletons,g_frameLimit #nparray: (skelNum, skelDim, frameNum)
    # global g_skeletons_GT #nparray: (skelNum, skelDim, frameNum)

    #If skel_list is not a list
    if isinstance(skel_list,list) == False and len(skel_list.shape)==2:
        skel_list = skel_list[np.newaxis,:]

    #add joint type
    if g_skeletons is None or bReset:
        g_skeletons =[]

    # if color is None:
    #     color = (255,0,0)
    #color can None

    for s in skel_list:
        g_skeletons.append({"skeleton":s, "color":colorRGB, "type":jointType})

    # else:# bisGT == False:      TODO: no need to have g_skeletons_GT anymore?
    #     g_skeletons_GT =[]
    #     for s in skel_list:
            
    #         g_skeletons_GT.append({"skeleton":s, "color":(255,0,0), "type":jointType})

    # if jointType =='smpl':
    #     print("Use smplcoco instead of smpl!")
    #     assert(False)
   

    # if bIsGT==False:
    #     #Add Skeleton Data

    #     # if len(skel_list)>1:
    #     #     lens = [len(l) for l in skel_list]
    #     #     minLeng=max(lens)

    #     #     for i in range(0,len(skel_list)):
    #     #         skel_list[i] = skel_list[i][:,:minLeng]

    #     #g_skeletons = np.asarray(skel_list)  #no effect if skel_list is already np.array
    #     g_skeletons = skel_list #List of 2dim np.array

    #     #g_frameLimit = g_skeletons.shape[2]
    #     # frameLens = [l.shape[1] for l in g_skeletons]
    #     # g_frameLimit = max(g_frameLimit,min(frameLens))

    # else:
    #      #Add Skeleton Data
    #     g_skeletons_GT = skel_list #List of 2dim np.array

    #     #g_frameLimit = g_skeletons.shape[2]
    #     # frameLens = [l.shape[1] for l in g_skeletons_GT]
    #     # g_frameLimit = max(g_frameLimit,min(frameLens))

    setFrameLimit()



#Input: skel_list (skelNum, dim, frames): nparray or list of arrays (dim, frames)
def showSkeleton(skel_list):

    #Add Skeleton Data
    global g_skeletons,g_frameLimit #nparray: (skelNum, skelDim, frameNum)

    # if len(skel_list)>1:
    #     lens = [len(l) for l in skel_list]
    #     minLeng=max(lens)

    #     for i in range(0,len(skel_list)):
    #         skel_list[i] = skel_list[i][:,:minLeng]

    #g_skeletons = np.asarray(skel_list)  #no effect if skel_list is already np.array
    g_skeletons = skel_list #List of 2dim np.array


    #g_frameLimit = g_skeletons.shape[2]
    frameLens = [l.shape[1] for l in g_skeletons]
    g_frameLimit = max(g_frameLimit,min(frameLens))

    # init_gl()


#Input: skel_list peopelNum x {'ver': vertexInfo, 'f': faceInfo}
#: vertexInfo should be (frames x vertexNum x 3 )
#: if vertexInfo has (vertexNum x 3 ), this function automatically changes it to (1 x vertexNum x 3)
#: faceInfo should be (vertexNum x 3 )
#'normal': if missing, draw mesh by wireframes
def setMeshData(mesh_list, bComputeNormal = False):

    global g_meshes

    #g_skeletons = np.asarray(skel_list)  #no effect if skel_list is already np.array

    ##
    g_meshes = [ d.copy() for d in mesh_list]

    if len(g_meshes)==0:
        return

    if len(g_meshes)>40:
        print("Warning: too many meshes ({})".format(len(g_meshes)))
        g_meshes =g_meshes[:40]


    if len(g_meshes)==0:
        return

    if len(g_meshes)>40:
        print("Warning: too many meshes ({})".format(len(g_meshes)))
        g_meshes =g_meshes[:40]


    for element in g_meshes:
        if len(element['ver'].shape) ==2:
            # print("## setMeshData: Warning: input size should be (N, verNum, 3). Current input is (verNum, 3). I am automatically fixing this.")
            element['ver'] = element['ver'][np.newaxis,:,:]
            if 'normal' in element.keys():
                element['normal'] = element['normal'][np.newaxis,:,:]

    #Auto computing normal
    if bComputeNormal:
        # print("## setMeshData: Computing face normals automatically.")
        for element in g_meshes:
            element['normal'] = ComputeNormal(element['ver'],element['f']) #output: (N, 18540, 3)

    #g_frameLimit = g_skeletons.shape[2]
    #mesh_list[0]['ver'].shape
    # frameLens = [l['ver'].shape[0] for l in g_meshes]
    # g_frameLimit = max(g_frameLimit,min(frameLens))

    setFrameLimit()


def setFrameLimit():
    global g_frameLimit
    g_frameLimit = 0

    if g_meshes is not None:
        frameLens = [1] + [l['ver'].shape[0] for l in g_meshes]
        g_frameLimit = max(g_frameLimit,min(frameLens))

    if g_skeletons is not None:
        # frameLens = [1] + [l['skeleton'].shape[1] for l in g_skeletons]
        frameLens = [l['skeleton'].shape[1] for l in g_skeletons]
        if len(frameLens)>0:
            g_frameLimit = max(g_frameLimit, min(frameLens))

    # if g_skeletons_GT is not None:
    #     frameLens = [1] + [l.shape[1] for l in g_skeletons_GT]
    #     g_frameLimit = max(g_frameLimit,min(frameLens))

    if g_faces is not None:
        frameLens = [1] + [l.shape[1] for l in g_faces]
        g_frameLimit = max(g_frameLimit,min(frameLens))

def resetFrameLimit():
    global g_frameLimit
    g_frameLimit =0


def resetMeshData(): 
    global g_meshes
    g_meshes =[]

def addMeshData(mesh_list, bComputeNormal = False):
    if g_meshes is not None:
        setMeshData(g_meshes + mesh_list, bComputeNormal) #TODO: no need to compute normal for already added one..
    else:
        setMeshData(mesh_list, bComputeNormal)




"""
    input:
     - list of faceParam, where each element has np array of (200 x frames)
"""
def setFaceParmData(faceParam_list, bComputeNormal = True):
    global g_faceModel

    if g_faceModel is None:
        import scipy.io as sio
        g_faceModel = sio.loadmat('/ssd/data/totalmodel/face_model_totalAligned.mat')

    #check the dimension. If dim <200, padding zeros
    for i,f in enumerate(faceParam_list):
        if (f.shape[0]<200):
            newData = np.zeros ( (200,f.shape[1]) )
            newData[:f.shape[0],:] = f
            faceParam_list[i] = newData

    faceParam_list = [ {'face_exp': f} for f in faceParam_list]

    faceMesh_list = GetFaceMesh(g_faceModel,faceParam_list, bComputeNormal= bComputeNormal)
    etMeshData( faceMesh_list)


"""
    input:
     - list of faceParam, where each element has np array of (200 x frames)
"""
def setFaceParmDataWithTrans(faceParam_list, bComputeNormal = True, trans= None, rot = None):
    global g_faceModel

    if g_faceModel is None:
        import scipy.io as sio
        g_faceModel = sio.loadmat('/ssd/data/totalmodel/face_model_totalAligned.mat')

    #check the dimension. If dim <200, padding zeros
    for i,f in enumerate(faceParam_list):
        if (f.shape[0]<200):
            newData = np.zeros ( (200,f.shape[1]) )
            newData[:f.shape[0],:] = f
            faceParam_list[i] = newData

    #faceParam_list = [ {'face_exp': f} for f in faceParam_list]
    faceParam_list_new =[]
    for i in range(len(faceParam_list)):
        data = dict()
        data['face_exp'] = faceParam_list[i]    #(200,frames)
        if trans is not None:
            data['trans'] = trans[i]        #(3,frames)

            frameLen = min( data['face_exp'].shape[1], data['trans'].shape[1] )
            data['face_exp'] = data['face_exp'][:,:frameLen]
            data['trans'] = data['trans'][:,:frameLen] *0.01

            ROT_PIVOT = np.array([0.003501, 0.475611, 0.115576])
            ROT_PIVOT[2] -=0.1
            #faceMassCenter = np.array([-0.002735  , -1.44728992,  0.2565446 ])#*100
            data['rot_pivot'] = data['trans']*0 + ROT_PIVOT[:,np.newaxis]




        if rot is not None:
            data['rot'] = rot[i]

            frameLen = min( data['face_exp'].shape[1], data['rot'].shape[1] )
            data['face_exp'] = data['face_exp'][:,:frameLen]
            if trans is not None:
                 data['trans'] = data['trans'][:,:frameLen]
            data['rot'] = data['rot'][:,:frameLen]

        faceParam_list_new.append(data)


    #faceMesh_list = GetFaceMesh(g_faceModel,faceParam_list, bComputeNormal= bComputeNormal)
    faceMesh_list = GetFaceMesh(g_faceModel,faceParam_list_new, bComputeNormal= bComputeNormal,bApplyTrans=(trans is not None),bApplyRot=(rot is not None), bApplyRotFlip=(rot is not None))
    etMeshData( faceMesh_list)


'''
    Input:
        face_list:  faceNum x element['face70'](210, frames)
    Output:
        faceNormal_list:
        face_list:  faceNum x (3, frames)
'''
def ComputeFaceNormal(face_list):
    ##Compute face normal
    faceNormal_list =[]
    for s in face_list:
        leftEye = s['face70'][(45*3):(45*3+3),:].transpose() #210xframes
        rightEye = s['face70'][(36*3):(36*3+3),:].transpose() #210xframes
        nose = s['face70'][(33*3):(33*3+3),:].transpose() #210xframes

        left2Right = rightEye - leftEye
        right2nose = nose - rightEye

        from sklearn.preprocessing import normalize
        left2Right = normalize(left2Right, axis=1)
        #Check: np.linalg.norm(left2Right,axis=1)
        right2nose = normalize(right2nose, axis=1)

        faceNormal = np.cross(left2Right,right2nose)
        faceNormal[:,1] = 0 #Project on x-z plane, ignoring y axis
        faceNormal = normalize(faceNormal, axis=1)
        faceNormal_list.append(faceNormal.transpose())

    return faceNormal_list


'''
    Input:
        body_list:  bodyNum x element['joints19'](21, frames)
    Output:
        faceNormal_list:
        face_list:  bodyNum x (3, frames)
'''
def ComputeBodyNormal_panoptic(body_list):
     #Compute Body Normal
    bodyNormal_list =[]
    for s in body_list:
        leftShoulder = s['joints19'][(3*3):(3*3+3),:].transpose() #210xframes
        rightShoulder = s['joints19'][(9*3):(9*3+3),:].transpose() #210xframes
        bodyCenter = s['joints19'][(2*3):(2*3+3),:].transpose() #210xframes

        left2Right = rightShoulder - leftShoulder
        right2center = bodyCenter - rightShoulder

        from sklearn.preprocessing import normalize
        left2Right = normalize(left2Right, axis=1)
        #Check: np.linalg.norm(left2Right,axis=1)
        right2center = normalize(right2center, axis=1)

        bodyNormal = np.cross(left2Right,right2center)
        bodyNormal[:,1] = 0 #Project on x-z plane, ignoring y axis
        bodyNormal = normalize(bodyNormal, axis=1)
        bodyNormal_list.append(bodyNormal.transpose())
    return bodyNormal_list

def getFaceRootCenter():
    global g_meshes

    speech_rootData = []

    for f in g_meshes:

        speech_rootData.append(f['centers'])
    return speech_rootData




#####################################################
### FaceOnly Model Visualization (todo. move this to another file)
# Input:
# - bApplyRot: if True, apply global face orientation
# Output:
# - 'ver': (frames, 11510, 3)
# - 'normal': (frames, 11510, 3)
# - 'f': (22800,3)
# - 'centers': (frames,3)


import scipy.io as sio
def GetFaceMesh(faceModel, faceParam_list, bComputeNormal = True, bApplyRot = False, bApplyTrans = False, bShowFaceId = False, bApplyRotFlip=False):

    MoshParam_list = []

    v_template = faceModel['v_template'] #11510 x 3
    v_template_flat = v_template.flatten()  #(34530,)
    v_template_flat = v_template_flat[:,np.newaxis] #(34530,1)  for broadcasting

    trifaces = faceModel['trifaces']  #22800 x 3

    U_id = faceModel['U_id'] #34530 x 150
    U_exp = faceModel['U_exp'] #34530 x 200


    for i, faceParam in enumerate(faceParam_list):

        print('processing: humanIdx{0}/{1}'.format(i, len(faceParam_list) ))
        #faceParam = faceParam_all[humanIdx]


        #Debug: only use the first 5
        faceParam['face_exp'][5:,:]=0

        """Computing face vertices for all frames simultaneously"""
        face_exp_component = np.matmul(U_exp, faceParam['face_exp']) #(34,530 x 200)  x  (200 x frames)

        v_face_allFrames = v_template_flat + face_exp_component  #(34530 x frames). Ignore face Identity information

        if bShowFaceId:
            face_id_component = np.matmul(U_id, faceParam['face_id']) #(34,530 x 150)  x  (150 x frames)
            v_face_allFrames += face_id_component #(34530 x frames)
            #v_face_allFrames = v_template_flat+ face_id_component +face_exp_component  #(34530 x frames)

        v_face_allFrames = v_face_allFrames.swapaxes(0,1) # (frames, 34530)
        v_face_allFrames = np.reshape(v_face_allFrames,[v_face_allFrames.shape[0], -1, 3]) # (frames, 11510, 3)


        faceMassCenter = np.zeros((3,faceParam['face_exp'].shape[1]))#*0.0 #(3, frames)
        if 'rot_pivot' in faceParam.keys():


            rot_pivot = np.swapaxes(faceParam['rot_pivot'],0,1) #(frames,3)
            rot_pivot = np.expand_dims(rot_pivot,1)  #(frames,1, 3)

            v_face_allFrames = v_face_allFrames - rot_pivot     # (frames, 11510, 3)

        if bApplyRot:
            assert False, "This code cannot be run."
            from modelViewer.batch_lbs import batch_rodrigues
            #Apply rotationsvf
            global_rot = None
            #computing global rotation
            global_rot =  batch_rodrigues(np.swapaxes(faceParam['rot'],0,1)) #input (Nx3), output: (N,3,3)



            #global_rot *( v_face_allFrames - rot_pivot)
            for f in range(v_face_allFrames.shape[0]):

                pts = np.swapaxes(v_face_allFrames[f,:,:],0,1) # (3,11510)

                if bApplyRotFlip:
                    #Flip
                    rot = np.array( [ 0, -1, 0,  1,  0,  0, 0, 0, 1])
                    rot = np.reshape(rot,(3,3))
                    pts =  np.matmul( rot, pts)  # (3,3)  x (11510,3) =>(3,11510)
                    pts =  np.matmul( rot, pts)  # (3,3)  x (11510,3) =>(3,11510)
                    pts *= 0.94


                # rot = np.array( [ 0, -1,  0, 1,  0,  0,0,  0,  1])
                # rot = np.reshape(rot,(3,3))
                rot = global_rot[f,:,:]
                pts =  np.matmul( rot, pts)  # (3,3)  x (11510,3) =>(3,11510)

                v_face_allFrames[f,:,:] = pts.transpose()
        else:   #Rotate 180 degrees to flip y axis
            #global_rot =  batch_rodrigues(np.swapaxes(faceParam['rot'],0,1)) #input (Nx3), output: (N,3,3)

            #global_rot *( v_face_allFrames - rot_pivot)
            for f in range(v_face_allFrames.shape[0]):

                pts = np.swapaxes(v_face_allFrames[f,:,:],0,1) # (3,11510)

                #rot = np.array( [ 1, 0, 0,  0, 1,  0,  0, 0, -1])
                #rot = np.array( [ 1, 0, 0,  0, 1,  0,  0, 0, -1])
                rot = np.array( [ 0, -1, 0,  1,  0,  0, 0, 0, 1])
                #rot = np.array( [ 0, 1, 0,  1,  0,  0, 0, 0, 1])
                rot = np.reshape(rot,(3,3))
                #rot = global_rot[f,:,:]
                pts =  np.matmul( rot, pts)  # (3,3)  x (11510,3) =>(3,11510)
                pts =  np.matmul( rot, pts)  # (3,3)  x (11510,3) =>(3,11510)

                #trans = np.array([[0.0,-1.5,0.2]])
                trans = np.array([[0.0,-1,0.2]])
                v_face_allFrames[f,:,:] = pts.transpose() +trans


        if bApplyTrans:
            trans = np.swapaxes(faceParam['trans'],0,1) #(frames,3)
            trans = np.expand_dims(trans,1)  #(frames,1, 3)
            v_face_allFrames = v_face_allFrames + trans     # (frames, 11510, 3)


            faceMassCenter += faceParam['trans'] #(3, frames)


        if bComputeNormal==False:
            #Debug. no normal
            faceMassCenter = v_face_allFrames[:,5885,:] #5885th vertex. around the head top.
            MoshParam = {'ver': v_face_allFrames, 'normal': [], 'f': trifaces, 'centers': faceMassCenter}  # support rendering two models together
            MoshParam_list.append(MoshParam)
            continue


        # # CPU version
        start = time.time()
        vertex_normals = ComputeNormal(v_face_allFrames, trifaces)
        print("CPU: normal computing time: {}".format( time.time() - start))

        # GPU version
        # start = time.time()
        # vertex_normals = ComputeNormal_gpu(v_face_allFrames, trifaces)
        # print("GPU: normal computing time: {}".format( time.time() - start))


        faceMassCenter = v_face_allFrames[:,5885,:] #5885th vertex. around the head top.
        #faceMassCenter = np.swapaxes(faceMassCenter,0,1) # (3,frames) - >(frames,3)

        MoshParam = {'ver': v_face_allFrames, 'normal': vertex_normals, 'f': trifaces, 'centers': faceMassCenter}  # support rendering two models together
        MoshParam_list.append(MoshParam)

    return MoshParam_list



#####################################################
### Haggling video rendering/generation

def LoadHagglingDataKeypoints(fileName):
    global g_hagglingseq_name
    g_hagglingseq_name = fileName[:-4]
    ##read face
    seqPath = '/ssd/codes/pytorch_motionSynth/motionsynth_data/data/processed_panoptic/panopticDB_face_pkl_hagglingProcessed/' +fileName
    motionData = pickle.load( open( seqPath, "rb" ) , encoding='latin1')
    if 'subjects' in motionData.keys() and len(motionData['subjects'])==3:
        setFace([motionData['subjects'][0]['face70'], motionData['subjects'][1]['face70'], motionData['subjects'][2]['face70']])
    else:
        setFace(None)


    ##read hand
    seqPath = '/ssd/codes/pytorch_motionSynth/motionsynth_data/data/processed_panoptic/panopticDB_hand_pkl_hagglingProcessed/' +fileName
    motionData = pickle.load( open( seqPath, "rb" ) , encoding='latin1')
    if 'hand_left' in motionData:
        if len(motionData['hand_left'])==3:
            setHand_left([motionData['hand_left'][0]['hand21'], motionData['hand_left'][1]['hand21'], motionData['hand_left'][2]['hand21']])
        else:
            setHand_left(None)


    if 'hand_right' in motionData:
        if len(motionData['hand_right'])==3:
            setHand_right([motionData['hand_right'][0]['hand21'], motionData['hand_right'][1]['hand21'], motionData['hand_right'][2]['hand21']])
        else:
            setHand_right(None)

    #read speech
    seqPath = '/ssd/codes/haggling_audio/panopticDB_pkl_speech_hagglingProcessed/' +fileName
    motionData = pickle.load( open( seqPath, "rb" ) , encoding='latin1')
    setSpeech([motionData['speechData'][0], motionData['speechData'][1], motionData['speechData'][2]])

    #read body
    seqPath = '/ssd/codes/pytorch_motionSynth/motionsynth_data/data/processed_panoptic/panopticDB_pkl_hagglingProcessed/' +fileName
    motionData = pickle.load( open( seqPath, "rb" ) , encoding='latin1')
    setSkeleton([motionData['subjects'][0]['joints19'], motionData['subjects'][1]['joints19'], motionData['subjects'][2]['joints19']])

    #Visualize Body Normal
    setBodyNormal([motionData['subjects'][0]['bodyNormal'], motionData['subjects'][1]['bodyNormal'], motionData['subjects'][2]['bodyNormal']])
    setFaceNormal([motionData['subjects'][0]['faceNormal'], motionData['subjects'][1]['faceNormal'], motionData['subjects'][2]['faceNormal']])


g_adamWrapper = None
def LoadHagglingData(fileName):
    global g_hagglingseq_name
    g_hagglingseq_name = fileName[:-4]

    bLoadFace = False
    bLoadAdam = True
    bLoadSkeleton = False
    bLoadKeypoints = False


    fileName_pkl = fileName.replace('npz','pkl')

    if bLoadFace :
        """Load Face data"""
        global g_faceModel
        if g_faceModel is None:
            import scipy.io as sio
            g_faceModel = sio.loadmat('/ssd/data/totalmodel/face_model_totalAligned.mat')

        ##read face mesh parameter
        seqPath = '/ssd/codes/pytorch_motionSynth/motionsynth_data/data/processed_panoptic/panopticDB_faceMesh_pkl_hagglingProcessed/' +fileName_pkl

        faceData = pickle.load( open( seqPath, "rb" ) , encoding='latin1')

        FaceParam_list = GetFaceMesh(g_faceModel,faceData['subjects'])
        setMeshData( FaceParam_list)

        # #read speech
        # seqPath = '/ssd/codes/haggling_audio/panopticDB_pkl_speech_hagglingProcessed/' +fileName
        # motionData = pickle.load( open( seqPath, "rb" ) )
        # setSpeech([motionData['speechData'][0], motionData['speechData'][1], motionData['speechData'][2]])

        #read speech
        seqPath = '/ssd/codes/haggling_audio/panopticDB_pkl_speech_hagglingProcessed/' +fileName_pkl
        motionData = pickle.load( open( seqPath, "rb" ) , encoding='latin1')
        speechData = [motionData['speechData'][0], motionData['speechData'][1], motionData['speechData'][2]]
        #setSpeech([motionData['speechData'][0], motionData['speechData'][1], motionData['speechData'][2]])
        speech_rootData = [FaceParam_list[0]['centers'], FaceParam_list[1]['centers'], FaceParam_list[2]['centers']]
        speech_rootData =np.multiply(speech_rootData, 100.0) #meter to cm
        setSpeech_withRoot(speechData,speech_rootData)

    if bLoadSkeleton:
        """Load Skeleton data (Holden's format)"""

        #Read body, holden's format
        #fileName_npz = fileName.replace('pkl','npz')
        X = np.load('/ssd/codes/pytorch_motionSynth/motionsynth_data/data/processed/panoptic_npz/' + fileName)['clips'] #(17944, 240, 73)
        X = np.swapaxes(X, 1, 2).astype(np.float32) #(17944, 73, 240)
        set_Holden_Data_73([ X[0,:,:], X[1,:,:], X[2,:,:] ], ignore_root=True)
    if bLoadKeypoints:
        LoadHagglingDataKeypoints(fileName_pkl)

    if bLoadAdam:

        from modelViewer.batch_adam import ADAM
        global g_adamWrapper

        """Load Adam data"""
        if g_adamWrapper==None:
            g_adamWrapper = ADAM()

        fileName_pkl = fileName.replace('npz','pkl')

        seqPath = '/ssd/codes/pytorch_motionSynth/motionsynth_data/data/processed_panoptic/panopticDB_adamMesh_pkl_hagglingProcessed_stage1/' +fileName_pkl
        adamParam_all = pickle.load( open( seqPath, "rb" ) , encoding='latin1')

         ##read face mesh parameter
        seqPath = '/ssd/codes/pytorch_motionSynth/motionsynth_data/data/processed_panoptic/panopticDB_faceMesh_pkl_hagglingProcessed/' +fileName_pkl
        faceData = pickle.load( open( seqPath, "rb" ) , encoding='latin1')

        start = time.time()
        meshes =[]
        frameStart = 0
        frameEnd = -1

        #for faceParam in faceParam_selected:
        for adamParam in adamParam_all['subjects']:

            betas = np.swapaxes(adamParam['betas'],0,1)[frameStart:frameEnd]  #Frames  x30
            faces = np.swapaxes(adamParam['faces'],0,1)[frameStart:frameEnd]  #Frames  200

            for faceParam in faceData['subjects']:
                if faceParam['humanId'] == adamParam['humanId']:
                    faces = faceParam['face_exp']
                    faces = np.swapaxes(faces,0,1)[frameStart:frameEnd]  #Frames  200

                    break

            pose = np.swapaxes(adamParam['pose'],0,1)[frameStart:frameEnd]  #Frames  186
            trans = np.swapaxes(adamParam['trans'],0,1)[frameStart:frameEnd]  #Frames  x3
            startTime = time.time()

            #Align the length
            frameNum = min([pose.shape[0], faces.shape[0], betas.shape[0], trans.shape[0]])
            pose = pose[:frameNum]
            faces = faces[:frameNum]
            betas = betas[:frameNum]
            trans = trans[:frameNum]


            v , j = g_adamWrapper(betas,pose,faces) #v:(frameNum, 18540, 3), j: (frameNum, 62, 3)
            print('time: {}'.format(time.time()-startTime))
            v += np.expand_dims(trans, axis=1)  # no translation in their LBS. trans: (humanNumn,30) ->(humanNumn,1, 30)

            v *=0.01

            normals = ComputeNormal(v, g_adamWrapper.f)
            meshes.append( {'ver': v, 'normal': normals, 'f': g_adamWrapper.f})  # support rendering two models together


        setMeshData( meshes)


    # global g_meshes,g_skeletons,g_speech
    # global g_frameLimit

    # g_frameLimit = min([l['ver'].shape[0] for l in g_meshes])
    # tempMinFrame = g_skeletons.shape[2]
    # g_frameLimit = min(g_frameLimit,tempMinFrame)

    # tempMinFrame = min([l['indicator'].shape[0] for l in g_speech])
    # tempMinFrame = min([l['root'].shape[0] for l in g_speech])

    # g_frameLimit = max(X.shape[2],FaceParam_list[0]['ver'].shape[0])


g_hagglingFileList = None
g_hagglingFileListIdx = -1
def LoadHagglingData_Caller():

    global g_hagglingFileList,g_hagglingFileListIdx,g_frameIdx

    if g_hagglingFileList==None: #If initial calling
        #directory = '/ssd/codes/pytorch_motionSynth/motionsynth_data/data/processed_panoptic/panopticDB_faceMesh_pkl_hagglingProcessed/'
        # directory = '/ssd/codes/pytorch_motionSynth/motionsynth_data/data/processed_panoptic/panopticDB_pkl_hagglingProcessed/'
        #directory = '/ssd/codes/pytorch_motionSynth/motionsynth_data/data/processed_panoptic/panopticDB_adamMesh_pkl_hagglingProcessed/'
        #g_hagglingFileList =  [f for f in sorted(list(os.listdir(directory))) if os.path.isfile(os.path.join(directory,f)) and f.endswith('.pkl')]

        #directory = '/ssd/codes/haggling_audio/panopticDB_pkl_speech_hagglingProcessed/'
        directory = '/ssd/codes/pytorch_motionSynth/motionsynth_data/data/processed/panoptic_npz/'
        g_hagglingFileList =  [f for f in sorted(list(os.listdir(directory))) if os.path.isfile(os.path.join(directory,f)) and f.endswith('.npz')]

        hagglingFileList_final = list()
        for f in g_hagglingFileList:
            seq_name = f[:-4]
            if not os.path.exists('/ssd/render_mesh/'+seq_name):
                hagglingFileList_final.append(f)

        g_hagglingFileList = hagglingFileList_final


        g_hagglingFileListIdx =-1


    g_hagglingFileListIdx +=1
    if g_hagglingFileListIdx>=len(g_hagglingFileList):
        print("g_hagglingFileListIdx>=len(g_hagglingFileList)")
        return False

    fileName = g_hagglingFileList[g_hagglingFileListIdx]
    print(fileName)
    LoadHagglingData(fileName)
    g_frameIdx =0

    return True

#Save all the frames, and exit opengl when all frames are saved
def setSaveOnlyMode(mode):
    global g_bSaveOnlyMode
    g_bSaveOnlyMode = mode

def setSave(mode):
    global g_bSaveToFile
    g_bSaveToFile = mode

def setSaveFolderName(folderName):
    global g_saveFolderName
    g_saveFolderName = folderName

def setSaveImgName(imgName):
    global g_saveImageName
    g_saveImageName = imgName



""" Using pywavefront"""
def LoadObjMesh(filename):

    import pywavefront
    mesh = pywavefront.Wavefront(filename, collect_faces=True)

    faces = np.array(mesh.mesh_list[0].faces)   # (#Faces, 3). Zero-based
    # faces = faces-1

    vertices = np.array(mesh.vertices)  # (#Vertices, 3)

    #Compute Normals
    vertices = vertices[np.newaxis,:]  #1 x meshVerNum x 3
    normal_all = ComputeNormal(vertices,faces)

    #Save as pkl
    meshModel =  {'vertices': vertices, 'normals':normal_all , 'faces': faces}
    return meshModel


def setupRotationView():
    global g_bShowFloor,g_zoom
    glutReshapeWindow(1000,1000)
    g_bShowFloor = False
    # g_zoom = 220
    g_zoom = 2220



def DrawPyramid(camWidth,camHeight,camDepth,lineWith=1):

	# glColorMaterial(GL_FRONT, GL_DIFFUSE);
	# glEnable(GL_COLOR_MATERIAL);
	# glColor4f(color.first.x,color.first.y,color.first.z,color.second);

    glLineWidth(lineWith)

    glBegin(GL_LINES)
    glVertex3f(0,0,0)
    glVertex3f(camWidth*0.5,camHeight*0.5,camDepth*1)
    glVertex3f(0,0,0)
    glVertex3f(camWidth*0.5,camHeight*-0.5,camDepth*1)
    glVertex3f(0,0,0)
    glVertex3f(camWidth*-0.5,camHeight*0.5,camDepth*1)
    glVertex3f(0,0,0)
    glVertex3f(camWidth*-0.5,camHeight*-0.5,camDepth*1)
    glEnd()
	

    glBegin( GL_LINE_STRIP)
    glVertex3f(camWidth*0.5,camHeight*0.5,camDepth*1)
    glVertex3f(camWidth*-0.5,camHeight*0.5,camDepth*1)
    glVertex3f(camWidth*-0.5,camHeight*-0.5,camDepth*1)
    glVertex3f(camWidth*0.5,camHeight*-0.5,camDepth*1) 
    glVertex3f(camWidth*0.5,camHeight*0.5,camDepth*1) 
    glEnd()
    glDisable(GL_COLOR_MATERIAL)

    glBegin( GL_QUADS)
    # glColor4f(color.first.x,color.first.y,color.first.z,color.second*0.1)
    glVertex3f(camWidth*0.5,camHeight*0.5,camDepth*1)
    glVertex3f(camWidth*-0.5,camHeight*0.5,camDepth*1)
    glVertex3f(camWidth*-0.5,camHeight*-0.5,camDepth*1)
    glVertex3f(camWidth*0.5,camHeight*-0.5,camDepth*1) 
    glVertex3f(camWidth*0.5,camHeight*0.5,camDepth*1) 
    glEnd()
    # glDisable(GL_COLOR_MATERIAL)

def DrawPtCloud():
    # glColor3f(0,1,0)

    glPointSize(g_ptSize); 
    glBegin(GL_POINTS)
    for i in range(g_ptCloud.shape[0]):
        glColor3f(g_ptCloudColor[i,0],g_ptCloudColor[i,1],g_ptCloudColor[i,2])
        glVertex3f(g_ptCloud[i,0],g_ptCloud[i,1],g_ptCloud[i,2])
    glEnd()
	


def DrawCameras():
    if g_cameraPoses is None:
        return
    
    # glColor3f(0,1,0)
    for i in range(len(g_cameraPoses)):

        glPushMatrix()


        # glTranslatef(g_cameraPoses[i,0],g_cameraPoses[i,1],g_cameraPoses[i,2])
        glTranslatef(g_cameraPoses[i][0],g_cameraPoses[i][1],g_cameraPoses[i][2])
        glMultMatrixd(g_cameraRots[i])
        glutSolidSphere(1, 10, 10)
        DrawPyramid(20, 20, 20)
            
        glPopMatrix()

    # glColor3f(g_meshColor[0],g_meshColor[1],g_meshColor[2])
    

def renderscene():
    global g_xRotate, g_rotateView_counter, g_saveFrameIdx

    start = timeit.default_timer()
    # global xrot
    # global yrot
    # global view_dist
    glClear(GL_COLOR_BUFFER_BIT|GL_DEPTH_BUFFER_BIT)

    #Some anti-aliasing code (seems not working, though)
    glEnable (GL_BLEND)
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
    glEnable (GL_LINE_SMOOTH)
    glHint (GL_LINE_SMOOTH_HINT, GL_NICEST)
    # glEnable(GL_POLYGON_SMOOTH)
    glEnable(GL_MULTISAMPLE)
    # glHint(GL_MULTISAMPLE_FILTER_HINT_NV, GL_NICEST)


    # Set up viewing transformation, looking down -Z axis
    glLoadIdentity()
    #gluLookAt(0, 0, -g_fViewDistance, 0, 0, 0, -.1, 0, 0)   #-.1,0,0
    gluLookAt(0,0,0, 0, 0, 1, 0, -1, 0)


    if g_viewMode=='camView':       #Set View Point in MTC Camera View
        # camidlist = ''.join(g_camid)
        # camid = int(camidlist)
        if g_bOrthoCam:
            setCameraViewOrth()
        else:
            setCameraView()

    else:#Free Mode
        # Set perspective (also zoom)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        #gluPerspective(zoom, float(g_Width)/float(g_Height), g_nearPlane, g_farPlane)
        gluPerspective(65, float(g_Width)/float(g_Height), g_nearPlane, g_farPlane)         # This should be called here (not in the reshpe)
        glMatrixMode(GL_MODELVIEW)
        # Render the scene

        # setFree3DView()
        glTranslatef(0,0,g_zoom)

        glRotatef( -g_yRotate, 1.0, 0.0, 0.0)
        glRotatef( -g_xRotate, 0.0, 1.0, 0.0)
        glRotatef( g_zRotate, 0.0, 0.0, 1.0)

        # glutSolidSphere(3, 10, 10)        #Draw Origin

        glTranslatef( g_xTrans,  0.0, 0.0 )
        glTranslatef(  0.0, g_yTrans, 0.0)
        glTranslatef(  0.0, 0, g_zTrans)
        
        glColor3f(0,1,0)



    #This should be drawn first, without depth test (it should be always back)
    if g_bShowBackground:
        if g_bOrthoCam:
            DrawBackgroundOrth()
        else:
            DrawBackground()

    glEnable(GL_LIGHTING)
    glEnable(GL_CULL_FACE)
    #drawbody(bodyDatas[cur_ind], connMat_coco19)
    #drawbody_haggling(m_landmarks[:, cur_ind], connMat_coco19)
    #if g_bSaveToFile:


    # #Debug
    # glColor3f(1,0,0)
    # glutSolidTeapot(100, 10, 10)
    # RenderDomeFloor()
    # glutSwapBuffers()
    # return

    # glUseProgram(0)
    glPolygonMode(GL_FRONT, GL_FILL)
    glPolygonMode(GL_BACK, GL_FILL)

    if g_bShowSkeleton:
        DrawSkeletons()
        # DrawSkeletonsGT()
    DrawTrajectory()

    DrawFaces()
    DrawHands()

    if g_bShowMesh:
        DrawMeshes()
    DrawPosOnly()


    glDisable(GL_LIGHTING)
    glDisable(GL_CULL_FACE)


    DrawCameras()
    if g_ptCloud is not None:
        DrawPtCloud()


    if g_bShowFloor:
        RenderDomeFloor()

    global g_frameIdx#, g_frameLimit
    global g_fps, g_show_fps
    if g_show_fps:
        RenderText("{0} fps".format(int(np.round(g_fps,0))))

    # swap the screen buffers for smooth animation
    glutSwapBuffers()

    if g_bRotateView:

        # g_rotateInnterval = 2.0
        g_xRotate += g_rotateInterval

        # print("{0}/rotview_{1:04d}.jpg".format("RENDER_DIR",g_rotateView_counter))

        g_saveFrameIdx = g_rotateView_counter
        g_rotateView_counter+=1


    if g_bSaveToFile:
        SaveScenesToFile()

    g_frameIdx +=1
    #time.sleep(1)
    if g_frameIdx>=g_frameLimit:
        #global g_bSaveOnlyMode
        if g_bSaveOnlyMode:
            #exit opengl
            global g_stopMainLoop
            g_stopMainLoop= True
            g_frameIdx = 0
        else:
            g_frameIdx =0

    if False:
        #Ensure 30fps
        stop = timeit.default_timer()
        time_sec = stop - start
        sleepTime = 0.03333*10 -time_sec
        if sleepTime>0:
            time.sleep(sleepTime)

    # #Ensure 60fps
    # stop = timeit.default_timer()
    # time_sec = stop - start
    # sleepTime = 0.06666 -time_sec
    # if sleepTime>0:
    #     time.sleep(sleepTime)
    # #print('Time: ', stop - start)


# This is to generate AMT example
# Set camera views, window size, and so on.
# Rotate scene, save as image
# Close glWindow
# def init_gl_rotationRendering(bSaveToFile = False, bShowSkeleton= False):
#     # Init_Haggling()
#     # global width
#     # global height
#     # Setup for double-buffered display and depth testing
#
#     global g_Width,g_Height
#     g_Width = 1000
#     g_Height = 1000
#
#     global g_viewMode, g_bShowBackground
#     g_viewMode = 'free' #'camView'
#     g_bShowBackground = False
#
#     # keyboard('R',0,0)
#     global g_bRotateView, g_rotateView_counter
#     g_bRotateView = True
#     g_rotateView_counter =0
#     global g_bSaveToFile, g_bShowSkeleton
#     g_bSaveToFile = bSaveToFile
#
#
#     # g_bShowSkeleton = False
#     g_bShowSkeleton = bShowSkeleton
#
#     global g_stopMainLoop
#     g_stopMainLoop=False
#
#
#     init_gl_util()
#     #Setup for rendering
#     keyboard('c',0,0)
#
#
#
#
#     setupRotationView()
#
#
#     burningCnt =5
#     # while True:
#     while g_rotateView_counter*g_rotateInterval<360:
#         glutPostRedisplay()
#         if bool(glutMainLoopEvent)==False:
#             continue
#         glutMainLoopEvent()
#         if g_stopMainLoop:
#             break
#
#         if burningCnt>=0:
#             # keyboard('c',0,0)
#             g_rotateView_counter =0
#             burningCnt -=1
#     # print("Escaped glut loop")
#
#     #Generate video from png
#     # import subprocess
#
#     # cmd = "cd /home/hjoo/temp/render_general; ffmpeg -r 30 -f image2 -i scene_%08d.png -vcodec libx264 -crf 25  -pix_fmt yuv420p ../test1.mp4"
#     # subprocess.call(cmd,shell=True)



# def init_gl_MTC_cameraView(bSaveToFile = False):
#     global g_Width,g_Height
#     g_Width = 1920
#     g_Height = 1080
#     show_SMPL(bSaveToFile = False, bResetSaveImgCnt = True, mode = 'init')

def show_SMPL_sideView(bSaveToFile = False, bResetSaveImgCnt=True, countImg = True, bReturnRendered= True):
    show_SMPL_cameraView(bSaveToFile, bResetSaveImgCnt, countImg, False)

    if bReturnRendered and g_saveImageName_last is not None:
        renderedImg = cv2.imread(g_saveImageName_last)
        return renderedImg

def show_SMPL_youtubeView(bSaveToFile = False, bResetSaveImgCnt=True, countImg = True, zoom = 230, bReturnRendered= False):
    show_SMPL(bSaveToFile = bSaveToFile, bResetSaveImgCnt = bResetSaveImgCnt, countImg = countImg, zoom = zoom, mode = 'youtube')

    if bReturnRendered and g_saveImageName_last is not None:
        renderedImg = cv2.imread(g_saveImageName_last)
        return renderedImg

def show_SMPL_cameraView(bSaveToFile = False, bResetSaveImgCnt=True, countImg = True, bShowBG = True, bReturnRendered= False):
    show_SMPL(bSaveToFile = bSaveToFile, bResetSaveImgCnt = bResetSaveImgCnt, countImg = countImg, bShowBG = bShowBG, mode = 'camera')

    if bReturnRendered and g_saveImageName_last is not None:
        renderedImg = cv2.imread(g_saveImageName_last)
        return renderedImg

# This is to render scene in camera view (especially MTC output)
def show_SMPL(bSaveToFile = False, bResetSaveImgCnt = True, countImg = True, bShowBG = True, zoom = 230, mode = 'camera'):
    init_gl_util()

    if mode == 'init':
        #Setup for rendering
        keyboard('c',0,0)

    global g_bSaveToFile, g_bSaveToFile_done, g_bShowSkeleton, g_bShowFloor, g_viewMode, g_saveFrameIdx
    global g_xTrans,  g_yTrans, g_zoom, g_xRotate, g_yRotate, g_zRotate

    g_bSaveToFile_done = False
    g_bSaveToFile = bSaveToFile
    # g_bShowSkeleton = False
    g_bShowFloor = False

    if mode == 'youtube':
        # bShowBG = True
        g_bShowFloor = True
        if False:   #Original
            g_xTrans=  -86.0
            g_yTrans= 0.0
            g_zoom= zoom
            g_xRotate = 34.0
            g_yRotate= -32.0
            g_zRotate= 0.0
            g_viewMode = 'free'
        elif True:   # cook
            g_xTrans=  170
            g_yTrans= 0.0
            # g_zoom= 1600          #Cook
            # g_zoom= 1000                #Comedian
            g_zoom= 800                #Bengio
            g_xRotate = 34.0
            g_yRotate= -32.0
            g_zRotate= 0.0
            g_viewMode = 'free'
        elif True:   # almost size
            g_xTrans=  0
            g_yTrans= 0
            g_zoom= zoom #230
            g_xRotate = 61
            g_yRotate= 3
            g_zRotate= 0.0
            g_viewMode = 'free'
    elif mode =="side":
        bShowBG = False
        g_bShowFloor = False
        # global g_xTrans,  g_yTrans, g_zoom, g_xRotate, g_yRotate, g_zRotate
        # g_xTrans=  170
        # g_yTrans= 0.0
        # g_zoom= 1600          #Cook
        # g_zoom= 1000                #Comedian
        g_zoom= 266                #Bengio
        g_xRotate = 90
        g_yRotate= 0
        g_zRotate= 0.0
        g_viewMode = 'free'

    elif mode == 'camera' or mode == 'init':
        g_viewMode = 'camView'

    global g_bShowBackground
    g_bShowBackground = bShowBG

    if bResetSaveImgCnt:
        g_saveFrameIdx = 0 #+= 1      #always save as: scene_00000000.png
    
    if mode == 'init':
        global g_stopMainLoop
        g_stopMainLoop=False
        # while True:
        while g_rotateView_counter*g_rotateInterval<360:
            glutPostRedisplay()
            if bool(glutMainLoopEvent)==False:
                continue
            glutMainLoopEvent()
            break
            if g_stopMainLoop:
                break
    else:
        if g_bSaveToFile:
            while g_bSaveToFile_done == False:
                glutPostRedisplay()
                if bool(glutMainLoopEvent)==False:
                    continue
                glutMainLoopEvent()
        else:
            for i in range(6):   ##Render more than one to be safer
                glutPostRedisplay()
                if bool(glutMainLoopEvent)==False:
                    continue
                glutMainLoopEvent()

    if countImg:
        g_saveFrameIdx +=1
    g_bSaveToFile = False



def render_on_image(saveDir, saveFileName, inputImage, scaleFactor=1, is_showBackground=True):
    # (dirName, fileName, rawImg)
    
    g_bShowBackground = is_showBackground
    setBackgroundTexture(inputImage)
    setWindowSize(inputImage.shape[1] *scaleFactor, inputImage.shape[0]*scaleFactor)
    setSaveFolderName(saveDir)
    setSaveImgName(saveFileName)
    show_SMPL(bSaveToFile = True, bResetSaveImgCnt = False, countImg = True, mode = 'camera')




#For easier use.
#skel should be a skeleton in a single frame
#Input: skel should be a vector (N,)
def VisSkeleton_single(skel):

    skel = skel[:,np.newaxis]       #(N, 1)
    setSkeleton( [skel] , jointType='smplcoco')#(skelNum, dim, frames)
    show()

def setNearPlane(p):
    global g_nearPlane
    g_nearPlane = p


#Aliasing since the "init_gl" is a bit ugly name
def show(maxIter=-10):
    init_gl(maxIter)


if __name__ == '__main__':

    import cv2
    imgPathFull =  '/run/media/hjoo/disk/data/posetrack/images/train/000001_bonn_train/000023.jpg'
    image = cv2.imread(imgPathFull)

    setBackgroundTexture(image)

    init_gl()
    #See demos.py

    #Ignore below
    # demo_holden_data()
    # demo_panoptic_data_haggling()
    #demo_holdenFormat_panoptic()
    # demo_panoptic_data_haggling_meshes()
    #demo_panoptic_faceModel_pkl_haggling()     #Visualize faceWarehouse model, driven by loaded face motion data
    # demo_panoptic_faceModel_pkl()
    # demo_loadObjMesh()
    #demo_panoptic_AdamModel_multiFrames()      Not implemented
imports/frankmocap/renderer/glViewer.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import cv2
import numpy as np

def draw_keypoints(image, kps, color=(0,0,255), radius=5, check_exist=False):
    # recover color 
    if color == 'red':
        color = (0, 0, 255)
    elif color == 'green':
        color = (0, 255, 0)
    elif color == 'blue':
        color = (255, 0, 0)
    else:
        assert isinstance(color, tuple) and len(color) == 3

    # draw keypoints
    res_img = image.copy()
    for i in range(kps.shape[0]):
        x, y = kps[i][:2].astype(np.int32)
        if check_exist:
            score = kps[i][2]
        else:
            score = 1.0
        # print(i, score)
        if score > 0.0:
            cv2.circle(res_img, (x,y), radius=radius, color=color, thickness=-1)
    return res_img.astype(np.uint8)


def draw_bbox(image, bbox, color=(0,0,255), thickness=3):
    x0, y0 = int(bbox[0]), int(bbox[1])
    x1, y1 = int(bbox[2]), int(bbox[3])
    res_img = cv2.rectangle(image.copy(), (x0,y0), (x1,y1), color=color, thickness=thickness)
    return res_img.astype(np.uint8)



def draw_raw_bbox(img, bboxes):
    img = img.copy()
    for bbox in bboxes:
        x0, y0, w, h = bbox
        bbox_xyxy = (x0, y0, x0+w, y0+h)
        img = draw_bbox(img, bbox_xyxy)
    return img


def draw_body_bbox(img, body_bbox_list):
    img = img.copy()
    for body_bbox in body_bbox_list:
        if body_bbox is not None:
            x0, y0, w, h = body_bbox
            img = draw_bbox(img, (x0, y0, x0+w, y0+h))
    return img


def draw_arm_pose(img, body_pose_list):
    img = img.copy()
    for body_pose in body_pose_list:
        # left & right arm
        img = draw_keypoints(
            img, body_pose[6:8, :], radius=10, color=(255, 0, 0))
        img = draw_keypoints(
            img, body_pose[3:5, :], radius=10, color=(0, 0, 255))
    return img


def draw_hand_bbox(img, hand_bbox_list):
    img = img.copy()
    for hand_bboxes in hand_bbox_list:
        if hand_bboxes is not None:
            for key in hand_bboxes:
                bbox = hand_bboxes[key]
                if bbox is not None:
                    x0, y0, w, h = bbox
                    bbox_new = (x0, y0, x0+w, y0+h)
                    color = (255, 0, 0) if key == 'left_hand' else (0, 255, 0)
                    img = draw_bbox(img, bbox_new, color=color)
    return img
imports/frankmocap/renderer/image_utils.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import numpy as np
from OpenGL.GLUT import *
from OpenGL.GLU import *
from renderer.shaders.framework import *

from renderer.glRenderer import glRenderer

# from renderer.render_utils import ComputeNormal

_glut_window = None

'''
#Usage:
    render.set_smpl_mesh(v)     #v for vertex locations in(6890,3)
    render.setBackgroundTexture(rawImg) #Optional BG texture
    render.setWindowSize(rawImg.shape[1], rawImg.shape[0])      #Optional: window size
    render.show_once()
'''

class meshRenderer(glRenderer):

    def __init__(self, width=1600, height=1200, name='GL Renderer',
                #  program_files=['renderer/shaders/simple140.fs', 'renderer/shaders/simple140.vs'],
                #  program_files=['renderer/shaders/normal140.fs', 'renderer/shaders/normal140.vs'],
                # program_files=['renderer/shaders/geo140.fs', 'renderer/shaders/geo140.vs'],
                render_mode ="normal",  #color, geo, normal
                color_size=1, ms_rate=1,
                frankmocap_path='./imports/frankmocap_2/'): ## ADDEd

        self.render_mode = render_mode
        self.program_files ={}
        self.program_files['color'] = ['renderer/shaders/simple140.fs', 'renderer/shaders/simple140.vs']
        self.program_files['normal'] = ['renderer/shaders/normal140.fs', 'renderer/shaders/normal140.vs']
        self.program_files['geo'] = ['renderer/shaders/colorgeo140.fs', 'renderer/shaders/colorgeo140.vs']

        for key in self.program_files:
            self.program_files[key]=[frankmocap_path+renderer_pth for renderer_pth in self.program_files[key]]

        glRenderer.__init__(self, width, height, name, self.program_files[render_mode], color_size, ms_rate)

    def setRenderMode(self, render_mode):
        """
        Set render mode among ['color', 'normal', 'geo']
        """
        if self.render_mode == render_mode:
            return
        
        self.render_mode = render_mode
        self.initShaderProgram(self.program_files[render_mode])


    def drawMesh(self):

        if self.vertex_dim is None:
            return
        # self.draw_init()

        glColor3f(1,1,0)
        glUseProgram(self.program)
        
        mvMat = glGetFloatv(GL_MODELVIEW_MATRIX)
        pMat = glGetFloatv(GL_PROJECTION_MATRIX)
        # mvpMat = pMat*mvMat

        self.model_view_matrix = mvMat
        self.projection_matrix = pMat

        # glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, self.model_view_matrix.transpose())
        # glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, self.projection_matrix.transpose())
        glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, self.model_view_matrix)
        glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, self.projection_matrix)

        # Handle vertex buffer
        glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer)
        glEnableVertexAttribArray(0)
        glVertexAttribPointer(0, self.vertex_dim, GL_DOUBLE, GL_FALSE, 0, None)

        # # Handle normal buffer
        glBindBuffer(GL_ARRAY_BUFFER, self.normal_buffer)
        glEnableVertexAttribArray(1)
        glVertexAttribPointer(1, 3, GL_DOUBLE, GL_FALSE, 0, None)

        # # Handle color buffer
        glBindBuffer(GL_ARRAY_BUFFER, self.color_buffer)
        glEnableVertexAttribArray(2)
        glVertexAttribPointer(2, 3, GL_DOUBLE, GL_FALSE, 0, None)
        

        if True:#self.meshindex_data:
            glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, self.index_buffer)           #Note "GL_ELEMENT_ARRAY_BUFFER" instead of GL_ARRAY_BUFFER
            glBufferData(GL_ELEMENT_ARRAY_BUFFER, self.meshindex_data, GL_STATIC_DRAW)

        # glDrawArrays(GL_TRIANGLES, 0, self.n_vertices)
        glDrawElements(GL_TRIANGLES, len(self.meshindex_data), GL_UNSIGNED_INT, None)       #For index array (mesh face data)
        glDisableVertexAttribArray(0)
        glBindBuffer(GL_ARRAY_BUFFER, 0)

        glUseProgram(0)
imports/frankmocap/renderer/meshRenderer.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

"""
Renders mesh using OpenDr / Pytorch-3d for visualization.
Part of code is modified from https://github.com/akanazawa/hmr
"""

import sys
import numpy as np
import cv2
import pdb
from PIL import Image, ImageDraw
from opendr.camera import ProjectPoints
from opendr.renderer import ColoredRenderer
from opendr.lighting import LambertianPointLight


class OpendrRenderer(object):
    def __init__(self,
                 img_size=224,
                 mesh_color=np.array([0.5, 0.5, 0.5]),):

        self.w = img_size
        self.h = img_size
        self.color = mesh_color
        self.img_size = img_size
        self.flength = 500.

    
    def render(self, verts, faces, bg_img):
        verts = verts.copy()
        faces = faces.copy()

        input_size = 500

        f = 10

        verts[:, 0] = (verts[:, 0] - input_size) / input_size
        verts[:, 1] = (verts[:, 1] - input_size) / input_size

        verts[:, 2] /= (5 * 112)
        verts[:, 2] += f

        cam_for_render = np.array([f, 1, 1]) * input_size

        rend_img = self.__call__(
            img=bg_img, cam=cam_for_render, 
            verts=verts, faces=faces, color=self.color)
        
        return rend_img
    

    def __call__(self,
                 verts,
                 faces,
                 cam=None,
                 img=None,
                 do_alpha=False,
                 far=None,
                 near=None,
                 color = np.array([0, 0, 255]),
                 img_size=None):
        """
        cam is 3D [f, px, py]
        """
        if img is not None:
            h, w = img.shape[:2]
        elif img_size is not None:
            h = img_size[0]
            w = img_size[1]
        else:
            h = self.h
            w = self.w

        if cam is None:
            cam = [self.flength, w / 2., h / 2.]

        use_cam = ProjectPoints(
            f=cam[0] * np.ones(2),
            rt=np.zeros(3),
            t=np.zeros(3),
            k=np.zeros(5),
            c=cam[1:3])

        if near is None:
            near = np.maximum(np.min(verts[:, 2]) - 25, 0.1)
        if far is None:
            far = np.maximum(np.max(verts[:, 2]) + 25, 25)

        return_value = render_model(
            verts,
            faces,
            w,
            h,
            use_cam,
            do_alpha=do_alpha,
            img=img,
            far=far,
            near=near,
            color=color)

        imtmp = return_value
        image = (imtmp * 255).astype('uint8')
        return image



def _create_renderer(w=640,
                     h=480,
                     rt=np.zeros(3),
                     t=np.zeros(3),
                     f=None,
                     c=None,
                     k=None,
                     near=.5,
                     far=10.):

    f = np.array([w, w]) / 2. if f is None else f
    c = np.array([w, h]) / 2. if c is None else c
    k = np.zeros(5) if k is None else k

    rn = ColoredRenderer()

    rn.camera = ProjectPoints(rt=rt, t=t, f=f, c=c, k=k)
    rn.frustum = {'near': near, 'far': far, 'height': h, 'width': w}
    return rn


def _rotateY(points, angle):
    """Rotate the points by a specified angle."""
    ry = np.array([[np.cos(angle), 0., np.sin(angle)], [0., 1., 0.],
                   [-np.sin(angle), 0., np.cos(angle)]])
    return np.dot(points, ry)


def simple_renderer(rn,
                    verts,
                    faces,
                    yrot=np.radians(70),
                    color=np.array([0, 0, 255])
                    ):

    # Rendered model color
    rn.set(v=verts, f=faces, vc=color, bgcolor=np.ones(3))
    albedo = rn.vc

    # Construct Back Light (on back right corner)
    rn.vc = LambertianPointLight(
        f=rn.f,
        v=rn.v,
        num_verts=len(rn.v),
        light_pos=_rotateY(np.array([-200, -100, -100]), yrot),
        vc=albedo,
        light_color=np.array([1, 1, 1]))

    # Construct Left Light
    rn.vc += LambertianPointLight(
        f=rn.f,
        v=rn.v,
        num_verts=len(rn.v),
        # light_pos=_rotateY(np.array([800, 10, 300]), yrot),
        light_pos=_rotateY(np.array([800, 10, 300]), yrot),
        vc=albedo,
        light_color=np.array([1, 1, 1]))

    # Construct Right Light
    rn.vc += LambertianPointLight(
        f=rn.f,
        v=rn.v,
        num_verts=len(rn.v),
        light_pos=_rotateY(np.array([-500, 500, 1000]), yrot),
        # light_pos=_rotateY(np.array([-500, 500, 1000]), yrot),
        vc=albedo,
        light_color=np.array([.7, .7, .7]))
 
    return rn.r


def get_alpha(imtmp, bgval=1.):
    h, w = imtmp.shape[:2]
    alpha = (~np.all(imtmp == bgval, axis=2)).astype(imtmp.dtype)

    b_channel, g_channel, r_channel = cv2.split(imtmp)

    im_RGBA = cv2.merge((b_channel, g_channel, r_channel, alpha.astype(
        imtmp.dtype)))
    return im_RGBA


def append_alpha(imtmp):
    alpha = np.ones_like(imtmp[:, :, 0]).astype(imtmp.dtype)
    if np.issubdtype(imtmp.dtype, np.uint8):
        alpha = alpha * 255
    b_channel, g_channel, r_channel = cv2.split(imtmp)
    im_RGBA = cv2.merge((b_channel, g_channel, r_channel, alpha))
    return im_RGBA


def render_model(verts,
                 faces,
                 w,
                 h,
                 cam,
                 near=0.5,
                 far=25,
                 img=None,
                 do_alpha=False,
                 color=None):
    rn = _create_renderer(
        w=w, h=h, near=near, far=far, rt=cam.rt, t=cam.t, f=cam.f, c=cam.c)

    # Uses img as background, otherwise white background.
    if img is not None:
        rn.background_image = img / 255. if img.max() > 1.1 else img

    imtmp = simple_renderer(rn, verts, faces, color=color)
    
    # If white bg, make transparent.
    if img is None and do_alpha:
        imtmp = get_alpha(imtmp)
    elif img is not None and do_alpha:
        imtmp = append_alpha(imtmp)
    
    return imtmp
imports/frankmocap/renderer/od_renderer.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

# Part of code is modified from https://github.com/facebookresearch/pytorch3d

import cv2
import os
import sys
import torch
import numpy as np

from pytorch3d.structures import Meshes
from pytorch3d.renderer.mesh import Textures
from pytorch3d.renderer import (
    PerspectiveCameras,
    FoVOrthographicCameras,
    PointLights, 
    RasterizationSettings, 
    MeshRenderer, 
    BlendParams,
    MeshRasterizer,  
    SoftPhongShader,
)

class Pytorch3dRenderer(object):

    def __init__(self, img_size, mesh_color):
        self.device = torch.device("cuda:0")
        # self.render_size = 1920

        self.img_size = img_size

        # mesh color
        mesh_color = np.array(mesh_color)[::-1]
        self.mesh_color = torch.from_numpy(
            mesh_color.copy()).view(1, 1, 3).float().to(self.device)

        # renderer for large objects, such as whole body.
        self.render_size_large = 700
        lights = PointLights(
            ambient_color = [[1.0, 1.0, 1.0],],
            diffuse_color = [[1.0, 1.0, 1.0],],
            device=self.device, location=[[1.0, 1.0, -30]])
        self.renderer_large = self.__get_renderer(self.render_size_large, lights)

        # renderer for small objects, such as whole body.
        self.render_size_medium = 400
        lights = PointLights(
            ambient_color = [[0.5, 0.5, 0.5],],
            diffuse_color = [[0.5, 0.5, 0.5],],
            device=self.device, location=[[1.0, 1.0, -30]])
        self.renderer_medium = self.__get_renderer(self.render_size_medium, lights)


        # renderer for small objects, such as whole body.
        self.render_size_small = 200
        lights = PointLights(
            ambient_color = [[0.5, 0.5, 0.5],],
            diffuse_color = [[0.5, 0.5, 0.5],],
            device=self.device, location=[[1.0, 1.0, -30]])
        self.renderer_small = self.__get_renderer(self.render_size_small, lights)



    def __get_renderer(self, render_size, lights):

        cameras = FoVOrthographicCameras(
            device = self.device,
            znear=0.1,
            zfar=10.0,
            max_y=1.0,
            min_y=-1.0,
            max_x=1.0,
            min_x=-1.0,
            scale_xyz=((1.0, 1.0, 1.0),),  # (1, 3)
        )

        raster_settings = RasterizationSettings(
            image_size = render_size,
            blur_radius = 0,
            faces_per_pixel = 1,
        )
        blend_params = BlendParams(sigma=1e-4, gamma=1e-4, background_color = (0,0,0))

        renderer = MeshRenderer(
            rasterizer=MeshRasterizer(
                cameras=cameras, 
                raster_settings=raster_settings
            ),
            shader=SoftPhongShader(
                device=self.device, 
                cameras=cameras,
                lights=lights,
                blend_params=blend_params
            )
        )

        return renderer
    

    def render(self, verts, faces, bg_img):
        verts = verts.copy()
        faces = faces.copy()

        # bbox for verts
        x0 = int(np.min(verts[:, 0]))
        x1 = int(np.max(verts[:, 0]))
        y0 = int(np.min(verts[:, 1]))
        y1 = int(np.max(verts[:, 1]))
        width = x1 - x0
        height = y1 - y0

        bbox_size = max(height, width)
        if bbox_size <= self.render_size_small:
            # print("Using small size renderer")
            render_size = self.render_size_small
            renderer = self.renderer_small
        else:
            if bbox_size <= self.render_size_medium:
                # print("Using medium size renderer")
                render_size = self.render_size_medium
                renderer = self.renderer_medium
            else:
                # print("Using large size renderer")
                render_size = self.render_size_large
                renderer = self.renderer_large
        
        # padding the tight bbox
        margin = int(max(width, height) * 0.1)
        x0 = max(0, x0-margin)
        y0 = max(0, y0-margin)
        x1 = min(self.img_size, x1+margin)
        y1 = min(self.img_size, y1+margin)

        # move verts to be in the bbox
        verts[:, 0] -= x0
        verts[:, 1] -= y0

        # normalize verts to (-1, 1)
        bbox_size = max(y1-y0, x1-x0)
        half_size = bbox_size / 2
        verts[:, 0] = (verts[:, 0] - half_size) / half_size
        verts[:, 1] = (verts[:, 1] - half_size) / half_size

        # the coords of pytorch-3d is (1, 1) for upper-left and (-1, -1) for lower-right
        # so need to multiple minus for vertices
        verts[:, :2] *= -1

        # shift verts along the z-axis
        verts[:, 2] /= 112
        verts[:, 2] += 5

        verts_tensor = torch.from_numpy(verts).float().unsqueeze(0).cuda()
        faces_tensor = torch.from_numpy(faces.copy()).long().unsqueeze(0).cuda()

        # set color
        mesh_color = self.mesh_color.repeat(1, verts.shape[0], 1)
        textures = Textures(verts_rgb = mesh_color)

        # rendering
        mesh = Meshes(verts=verts_tensor, faces=faces_tensor, textures=textures)

        # blending rendered mesh with background image
        rend_img = renderer(mesh)
        rend_img = rend_img[0].cpu().numpy()


        scale_ratio = render_size / bbox_size
        img_size_new = int(self.img_size * scale_ratio)
        bg_img_new = cv2.resize(bg_img, (img_size_new, img_size_new))

        x0 = max(int(x0 * scale_ratio), 0)
        y0 = max(int(y0 * scale_ratio), 0)
        x1 = min(int(x1 * scale_ratio), img_size_new)
        y1 = min(int(y1 * scale_ratio), img_size_new)

        h0 = min(y1-y0, render_size)
        w0 = min(x1-x0, render_size)

        y1 = y0 + h0
        x1 = x0 + w0

        rend_img_new = np.zeros((img_size_new, img_size_new, 4))
        rend_img_new[y0:y1, x0:x1, :] = rend_img[:h0, :w0, :]
        rend_img = rend_img_new

        alpha = rend_img[:, :, 3:4]
        alpha[alpha>0] = 1.0
        

        rend_img = rend_img[:, :, :3] 
        maxColor = rend_img.max()
        rend_img *= 255 /maxColor #Make sure <1.0
        rend_img = rend_img[:, :, ::-1]

        res_img = alpha * rend_img + (1.0 - alpha) * bg_img_new

        res_img = cv2.resize(res_img, (self.img_size, self.img_size))

        return res_img
imports/frankmocap/renderer/p3d_renderer.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

import numpy as np


# vertices: frames x meshVerNum x 3
# trifaces: facePolygonNum x 3 = 22800 x 3
def ComputeNormal(vertices, trifaces):

    if vertices.shape[0] > 5000:
        print('ComputeNormal: Warning: too big to compute {0}'.format(vertices.shape) )
        return

    #compute vertex Normals for all frames
    U = vertices[:,trifaces[:,1],:] - vertices[:,trifaces[:,0],:]  #frames x faceNum x 3
    V = vertices[:,trifaces[:,2],:] - vertices[:,trifaces[:,1],:]  #frames x faceNum x 3
    originalShape = U.shape  #remember: frames x faceNum x 3

    U = np.reshape(U, [-1,3])
    V = np.reshape(V, [-1,3])
    faceNormals = np.cross(U,V)     #frames x 13776 x 3
    from sklearn.preprocessing import normalize

    if np.isnan(np.max(faceNormals)):
        print('ComputeNormal: Warning nan is detected {0}')
        return
    faceNormals = normalize(faceNormals)

    faceNormals = np.reshape(faceNormals, originalShape)

    if False:        #Slow version
        vertex_normals = np.zeros(vertices.shape) #(frames x 11510) x 3
        for fIdx, vIdx in enumerate(trifaces[:,0]):
            vertex_normals[:,vIdx,:] += faceNormals[:,fIdx,:]
        for fIdx, vIdx in enumerate(trifaces[:,1]):
            vertex_normals[:,vIdx,:] += faceNormals[:,fIdx,:]
        for fIdx, vIdx in enumerate(trifaces[:,2]):
            vertex_normals[:,vIdx,:] += faceNormals[:,fIdx,:]
    else:   #Faster version
        # Computing vertex normals, much faster (and obscure) replacement
        index = np.vstack((np.ravel(trifaces), np.repeat(np.arange(len(trifaces)), 3))).T
        index_sorted = index[index[:,0].argsort()]
        vertex_normals = np.add.reduceat(faceNormals[:,index_sorted[:, 1],:][0],
            np.concatenate(([0], np.cumsum(np.unique(index_sorted[:, 0],
            return_counts=True)[1])[:-1])))[None, :]
        vertex_normals = vertex_normals.astype(np.float64)

    originalShape = vertex_normals.shape
    vertex_normals = np.reshape(vertex_normals, [-1,3])
    vertex_normals = normalize(vertex_normals)
    vertex_normals = np.reshape(vertex_normals,originalShape)

    return vertex_normals



def ComputeNormal_gpu(vertices, trifaces):
    import torch
    import torch.nn.functional as F

    if vertices.shape[0] > 5000:
        print('ComputeNormal: Warning: too big to compute {0}'.format(vertices.shape) )
        return

    #compute vertex Normals for all frames
    #trifaces_cuda = torch.from_numpy(trifaces.astype(np.long)).cuda()
    vertices_cuda = torch.from_numpy(vertices.astype(np.float32)).cuda()

    U_cuda = vertices_cuda[:,trifaces[:,1],:] - vertices_cuda[:,trifaces[:,0],:]  #frames x faceNum x 3
    V_cuda = vertices_cuda[:,trifaces[:,2],:] - vertices_cuda[:,trifaces[:,1],:]  #frames x faceNum x 3
    originalShape = list(U_cuda.size())  #remember: frames x faceNum x 3

    U_cuda = torch.reshape(U_cuda, [-1,3])#.astype(np.float32)
    V_cuda = torch.reshape(V_cuda, [-1,3])#.astype(np.float32)

    faceNormals = U_cuda.cross(V_cuda)
    faceNormals = F.normalize(faceNormals,dim=1)

    faceNormals = torch.reshape(faceNormals, originalShape)

    # trifaces has duplicated vertex index, so cannot be parallazied
    # vertex_normals = torch.zeros(vertices.shape,dtype=torch.float32).cuda() #(frames x 11510) x 3
    # for fIdx, vIdx in enumerate(trifaces[:,0]):
    #    vertex_normals[:,vIdx,:] += faceNormals[:,fIdx,:]
    # for fIdx, vIdx in enumerate(trifaces[:,1]):
    #     vertex_normals[:,vIdx,:] += faceNormals[:,fIdx,:]
    # for fIdx, vIdx in enumerate(trifaces[:,2]):
    #     vertex_normals[:,vIdx,:] += faceNormals[:,fIdx,:]

    # Computing vertex normals, much faster (and obscure) replacement
    index = np.vstack((np.ravel(trifaces), np.repeat(np.arange(len(trifaces)), 3))).T
    index_sorted = index[index[:,0].argsort()]
    vertex_normals = np.add.reduceat(faceNormals[:,index_sorted[:, 1],:][0],
        np.concatenate(([0], np.cumsum(np.unique(index_sorted[:, 0],
        return_counts=True)[1])[:-1])))[None, :]
    vertex_normals = torch.from_numpy(vertex_normals).float().cuda()

    vertex_normals = F.normalize(vertex_normals,dim=2)
    vertex_normals = vertex_normals.data.cpu().numpy()  #(batch, chunksize, dim)

    return vertex_normals
imports/frankmocap/renderer/render_utils.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

"""
Renders mesh using OpenDr / Pytorch-3D for visualization.
"""

import sys
import renderer.p3d_renderer as p3d_renderer
import renderer.od_renderer as od_renderer

import numpy as np
import cv2
import pdb
from .image_utils import draw_raw_bbox, draw_hand_bbox, draw_body_bbox, draw_arm_pose

# To use screen_free visualizer. Either OpenDR or Pytorch3D should be installed.
# g_valid_visualize = False
# try:
#     from Ego4d_Codes.github_codes.frankmocap.renderer.od_renderer import OpendrRenderer
#     g_valid_visualize = True
# except ImportError:
#     print("Cannot import OpenDR Renderer")
# try:
#     from Ego4d_Codes.github_codes.frankmocap.renderer.p3d_renderer import Pytorch3dRenderer
#     g_valid_visualize = True
# except ImportError:
#     print("Cannot import Pytorch3D Renderer")
# assert g_valid_visualize, "You should import either OpenDR or Pytorch3D"


class Visualizer(object):

    def __init__(self, renderer_backend):
        colors = {
            # colorbline/print/copy safe:
            'light_gray':  [0.9, 0.9, 0.9],
            'light_purple':  [0.8, 0.53, 0.53],
            'light_green': [166/255.0, 178/255.0, 30/255.0],
            'light_blue': [0.65098039, 0.74117647, 0.85882353],
        }

        self.input_size = 1920

        # set-up renderer
        assert renderer_backend in ['opendr', 'pytorch3d']
        if renderer_backend == 'opendr':
            self.renderer = od_renderer.OpendrRenderer(
                img_size=self.input_size, 
                mesh_color=colors['light_purple'])
        else:
            self.renderer = p3d_renderer.Pytorch3dRenderer(
                img_size=self.input_size, 
                mesh_color=colors['light_purple'])


    def __render_pred_verts(self, img_original, pred_mesh_list):
        assert max(img_original.shape) <= self.input_size, \
            f"Currently, we donlt support images size larger than:{self.input_size}"

        res_img = img_original.copy()
        rend_img = np.ones((self.input_size, self.input_size, 3))
        h, w = img_original.shape[:2]
        rend_img[:h, :w, :] = img_original

        for mesh in pred_mesh_list:
            verts = mesh['vertices']
            faces = mesh['faces']
            rend_img = self.renderer.render(verts, faces, rend_img)

        res_img = rend_img[:h, :w, :]
        return res_img


    def visualize(self, 
        input_img, 
        hand_bbox_list = None, 
        body_bbox_list = None,
        body_pose_list = None,
        raw_hand_bboxes = None,
        pred_mesh_list = None,
        vis_raw_hand_bbox = True,
        vis_body_pose = True,
        vis_hand_bbox = True,
    ):
        # init
        res_img = input_img.copy()

        # draw raw hand bboxes
        if raw_hand_bboxes is not None and vis_raw_hand_bbox:
            res_img = draw_raw_bbox(input_img, raw_hand_bboxes)
            # res_img = np.concatenate((res_img, raw_bbox_img), axis=1)

        # draw 2D Pose
        if body_pose_list is not None and vis_body_pose:
            res_img = draw_arm_pose(res_img, body_pose_list)

        # draw body bbox
        if body_bbox_list is not None:
            body_bbox_img = draw_body_bbox(input_img, body_bbox_list)
            res_img = body_bbox_img

        # draw hand bbox
        if hand_bbox_list is not None:
            res_img = draw_hand_bbox(res_img, hand_bbox_list)
        
        # render predicted meshes
        if pred_mesh_list is not None:
            rend_img = self.__render_pred_verts(input_img, pred_mesh_list)
            res_img = np.concatenate((res_img, rend_img), axis=1)
            # res_img = rend_img
        
        return res_img
imports/frankmocap/renderer/screen_free_visualizer.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

#Visualization Function

import cv2

import numpy as np
import PIL
from PIL.Image import Image

def __ValidateNumpyImg(inputImg):
    if isinstance(inputImg, Image):
        # inputImg = cv2.cvtColor(np.array(inputImg), cv2.COLOR_RGB2BGR)
        inputImg = np.array(inputImg)

    return inputImg     #Q? is this copying someting (wasting memory or time?)?

veryFirstImShow = True
def ImShow(inputImg, waitTime=1, bConvRGB2BGR=False,name='image', scale=1.0):

    inputImg = __ValidateNumpyImg(inputImg)

    if scale!=1.0:
        inputImg = cv2.resize(inputImg, (inputImg.shape[0]*int(scale), inputImg.shape[1]*int(scale)))


    if bConvRGB2BGR:
        inputImg = cv2.cvtColor(inputImg, cv2.COLOR_RGB2BGR)

    cv2.imshow(name,inputImg)

    global veryFirstImShow
    if False:#veryFirstImShow:
        print(">> Press any key to move on")
        cv2.waitKey(0)      #the initial one is always blank... why?
        veryFirstImShow = 0
    else:
        cv2.waitKey(waitTime)

def ImgSC(inputImg, waitTime=1, bConvRGB2BGR=False,name='image', scale=1.0):

    inputImg = __ValidateNumpyImg(inputImg)

    minVal = np.min(inputImg)
    maxVal = np.max(inputImg)

    #rescale 
    inputImg = (inputImg-minVal)/ (maxVal-minVal)*255

    if scale!=1.0:
        inputImg = cv2.resize(inputImg, (inputImg.shape[0]*int(scale), inputImg.shape[1]*int(scale)))


    if bConvRGB2BGR:
        inputImg = cv2.cvtColor(inputImg, cv2.COLOR_RGB2BGR)

    cv2.imshow(name,inputImg)

    global veryFirstImShow
    if veryFirstImShow:
        print(">> Press any key to move on")
        cv2.waitKey(0)      #the initial one is always blank... why?
        veryFirstImShow = 0
    else:
        cv2.waitKey(waitTime)

#bbe: min_pt, max_pt
def Vis_Bbox_minmaxPt(inputImg, min_pt, max_pt, color=None):

    bbr = [min_pt[0],min_pt[1], max_pt[0]- min_pt[0], max_pt[1]- min_pt[1]]
    return Vis_Bbox(inputImg, bbr, color)


def Vis_Bbox_XYXY(inputImg, bbox_xyxy, color=None):

    #draw biggest bbox
    pt1 = ( int(bbox_xyxy[0]),int(bbox_xyxy[1]) )
    pt2 = (int(bbox_xyxy[2]),int(bbox_xyxy[3]) )

    if color is None:
        color = (0,0,255)
    cv2.rectangle(inputImg, pt1, pt2,color, 3)

    return inputImg

def Vis_Bbox(inputImg, bbox_xyhw, color= None):
    return Vis_Bbox_XYWH(inputImg, bbox_xyhw, color)

#bbe: [leftTop_x,leftTop_y,width,height]
def Vis_Bbox_XYWH(inputImg, bbox_xyhw, color= None):

    inputImg = __ValidateNumpyImg(inputImg)

    #draw biggest bbox
    pt1 = ( int(bbox_xyhw[0]),int(bbox_xyhw[1]) )
    pt2 = (int(bbox_xyhw[0]  + bbox_xyhw[2]),int(bbox_xyhw[1] + bbox_xyhw[3]) )

    if color is None:
        color = (0,0,255)
    cv2.rectangle(inputImg, pt1, pt2,color, 3)

    return inputImg

def Vis_CocoBbox(inputImg, coco_annot):

    inputImg = __ValidateNumpyImg(inputImg)

    bbr =  np.round(coco_annot['bbox'])  #[leftTop_x,leftTop_y,width,height]

    #draw biggest bbox
    pt1 = ( int(bbr[0]),int(bbr[1]) )
    pt2 = (int(bbr[0]  + bbr[2]),int(bbr[1] + bbr[3]) )
    cv2.rectangle(inputImg, pt1, pt2,(255,255,255), 3)

    return inputImg

def Vis_CocoSkeleton(keypoints, image=None):
# def Vis_CocoSkeleton(inputImg, coco_annot):

    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    image = __ValidateNumpyImg(image)

    #COCO17 original annotation ordering
    link2D = [ [0, 1], [1,3],		#nose(0), leftEye(1), leftEar(3)
                [0,5], [5, 7], [7, 9],			#leftShoulder(5), leftArm(7), leftWrist(9)
        	    [0, 11], [11, 13], [13, 15],  #leftHip(11), leftKnee(13), leftAnkle(15)

                [0,2], [2,4],     #nose(0), rightEye(2), rightEar(4)
                [0,6], [6, 8], [8, 10],			#rightShoulder(6), rightArm(8), rightWrist(10)
                [0, 12], [12, 14], [14, 16]  #rightHip(12), rightKnee(14), rightAnkle(16)
                ]

    bLeft = [ 1,1,
             1, 1, 1,
              1,1,1,
              0,0,
              0,0,0,
               0,0,0]

    # keypoints =  np.round(coco_annot['keypoints'])  #coco_annot['keypoints']: list with length 51
    if keypoints.shape[0] == 51:
        keypoints = np.reshape(keypoints, (-1,3))   #(17,3): (X, Y, Label)
    else:
        keypoints = np.reshape(keypoints, (-1,2))   #(17,3): (X, Y, Label)

    radius = 4

    for k in np.arange( len(keypoints) ):
        cv2.circle(image, (int(keypoints[k][0]), int(keypoints[k][1]) ), radius,(0,0,255),-1)

    for k in np.arange( len(link2D) ):
        parent = link2D[k][0]
        child = link2D[k][1]
        if bLeft[k]:
            c = (0,0,255)#BGR, RED
        else:
            c = (200,200,200)

        if keypoints[parent][0] ==0 or keypoints[child][0]==0: #		//not annotated one
            continue

        cv2.line(image, (int(keypoints[parent][0]), int(keypoints[parent][1])), (int(keypoints[child][0]), int(keypoints[child][1])), c, radius - 2)

    return image

DP_partIdx ={
    'Torso_Back': 1,
    'Torso_Front': 2,
    'RHand': 3,
    'LHand': 4,
    'LFoot': 5,
    'RFoot': 6,

    'R_upperLeg_back': 7,
    'L_upperLeg_back': 8,
    'R_upperLeg_front': 9,
    'L_upperLeg_front': 10,

    'R_lowerLeg_back': 11,
    'L_lowerLeg_back': 12,
    'R_lowerLeg_front': 13,
    'L_lowerLeg_front': 14,

    'L_upperArm_front': 15,
    'R_upperArm_front': 16,
    'L_upperArm_back': 17,
    'R_upperArm_back': 18,

    'L_lowerArm_back': 19,
    'R_lowerArm_back': 20,
    'L_lowerArm_front': 21,
    'R_lowerArm_front': 22,

    'RFace': 23,
    'LFace': 24
}
def Vis_Densepose(inputImg, coco_annot):

    inputImg = __ValidateNumpyImg(inputImg)

    import densepose_methods as dp_utils
    DP = dp_utils.DensePoseMethods()

    if('dp_x' not in coco_annot.keys()):
        print("## Warning: No Densepose coco_annotation")
        return inputImg

    bbr =  np.round(coco_annot['bbox'])  #[leftTop_x,leftTop_y,width,height]

    Point_x = np.array(coco_annot['dp_x'])/ 255. * bbr[2] + bbr[0] # Strech the points to current box. from 255x255 -> [bboxWidth,bboxheight]
    Point_y = np.array(coco_annot['dp_y'])/ 255. * bbr[3] + bbr[1] # Strech the points to current box.
    # part_seg_index = np.array(coco_annot['dp_I'])   # part segment info

    #coco_annot['dp_I']: indexing
    # Torso Back: 1
    # Torso front: 2
    # RHand: 3
    # LHand: 4
    # LFoot: 5
    # RFoot: 6

    # R_upperLeg_back 7
    # L_upperLeg_back 8
    # R_upperLeg_front 9
    # L_upperLeg_front 10

    # R_lowerLeg_back 11
    # L_lowerLeg_back 12
    # R_lowerLeg_front 13
    # L_lowerLeg_front 14

    # L_upperArm_front 15
    # R_upperArm_front 16
    # L_upperArm_back 17
    # R_upperArm_back 18

    # L_lowerArm_back 19
    # R_lowerArm_back 20
    # L_lowerArm_front 21
    # R_lowerArm_front 22

    # RFace: 23
    # LFace: 24

    #Found BBoxes for rhand, lhand, and face using DensePose Data
    RHandIdx = [i for i,x in enumerate(coco_annot['dp_I']) if x == DP_partIdx['RHand'] ] #3.0]
    if len(RHandIdx)>0:

        minX = min(Point_x[RHandIdx])
        maxX = max(Point_x[RHandIdx])
        minY = min(Point_y[RHandIdx])
        maxY = max(Point_y[RHandIdx])
        RhandBBox = [minX, minY, maxX-minX, maxY-minY]
    else:
        RhandBBox = [-1,-1,-1,-1]

    LHandIdx = [i for i,x in enumerate(coco_annot['dp_I']) if x == DP_partIdx['LHand'] ]#4.0]
    if len(LHandIdx)>0:
        minX = min(Point_x[LHandIdx])
        maxX = max(Point_x[LHandIdx])
        minY = min(Point_y[LHandIdx])
        maxY = max(Point_y[LHandIdx])
        LhandBBox = [minX, minY, maxX-minX, maxY-minY]
    else:
        LhandBBox = [-1,-1,-1,-1]

    FaceIdx = [i for i,x in enumerate(coco_annot['dp_I']) if x == DP_partIdx['RFace'] or x == DP_partIdx['LFace'] ]   #23.0 or x == 24.0]
    if len(FaceIdx)>0:
        minX = min(Point_x[FaceIdx])
        maxX = max(Point_x[FaceIdx])
        minY = min(Point_y[FaceIdx])
        maxY = max(Point_y[FaceIdx])
        FaceBBox = [minX, minY, maxX-minX, maxY-minY]
    else:
        FaceBBox = [-1,-1,-1,-1]

    # #U,V,I -> Adam vertex  (Todo: should be reverified)
    # adamVerIdx_vec = np.zeros(len(coco_annot['dp_I']))
    # for i, (ii,uu,vv) in enumerate(zip(coco_annot['dp_I'],coco_annot['dp_U'],coco_annot['dp_V'])):

    #     vertexId = DP.IUV2VertexId(ii,uu,vv)
    #     adamVerIdx_vec[i] = vertexId

    # #draw biggest bbox
    # pt1 = ( int(bbr[0]),int(bbr[1]) )
    # pt2 = (int(bbr[0]  + bbr[2]),int(bbr[1] + bbr[3]) )
    # cv2.rectangle(inputImg, pt1, pt2,(0,0,0),1)

    #draw RHand bbox
    pt1 = ( int(RhandBBox[0]),int(RhandBBox[1]) )
    pt2 = (int(RhandBBox[0]  + RhandBBox[2]),int(RhandBBox[1] + RhandBBox[3]) )
    cv2.rectangle(inputImg, pt1, pt2,(0,0,255),2)

    #draw lHand bbox
    pt1 = ( int(LhandBBox[0]),int(LhandBBox[1]) )
    pt2 = (int(LhandBBox[0]  + LhandBBox[2]),int(LhandBBox[1] + LhandBBox[3]) )
    cv2.rectangle(inputImg, pt1, pt2,(0,255,0),2)

    #draw Face bbox
    pt1 = ( int(FaceBBox[0]),int(FaceBBox[1]) )
    pt2 = (int(FaceBBox[0]  + FaceBBox[2]),int(FaceBBox[1] + FaceBBox[3]) )
    cv2.rectangle(inputImg, pt1, pt2,(255,0,0),2)

    # Draw Densepose Keypoints
    tempColorIdx  = np.array(coco_annot['dp_I'])/ 24 *255
    #tempColorIdx  = np.array(coco_annot['dp_U']) *255
    #tempColorIdx  = np.array(coco_annot['dp_V']) *255
    tempColorIdx = np.uint8(tempColorIdx)
    tempColorIdx = cv2.applyColorMap(tempColorIdx, cv2.COLORMAP_JET)
    for cnt, pt in enumerate(zip(Point_x,Point_y,tempColorIdx, coco_annot['dp_I'])):
        # if pt[3] != DP_partIdx['Torso_Front']: #Uncomment this if you want to draw specific part
        #     continue
        #tempColorIdx  = coco_annot['dp_I']
        tempColor = pt[2][0].astype(np.int32).tolist()
        cv2.circle(inputImg,(int(pt[0]),int(pt[1])), 5,tempColor, -1)

    return inputImg


#H36m skeleton32
#skel can be
# : (17,2)       #with Nose
# : (16,2)       #without NOse
# : (32,2)       #original
def Vis_Skeleton_2D_H36m(pt2d, image = None, color=None):
    pt2d = np.reshape(pt2d,[-1,2])          #Just in case. Make sure (32, 2)


    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4

    if(pt2d.shape[0]==16):
        print("Vis_Skeleton_2D_H36m: {} joints".format(16))
        #Without Nose
        link2D = [ [0,1],[1,2],[2,3],#root(0), rHip(1), rKnee(2), rAnkle(3)
                            [0,4],[4,5],[5,6],#root(0, lHip(4), lKnee(5), lAnkle(6)
                            [0,7], [7,8], [8,9], #root(0, spineMid(7), neck(8), head(9)
                            [8,10], [10,11], [11,12], #Left Arms. neck(8). lshoulder(10),  lElbow(11), lWrist (12)
                            [8,13], [13,14], [14,15] #Right Arm, neck(8), rshoulder(13),  rElbow(14), rWrist (15)
                            ]
        bLeft = [ 0,0,0,
            1, 1, 1,
            1,1,1,
            1,1,1,
            0,0,0]
    elif pt2d.shape[0]==17:
        print("Vis_Skeleton_2D_H36m: {} joints".format(17))
        #With Nose
        link2D = [ [0,1],[1,2],[2,3],#root(0), rHip(1), rKnee(2), rAnkle(3)
                            [0,4],[4,5],[5,6],#root(0, lHip(4), lKnee(5), lAnkle(6)
                            [0,7], [7,8], [8,9], [9,10], #root(0, spineMid(7), neck(8), nose(9), head(9)
                            [8,11], [11,12], [12,13], #Left Arms. neck(8). lshoulder(11),  lElbow(12), lWrist (13)
                            [8,14], [14,15], [15,16] #Right Arm, neck(8), rshoulder(14),  rElbow(15), rWrist (16)
                            ]
        bLeft = [ 0,0,0,
            1, 1, 1,
            1,1,1, 1,
            1,1,1,
            0,0,0]
    else:
        print("Vis_Skeleton_2D_H36m: {} joints".format(32))
        #Human 36m DB's mocap data. 32 joints
        link2D = [ [0,1],[1,2],[2,3],[3,4],[4,5], #RightLeg: root(0), rHip(1), rKnee(2), rAnkle(3), rFootMid(4), rFootEnd(5)
                     [0,6],[6,7],[7,8],[8,9], [9,10], #LeftLeg: root, lHip(6), lKnee(7), lAnkle(8), lFootMid(9), lFootEnd(10)
                     [11,12], [12,13], [13,14], [14,15], #root2(11), spineMid(12), neck(13), nose(14), head(15) #0,11 are the same points?
                     [16,17], [17,18], [18,19], [20,21], [20,22],   #Left Arms. neck(16==13), lshoulder(17),  lElbow(18), lWrist (19=20), lThumb(21), lMiddleFinger(22)
                     [24,25], [25,26], [26,27], [27,29], [27,30]   #Right Arm, neck(24==13), rshoulder(25),  rElbow(26), rWrist (27=28), rThumb(29), rMiddleFinger(30)
                     ]
        bLeft = [0 ,0, 0, 0, 0,
                    1, 1, 1, 1, 1,
                    1, 1, 1, 1,
                    1, 1, 1, 1, 1,
                    0, 0, 0, 0, 0] #To draw left as different color. Torso is treated as left


    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    for k in np.arange( len(link2D) ):
        parent = link2D[k][0]
        child = link2D[k][1]
        if color is not None:
            c = color
        else:
            if bLeft[k]:
                c = (0,0,255)#BGR, RED
            else:
                c = (0,0,0)

        cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    return image




#Panoptic Studio SMC19 ordering

def Vis_Skeleton_2D_SMC19(pt2d, image = None, color=None):
    pt2d = np.reshape(pt2d,[-1,2])          #Just in case. Make sure (32, 2)


    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4

    assert pt2d.shape[0]==19
    print("Vis_Skeleton_2D_H36m: {} joints".format(16))
    #Without Nose
    link2D = [ [0,1], [0,2],  #neck -> nose, neck-> bodyCenter
                 [0,3], [3,4], [4,5],   #Left Arm

                 [2,6], [6,7], [7,8],   #left leg
                [2,12],[12,13], [13,14], #Right leg
                [0,9], [9, 10], [10, 11], #Right Arm
                [1, 15], [15, 16], #left eye
                [1, 17], [17, 18]] #right eye

    bLeft = [ 1,1,
        1, 1, 1,
        1,1,1,
        0,0,0,
        0,0,0,
        1,1,
        0,0]


    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    for k in np.arange( len(link2D) ):
        parent = link2D[k][0]
        child = link2D[k][1]
        if color is not None:
            c = color
        else:
            if bLeft[k]:
                c = (0,0,255)#BGR, RED
            else:
                c = (0,0,0)

        cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    return image




#Panoptic Studio SMC19 ordering
 
def Vis_Skeleton_2D_SMC19(pt2d, image = None, color=None):
    pt2d = np.reshape(pt2d,[-1,2])          #Just in case. Make sure (32, 2)


    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255        

    radius = 4

    assert pt2d.shape[0]==19
    print("Vis_Skeleton_2D_H36m: {} joints".format(16))
    #Without Nose
    link2D = [ [0,1], [0,2],  #neck -> nose, neck-> bodyCenter
                 [0,3], [3,4], [4,5],   #Left Arm
                
                 [2,6], [6,7], [7,8],   #left leg
                [2,12],[12,13], [13,14], #Right leg
                [0,9], [9, 10], [10, 11], #Right Arm
                [1, 15], [15, 16], #left eye
                [1, 17], [17, 18]] #right eye
    
    bLeft = [ 1,1,
        1, 1, 1,
        1,1,1,
        0,0,0,
        0,0,0,
        1,1,
        0,0]

            
    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    for k in np.arange( len(link2D) ):
        parent = link2D[k][0]
        child = link2D[k][1]
        if color is not None:
            c = color
        else:
            if bLeft[k]:
                c = (0,0,255)#BGR, RED
            else:
                c = (0,0,0)

        cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    return image


def Vis_Skeleton_2D_Hand(pt2d, image = None, color=None):
    pt2d = np.reshape(pt2d,[-1,2])          #Just in case. Make sure (32, 2)

    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255        

    radius = 4

    # assert pt2d.shape[0]==19
    print("Vis_Skeleton_2D_H36m: {} joints".format(16))
    #Without Nose
    link2D = [ [0,1], [1,2], [2,3], [3,4],  #thumb
                [0,5], [5,6],[6,7],[7,8],   #index
                [0,9],[9,10],[10,11],[11,12],
                [0,13],[13,14],[14,15],[15,16],
                [0,17],[17,18],[18,19],[19,20]
                ]
            
    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    for k in np.arange( len(link2D) ):
        parent = link2D[k][0]
        child = link2D[k][1]
        if color is not None:
            c = color
        else:
            if True:#bLeft[k]:
                c = (0,255,255)#BGR, RED
            # else:
            #     c = (0,0,0)

        cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    return image





#smplCOCO19 or MTC20
#smplCOCO: (19,2)
#MTC: (20,2)
#pt2d: (36,)
#pt2d_visibility: (18,)
def Vis_Skeleton_2D_smplCOCO(pt2d, pt2d_visibility=None, image = None, blankImSize = 1000, bVis = False, color=None, bBoxWidth=None):
    pt2d = np.reshape(pt2d,[-1,2])          #Just in case. Make sure (32, 2)


    if pt2d_visibility is not None and len(pt2d_visibility) == len(pt2d)*2:
        pt2d_visibility = pt2d_visibility[::2]
    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((blankImSize,blankImSize,3),np.uint8) *255

    radius = 4

    if(pt2d.shape[0]==19 or pt2d.shape[0]==20):
        # print("Vis_Skeleton_2D_smplCOCO: {} joints".format(16))
        #Without Nose
        link2D = [ [12,2], [2,1], [1,0], #Right leg
                     [12,3], [3,4], [4,5], #Left leg
                     [12,9], [9,10], [10,11], #Left Arm
                     [12,8], [8,7], [7,6], #Right shoulder
                      [12,14],[14,16],[16,18],  #Neck(12)->Nose(14)->rightEye(16)->rightEar(18)
                      [14,15],[15,17],   #Nose(14)->leftEye(15)->leftEar(17).
                      [14,13] #Nose->headTop(13)
                     ]
        bLeft = [ 0,0,0,
            1, 1, 1,
            0,0,0,
            1,1,1,
            1,1,1,
            11,0,11,0]
    elif(pt2d.shape[0]==18): #No head (13)
        # print("Vis_Skeleton_2D_smplCOCO: {} joints".format(16))
        #Without Nose
        link2D = [ [12,2], [2,1], [1,0], #Right leg
                     [12,3], [3,4], [4,5], #Left leg
                     [12,9], [9,10], [10,11], #Left Arm
                     [12,8], [8,7], [7,6], #Right shoulder
                      [12,13],[13,15],[15,17],  #Neck(12)->Nose(14)->rightEye(16)->rightEar(18)
                      [13,14],[14,16]   #Nose(14)->leftEye(15)->leftEar(17).
                    #   [14,13] #Nose->headTop(13)
                     ]
        bLeft = [ 0,0,0,
            1, 1, 1,
            1,1,1,
            0,0,0,
            1,0,0,
            1,1]
    elif(pt2d.shape[0]==26): #SMPLCOCO totalCpa26
        #Without Nose
        link2D = [ [12,2], [2,1], [1,0], #Right leg
                    [12,3], [3,4], [4,5], #Left leg
                    [12,9], [9,10], [10,11], #Left Arm
                    [12,8], [8,7], [7,6], #Right shoulder
                    [12,14],[14,16],[16,18],  #Neck(12)->Nose(14)->rightEye(16)->rightEar(18)
                    [14,15],[15,17],   #Nose(14)->leftEye(15)->leftEar(17).
                    # [14,13], #Nose->headMidle(13)
                    [12,19],       #headTop19
                    [5,20], [5,21], [5,22],       #leftFoot
                    [0,23], [0,24], [0,25]       #rightFoot
                    ]
        bLeft = [ 0,0,0,
            1, 1, 1,
            1,1,1,
            0,0,0,
            1,0,0,
            1,1,
            1,
            1,1,1,
            0,0,0]

    else:
        assert False

    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        if pt2d_visibility is None:
            cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)
        else:
            if pt2d_visibility[k]:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    for k in np.arange( len(link2D) ):
        parent = link2D[k][0]
        child = link2D[k][1]
        if color is not None:
            c = color
        else:
            if bLeft[k]:
                c = (0,0,255)#BGR, RED
            else:
                c = (0,0,0)
        if pt2d_visibility is None:
            cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)
        else:
            if pt2d_visibility[parent] and pt2d_visibility[child]:
                cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)


    if bBoxWidth is not None:
        image = Vis_Bbox_minmaxPt(image, [0,0], [bBoxWidth,bBoxWidth])

    if bVis:
        ImShow(image,name='Vis_Skeleton_2D_smplCOCO')
    return image


def Vis_Skeleton_2D_smpl24(pt2d, image = None, bVis = False, color=None):
    pt2d = np.reshape(pt2d,(-1,2))          #Just in case. Make sure (32, 2)

    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4


    #SMPL 24 joints used for LBS
    link2D = [ [0,3],[3,6],[6,9],[9,12],[12,15],  #root-> torso -> head
                    [9,13],[13,16],[16,18],[18,20],[20,22], #Nect-> left hand
                    [9,14], [14,17], [17,19], [19,21], [21,23],  #Nect-> right hand
                    [0,1], [1,4], [4,7], [7,10], # left Leg
                    [0,2], [2,5], [5,8], [8,11] #right leg
                    ]

    bLeft = [ 0,0,0,
        1, 1, 1,
        0,0,0,
        1,1,1,
        1,1,1,
        11,0,11,0]

    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    for k in np.arange( len(link2D) ):
        parent = link2D[k][0]
        child = link2D[k][1]
        if color is not None:
            c = color
        else:
            if True:#bLeft[k]:
                c = (0,0,255)#BGR, RED
            else:
                c = (0,0,0)

        cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    if bVis:
        ImShow(image)
    return image


def Vis_Skeleton_2D_smpl45(pt2d, image = None, bVis = False, color=None):
    pt2d = np.reshape(pt2d,(-1,2))          #Just in case. Make sure (32, 2)

    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4


    #SMPL 24 joints used for LBS
    link2D = [ [0,3],[3,6],[6,9],[9,12],[12,15],  #root-> torso -> head
                    [9,13],[13,16],[16,18],[18,20],[20,22], #Nect-> left hand
                    [9,14], [14,17], [17,19], [19,21], [21,23],  #Nect-> right hand
                    [0,1], [1,4], [4,7], [7,10], # left Leg
                    [0,2], [2,5], [5,8], [8,11] #right leg
                    ]

    bLeft = [ 0,0,0,
        1, 1, 1,
        0,0,0,
        1,1,1,
        1,1,1,
        11,0,11,0]

    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    for k in np.arange( len(link2D) ):
        parent = link2D[k][0]
        child = link2D[k][1]
        if color is not None:
            c = color
        else:
            if True:#bLeft[k]:
                c = (0,0,255)#BGR, RED
            else:
                c = (0,0,0)

        cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    if bVis:
        ImShow(image)
    return image




def Vis_Skeleton_2D_MPII(pt2d, pt2d_visibility = None, image = None, bVis = False, color=None):
    pt2d = np.reshape(pt2d,(-1,2))          #Just in case. Make sure (32, 2)

    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4
     #SMPL 24 joints used for LBS
    link2D = [ [6,7],[7,8],[8,9],  #root-> torso -> head
                    [7,12], [12,11],[11,10], #right arm
                    [7,13], [13,14], [14,15],  #left arm
                    [6,2],[2,1], [1,0], #right leg
                    [6,3], [3,4], [4,5] #left leg
                    ]

    bLeft = [ 1,1,1,
        0, 0, 0,
        1,1,1,
        0,0,0,
        1,1,1]
    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        if pt2d_visibility is None:
            cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)
        else:
            if pt2d_visibility[k]:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    for k in np.arange( len(link2D) ):
        parent = link2D[k][0]
        child = link2D[k][1]
        if color is not None:
            c = color
        else:
            if bLeft[k]:
                c = (0,0,255)#BGR, RED
            else:
                c = (0,0,0)
        if pt2d_visibility is None:
            cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)
        else:
            if pt2d_visibility[parent] and pt2d_visibility[child]:
                cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    if bVis:
        ImShow(image)
    return image



def Vis_Skeleton_2D_foot(pt2d, pt2d_visibility = None, image = None, bVis = False, color=None):
    pt2d = np.reshape(pt2d,(-1,2))          #Just in case. Make sure (32, 2)

    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4
     #SMPL 24 joints used for LBS
    link2D = [ [0,1],[1,2],  #root-> torso -> head
                [3,4], [4,5] ]

    bLeft = [ 1,1,
        0, 0]
    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        if pt2d_visibility is None:
            cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)
        else:
            if pt2d_visibility[k]:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    for k in np.arange( len(link2D) ):
        parent = link2D[k][0]
        child = link2D[k][1]
        if color is not None:
            c = color
        else:
            if bLeft[k]:
                c = (0,0,255)#BGR, RED
            else:
                c = (0,0,0)
        if pt2d_visibility is None:
            cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)
        else:
            if pt2d_visibility[parent] and pt2d_visibility[child]:
                cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    if bVis:
        ImShow(image)
    return image





def Vis_Skeleton_2D_Openpose25(pt2d, pt2d_visibility = None, image = None, bVis = False, color=None):
    pt2d = np.reshape(pt2d,(-1,2))          #Just in case. Make sure (32, 2)

    if pt2d.shape[0]==49:        #SPIN 25 (openpose) + 24 (superset) joint
        return Vis_Skeleton_2D_SPIN49(pt2d, pt2d_visibility, image, bVis, color)



    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4

    #Openpose25
    link_openpose = [  [8,1], [1,0] , [0,16] , [16,18] , [0,15], [15,17],
            [1,2],[2,3],[3,4],      #Right Arm
            [1,5], [5,6], [6,7],       #Left Arm
            [8,12], [12,13], [13,14], [14,21], [14,19], [14,20],
            [8,9], [9,10], [10,11], [11,24], [11,22], [11,23]
            ]

    bLeft = [ 1,1,1,1,0,0,
        0,0,0,
        1,1,1,
        1,1,1,1,1,1,
        0,0,0,0,0,0]



    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        if pt2d_visibility is None:
            cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)
        else:
            if pt2d_visibility[k]:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    #Openpose joint drawn as blue
    for k in np.arange( len(link_openpose) ):
        parent = link_openpose[k][0]
        child = link_openpose[k][1]
        if color is not None:
            c = color
        else:
            if bLeft[k]:
                c = (255,0,0)#BGR, Blue
            else:
                c = (0,0,0) #Right Black
        if pt2d_visibility is None:
            cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)
        else:
            if pt2d_visibility[parent] and pt2d_visibility[child]:
                cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    return image


def Vis_Skeleton_2D_Openpose_hand(pt2d, pt2d_visibility = None, image = None, bVis = False, color=None):
    pt2d = np.reshape(pt2d,(-1,2))          #Just in case. Make sure (32, 2)

    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4

    #Openpose25
    link_openpose =  [ [0,1], [1,2], [2,3], [3,4],  #thumb
                    [0,5], [5,6],[6,7],[7,8],   #index
                    [0,9],[9,10],[10,11],[11,12],
                    [0,13],[13,14],[14,15],[15,16],
                    [0,17],[17,18],[18,19],[19,20]
                    ]
    link_openpose = np.array(link_openpose)


    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        if pt2d_visibility is None:
            cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)
        else:
            if pt2d_visibility[k]:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    #Openpose joint drawn as blue
    for k in np.arange( len(link_openpose) ):
        parent = link_openpose[k][0]
        child = link_openpose[k][1]
        if color is not None:
            c = color
        else:
            c = (255,0,0)#BGR, Blue
        if pt2d_visibility is None:
            cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)
        else:
            if pt2d_visibility[parent] and pt2d_visibility[child]:
                cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    return image



def Vis_Skeleton_2D_Openpose18(pt2d, pt2d_visibility = None, image = None, bVis = False, color=None):
    pt2d = np.reshape(pt2d,(-1,2))          #Just in case. Make sure (32, 2)

    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4

    #Openpose18
    link_openpose = [    [1,0] , [0,14] , [14,16] , [0,15], [15,17],
                [1,2],[2,3],[3,4],      #Right Arm
                [1,5], [5,6], [6,7],       #Left Arm
                [1,11], [11,12], [12,13],       #Left Leg
                [8,1], [8,9], [9,10]       #Right Leg
                ]

    bLeft = [ 1,1,1,1,1,
        0,0,0,
        1,1,1,
        1,1,1,
        0,0,0]



    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        if pt2d_visibility is None:
            cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)
        else:
            if pt2d_visibility[k]:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    #Openpose joint drawn as blue
    for k in np.arange( len(link_openpose) ):
        parent = link_openpose[k][0]
        child = link_openpose[k][1]
        if color is not None:
            c = color
        else:
            if bLeft[k]:
                c = (255,0,0)#BGR, Blue
            else:
                c = (0,0,0)
        if pt2d_visibility is None:
            cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)
        else:
            if pt2d_visibility[parent] and pt2d_visibility[child]:
                cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    return image



def Vis_Skeleton_2D_SPIN24(pt2d, pt2d_visibility = None, image = None, bVis = False, color=None):
    pt2d = np.reshape(pt2d,(-1,2))          #Just in case. Make sure (32, 2)

    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4



    #Openpose25 in Spin Defition + SPIN global 24
    # 'OP Nose', 'OP Neck', 'OP RShoulder',           #0,1,2
    # 'OP RElbow', 'OP RWrist', 'OP LShoulder',       #3,4,5
    # 'OP LElbow', 'OP LWrist', 'OP MidHip',          #6, 7,8
    # 'OP RHip', 'OP RKnee', 'OP RAnkle',             #9,10,11
    # 'OP LHip', 'OP LKnee', 'OP LAnkle',             #12,13,14
    # 'OP REye', 'OP LEye', 'OP REar',                #15,16,17
    # 'OP LEar', 'OP LBigToe', 'OP LSmallToe',        #18,19,20
    # 'OP LHeel', 'OP RBigToe', 'OP RSmallToe', 'OP RHeel',  #21, 22, 23, 24  ##Total 25 joints  for openpose
    link_openpose = [  [8,1], [1,0] , [0,16] , [16,18] , [0,15], [15,17],
                [1,2],[2,3],[3,4],      #Right Arm
                [1,5], [5,6], [6,7],       #Left Arm
                [8,12], [12,13], [13,14], [14,19], [19,20], [20,21],    #Left Leg
                [8,9], [9,10], [10,11], [11,22], [22,23], [23,24]       #Right left
                ]

    link_spin24 =[  [14,16], [16,12], [12,17] , [17,18] ,
                [12,9],[9,10],[10,11],      #Right Arm
                [12,8], [8,7], [7,6],       #Left Arm
                [14,3], [3,4], [4,5],
                [14,2], [2,1], [1,0]]


    link_spin24 = np.array(link_spin24) + 25

    # bLeft = [ 1,1,1,1,0,0,
    #     0,0,0,
    #     1,1,1,
    #     1,1,1,1,1,1,
    #     0,0,0,0,0,0]
    bLeft = [ 0,0,0,0,
        1,1,1,
        0,0,0,
        1,1,1,
        0,0,0]



    # for i in np.arange( len(link) ):
    for k in np.arange( 25,len(pt2d) ):
        if color is not None:
            if pt2d_visibility is None:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, color,-1)
            else:
                if pt2d_visibility[k]:
                    cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, color,-1)
        else:
            if pt2d_visibility is None:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,0,255),-1)
            else:
                if pt2d_visibility[k]:
                    cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,0,255),-1)

    # # #Openpose joint drawn as blue
    # for k in np.arange( len(link_openpose) ):
    #     parent = link_openpose[k][0]
    #     child = link_openpose[k][1]
    #     if color is not None:
    #         c = color
    #     else:
    #         if True:#bLeft[k]:
    #             c = (255,0,0)#BGR, Blue
    #         else:
    #             c = (0,0,0)
    #     if pt2d_visibility is None:
    #         cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)
    #     else:
    #         if pt2d_visibility[parent] and pt2d_visibility[child]:
    #             cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)


    #SPIN24 joint drawn as red
    for k in np.arange( len(link_spin24) ):
        parent = link_spin24[k][0]
        child = link_spin24[k][1]
        if color is not None:
            c = color
        else:
            if True:#bLeft[k]:
                c = (0,0,255)#BGR, RED
            else:
                c = (0,0,0)
        if pt2d_visibility is None:
            cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)
        else:
            if pt2d_visibility[parent] and pt2d_visibility[child]:
                cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)



    if bVis:
        ImShow(image)
    return image




def Vis_Skeleton_2D_SPIN49(pt2d, pt2d_visibility = None, image = None, bVis = False, color=None):
    pt2d = np.reshape(pt2d,(-1,2))          #Just in case. Make sure (32, 2)

    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4



    #Openpose25 in Spin Defition + SPIN global 24
    # 'OP Nose', 'OP Neck', 'OP RShoulder',           #0,1,2
    # 'OP RElbow', 'OP RWrist', 'OP LShoulder',       #3,4,5
    # 'OP LElbow', 'OP LWrist', 'OP MidHip',          #6, 7,8
    # 'OP RHip', 'OP RKnee', 'OP RAnkle',             #9,10,11
    # 'OP LHip', 'OP LKnee', 'OP LAnkle',             #12,13,14
    # 'OP REye', 'OP LEye', 'OP REar',                #15,16,17
    # 'OP LEar', 'OP LBigToe', 'OP LSmallToe',        #18,19,20
    # 'OP LHeel', 'OP RBigToe', 'OP RSmallToe', 'OP RHeel',  #21, 22, 23, 24  ##Total 25 joints  for openpose
    link_openpose = [  [8,1], [1,0] , [0,16] , [16,18] , [0,15], [15,17],
                [1,2],[2,3],[3,4],      #Right Arm
                [1,5], [5,6], [6,7],       #Left Arm
                [8,12], [12,13], [13,14], [14,19], [19,20], [20,21],    #Left Leg
                [8,9], [9,10], [10,11], [11,22], [22,23], [23,24]       #Right left
                ]

    link_spin24 =[  [14,16], [16,12], [12,17] , [17,18] ,
                [12,9],[9,10],[10,11],      #Right Arm
                [12,8], [8,7], [7,6],       #Left Arm
                [14,3], [3,4], [4,5],
                [14,2], [2,1], [1,0]]


    link_spin24 = np.array(link_spin24) + 25

    # bLeft = [ 1,1,1,1,0,0,
    #     0,0,0,
    #     1,1,1,
    #     1,1,1,1,1,1,
    #     0,0,0,0,0,0]
    bLeft = [ 0,0,0,0,
        1,1,1,
        0,0,0,
        1,1,1,
        0,0,0]



    # for i in np.arange( len(link) ):
    for k in np.arange( 25,len(pt2d) ):
        if color is not None:
            if pt2d_visibility is None:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, color,-1)
            else:
                if pt2d_visibility[k]:
                    cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, color,-1)
        else:
            if pt2d_visibility is None:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,0,255),-1)
            else:
                if pt2d_visibility[k]:
                    cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,0,255),-1)

    # #Openpose joint drawn as blue
    for k in np.arange( len(link_openpose) ):
        parent = link_openpose[k][0]
        child = link_openpose[k][1]
        if color is not None:
            c = color
        else:
            if True:#bLeft[k]:
                c = (255,0,0)#BGR, Blue
            else:
                c = (0,0,0)
        if pt2d_visibility is None:
            cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)
        else:
            if pt2d_visibility[parent] and pt2d_visibility[child]:
                cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)


    #SPIN24 joint drawn as red
    for k in np.arange( len(link_spin24) ):
        parent = link_spin24[k][0]
        child = link_spin24[k][1]
        if color is not None:
            c = color
        else:
            if True:#bLeft[k]:
                c = (0,0,255)#BGR, RED
            else:
                c = (0,0,0)
        if pt2d_visibility is None:
            cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)
        else:
            if pt2d_visibility[parent] and pt2d_visibility[child]:
                cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)



    if bVis:
        ImShow(image)
    return image




def Vis_Skeleton_2D_coco(pt2d, pt2d_visibility = None, image = None,  bVis = False, color=None , offsetXY =None):
    pt2d = np.reshape(pt2d,(-1,2))          #Just in case. Make sure (32, 2)

    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4

    # 'OP RHip', 'OP RKnee', 'OP RAnkle',             #9,10,11
    # 'OP LHip', 'OP LKnee', 'OP LAnkle',             #12,13,14
    # 'OP REye', 'OP LEye', 'OP REar',                #15,16,17
    # 'OP LEar', 'OP LBigToe', 'OP LSmallToe',        #18,19,20
    # 'OP LHeel', 'OP RBigToe', 'OP RSmallToe', 'OP RHeel',  #21, 22, 23, 24  ##Total 25 joints  for openpose
    link_coco = [  [0,1], [1,3] , [0,2] , [2,4],
                [6,8],[8,10],      #Right Arm
                [5,7], [7,9],       #Left Arm
                [15,13], [13,11], [11,5],    #Left Leg
                [16,14], [14,12], [12,6], #Right left
                ]

    for k in np.arange( len(pt2d) ):
        if color is not None:
            if pt2d_visibility is None:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, color,-1)
            else:
                if pt2d_visibility[k]:
                    cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, color,-1)
        else:
            if pt2d_visibility is None:
                cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,0,255),-1)
            else:
                if pt2d_visibility[k]:
                    cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,0,255),-1)

    # # #Openpose joint drawn as blue
    for k in np.arange( len(link_coco) ):
        parent = link_coco[k][0]
        child = link_coco[k][1]
        if color is not None:
            c = color
        else:
            if True:#bLeft[k]:
                c = (255,0,0)#BGR, Blue
            else:
                c = (0,0,0)
        if pt2d_visibility is None:
            cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)
        else:
            if pt2d_visibility[parent] and pt2d_visibility[child]:
                cv2.line(image, (int(pt2d[parent][0]), int(pt2d[parent][1])), (int(pt2d[child][0]), int(pt2d[child][1])), c, radius - 2)

    if bVis:
        ImShow(image)
    return image




    

def Vis_Skeleton_2D_general(pt2d, pt2d_visibility = None, image = None,  bVis = False, color=None , offsetXY =None):
    pt2d = np.reshape(pt2d,(-1,2))          #Just in case. Make sure (32, 2)

    if offsetXY is not None:
        pt2d = pt2d + np.array(offsetXY)

    if pt2d.shape[0]==49:        #SPIN 25 (openpose) + 24 (superset) joint
        return Vis_Skeleton_2D_SPIN49(pt2d, pt2d_visibility, image, bVis, color)



    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4

    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    if bVis:
        ImShow(image)
    return image


def Vis_Skeleton_3Dto2D_general(pt2d, pt2d_visibility = None, image = None,  bVis = False, color=None, offsetXY =None):
    pt2d = np.reshape(pt2d,(-1,3))          #Just in case. Make sure (32, 2)


    if pt2d.shape[0]==49:        #SPIN 25 (openpose) + 24 (superset) joint
        return Vis_Skeleton_2D_SPIN49(pt2d, pt2d_visibility, image, bVis, color)

    #Draw via opencv
    if not isinstance(image, np.ndarray):#not image: #If no image is given, generate Blank image
        image = np.ones((1000,1000,3),np.uint8) *255

    radius = 4

    # for i in np.arange( len(link) ):
    for k in np.arange( len(pt2d) ):
        cv2.circle(image, (int(pt2d[k][0]), int(pt2d[k][1]) ), radius, (0,255,0),-1)

    if bVis:
        ImShow(image)
    return image

#H36m skeleton32
# def Vis_Skeleton_H36m16(inputImg, coco_annot):
imports/frankmocap/renderer/viewer2D.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.

"""
Visualizing 3D humans via Opengl
- Options:
    GUI mode: a screen is required
    Scnreenless mode: xvfb-run can be used to avoid screen requirement

"""
import numpy as np
import cv2
import torch
from renderer import viewer2D#, glViewer, glRenderer
from renderer import glViewer
from renderer import meshRenderer #glRenderer

from mocap_utils.coordconv import convert_smpl_to_bbox, convert_bbox_to_oriIm

from renderer.image_utils import draw_raw_bbox, draw_hand_bbox, draw_body_bbox, draw_arm_pose

class Visualizer(object):
    """
   Visualizer to visuzlie SMPL reconstruction output from HMR family (HMR, SPIN, EFT)

    Args:
        reconstruction output
        rawImg, bbox, 
        smpl_params (shape, pose, cams )
    """

    def __init__(
        self,
        rendererType ='opengl_gui'          #nongui or gui
    ):
        self.rendererType = rendererType
        if rendererType != "opengl_gui" and rendererType!= "opengl":
            print("Wrong rendererType: {rendererType}")
            assert False

        self.cam_all = []
        self.vert_all = []
        self.bboxXYXY_all = []

        self.bg_image = None

        #Screenless rendering
        if rendererType =='opengl':
            self.renderer = meshRenderer.meshRenderer()
            self.renderer.setRenderMode('geo')
            self.renderer.offscreenMode(True)
        else:
            from renderer import glViewer #glRenderer
            self.renderer = None

        #Output rendering
        self.renderout = None

    def visualize(self,
        input_img, 
        hand_bbox_list = None, 
        body_bbox_list = None,
        body_pose_list = None,
        raw_hand_bboxes = None,
        pred_mesh_list = None,
        vis_raw_hand_bbox = True,
        vis_body_pose = True,
        vis_hand_bbox = True,
        concatenate=True
    ):
         # init
        res_img = input_img.copy()

        # draw raw hand bboxes
        if raw_hand_bboxes is not None and vis_raw_hand_bbox:
            res_img = draw_raw_bbox(input_img, raw_hand_bboxes)
            # res_img = np.concatenate((res_img, raw_bbox_img), axis=1)

        # draw 2D Pose
        if body_pose_list is not None and vis_body_pose:
            res_img = draw_arm_pose(res_img, body_pose_list)

        # draw body bbox
        if body_bbox_list is not None:
            body_bbox_img = draw_body_bbox(input_img, body_bbox_list)
            res_img = body_bbox_img

        # draw hand bbox
        if hand_bbox_list is not None:
            res_img = draw_hand_bbox(res_img, hand_bbox_list)

        # render predicted meshes
        if pred_mesh_list is not None:
            rend_img = self.__render_pred_verts(input_img, pred_mesh_list)
            if rend_img is not None:
                if concatenate:
                    res_img = np.concatenate((res_img, rend_img), axis=1) # Concatenate (H, W, C) & (H, W, C) images to (H, 2W, C) image
                else:
                    res_img = draw_hand_bbox(rend_img, hand_bbox_list)
            # res_img = rend_img
        
        return res_img
        
    def __render_pred_verts(self, img_original, pred_mesh_list):

        res_img = img_original.copy()

        pred_mesh_list_offset =[]
        for mesh in pred_mesh_list:

            # Mesh vertices have in image cooridnate (left, top origin)
            # Move the X-Y origin in image center
            mesh_offset = mesh['vertices'].copy()
            mesh_offset[:,0] -= img_original.shape[1]*0.5
            mesh_offset[:,1] -= img_original.shape[0]*0.5
            pred_mesh_list_offset.append( {'ver': mesh_offset, 'f':mesh['faces'] })# verts = mesh['vertices']
            # faces = mesh['faces']
        if self.rendererType =="opengl_gui":
            self._visualize_gui_naive(pred_mesh_list_offset, img_original=res_img)
            overlaidImg = None
        else:
            self._visualize_screenless_naive(pred_mesh_list_offset, img_original=res_img, maxHeight = 100000)
            overlaidImg = self.renderout['render_camview']

        return overlaidImg


    def _visualize_screenless_naive(self, meshList, skelList=None, body_bbox_list=None, img_original=None, show_side = False, vis=False, maxHeight = 1080):
        """
            args:
                meshList: list of {'ver': pred_vertices, 'f': smpl.faces}
                skelList: list of [JointNum*3, 1]       (where 1 means num. of frames in glviewer)
                bbr_list: list of [x,y,w,h] 
            output:
                #Rendered images are saved in 
                self.renderout['render_camview']
                self.renderout['render_sideview']

            #Note: The size of opengl rendering is restricted by the current screen size. Set the maxHeight accordingly

        """
        assert self.renderer is not None

        if len(meshList)==0:
               # sideImg = cv2.resize(sideImg, (renderImg.shape[1], renderImg.shape[0]) )
            self.renderout  ={}
            self.renderout['render_camview'] = img_original.copy()

            blank = np.ones(img_original.shape, dtype=np.uint8)*255       #generate blank image
            self.renderout['render_sideview'] = blank
            
            return
        
        if body_bbox_list is not None:
            for bbr in body_bbox_list:
                viewer2D.Vis_Bbox(img_original,bbr)
        # viewer2D.ImShow(img_original)

        #Check image height
        imgHeight, imgWidth = img_original.shape[0], img_original.shape[1]
        
        
        if maxHeight <imgHeight:        #Resize
            ratio = maxHeight/imgHeight
        
            #Resize Img
            newWidth = int(imgWidth*ratio)
            newHeight = int(imgHeight*ratio)
            img_original_resized = cv2.resize(img_original, (newWidth,newHeight))

            #Resize skeleton
            for m in meshList:
                m['ver'] *=ratio

            if skelList is not None:
                for s in skelList:
                    s *=ratio
        else:
            img_original_resized = img_original

        self.renderer.setWindowSize(img_original_resized.shape[1], img_original_resized.shape[0])
        self.renderer.setBackgroundTexture(img_original_resized)
        self.renderer.setViewportSize(img_original_resized.shape[1], img_original_resized.shape[0])

        # self.renderer.add_mesh(meshList[0]['ver'],meshList[0]['f'])
        self.renderer.clear_mesh()
        for mesh in meshList:
            self.renderer.add_mesh(mesh['ver'],mesh['f'])
        self.renderer.showBackground(True)
        self.renderer.setWorldCenterBySceneCenter()
        self.renderer.setCameraViewMode("cam")
        # self.renderer.setViewportSize(img_original_resized.shape[1], img_original_resized.shape[0])
                
        self.renderer.display()
        renderImg = self.renderer.get_screen_color_ibgr()

        if vis:
            viewer2D.ImShow(renderImg,waitTime=1,name="rendered")

        ###Render Side View
        if show_side:
            self.renderer.setCameraViewMode("free")     
            self.renderer.setViewAngle(90,20)
            self.renderer.showBackground(False)
            self.renderer.setViewportSize(img_original_resized.shape[1], img_original_resized.shape[0])
            self.renderer.display()
            sideImg = self.renderer.get_screen_color_ibgr()        #Overwite on rawImg

            if vis:
                viewer2D.ImShow(sideImg,waitTime=0,name="sideview")
        
        # sideImg = cv2.resize(sideImg, (renderImg.shape[1], renderImg.shape[0]) )
        self.renderout  ={}
        self.renderout['render_camview'] = renderImg

        if show_side:
            self.renderout['render_sideview'] = sideImg


    def _visualize_gui_naive(self, meshList, skelList=None, body_bbox_list=None, img_original=None, normal_compute=True):
        """
            args:
                meshList: list of {'ver': pred_vertices, 'f': smpl.faces}
                skelList: list of [JointNum*3, 1]       (where 1 means num. of frames in glviewer)
                bbr_list: list of [x,y,w,h] 
        """
        if body_bbox_list is not None:
            for bbr in body_bbox_list:
                viewer2D.Vis_Bbox(img_original, bbr)
        # viewer2D.ImShow(img_original)

        glViewer.setWindowSize(img_original.shape[1], img_original.shape[0])
        # glViewer.setRenderOutputSize(inputImg.shape[1],inputImg.shape[0])
        glViewer.setBackgroundTexture(img_original)
        glViewer.SetOrthoCamera(True)
        glViewer.setMeshData(meshList, bComputeNormal= normal_compute)        # meshes = {'ver': pred_vertices, 'f': smplWrapper.f}

        if skelList is not None:
            glViewer.setSkeleton(skelList)

        if True:   #Save to File
            if True:        #Cam view rendering
                # glViewer.setSaveFolderName(overlaidImageFolder)
                glViewer.setNearPlane(50)
                glViewer.setWindowSize(img_original.shape[1], img_original.shape[0])
                # glViewer.show_SMPL(bSaveToFile = True, bResetSaveImgCnt = False, countImg = False, mode = 'camera')
                glViewer.show(1)

            if False:    #Side view rendering
                # glViewer.setSaveFolderName(sideImageFolder)
                glViewer.setNearPlane(50)
                glViewer.setWindowSize(img_original.shape[1], img_original.shape[0])
                glViewer.show_SMPL(bSaveToFile = True, bResetSaveImgCnt = False, countImg = True, zoom=1108, mode = 'youtube')
                # glViewer.setSaveFolderName(g_renderDir)
                # glViewer.show(0)
    

    def _visualize_gui_smplpose_basic(self, smpl, poseParamList, shapeParamList =None,  colorList = None, isRotMat = False, scalingFactor=300, waittime =1):
        '''
            Visualize SMPL vertices from SMPL pose parameters
            This can be used as a quick visualize function if you have a pose parameters

            args: 
                poseParamList: list of pose parameters (numpy array) in angle axis (72,) by default  or rot matrix (24,3,3) with isRotMat==True
                shapeParamList: (optional) list of shape parameters (numpy array) (10,). If not provided, use a zero vector
                colorList: (optional) list of color RGB values e.g., (255,0,0) for red
        '''
        zero_betas = torch.from_numpy(np.zeros( (1,10), dtype=np.float32))
        default_color = glViewer.g_colorSet['eft']
        meshList =[]
        for i, poseParam in enumerate(poseParamList):
            
            if shapeParamList is not None:
                shapeParam = torch.from_numpy(shapeParamList[i][np.newaxis,:])
            else:
                shapeParam = zero_betas#.copy()

            if colorList is not None:
                color = colorList[i]
            else:
                color = default_color

            poseParam_tensor = torch.from_numpy( poseParam[np.newaxis,:]).float()
            if isRotMat:        #rot mat
                pred_output = smpl(betas=shapeParam, body_pose=poseParam_tensor[:,1:], global_orient=poseParam_tensor[:,[0]], pose2rot=False)

            else:  #angle axis
                pred_output = smpl(betas=shapeParam, body_pose=poseParam_tensor[:,3:], global_orient=poseParam_tensor[:,:3], pose2rot=True)
        
            nn_vertices = pred_output.vertices.detach()[0].numpy() * scalingFactor
            tempMesh = {'ver': nn_vertices, 'f':  smpl.faces, 'color':color}
            meshList.append(tempMesh)
            
        
        # glViewer.setRenderOutputSize(inputImg.shape[1],inputImg.shape[0])
        # glViewer.setBackgroundTexture(img_original)
        glViewer.g_bShowFloor = True
        # glViewer.SetOrthoCamera(True)
        glViewer.g_viewMode = 'free'
        glViewer.setMeshData(meshList, bComputeNormal= True)        # meshes = {'ver': pred_vertices, 'f': smplWrapper.f}
        # glViewer.PuttingObjectCenter()
        # glViewer.setSkeleton(skelList)
        glViewer.show(waittime)        #Press q to escape
imports/frankmocap/renderer/visualizer.py
jellyheadandrew/CHORUS
from setuptools import setup, find_packages

setup(
    name='frankmocap4chorus',
    version='0.0.1',
    description='',
    packages=find_packages(),
    install_requires=[],
)
imports/frankmocap/setup.py
jellyheadandrew/CHORUS
import torch
import torch.nn.functional as F
import math


class NoiseScheduleVP:
    def __init__(
            self,
            schedule='discrete',
            betas=None,
            alphas_cumprod=None,
            continuous_beta_0=0.1,
            continuous_beta_1=20.,
        ):
        """Create a wrapper class for the forward SDE (VP type).

        ***
        Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
                We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
        ***

        The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
        We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
        Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:

            log_alpha_t = self.marginal_log_mean_coeff(t)
            sigma_t = self.marginal_std(t)
            lambda_t = self.marginal_lambda(t)

        Moreover, as lambda(t) is an invertible function, we also support its inverse function:

            t = self.inverse_lambda(lambda_t)

        ===============================================================

        We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).

        1. For discrete-time DPMs:

            For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
                t_i = (i + 1) / N
            e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
            We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.

            Args:
                betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
                alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)

            Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.

            **Important**:  Please pay special attention for the args for `alphas_cumprod`:
                The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
                    q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
                Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
                    alpha_{t_n} = \sqrt{\hat{alpha_n}},
                and
                    log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).


        2. For continuous-time DPMs:

            We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
            schedule are the default settings in DDPM and improved-DDPM:

            Args:
                beta_min: A `float` number. The smallest beta for the linear schedule.
                beta_max: A `float` number. The largest beta for the linear schedule.
                cosine_s: A `float` number. The hyperparameter in the cosine schedule.
                cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
                T: A `float` number. The ending time of the forward process.

        ===============================================================

        Args:
            schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
                    'linear' or 'cosine' for continuous-time DPMs.
        Returns:
            A wrapper object of the forward SDE (VP type).
        
        ===============================================================

        Example:

        # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
        >>> ns = NoiseScheduleVP('discrete', betas=betas)

        # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
        >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)

        # For continuous-time DPMs (VPSDE), linear schedule:
        >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)

        """

        if schedule not in ['discrete', 'linear', 'cosine']:
            raise ValueError("Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(schedule))

        self.schedule = schedule
        if schedule == 'discrete':
            if betas is not None:
                log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
            else:
                assert alphas_cumprod is not None
                log_alphas = 0.5 * torch.log(alphas_cumprod)
            self.total_N = len(log_alphas)
            self.T = 1.
            self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))
            self.log_alpha_array = log_alphas.reshape((1, -1,))
        else:
            self.total_N = 1000
            self.beta_0 = continuous_beta_0
            self.beta_1 = continuous_beta_1
            self.cosine_s = 0.008
            self.cosine_beta_max = 999.
            self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
            self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
            self.schedule = schedule
            if schedule == 'cosine':
                # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
                # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
                self.T = 0.9946
            else:
                self.T = 1.

    def marginal_log_mean_coeff(self, t):
        """
        Compute log(alpha_t) of a given continuous-time label t in [0, T].
        """
        if self.schedule == 'discrete':
            return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device), self.log_alpha_array.to(t.device)).reshape((-1))
        elif self.schedule == 'linear':
            return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
        elif self.schedule == 'cosine':
            log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
            log_alpha_t =  log_alpha_fn(t) - self.cosine_log_alpha_0
            return log_alpha_t

    def marginal_alpha(self, t):
        """
        Compute alpha_t of a given continuous-time label t in [0, T].
        """
        return torch.exp(self.marginal_log_mean_coeff(t))

    def marginal_std(self, t):
        """
        Compute sigma_t of a given continuous-time label t in [0, T].
        """
        return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))

    def marginal_lambda(self, t):
        """
        Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
        """
        log_mean_coeff = self.marginal_log_mean_coeff(t)
        log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
        return log_mean_coeff - log_std

    def inverse_lambda(self, lamb):
        """
        Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
        """
        if self.schedule == 'linear':
            tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
            Delta = self.beta_0**2 + tmp
            return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
        elif self.schedule == 'discrete':
            log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
            t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]), torch.flip(self.t_array.to(lamb.device), [1]))
            return t.reshape((-1,))
        else:
            log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
            t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
            t = t_fn(log_alpha)
            return t


def model_wrapper(
    model,
    noise_schedule,
    model_type="noise",
    model_kwargs={},
    guidance_type="uncond",
    condition=None,
    unconditional_condition=None,
    guidance_scale=1.,
    classifier_fn=None,
    classifier_kwargs={},
):
    """Create a wrapper function for the noise prediction model.

    DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
    firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.

    We support four types of the diffusion model by setting `model_type`:

        1. "noise": noise prediction model. (Trained by predicting noise).

        2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).

        3. "v": velocity prediction model. (Trained by predicting the velocity).
            The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].

            [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
                arXiv preprint arXiv:2202.00512 (2022).
            [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
                arXiv preprint arXiv:2210.02303 (2022).
    
        4. "score": marginal score function. (Trained by denoising score matching).
            Note that the score function and the noise prediction model follows a simple relationship:
            ```
                noise(x_t, t) = -sigma_t * score(x_t, t)
            ```

    We support three types of guided sampling by DPMs by setting `guidance_type`:
        1. "uncond": unconditional sampling by DPMs.
            The input `model` has the following format:
            ``
                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
            ``

        2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
            The input `model` has the following format:
            ``
                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
            `` 

            The input `classifier_fn` has the following format:
            ``
                classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
            ``

            [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
                in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.

        3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
            The input `model` has the following format:
            ``
                model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
            `` 
            And if cond == `unconditional_condition`, the model output is the unconditional DPM output.

            [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
                arXiv preprint arXiv:2207.12598 (2022).
        

    The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
    or continuous-time labels (i.e. epsilon to T).

    We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
    ``
        def model_fn(x, t_continuous) -> noise:
            t_input = get_model_input_time(t_continuous)
            return noise_pred(model, x, t_input, **model_kwargs)         
    ``
    where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.

    ===============================================================

    Args:
        model: A diffusion model with the corresponding format described above.
        noise_schedule: A noise schedule object, such as NoiseScheduleVP.
        model_type: A `str`. The parameterization type of the diffusion model.
                    "noise" or "x_start" or "v" or "score".
        model_kwargs: A `dict`. A dict for the other inputs of the model function.
        guidance_type: A `str`. The type of the guidance for sampling.
                    "uncond" or "classifier" or "classifier-free".
        condition: A pytorch tensor. The condition for the guided sampling.
                    Only used for "classifier" or "classifier-free" guidance type.
        unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
                    Only used for "classifier-free" guidance type.
        guidance_scale: A `float`. The scale for the guided sampling.
        classifier_fn: A classifier function. Only used for the classifier guidance.
        classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
    Returns:
        A noise prediction model that accepts the noised data and the continuous time as the inputs.
    """

    def get_model_input_time(t_continuous):
        """
        Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
        For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
        For continuous-time DPMs, we just use `t_continuous`.
        """
        if noise_schedule.schedule == 'discrete':
            return (t_continuous - 1. / noise_schedule.total_N) * 1000.
        else:
            return t_continuous

    def noise_pred_fn(x, t_continuous, cond=None):
        if t_continuous.reshape((-1,)).shape[0] == 1:
            t_continuous = t_continuous.expand((x.shape[0]))
        t_input = get_model_input_time(t_continuous)
        if cond is None:
            output = model(x, t_input, **model_kwargs)
        else:
            output = model(x, t_input, cond, **model_kwargs)
        if model_type == "noise":
            return output
        elif model_type == "x_start":
            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
            dims = x.dim()
            return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
        elif model_type == "v":
            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
            dims = x.dim()
            return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
        elif model_type == "score":
            sigma_t = noise_schedule.marginal_std(t_continuous)
            dims = x.dim()
            return -expand_dims(sigma_t, dims) * output

    def cond_grad_fn(x, t_input):
        """
        Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
        """
        with torch.enable_grad():
            x_in = x.detach().requires_grad_(True)
            log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
            return torch.autograd.grad(log_prob.sum(), x_in)[0]

    def model_fn(x, t_continuous):
        """
        The noise predicition model function that is used for DPM-Solver.
        """
        if t_continuous.reshape((-1,)).shape[0] == 1:
            t_continuous = t_continuous.expand((x.shape[0]))
        if guidance_type == "uncond":
            return noise_pred_fn(x, t_continuous)
        elif guidance_type == "classifier":
            assert classifier_fn is not None
            t_input = get_model_input_time(t_continuous)
            cond_grad = cond_grad_fn(x, t_input)
            sigma_t = noise_schedule.marginal_std(t_continuous)
            noise = noise_pred_fn(x, t_continuous)
            return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
        elif guidance_type == "classifier-free":
            if guidance_scale == 1. or unconditional_condition is None:
                return noise_pred_fn(x, t_continuous, cond=condition)
            else:
                x_in = torch.cat([x] * 2)
                t_in = torch.cat([t_continuous] * 2)
                c_in = torch.cat([unconditional_condition, condition])
                noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
                return noise_uncond + guidance_scale * (noise - noise_uncond)

    assert model_type in ["noise", "x_start", "v"]
    assert guidance_type in ["uncond", "classifier", "classifier-free"]
    return model_fn


class DPM_Solver:
    def __init__(self, model_fn, noise_schedule, predict_x0=False, thresholding=False, max_val=1.):
        """Construct a DPM-Solver. 

        We support both the noise prediction model ("predicting epsilon") and the data prediction model ("predicting x0").
        If `predict_x0` is False, we use the solver for the noise prediction model (DPM-Solver).
        If `predict_x0` is True, we use the solver for the data prediction model (DPM-Solver++).
            In such case, we further support the "dynamic thresholding" in [1] when `thresholding` is True.
            The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales.

        Args:
            model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]):
                ``
                def model_fn(x, t_continuous):
                    return noise
                ``
            noise_schedule: A noise schedule object, such as NoiseScheduleVP.
            predict_x0: A `bool`. If true, use the data prediction model; else, use the noise prediction model.
            thresholding: A `bool`. Valid when `predict_x0` is True. Whether to use the "dynamic thresholding" in [1].
            max_val: A `float`. Valid when both `predict_x0` and `thresholding` are True. The max value for thresholding.
        
        [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b.
        """
        self.model = model_fn
        self.noise_schedule = noise_schedule
        self.predict_x0 = predict_x0
        self.thresholding = thresholding
        self.max_val = max_val

    def noise_prediction_fn(self, x, t):
        """
        Return the noise prediction model.
        """
        return self.model(x, t)

    def data_prediction_fn(self, x, t):
        """
        Return the data prediction model (with thresholding).
        """
        noise = self.noise_prediction_fn(x, t)
        dims = x.dim()
        alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
        x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
        if self.thresholding:
            p = 0.995   # A hyperparameter in the paper of "Imagen" [1].
            s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
            s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
            x0 = torch.clamp(x0, -s, s) / s
        return x0

    def model_fn(self, x, t):
        """
        Convert the model to the noise prediction model or the data prediction model. 
        """
        if self.predict_x0:
            return self.data_prediction_fn(x, t)
        else:
            return self.noise_prediction_fn(x, t)

    def get_time_steps(self, skip_type, t_T, t_0, N, device):
        """Compute the intermediate time steps for sampling.

        Args:
            skip_type: A `str`. The type for the spacing of the time steps. We support three types:
                - 'logSNR': uniform logSNR for the time steps.
                - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
                - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
            t_T: A `float`. The starting time of the sampling (default is T).
            t_0: A `float`. The ending time of the sampling (default is epsilon).
            N: A `int`. The total number of the spacing of the time steps.
            device: A torch device.
        Returns:
            A pytorch tensor of the time steps, with the shape (N + 1,).
        """
        if skip_type == 'logSNR':
            lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
            lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
            logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
            return self.noise_schedule.inverse_lambda(logSNR_steps)
        elif skip_type == 'time_uniform':
            return torch.linspace(t_T, t_0, N + 1).to(device)
        elif skip_type == 'time_quadratic':
            t_order = 2
            t = torch.linspace(t_T**(1. / t_order), t_0**(1. / t_order), N + 1).pow(t_order).to(device)
            return t
        else:
            raise ValueError("Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))

    def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
        """
        Get the order of each step for sampling by the singlestep DPM-Solver.

        We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast".
        Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is:
            - If order == 1:
                We take `steps` of DPM-Solver-1 (i.e. DDIM).
            - If order == 2:
                - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling.
                - If steps % 2 == 0, we use K steps of DPM-Solver-2.
                - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1.
            - If order == 3:
                - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
                - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1.
                - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1.
                - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2.

        ============================================
        Args:
            order: A `int`. The max order for the solver (2 or 3).
            steps: A `int`. The total number of function evaluations (NFE).
            skip_type: A `str`. The type for the spacing of the time steps. We support three types:
                - 'logSNR': uniform logSNR for the time steps.
                - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
                - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
            t_T: A `float`. The starting time of the sampling (default is T).
            t_0: A `float`. The ending time of the sampling (default is epsilon).
            device: A torch device.
        Returns:
            orders: A list of the solver order of each step.
        """
        if order == 3:
            K = steps // 3 + 1
            if steps % 3 == 0:
                orders = [3,] * (K - 2) + [2, 1]
            elif steps % 3 == 1:
                orders = [3,] * (K - 1) + [1]
            else:
                orders = [3,] * (K - 1) + [2]
        elif order == 2:
            if steps % 2 == 0:
                K = steps // 2
                orders = [2,] * K
            else:
                K = steps // 2 + 1
                orders = [2,] * (K - 1) + [1]
        elif order == 1:
            K = 1
            orders = [1,] * steps
        else:
            raise ValueError("'order' must be '1' or '2' or '3'.")
        if skip_type == 'logSNR':
            # To reproduce the results in DPM-Solver paper
            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
        else:
            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[torch.cumsum(torch.tensor([0,] + orders)).to(device)]
        return timesteps_outer, orders

    def denoise_to_zero_fn(self, x, s):
        """
        Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization. 
        """
        return self.data_prediction_fn(x, s)

    def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False):
        """
        DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`.

        Args:
            x: A pytorch tensor. The initial value at time `s`.
            s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
            model_s: A pytorch tensor. The model function evaluated at time `s`.
                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
            return_intermediate: A `bool`. If true, also return the model value at time `s`.
        Returns:
            x_t: A pytorch tensor. The approximated solution at time `t`.
        """
        ns = self.noise_schedule
        dims = x.dim()
        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
        h = lambda_t - lambda_s
        log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t)
        sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t)
        alpha_t = torch.exp(log_alpha_t)

        if self.predict_x0:
            phi_1 = torch.expm1(-h)
            if model_s is None:
                model_s = self.model_fn(x, s)
            x_t = (
                expand_dims(sigma_t / sigma_s, dims) * x
                - expand_dims(alpha_t * phi_1, dims) * model_s
            )
            if return_intermediate:
                return x_t, {'model_s': model_s}
            else:
                return x_t
        else:
            phi_1 = torch.expm1(h)
            if model_s is None:
                model_s = self.model_fn(x, s)
            x_t = (
                expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
                - expand_dims(sigma_t * phi_1, dims) * model_s
            )
            if return_intermediate:
                return x_t, {'model_s': model_s}
            else:
                return x_t

    def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s=None, return_intermediate=False, solver_type='dpm_solver'):
        """
        Singlestep solver DPM-Solver-2 from time `s` to time `t`.

        Args:
            x: A pytorch tensor. The initial value at time `s`.
            s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
            r1: A `float`. The hyperparameter of the second-order solver.
            model_s: A pytorch tensor. The model function evaluated at time `s`.
                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
            return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time).
            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
        Returns:
            x_t: A pytorch tensor. The approximated solution at time `t`.
        """
        if solver_type not in ['dpm_solver', 'taylor']:
            raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
        if r1 is None:
            r1 = 0.5
        ns = self.noise_schedule
        dims = x.dim()
        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
        h = lambda_t - lambda_s
        lambda_s1 = lambda_s + r1 * h
        s1 = ns.inverse_lambda(lambda_s1)
        log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(t)
        sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t)
        alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t)

        if self.predict_x0:
            phi_11 = torch.expm1(-r1 * h)
            phi_1 = torch.expm1(-h)

            if model_s is None:
                model_s = self.model_fn(x, s)
            x_s1 = (
                expand_dims(sigma_s1 / sigma_s, dims) * x
                - expand_dims(alpha_s1 * phi_11, dims) * model_s
            )
            model_s1 = self.model_fn(x_s1, s1)
            if solver_type == 'dpm_solver':
                x_t = (
                    expand_dims(sigma_t / sigma_s, dims) * x
                    - expand_dims(alpha_t * phi_1, dims) * model_s
                    - (0.5 / r1) * expand_dims(alpha_t * phi_1, dims) * (model_s1 - model_s)
                )
            elif solver_type == 'taylor':
                x_t = (
                    expand_dims(sigma_t / sigma_s, dims) * x
                    - expand_dims(alpha_t * phi_1, dims) * model_s
                    + (1. / r1) * expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * (model_s1 - model_s)
                )
        else:
            phi_11 = torch.expm1(r1 * h)
            phi_1 = torch.expm1(h)

            if model_s is None:
                model_s = self.model_fn(x, s)
            x_s1 = (
                expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
                - expand_dims(sigma_s1 * phi_11, dims) * model_s
            )
            model_s1 = self.model_fn(x_s1, s1)
            if solver_type == 'dpm_solver':
                x_t = (
                    expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
                    - expand_dims(sigma_t * phi_1, dims) * model_s
                    - (0.5 / r1) * expand_dims(sigma_t * phi_1, dims) * (model_s1 - model_s)
                )
            elif solver_type == 'taylor':
                x_t = (
                    expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
                    - expand_dims(sigma_t * phi_1, dims) * model_s
                    - (1. / r1) * expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * (model_s1 - model_s)
                )
        if return_intermediate:
            return x_t, {'model_s': model_s, 'model_s1': model_s1}
        else:
            return x_t

    def singlestep_dpm_solver_third_update(self, x, s, t, r1=1./3., r2=2./3., model_s=None, model_s1=None, return_intermediate=False, solver_type='dpm_solver'):
        """
        Singlestep solver DPM-Solver-3 from time `s` to time `t`.

        Args:
            x: A pytorch tensor. The initial value at time `s`.
            s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
            r1: A `float`. The hyperparameter of the third-order solver.
            r2: A `float`. The hyperparameter of the third-order solver.
            model_s: A pytorch tensor. The model function evaluated at time `s`.
                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
            model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`).
                If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it.
            return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
        Returns:
            x_t: A pytorch tensor. The approximated solution at time `t`.
        """
        if solver_type not in ['dpm_solver', 'taylor']:
            raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
        if r1 is None:
            r1 = 1. / 3.
        if r2 is None:
            r2 = 2. / 3.
        ns = self.noise_schedule
        dims = x.dim()
        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
        h = lambda_t - lambda_s
        lambda_s1 = lambda_s + r1 * h
        lambda_s2 = lambda_s + r2 * h
        s1 = ns.inverse_lambda(lambda_s1)
        s2 = ns.inverse_lambda(lambda_s2)
        log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t)
        sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(s2), ns.marginal_std(t)
        alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t)

        if self.predict_x0:
            phi_11 = torch.expm1(-r1 * h)
            phi_12 = torch.expm1(-r2 * h)
            phi_1 = torch.expm1(-h)
            phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1.
            phi_2 = phi_1 / h + 1.
            phi_3 = phi_2 / h - 0.5

            if model_s is None:
                model_s = self.model_fn(x, s)
            if model_s1 is None:
                x_s1 = (
                    expand_dims(sigma_s1 / sigma_s, dims) * x
                    - expand_dims(alpha_s1 * phi_11, dims) * model_s
                )
                model_s1 = self.model_fn(x_s1, s1)
            x_s2 = (
                expand_dims(sigma_s2 / sigma_s, dims) * x
                - expand_dims(alpha_s2 * phi_12, dims) * model_s
                + r2 / r1 * expand_dims(alpha_s2 * phi_22, dims) * (model_s1 - model_s)
            )
            model_s2 = self.model_fn(x_s2, s2)
            if solver_type == 'dpm_solver':
                x_t = (
                    expand_dims(sigma_t / sigma_s, dims) * x
                    - expand_dims(alpha_t * phi_1, dims) * model_s
                    + (1. / r2) * expand_dims(alpha_t * phi_2, dims) * (model_s2 - model_s)
                )
            elif solver_type == 'taylor':
                D1_0 = (1. / r1) * (model_s1 - model_s)
                D1_1 = (1. / r2) * (model_s2 - model_s)
                D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
                D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
                x_t = (
                    expand_dims(sigma_t / sigma_s, dims) * x
                    - expand_dims(alpha_t * phi_1, dims) * model_s
                    + expand_dims(alpha_t * phi_2, dims) * D1
                    - expand_dims(alpha_t * phi_3, dims) * D2
                )
        else:
            phi_11 = torch.expm1(r1 * h)
            phi_12 = torch.expm1(r2 * h)
            phi_1 = torch.expm1(h)
            phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1.
            phi_2 = phi_1 / h - 1.
            phi_3 = phi_2 / h - 0.5

            if model_s is None:
                model_s = self.model_fn(x, s)
            if model_s1 is None:
                x_s1 = (
                    expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
                    - expand_dims(sigma_s1 * phi_11, dims) * model_s
                )
                model_s1 = self.model_fn(x_s1, s1)
            x_s2 = (
                expand_dims(torch.exp(log_alpha_s2 - log_alpha_s), dims) * x
                - expand_dims(sigma_s2 * phi_12, dims) * model_s
                - r2 / r1 * expand_dims(sigma_s2 * phi_22, dims) * (model_s1 - model_s)
            )
            model_s2 = self.model_fn(x_s2, s2)
            if solver_type == 'dpm_solver':
                x_t = (
                    expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
                    - expand_dims(sigma_t * phi_1, dims) * model_s
                    - (1. / r2) * expand_dims(sigma_t * phi_2, dims) * (model_s2 - model_s)
                )
            elif solver_type == 'taylor':
                D1_0 = (1. / r1) * (model_s1 - model_s)
                D1_1 = (1. / r2) * (model_s2 - model_s)
                D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
                D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
                x_t = (
                    expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
                    - expand_dims(sigma_t * phi_1, dims) * model_s
                    - expand_dims(sigma_t * phi_2, dims) * D1
                    - expand_dims(sigma_t * phi_3, dims) * D2
                )

        if return_intermediate:
            return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2}
        else:
            return x_t

    def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpm_solver"):
        """
        Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`.

        Args:
            x: A pytorch tensor. The initial value at time `s`.
            model_prev_list: A list of pytorch tensor. The previous computed model values.
            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
        Returns:
            x_t: A pytorch tensor. The approximated solution at time `t`.
        """
        if solver_type not in ['dpm_solver', 'taylor']:
            raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
        ns = self.noise_schedule
        dims = x.dim()
        model_prev_1, model_prev_0 = model_prev_list
        t_prev_1, t_prev_0 = t_prev_list
        lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t)
        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
        alpha_t = torch.exp(log_alpha_t)

        h_0 = lambda_prev_0 - lambda_prev_1
        h = lambda_t - lambda_prev_0
        r0 = h_0 / h
        D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
        if self.predict_x0:
            if solver_type == 'dpm_solver':
                x_t = (
                    expand_dims(sigma_t / sigma_prev_0, dims) * x
                    - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
                    - 0.5 * expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * D1_0
                )
            elif solver_type == 'taylor':
                x_t = (
                    expand_dims(sigma_t / sigma_prev_0, dims) * x
                    - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
                    + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1_0
                )
        else:
            if solver_type == 'dpm_solver':
                x_t = (
                    expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
                    - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
                    - 0.5 * expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * D1_0
                )
            elif solver_type == 'taylor':
                x_t = (
                    expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
                    - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
                    - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1_0
                )
        return x_t

    def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpm_solver'):
        """
        Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`.

        Args:
            x: A pytorch tensor. The initial value at time `s`.
            model_prev_list: A list of pytorch tensor. The previous computed model values.
            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
        Returns:
            x_t: A pytorch tensor. The approximated solution at time `t`.
        """
        ns = self.noise_schedule
        dims = x.dim()
        model_prev_2, model_prev_1, model_prev_0 = model_prev_list
        t_prev_2, t_prev_1, t_prev_0 = t_prev_list
        lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda(t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t)
        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
        alpha_t = torch.exp(log_alpha_t)

        h_1 = lambda_prev_1 - lambda_prev_2
        h_0 = lambda_prev_0 - lambda_prev_1
        h = lambda_t - lambda_prev_0
        r0, r1 = h_0 / h, h_1 / h
        D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
        D1_1 = expand_dims(1. / r1, dims) * (model_prev_1 - model_prev_2)
        D1 = D1_0 + expand_dims(r0 / (r0 + r1), dims) * (D1_0 - D1_1)
        D2 = expand_dims(1. / (r0 + r1), dims) * (D1_0 - D1_1)
        if self.predict_x0:
            x_t = (
                expand_dims(sigma_t / sigma_prev_0, dims) * x
                - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
                + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1
                - expand_dims(alpha_t * ((torch.exp(-h) - 1. + h) / h**2 - 0.5), dims) * D2
            )
        else:
            x_t = (
                expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
                - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
                - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1
                - expand_dims(sigma_t * ((torch.exp(h) - 1. - h) / h**2 - 0.5), dims) * D2
            )
        return x_t

    def singlestep_dpm_solver_update(self, x, s, t, order, return_intermediate=False, solver_type='dpm_solver', r1=None, r2=None):
        """
        Singlestep DPM-Solver with the order `order` from time `s` to time `t`.

        Args:
            x: A pytorch tensor. The initial value at time `s`.
            s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
            order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
            return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
            r1: A `float`. The hyperparameter of the second-order or third-order solver.
            r2: A `float`. The hyperparameter of the third-order solver.
        Returns:
            x_t: A pytorch tensor. The approximated solution at time `t`.
        """
        if order == 1:
            return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate)
        elif order == 2:
            return self.singlestep_dpm_solver_second_update(x, s, t, return_intermediate=return_intermediate, solver_type=solver_type, r1=r1)
        elif order == 3:
            return self.singlestep_dpm_solver_third_update(x, s, t, return_intermediate=return_intermediate, solver_type=solver_type, r1=r1, r2=r2)
        else:
            raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))

    def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpm_solver'):
        """
        Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`.

        Args:
            x: A pytorch tensor. The initial value at time `s`.
            model_prev_list: A list of pytorch tensor. The previous computed model values.
            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
            order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
        Returns:
            x_t: A pytorch tensor. The approximated solution at time `t`.
        """
        if order == 1:
            return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1])
        elif order == 2:
            return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
        elif order == 3:
            return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
        else:
            raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))

    def dpm_solver_adaptive(self, x, order, t_T, t_0, h_init=0.05, atol=0.0078, rtol=0.05, theta=0.9, t_err=1e-5, solver_type='dpm_solver'):
        """
        The adaptive step size solver based on singlestep DPM-Solver.

        Args:
            x: A pytorch tensor. The initial value at time `t_T`.
            order: A `int`. The (higher) order of the solver. We only support order == 2 or 3.
            t_T: A `float`. The starting time of the sampling (default is T).
            t_0: A `float`. The ending time of the sampling (default is epsilon).
            h_init: A `float`. The initial step size (for logSNR).
            atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1].
            rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05.
            theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1].
            t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the 
                current time and `t_0` is less than `t_err`. The default setting is 1e-5.
            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
        Returns:
            x_0: A pytorch tensor. The approximated solution at time `t_0`.

        [1] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, "Gotta go fast when generating data with score-based models," arXiv preprint arXiv:2105.14080, 2021.
        """
        ns = self.noise_schedule
        s = t_T * torch.ones((x.shape[0],)).to(x)
        lambda_s = ns.marginal_lambda(s)
        lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x))
        h = h_init * torch.ones_like(s).to(x)
        x_prev = x
        nfe = 0
        if order == 2:
            r1 = 0.5
            lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True)
            higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1, solver_type=solver_type, **kwargs)
        elif order == 3:
            r1, r2 = 1. / 3., 2. / 3.
            lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1, return_intermediate=True, solver_type=solver_type)
            higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(x, s, t, r1=r1, r2=r2, solver_type=solver_type, **kwargs)
        else:
            raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order))
        while torch.abs((s - t_0)).mean() > t_err:
            t = ns.inverse_lambda(lambda_s + h)
            x_lower, lower_noise_kwargs = lower_update(x, s, t)
            x_higher = higher_update(x, s, t, **lower_noise_kwargs)
            delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev)))
            norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True))
            E = norm_fn((x_higher - x_lower) / delta).max()
            if torch.all(E <= 1.):
                x = x_higher
                s = t
                x_prev = x_lower
                lambda_s = ns.marginal_lambda(s)
            h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s)
            nfe += order
        print('adaptive solver nfe', nfe)
        return x

    def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform',
        method='singlestep', lower_order_final=True, denoise_to_zero=False, solver_type='dpm_solver',
        atol=0.0078, rtol=0.05,
    ):
        """
        Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`.

        =====================================================

        We support the following algorithms for both noise prediction model and data prediction model:
            - 'singlestep':
                Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver. 
                We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps).
                The total number of function evaluations (NFE) == `steps`.
                Given a fixed NFE == `steps`, the sampling procedure is:
                    - If `order` == 1:
                        - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM).
                    - If `order` == 2:
                        - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling.
                        - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2.
                        - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
                    - If `order` == 3:
                        - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
                        - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, and 1 step of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
                        - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1.
                        - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2.
            - 'multistep':
                Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`.
                We initialize the first `order` values by lower order multistep solvers.
                Given a fixed NFE == `steps`, the sampling procedure is:
                    Denote K = steps.
                    - If `order` == 1:
                        - We use K steps of DPM-Solver-1 (i.e. DDIM).
                    - If `order` == 2:
                        - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2.
                    - If `order` == 3:
                        - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, then (K - 2) step of multistep DPM-Solver-3.
            - 'singlestep_fixed':
                Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3).
                We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE.
            - 'adaptive':
                Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper).
                We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`.
                You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs
                (NFE) and the sample quality.
                    - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2.
                    - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3.

        =====================================================

        Some advices for choosing the algorithm:
            - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:
                Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with `order = 3`.
                e.g.
                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False)
                    >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,
                            skip_type='time_uniform', method='singlestep')
            - For **guided sampling with large guidance scale** by DPMs:
                Use multistep DPM-Solver with `predict_x0 = True` and `order = 2`.
                e.g.
                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True)
                    >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2,
                            skip_type='time_uniform', method='multistep')

        We support three types of `skip_type`:
            - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images**
            - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**.
            - 'time_quadratic': quadratic time for the time steps.

        =====================================================
        Args:
            x: A pytorch tensor. The initial value at time `t_start`
                e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution.
            steps: A `int`. The total number of function evaluations (NFE).
            t_start: A `float`. The starting time of the sampling.
                If `T` is None, we use self.noise_schedule.T (default is 1.0).
            t_end: A `float`. The ending time of the sampling.
                If `t_end` is None, we use 1. / self.noise_schedule.total_N.
                e.g. if total_N == 1000, we have `t_end` == 1e-3.
                For discrete-time DPMs:
                    - We recommend `t_end` == 1. / self.noise_schedule.total_N.
                For continuous-time DPMs:
                    - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15.
            order: A `int`. The order of DPM-Solver.
            skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'.
            method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'.
            denoise_to_zero: A `bool`. Whether to denoise to time 0 at the final step.
                Default is `False`. If `denoise_to_zero` is `True`, the total NFE is (`steps` + 1).

                This trick is firstly proposed by DDPM (https://arxiv.org/abs/2006.11239) and
                score_sde (https://arxiv.org/abs/2011.13456). Such trick can improve the FID
                for diffusion models sampling by diffusion SDEs for low-resolutional images
                (such as CIFAR-10). However, we observed that such trick does not matter for
                high-resolutional images. As it needs an additional NFE, we do not recommend
                it for high-resolutional images.
            lower_order_final: A `bool`. Whether to use lower order solvers at the final steps.
                Only valid for `method=multistep` and `steps < 15`. We empirically find that
                this trick is a key to stabilizing the sampling by DPM-Solver with very few steps
                (especially for steps <= 10). So we recommend to set it to be `True`.
            solver_type: A `str`. The taylor expansion type for the solver. `dpm_solver` or `taylor`. We recommend `dpm_solver`.
            atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
            rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
        Returns:
            x_end: A pytorch tensor. The approximated solution at time `t_end`.

        """
        t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
        t_T = self.noise_schedule.T if t_start is None else t_start
        device = x.device
        if method == 'adaptive':
            with torch.no_grad():
                x = self.dpm_solver_adaptive(x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol, solver_type=solver_type)
        elif method == 'multistep':
            assert steps >= order
            timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
            assert timesteps.shape[0] - 1 == steps
            with torch.no_grad():
                vec_t = timesteps[0].expand((x.shape[0]))
                model_prev_list = [self.model_fn(x, vec_t)]
                t_prev_list = [vec_t]
                # Init the first `order` values by lower order multistep DPM-Solver.
                for init_order in range(1, order):
                    vec_t = timesteps[init_order].expand(x.shape[0])
                    x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, init_order, solver_type=solver_type)
                    model_prev_list.append(self.model_fn(x, vec_t))
                    t_prev_list.append(vec_t)
                # Compute the remaining values by `order`-th order multistep DPM-Solver.
                for step in range(order, steps + 1):
                    vec_t = timesteps[step].expand(x.shape[0])
                    if lower_order_final and steps < 15:
                        step_order = min(order, steps + 1 - step)
                    else:
                        step_order = order
                    x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, step_order, solver_type=solver_type)
                    for i in range(order - 1):
                        t_prev_list[i] = t_prev_list[i + 1]
                        model_prev_list[i] = model_prev_list[i + 1]
                    t_prev_list[-1] = vec_t
                    # We do not need to evaluate the final model value.
                    if step < steps:
                        model_prev_list[-1] = self.model_fn(x, vec_t)
        elif method in ['singlestep', 'singlestep_fixed']:
            if method == 'singlestep':
                timesteps_outer, orders = self.get_orders_and_timesteps_for_singlestep_solver(steps=steps, order=order, skip_type=skip_type, t_T=t_T, t_0=t_0, device=device)
            elif method == 'singlestep_fixed':
                K = steps // order
                orders = [order,] * K
                timesteps_outer = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=K, device=device)
            for i, order in enumerate(orders):
                t_T_inner, t_0_inner = timesteps_outer[i], timesteps_outer[i + 1]
                timesteps_inner = self.get_time_steps(skip_type=skip_type, t_T=t_T_inner.item(), t_0=t_0_inner.item(), N=order, device=device)
                lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner)
                vec_s, vec_t = t_T_inner.tile(x.shape[0]), t_0_inner.tile(x.shape[0])
                h = lambda_inner[-1] - lambda_inner[0]
                r1 = None if order <= 1 else (lambda_inner[1] - lambda_inner[0]) / h
                r2 = None if order <= 2 else (lambda_inner[2] - lambda_inner[0]) / h
                x = self.singlestep_dpm_solver_update(x, vec_s, vec_t, order, solver_type=solver_type, r1=r1, r2=r2)
        if denoise_to_zero:
            x = self.denoise_to_zero_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
        return x



#############################################################
# other utility functions
#############################################################

def interpolate_fn(x, xp, yp):
    """
    A piecewise linear function y = f(x), using xp and yp as keypoints.
    We implement f(x) in a differentiable way (i.e. applicable for autograd).
    The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)

    Args:
        x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).
        xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
        yp: PyTorch tensor with shape [C, K].
    Returns:
        The function values f(x), with shape [N, C].
    """
    N, K = x.shape[0], xp.shape[1]
    all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
    sorted_all_x, x_indices = torch.sort(all_x, dim=2)
    x_idx = torch.argmin(x_indices, dim=2)
    cand_start_idx = x_idx - 1
    start_idx = torch.where(
        torch.eq(x_idx, 0),
        torch.tensor(1, device=x.device),
        torch.where(
            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
        ),
    )
    end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
    start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
    end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
    start_idx2 = torch.where(
        torch.eq(x_idx, 0),
        torch.tensor(0, device=x.device),
        torch.where(
            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
        ),
    )
    y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
    start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
    end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
    cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
    return cand


def expand_dims(v, dims):
    """
    Expand the tensor `v` to the dim `dims`.

    Args:
        `v`: a PyTorch tensor with shape [N].
        `dim`: a `int`.
    Returns:
        a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
    """
    return v[(...,) + (None,)*(dims - 1)]
imports/ldm/ldm/models/diffusion/dpm_solver/dpm_solver.py
jellyheadandrew/CHORUS
"""SAMPLING ONLY."""

import torch

from .dpm_solver import NoiseScheduleVP, model_wrapper, DPM_Solver


class DPMSolverSampler(object):
    def __init__(self, model, **kwargs):
        super().__init__()
        self.model = model
        to_torch = lambda x: x.clone().detach().to(torch.float32).to(model.device)
        self.register_buffer('alphas_cumprod', to_torch(model.alphas_cumprod))

    def register_buffer(self, name, attr):
        if type(attr) == torch.Tensor:
            if attr.device != torch.device("cuda"):
                attr = attr.to(torch.device("cuda"))
        setattr(self, name, attr)

    @torch.no_grad()
    def sample(self,
               S,
               batch_size,
               shape,
               conditioning=None,
               callback=None,
               normals_sequence=None,
               img_callback=None,
               quantize_x0=False,
               eta=0.,
               mask=None,
               x0=None,
               temperature=1.,
               noise_dropout=0.,
               score_corrector=None,
               corrector_kwargs=None,
               verbose=True,
               x_T=None,
               log_every_t=100,
               unconditional_guidance_scale=1.,
               unconditional_conditioning=None,
               # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
               **kwargs
               ):
        if conditioning is not None:
            if isinstance(conditioning, dict):
                cbs = conditioning[list(conditioning.keys())[0]].shape[0]
                if cbs != batch_size:
                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
            else:
                if conditioning.shape[0] != batch_size:
                    print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")

        # sampling
        C, H, W = shape
        size = (batch_size, C, H, W)

        # print(f'Data shape for DPM-Solver sampling is {size}, sampling steps {S}')

        device = self.model.betas.device
        if x_T is None:
            img = torch.randn(size, device=device)
        else:
            img = x_T

        ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod)

        model_fn = model_wrapper(
            lambda x, t, c: self.model.apply_model(x, t, c),
            ns,
            model_type="noise",
            guidance_type="classifier-free",
            condition=conditioning,
            unconditional_condition=unconditional_conditioning,
            guidance_scale=unconditional_guidance_scale,
        )

        dpm_solver = DPM_Solver(model_fn, ns, predict_x0=True, thresholding=False)
        x = dpm_solver.sample(img, steps=S, skip_type="time_uniform", method="multistep", order=2, lower_order_final=True)

        return x.to(device), None

imports/ldm/ldm/models/diffusion/dpm_solver/sampler.py
jellyheadandrew/CHORUS
"""SAMPLING ONLY."""

import torch
import numpy as np
from tqdm import tqdm
from functools import partial

from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like


class PLMSSampler(object):
    def __init__(self, model, schedule="linear", **kwargs):
        super().__init__()
        self.model = model
        self.ddpm_num_timesteps = model.num_timesteps
        self.schedule = schedule

    def register_buffer(self, name, attr):
        if type(attr) == torch.Tensor:
            if attr.device != torch.device("cuda"):
                attr = attr.to(torch.device("cuda"))
        setattr(self, name, attr)

    def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
        if ddim_eta != 0:
            raise ValueError('ddim_eta must be 0 for PLMS')
        self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
                                                  num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
        alphas_cumprod = self.model.alphas_cumprod
        assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
        to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)

        self.register_buffer('betas', to_torch(self.model.betas))
        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
        self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))

        # calculations for diffusion q(x_t | x_{t-1}) and others
        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))

        # ddim sampling parameters
        ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
                                                                                   ddim_timesteps=self.ddim_timesteps,
                                                                                   eta=ddim_eta,verbose=verbose)
        self.register_buffer('ddim_sigmas', ddim_sigmas)
        self.register_buffer('ddim_alphas', ddim_alphas)
        self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
        self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
        sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
            (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
                        1 - self.alphas_cumprod / self.alphas_cumprod_prev))
        self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)

    @torch.no_grad()
    def sample(self,
               S,
               batch_size,
               shape,
               conditioning=None,
               callback=None,
               normals_sequence=None,
               img_callback=None,
               quantize_x0=False,
               eta=0.,
               mask=None,
               x0=None,
               temperature=1.,
               noise_dropout=0.,
               score_corrector=None,
               corrector_kwargs=None,
               verbose=True,
               x_T=None,
               log_every_t=100,
               unconditional_guidance_scale=1.,
               unconditional_conditioning=None,
               # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
               **kwargs
               ):
        if conditioning is not None:
            if isinstance(conditioning, dict):
                cbs = conditioning[list(conditioning.keys())[0]].shape[0]
                if cbs != batch_size:
                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
            else:
                if conditioning.shape[0] != batch_size:
                    print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")

        self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
        # sampling
        C, H, W = shape
        size = (batch_size, C, H, W)

        if verbose:        
            print(f'Data shape for PLMS sampling is {size}')

        samples, intermediates = self.plms_sampling(conditioning, size,
                                                    callback=callback,
                                                    img_callback=img_callback,
                                                    quantize_denoised=quantize_x0,
                                                    mask=mask, x0=x0,
                                                    ddim_use_original_steps=False,
                                                    noise_dropout=noise_dropout,
                                                    temperature=temperature,
                                                    score_corrector=score_corrector,
                                                    corrector_kwargs=corrector_kwargs,
                                                    x_T=x_T,
                                                    log_every_t=log_every_t,
                                                    unconditional_guidance_scale=unconditional_guidance_scale,
                                                    unconditional_conditioning=unconditional_conditioning,
                                                    verbose=verbose
                                                    )
        return samples, intermediates

    @torch.no_grad()
    def plms_sampling(self, cond, shape,
                      x_T=None, ddim_use_original_steps=False,
                      callback=None, timesteps=None, quantize_denoised=False,
                      mask=None, x0=None, img_callback=None, log_every_t=100,
                      temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
                      unconditional_guidance_scale=1., unconditional_conditioning=None,verbose=False):
        device = self.model.betas.device
        b = shape[0]
        if x_T is None:
            img = torch.randn(shape, device=device)
        else:
            img = x_T

        if timesteps is None:
            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
        elif timesteps is not None and not ddim_use_original_steps:
            subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
            timesteps = self.ddim_timesteps[:subset_end]

        intermediates = {'x_inter': [img], 'pred_x0': [img]}
        time_range = list(reversed(range(0,timesteps))) if ddim_use_original_steps else np.flip(timesteps)
        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
        if verbose:
            print(f"Running PLMS Sampling with {total_steps} timesteps")
            iterator = tqdm(time_range, desc='PLMS Sampler', total=total_steps)
        else:
            iterator = time_range
        old_eps = []

        for i, step in enumerate(iterator):
            index = total_steps - i - 1
            ts = torch.full((b,), step, device=device, dtype=torch.long)
            ts_next = torch.full((b,), time_range[min(i + 1, len(time_range) - 1)], device=device, dtype=torch.long)

            if mask is not None:
                assert x0 is not None
                img_orig = self.model.q_sample(x0, ts)  # TODO: deterministic forward pass?
                img = img_orig * mask + (1. - mask) * img

            outs = self.p_sample_plms(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
                                      quantize_denoised=quantize_denoised, temperature=temperature,
                                      noise_dropout=noise_dropout, score_corrector=score_corrector,
                                      corrector_kwargs=corrector_kwargs,
                                      unconditional_guidance_scale=unconditional_guidance_scale,
                                      unconditional_conditioning=unconditional_conditioning,
                                      old_eps=old_eps, t_next=ts_next)
            img, pred_x0, e_t = outs
            old_eps.append(e_t)
            if len(old_eps) >= 4:
                old_eps.pop(0)
            if callback: callback(i)
            if img_callback: img_callback(pred_x0, i)

            if index % log_every_t == 0 or index == total_steps - 1:
                intermediates['x_inter'].append(img)
                intermediates['pred_x0'].append(pred_x0)

        return img, intermediates

    @torch.no_grad()
    def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
                      temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
                      unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None):
        b, *_, device = *x.shape, x.device

        def get_model_output(x, t):
            if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
                e_t = self.model.apply_model(x, t, c)
            else:
                x_in = torch.cat([x] * 2)
                t_in = torch.cat([t] * 2)
                c_in = torch.cat([unconditional_conditioning, c])
                e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
                e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)

            if score_corrector is not None:
                assert self.model.parameterization == "eps"
                e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)

            return e_t

        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
        alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
        sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas

        def get_x_prev_and_pred_x0(e_t, index):
            # select parameters corresponding to the currently considered timestep
            a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
            a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
            sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
            sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)

            # current prediction for x_0
            pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
            if quantize_denoised:
                pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
            # direction pointing to x_t
            dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
            noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
            if noise_dropout > 0.:
                noise = torch.nn.functional.dropout(noise, p=noise_dropout)
            x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
            return x_prev, pred_x0

        e_t = get_model_output(x, t)
        if len(old_eps) == 0:
            # Pseudo Improved Euler (2nd order)
            x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
            e_t_next = get_model_output(x_prev, t_next)
            e_t_prime = (e_t + e_t_next) / 2
        elif len(old_eps) == 1:
            # 2nd order Pseudo Linear Multistep (Adams-Bashforth)
            e_t_prime = (3 * e_t - old_eps[-1]) / 2
        elif len(old_eps) == 2:
            # 3nd order Pseudo Linear Multistep (Adams-Bashforth)
            e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
        elif len(old_eps) >= 3:
            # 4nd order Pseudo Linear Multistep (Adams-Bashforth)
            e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24

        x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)

        return x_prev, pred_x0, e_t

imports/ldm/ldm/models/diffusion/plms.py
jellyheadandrew/CHORUS
import torch
import torch.nn as nn
from functools import partial
import clip
from einops import rearrange, repeat
from transformers import CLIPTokenizer, CLIPTextModel
import kornia

from ldm.modules.x_transformer import Encoder, TransformerWrapper  # TODO: can we directly rely on lucidrains code and simply add this as a reuirement? --> test


class AbstractEncoder(nn.Module):
    def __init__(self):
        super().__init__()

    def encode(self, *args, **kwargs):
        raise NotImplementedError



class ClassEmbedder(nn.Module):
    def __init__(self, embed_dim, n_classes=1000, key='class'):
        super().__init__()
        self.key = key
        self.embedding = nn.Embedding(n_classes, embed_dim)

    def forward(self, batch, key=None):
        if key is None:
            key = self.key
        # this is for use in crossattn
        c = batch[key][:, None]
        c = self.embedding(c)
        return c


class TransformerEmbedder(AbstractEncoder):
    """Some transformer encoder layers"""
    def __init__(self, n_embed, n_layer, vocab_size, max_seq_len=77, device="cuda"):
        super().__init__()
        self.device = device
        self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len,
                                              attn_layers=Encoder(dim=n_embed, depth=n_layer))

    def forward(self, tokens):
        tokens = tokens.to(self.device)  # meh
        z = self.transformer(tokens, return_embeddings=True)
        return z

    def encode(self, x):
        return self(x)


class BERTTokenizer(AbstractEncoder):
    """ Uses a pretrained BERT tokenizer by huggingface. Vocab size: 30522 (?)"""
    def __init__(self, device="cuda", vq_interface=True, max_length=77):
        super().__init__()
        from transformers import BertTokenizerFast  # TODO: add to reuquirements
        self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
        self.device = device
        self.vq_interface = vq_interface
        self.max_length = max_length

    def forward(self, text):
        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
        tokens = batch_encoding["input_ids"].to(self.device)
        return tokens

    @torch.no_grad()
    def encode(self, text):
        tokens = self(text)
        if not self.vq_interface:
            return tokens
        return None, None, [None, None, tokens]

    def decode(self, text):
        return text


class BERTEmbedder(AbstractEncoder):
    """Uses the BERT tokenizr model and add some transformer encoder layers"""
    def __init__(self, n_embed, n_layer, vocab_size=30522, max_seq_len=77,
                 device="cuda",use_tokenizer=True, embedding_dropout=0.0):
        super().__init__()
        self.use_tknz_fn = use_tokenizer
        if self.use_tknz_fn:
            self.tknz_fn = BERTTokenizer(vq_interface=False, max_length=max_seq_len)
        self.device = device
        self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len,
                                              attn_layers=Encoder(dim=n_embed, depth=n_layer),
                                              emb_dropout=embedding_dropout)

    def forward(self, text):
        if self.use_tknz_fn:
            tokens = self.tknz_fn(text)#.to(self.device)
        else:
            tokens = text
        z = self.transformer(tokens, return_embeddings=True)
        return z

    def encode(self, text):
        # output of length 77
        return self(text)


class SpatialRescaler(nn.Module):
    def __init__(self,
                 n_stages=1,
                 method='bilinear',
                 multiplier=0.5,
                 in_channels=3,
                 out_channels=None,
                 bias=False):
        super().__init__()
        self.n_stages = n_stages
        assert self.n_stages >= 0
        assert method in ['nearest','linear','bilinear','trilinear','bicubic','area']
        self.multiplier = multiplier
        self.interpolator = partial(torch.nn.functional.interpolate, mode=method)
        self.remap_output = out_channels is not None
        if self.remap_output:
            print(f'Spatial Rescaler mapping from {in_channels} to {out_channels} channels after resizing.')
            self.channel_mapper = nn.Conv2d(in_channels,out_channels,1,bias=bias)

    def forward(self,x):
        for stage in range(self.n_stages):
            x = self.interpolator(x, scale_factor=self.multiplier)


        if self.remap_output:
            x = self.channel_mapper(x)
        return x

    def encode(self, x):
        return self(x)

class FrozenCLIPEmbedder(AbstractEncoder):
    """Uses the CLIP transformer encoder for text (from Hugging Face)"""
    def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77):
        super().__init__()
        self.tokenizer = CLIPTokenizer.from_pretrained(version)
        self.transformer = CLIPTextModel.from_pretrained(version)
        self.device = device
        self.max_length = max_length
        self.freeze()

    def freeze(self):
        self.transformer = self.transformer.eval()
        for param in self.parameters():
            param.requires_grad = False

    def forward(self, text):
        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
        tokens = batch_encoding["input_ids"].to(self.device)
        outputs = self.transformer(input_ids=tokens)

        z = outputs.last_hidden_state
        return z

    def encode(self, text):
        return self(text)


class FrozenCLIPTextEmbedder(nn.Module):
    """
    Uses the CLIP transformer encoder for text.
    """
    def __init__(self, version='ViT-L/14', device="cuda", max_length=77, n_repeat=1, normalize=True):
        super().__init__()
        self.model, _ = clip.load(version, jit=False, device="cpu")
        self.device = device
        self.max_length = max_length
        self.n_repeat = n_repeat
        self.normalize = normalize

    def freeze(self):
        self.model = self.model.eval()
        for param in self.parameters():
            param.requires_grad = False

    def forward(self, text):
        tokens = clip.tokenize(text).to(self.device)
        z = self.model.encode_text(tokens)
        if self.normalize:
            z = z / torch.linalg.norm(z, dim=1, keepdim=True)
        return z

    def encode(self, text):
        z = self(text)
        if z.ndim==2:
            z = z[:, None, :]
        z = repeat(z, 'b 1 d -> b k d', k=self.n_repeat)
        return z


class FrozenClipImageEmbedder(nn.Module):
    """
        Uses the CLIP image encoder.
        """
    def __init__(
            self,
            model,
            jit=False,
            device='cuda' if torch.cuda.is_available() else 'cpu',
            antialias=False,
        ):
        super().__init__()
        self.model, _ = clip.load(name=model, device=device, jit=jit)

        self.antialias = antialias

        self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
        self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)

    def preprocess(self, x):
        # normalize to [0,1]
        x = kornia.geometry.resize(x, (224, 224),
                                   interpolation='bicubic',align_corners=True,
                                   antialias=self.antialias)
        x = (x + 1.) / 2.
        # renormalize according to clip
        x = kornia.enhance.normalize(x, self.mean, self.std)
        return x

    def forward(self, x):
        # x is assumed to be in range [-1,1]
        return self.model.encode_image(self.preprocess(x))


if __name__ == "__main__":
    from ldm.util import count_params
    model = FrozenCLIPEmbedder()
    count_params(model, verbose=True)
imports/ldm/ldm/modules/encoders/modules.py
jellyheadandrew/CHORUS
import importlib

import torch
import numpy as np
from collections import abc
from einops import rearrange
from functools import partial

import multiprocessing as mp
from threading import Thread
from queue import Queue

from inspect import isfunction
from PIL import Image, ImageDraw, ImageFont


def log_txt_as_img(wh, xc, size=10):
    # wh a tuple of (width, height)
    # xc a list of captions to plot
    b = len(xc)
    txts = list()
    for bi in range(b):
        txt = Image.new("RGB", wh, color="white")
        draw = ImageDraw.Draw(txt)
        font = ImageFont.truetype('data/DejaVuSans.ttf', size=size)
        nc = int(40 * (wh[0] / 256))
        lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))

        try:
            draw.text((0, 0), lines, fill="black", font=font)
        except UnicodeEncodeError:
            print("Cant encode string for logging. Skipping.")

        txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
        txts.append(txt)
    txts = np.stack(txts)
    txts = torch.tensor(txts)
    return txts


def ismap(x):
    if not isinstance(x, torch.Tensor):
        return False
    return (len(x.shape) == 4) and (x.shape[1] > 3)


def isimage(x):
    if not isinstance(x, torch.Tensor):
        return False
    return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)


def exists(x):
    return x is not None


def default(val, d):
    if exists(val):
        return val
    return d() if isfunction(d) else d


def mean_flat(tensor):
    """
    https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86
    Take the mean over all non-batch dimensions.
    """
    return tensor.mean(dim=list(range(1, len(tensor.shape))))


def count_params(model, verbose=False):
    total_params = sum(p.numel() for p in model.parameters())
    if verbose:
        print(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")
    return total_params


def instantiate_from_config(config):
    if not "target" in config:
        if config == '__is_first_stage__':
            return None
        elif config == "__is_unconditional__":
            return None
        raise KeyError("Expected key `target` to instantiate.")
    return get_obj_from_str(config["target"])(**config.get("params", dict()))


def get_obj_from_str(string, reload=False):
    module, cls = string.rsplit(".", 1)
    if reload:
        module_imp = importlib.import_module(module)
        importlib.reload(module_imp)
    return getattr(importlib.import_module(module, package=None), cls)


def _do_parallel_data_prefetch(func, Q, data, idx, idx_to_fn=False):
    # create dummy dataset instance

    # run prefetching
    if idx_to_fn:
        res = func(data, worker_id=idx)
    else:
        res = func(data)
    Q.put([idx, res])
    Q.put("Done")


def parallel_data_prefetch(
        func: callable, data, n_proc, target_data_type="ndarray", cpu_intensive=True, use_worker_id=False
):
    # if target_data_type not in ["ndarray", "list"]:
    #     raise ValueError(
    #         "Data, which is passed to parallel_data_prefetch has to be either of type list or ndarray."
    #     )
    if isinstance(data, np.ndarray) and target_data_type == "list":
        raise ValueError("list expected but function got ndarray.")
    elif isinstance(data, abc.Iterable):
        if isinstance(data, dict):
            print(
                f'WARNING:"data" argument passed to parallel_data_prefetch is a dict: Using only its values and disregarding keys.'
            )
            data = list(data.values())
        if target_data_type == "ndarray":
            data = np.asarray(data)
        else:
            data = list(data)
    else:
        raise TypeError(
            f"The data, that shall be processed parallel has to be either an np.ndarray or an Iterable, but is actually {type(data)}."
        )

    if cpu_intensive:
        Q = mp.Queue(1000)
        proc = mp.Process
    else:
        Q = Queue(1000)
        proc = Thread
    # spawn processes
    if target_data_type == "ndarray":
        arguments = [
            [func, Q, part, i, use_worker_id]
            for i, part in enumerate(np.array_split(data, n_proc))
        ]
    else:
        step = (
            int(len(data) / n_proc + 1)
            if len(data) % n_proc != 0
            else int(len(data) / n_proc)
        )
        arguments = [
            [func, Q, part, i, use_worker_id]
            for i, part in enumerate(
                [data[i: i + step] for i in range(0, len(data), step)]
            )
        ]
    processes = []
    for i in range(n_proc):
        p = proc(target=_do_parallel_data_prefetch, args=arguments[i])
        processes += [p]

    # start processes
    print(f"Start prefetching...")
    import time

    start = time.time()
    gather_res = [[] for _ in range(n_proc)]
    try:
        for p in processes:
            p.start()

        k = 0
        while k < n_proc:
            # get result
            res = Q.get()
            if res == "Done":
                k += 1
            else:
                gather_res[res[0]] = res[1]

    except Exception as e:
        print("Exception: ", e)
        for p in processes:
            p.terminate()

        raise e
    finally:
        for p in processes:
            p.join()
        print(f"Prefetching complete. [{time.time() - start} sec.]")

    if target_data_type == 'ndarray':
        if not isinstance(gather_res[0], np.ndarray):
            return np.concatenate([np.asarray(r) for r in gather_res], axis=0)

        # order outputs
        return np.concatenate(gather_res, axis=0)
    elif target_data_type == 'list':
        out = []
        for r in gather_res:
            out.extend(r)
        return out
    else:
        return gather_res

def load_model_from_config(config, ckpt, verbose=False):
    print(f"Loading model from {ckpt}")
    pl_sd = torch.load(ckpt, map_location="cpu")
    if "global_step" in pl_sd:
        print(f"Global Step: {pl_sd['global_step']}")
    sd = pl_sd["state_dict"]
    model = instantiate_from_config(config.model)
    m, u = model.load_state_dict(sd, strict=False)
    if len(m) > 0 and verbose:
        print("missing keys:")
        print(m)
    if len(u) > 0 and verbose:
        print("unexpected keys:")
        print(u)

    model.cuda()
    model.eval()
    return model
imports/ldm/ldm/util.py
jellyheadandrew/CHORUS
import argparse, os, sys, datetime, glob, importlib, csv
import numpy as np
import time
import torch
import torchvision
import pytorch_lightning as pl

from packaging import version
from omegaconf import OmegaConf
from torch.utils.data import random_split, DataLoader, Dataset, Subset
from functools import partial
from PIL import Image

from pytorch_lightning import seed_everything
from pytorch_lightning.trainer import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint, Callback, LearningRateMonitor
from pytorch_lightning.utilities.distributed import rank_zero_only
from pytorch_lightning.utilities import rank_zero_info

from ldm.data.base import Txt2ImgIterableBaseDataset
from ldm.util import instantiate_from_config


def get_parser(**parser_kwargs):
    def str2bool(v):
        if isinstance(v, bool):
            return v
        if v.lower() in ("yes", "true", "t", "y", "1"):
            return True
        elif v.lower() in ("no", "false", "f", "n", "0"):
            return False
        else:
            raise argparse.ArgumentTypeError("Boolean value expected.")

    parser = argparse.ArgumentParser(**parser_kwargs)
    parser.add_argument(
        "-n",
        "--name",
        type=str,
        const=True,
        default="",
        nargs="?",
        help="postfix for logdir",
    )
    parser.add_argument(
        "-r",
        "--resume",
        type=str,
        const=True,
        default="",
        nargs="?",
        help="resume from logdir or checkpoint in logdir",
    )
    parser.add_argument(
        "-b",
        "--base",
        nargs="*",
        metavar="base_config.yaml",
        help="paths to base configs. Loaded from left-to-right. "
             "Parameters can be overwritten or added with command-line options of the form `--key value`.",
        default=list(),
    )
    parser.add_argument(
        "-t",
        "--train",
        type=str2bool,
        const=True,
        default=False,
        nargs="?",
        help="train",
    )
    parser.add_argument(
        "--no-test",
        type=str2bool,
        const=True,
        default=False,
        nargs="?",
        help="disable test",
    )
    parser.add_argument(
        "-p",
        "--project",
        help="name of new or path to existing project"
    )
    parser.add_argument(
        "-d",
        "--debug",
        type=str2bool,
        nargs="?",
        const=True,
        default=False,
        help="enable post-mortem debugging",
    )
    parser.add_argument(
        "-s",
        "--seed",
        type=int,
        default=23,
        help="seed for seed_everything",
    )
    parser.add_argument(
        "-f",
        "--postfix",
        type=str,
        default="",
        help="post-postfix for default name",
    )
    parser.add_argument(
        "-l",
        "--logdir",
        type=str,
        default="logs",
        help="directory for logging",
    )
    parser.add_argument(
        "--scale_lr",
        type=str2bool,
        nargs="?",
        const=True,
        default=True,
        help="scale base-lr by ngpu * batch_size * n_accumulate",
    )
    return parser


def nondefault_trainer_args(opt):
    parser = argparse.ArgumentParser()
    parser = Trainer.add_argparse_args(parser)
    args = parser.parse_args([])
    return sorted(k for k in vars(args) if getattr(opt, k) != getattr(args, k))


class WrappedDataset(Dataset):
    """Wraps an arbitrary object with __len__ and __getitem__ into a pytorch dataset"""

    def __init__(self, dataset):
        self.data = dataset

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


def worker_init_fn(_):
    worker_info = torch.utils.data.get_worker_info()

    dataset = worker_info.dataset
    worker_id = worker_info.id

    if isinstance(dataset, Txt2ImgIterableBaseDataset):
        split_size = dataset.num_records // worker_info.num_workers
        # reset num_records to the true number to retain reliable length information
        dataset.sample_ids = dataset.valid_ids[worker_id * split_size:(worker_id + 1) * split_size]
        current_id = np.random.choice(len(np.random.get_state()[1]), 1)
        return np.random.seed(np.random.get_state()[1][current_id] + worker_id)
    else:
        return np.random.seed(np.random.get_state()[1][0] + worker_id)


class DataModuleFromConfig(pl.LightningDataModule):
    def __init__(self, batch_size, train=None, validation=None, test=None, predict=None,
                 wrap=False, num_workers=None, shuffle_test_loader=False, use_worker_init_fn=False,
                 shuffle_val_dataloader=False):
        super().__init__()
        self.batch_size = batch_size
        self.dataset_configs = dict()
        self.num_workers = num_workers if num_workers is not None else batch_size * 2
        self.use_worker_init_fn = use_worker_init_fn
        if train is not None:
            self.dataset_configs["train"] = train
            self.train_dataloader = self._train_dataloader
        if validation is not None:
            self.dataset_configs["validation"] = validation
            self.val_dataloader = partial(self._val_dataloader, shuffle=shuffle_val_dataloader)
        if test is not None:
            self.dataset_configs["test"] = test
            self.test_dataloader = partial(self._test_dataloader, shuffle=shuffle_test_loader)
        if predict is not None:
            self.dataset_configs["predict"] = predict
            self.predict_dataloader = self._predict_dataloader
        self.wrap = wrap

    def prepare_data(self):
        for data_cfg in self.dataset_configs.values():
            instantiate_from_config(data_cfg)

    def setup(self, stage=None):
        self.datasets = dict(
            (k, instantiate_from_config(self.dataset_configs[k]))
            for k in self.dataset_configs)
        if self.wrap:
            for k in self.datasets:
                self.datasets[k] = WrappedDataset(self.datasets[k])

    def _train_dataloader(self):
        is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset)
        if is_iterable_dataset or self.use_worker_init_fn:
            init_fn = worker_init_fn
        else:
            init_fn = None
        return DataLoader(self.datasets["train"], batch_size=self.batch_size,
                          num_workers=self.num_workers, shuffle=False if is_iterable_dataset else True,
                          worker_init_fn=init_fn)

    def _val_dataloader(self, shuffle=False):
        if isinstance(self.datasets['validation'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn:
            init_fn = worker_init_fn
        else:
            init_fn = None
        return DataLoader(self.datasets["validation"],
                          batch_size=self.batch_size,
                          num_workers=self.num_workers,
                          worker_init_fn=init_fn,
                          shuffle=shuffle)

    def _test_dataloader(self, shuffle=False):
        is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset)
        if is_iterable_dataset or self.use_worker_init_fn:
            init_fn = worker_init_fn
        else:
            init_fn = None

        # do not shuffle dataloader for iterable dataset
        shuffle = shuffle and (not is_iterable_dataset)

        return DataLoader(self.datasets["test"], batch_size=self.batch_size,
                          num_workers=self.num_workers, worker_init_fn=init_fn, shuffle=shuffle)

    def _predict_dataloader(self, shuffle=False):
        if isinstance(self.datasets['predict'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn:
            init_fn = worker_init_fn
        else:
            init_fn = None
        return DataLoader(self.datasets["predict"], batch_size=self.batch_size,
                          num_workers=self.num_workers, worker_init_fn=init_fn)


class SetupCallback(Callback):
    def __init__(self, resume, now, logdir, ckptdir, cfgdir, config, lightning_config):
        super().__init__()
        self.resume = resume
        self.now = now
        self.logdir = logdir
        self.ckptdir = ckptdir
        self.cfgdir = cfgdir
        self.config = config
        self.lightning_config = lightning_config

    def on_keyboard_interrupt(self, trainer, pl_module):
        if trainer.global_rank == 0:
            print("Summoning checkpoint.")
            ckpt_path = os.path.join(self.ckptdir, "last.ckpt")
            trainer.save_checkpoint(ckpt_path)

    def on_pretrain_routine_start(self, trainer, pl_module):
        if trainer.global_rank == 0:
            # Create logdirs and save configs
            os.makedirs(self.logdir, exist_ok=True)
            os.makedirs(self.ckptdir, exist_ok=True)
            os.makedirs(self.cfgdir, exist_ok=True)

            if "callbacks" in self.lightning_config:
                if 'metrics_over_trainsteps_checkpoint' in self.lightning_config['callbacks']:
                    os.makedirs(os.path.join(self.ckptdir, 'trainstep_checkpoints'), exist_ok=True)
            print("Project config")
            print(OmegaConf.to_yaml(self.config))
            OmegaConf.save(self.config,
                           os.path.join(self.cfgdir, "{}-project.yaml".format(self.now)))

            print("Lightning config")
            print(OmegaConf.to_yaml(self.lightning_config))
            OmegaConf.save(OmegaConf.create({"lightning": self.lightning_config}),
                           os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now)))

        else:
            # ModelCheckpoint callback created log directory --- remove it
            if not self.resume and os.path.exists(self.logdir):
                dst, name = os.path.split(self.logdir)
                dst = os.path.join(dst, "child_runs", name)
                os.makedirs(os.path.split(dst)[0], exist_ok=True)
                try:
                    os.rename(self.logdir, dst)
                except FileNotFoundError:
                    pass


class ImageLogger(Callback):
    def __init__(self, batch_frequency, max_images, clamp=True, increase_log_steps=True,
                 rescale=True, disabled=False, log_on_batch_idx=False, log_first_step=False,
                 log_images_kwargs=None):
        super().__init__()
        self.rescale = rescale
        self.batch_freq = batch_frequency
        self.max_images = max_images
        self.logger_log_images = {
            pl.loggers.TestTubeLogger: self._testtube,
        }
        self.log_steps = [2 ** n for n in range(int(np.log2(self.batch_freq)) + 1)]
        if not increase_log_steps:
            self.log_steps = [self.batch_freq]
        self.clamp = clamp
        self.disabled = disabled
        self.log_on_batch_idx = log_on_batch_idx
        self.log_images_kwargs = log_images_kwargs if log_images_kwargs else {}
        self.log_first_step = log_first_step

    @rank_zero_only
    def _testtube(self, pl_module, images, batch_idx, split):
        for k in images:
            grid = torchvision.utils.make_grid(images[k])
            grid = (grid + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w

            tag = f"{split}/{k}"
            pl_module.logger.experiment.add_image(
                tag, grid,
                global_step=pl_module.global_step)

    @rank_zero_only
    def log_local(self, save_dir, split, images,
                  global_step, current_epoch, batch_idx):
        root = os.path.join(save_dir, "images", split)
        for k in images:
            grid = torchvision.utils.make_grid(images[k], nrow=4)
            if self.rescale:
                grid = (grid + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w
            grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1)
            grid = grid.numpy()
            grid = (grid * 255).astype(np.uint8)
            filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(
                k,
                global_step,
                current_epoch,
                batch_idx)
            path = os.path.join(root, filename)
            os.makedirs(os.path.split(path)[0], exist_ok=True)
            Image.fromarray(grid).save(path)

    def log_img(self, pl_module, batch, batch_idx, split="train"):
        check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step
        if (self.check_frequency(check_idx) and  # batch_idx % self.batch_freq == 0
                hasattr(pl_module, "log_images") and
                callable(pl_module.log_images) and
                self.max_images > 0):
            logger = type(pl_module.logger)

            is_train = pl_module.training
            if is_train:
                pl_module.eval()

            with torch.no_grad():
                images = pl_module.log_images(batch, split=split, **self.log_images_kwargs)

            for k in images:
                N = min(images[k].shape[0], self.max_images)
                images[k] = images[k][:N]
                if isinstance(images[k], torch.Tensor):
                    images[k] = images[k].detach().cpu()
                    if self.clamp:
                        images[k] = torch.clamp(images[k], -1., 1.)

            self.log_local(pl_module.logger.save_dir, split, images,
                           pl_module.global_step, pl_module.current_epoch, batch_idx)

            logger_log_images = self.logger_log_images.get(logger, lambda *args, **kwargs: None)
            logger_log_images(pl_module, images, pl_module.global_step, split)

            if is_train:
                pl_module.train()

    def check_frequency(self, check_idx):
        if ((check_idx % self.batch_freq) == 0 or (check_idx in self.log_steps)) and (
                check_idx > 0 or self.log_first_step):
            try:
                self.log_steps.pop(0)
            except IndexError as e:
                print(e)
                pass
            return True
        return False

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        if not self.disabled and (pl_module.global_step > 0 or self.log_first_step):
            self.log_img(pl_module, batch, batch_idx, split="train")

    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        if not self.disabled and pl_module.global_step > 0:
            self.log_img(pl_module, batch, batch_idx, split="val")
        if hasattr(pl_module, 'calibrate_grad_norm'):
            if (pl_module.calibrate_grad_norm and batch_idx % 25 == 0) and batch_idx > 0:
                self.log_gradients(trainer, pl_module, batch_idx=batch_idx)


class CUDACallback(Callback):
    # see https://github.com/SeanNaren/minGPT/blob/master/mingpt/callback.py
    def on_train_epoch_start(self, trainer, pl_module):
        # Reset the memory use counter
        torch.cuda.reset_peak_memory_stats(trainer.root_gpu)
        torch.cuda.synchronize(trainer.root_gpu)
        self.start_time = time.time()

    def on_train_epoch_end(self, trainer, pl_module, outputs):
        torch.cuda.synchronize(trainer.root_gpu)
        max_memory = torch.cuda.max_memory_allocated(trainer.root_gpu) / 2 ** 20
        epoch_time = time.time() - self.start_time

        try:
            max_memory = trainer.training_type_plugin.reduce(max_memory)
            epoch_time = trainer.training_type_plugin.reduce(epoch_time)

            rank_zero_info(f"Average Epoch time: {epoch_time:.2f} seconds")
            rank_zero_info(f"Average Peak memory {max_memory:.2f}MiB")
        except AttributeError:
            pass


if __name__ == "__main__":
    # custom parser to specify config files, train, test and debug mode,
    # postfix, resume.
    # `--key value` arguments are interpreted as arguments to the trainer.
    # `nested.key=value` arguments are interpreted as config parameters.
    # configs are merged from left-to-right followed by command line parameters.

    # model:
    #   base_learning_rate: float
    #   target: path to lightning module
    #   params:
    #       key: value
    # data:
    #   target: main.DataModuleFromConfig
    #   params:
    #      batch_size: int
    #      wrap: bool
    #      train:
    #          target: path to train dataset
    #          params:
    #              key: value
    #      validation:
    #          target: path to validation dataset
    #          params:
    #              key: value
    #      test:
    #          target: path to test dataset
    #          params:
    #              key: value
    # lightning: (optional, has sane defaults and can be specified on cmdline)
    #   trainer:
    #       additional arguments to trainer
    #   logger:
    #       logger to instantiate
    #   modelcheckpoint:
    #       modelcheckpoint to instantiate
    #   callbacks:
    #       callback1:
    #           target: importpath
    #           params:
    #               key: value

    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")

    # add cwd for convenience and to make classes in this file available when
    # running as `python main.py`
    # (in particular `main.DataModuleFromConfig`)
    sys.path.append(os.getcwd())

    parser = get_parser()
    parser = Trainer.add_argparse_args(parser)

    opt, unknown = parser.parse_known_args()
    if opt.name and opt.resume:
        raise ValueError(
            "-n/--name and -r/--resume cannot be specified both."
            "If you want to resume training in a new log folder, "
            "use -n/--name in combination with --resume_from_checkpoint"
        )
    if opt.resume:
        if not os.path.exists(opt.resume):
            raise ValueError("Cannot find {}".format(opt.resume))
        if os.path.isfile(opt.resume):
            paths = opt.resume.split("/")
            # idx = len(paths)-paths[::-1].index("logs")+1
            # logdir = "/".join(paths[:idx])
            logdir = "/".join(paths[:-2])
            ckpt = opt.resume
        else:
            assert os.path.isdir(opt.resume), opt.resume
            logdir = opt.resume.rstrip("/")
            ckpt = os.path.join(logdir, "checkpoints", "last.ckpt")

        opt.resume_from_checkpoint = ckpt
        base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*.yaml")))
        opt.base = base_configs + opt.base
        _tmp = logdir.split("/")
        nowname = _tmp[-1]
    else:
        if opt.name:
            name = "_" + opt.name
        elif opt.base:
            cfg_fname = os.path.split(opt.base[0])[-1]
            cfg_name = os.path.splitext(cfg_fname)[0]
            name = "_" + cfg_name
        else:
            name = ""
        nowname = now + name + opt.postfix
        logdir = os.path.join(opt.logdir, nowname)

    ckptdir = os.path.join(logdir, "checkpoints")
    cfgdir = os.path.join(logdir, "configs")
    seed_everything(opt.seed)

    try:
        # init and save configs
        configs = [OmegaConf.load(cfg) for cfg in opt.base]
        cli = OmegaConf.from_dotlist(unknown)
        config = OmegaConf.merge(*configs, cli)
        lightning_config = config.pop("lightning", OmegaConf.create())
        # merge trainer cli with config
        trainer_config = lightning_config.get("trainer", OmegaConf.create())
        # default to ddp
        trainer_config["accelerator"] = "ddp"
        for k in nondefault_trainer_args(opt):
            trainer_config[k] = getattr(opt, k)
        if not "gpus" in trainer_config:
            del trainer_config["accelerator"]
            cpu = True
        else:
            gpuinfo = trainer_config["gpus"]
            print(f"Running on GPUs {gpuinfo}")
            cpu = False
        trainer_opt = argparse.Namespace(**trainer_config)
        lightning_config.trainer = trainer_config

        # model
        model = instantiate_from_config(config.model)

        # trainer and callbacks
        trainer_kwargs = dict()

        # default logger configs
        default_logger_cfgs = {
            "wandb": {
                "target": "pytorch_lightning.loggers.WandbLogger",
                "params": {
                    "name": nowname,
                    "save_dir": logdir,
                    "offline": opt.debug,
                    "id": nowname,
                }
            },
            "testtube": {
                "target": "pytorch_lightning.loggers.TestTubeLogger",
                "params": {
                    "name": "testtube",
                    "save_dir": logdir,
                }
            },
        }
        default_logger_cfg = default_logger_cfgs["testtube"]
        if "logger" in lightning_config:
            logger_cfg = lightning_config.logger
        else:
            logger_cfg = OmegaConf.create()
        logger_cfg = OmegaConf.merge(default_logger_cfg, logger_cfg)
        trainer_kwargs["logger"] = instantiate_from_config(logger_cfg)

        # modelcheckpoint - use TrainResult/EvalResult(checkpoint_on=metric) to
        # specify which metric is used to determine best models
        default_modelckpt_cfg = {
            "target": "pytorch_lightning.callbacks.ModelCheckpoint",
            "params": {
                "dirpath": ckptdir,
                "filename": "{epoch:06}",
                "verbose": True,
                "save_last": True,
            }
        }
        if hasattr(model, "monitor"):
            print(f"Monitoring {model.monitor} as checkpoint metric.")
            default_modelckpt_cfg["params"]["monitor"] = model.monitor
            default_modelckpt_cfg["params"]["save_top_k"] = 3

        if "modelcheckpoint" in lightning_config:
            modelckpt_cfg = lightning_config.modelcheckpoint
        else:
            modelckpt_cfg =  OmegaConf.create()
        modelckpt_cfg = OmegaConf.merge(default_modelckpt_cfg, modelckpt_cfg)
        print(f"Merged modelckpt-cfg: \n{modelckpt_cfg}")
        if version.parse(pl.__version__) < version.parse('1.4.0'):
            trainer_kwargs["checkpoint_callback"] = instantiate_from_config(modelckpt_cfg)

        # add callback which sets up log directory
        default_callbacks_cfg = {
            "setup_callback": {
                "target": "main.SetupCallback",
                "params": {
                    "resume": opt.resume,
                    "now": now,
                    "logdir": logdir,
                    "ckptdir": ckptdir,
                    "cfgdir": cfgdir,
                    "config": config,
                    "lightning_config": lightning_config,
                }
            },
            "image_logger": {
                "target": "main.ImageLogger",
                "params": {
                    "batch_frequency": 750,
                    "max_images": 4,
                    "clamp": True
                }
            },
            "learning_rate_logger": {
                "target": "main.LearningRateMonitor",
                "params": {
                    "logging_interval": "step",
                    # "log_momentum": True
                }
            },
            "cuda_callback": {
                "target": "main.CUDACallback"
            },
        }
        if version.parse(pl.__version__) >= version.parse('1.4.0'):
            default_callbacks_cfg.update({'checkpoint_callback': modelckpt_cfg})

        if "callbacks" in lightning_config:
            callbacks_cfg = lightning_config.callbacks
        else:
            callbacks_cfg = OmegaConf.create()

        if 'metrics_over_trainsteps_checkpoint' in callbacks_cfg:
            print(
                'Caution: Saving checkpoints every n train steps without deleting. This might require some free space.')
            default_metrics_over_trainsteps_ckpt_dict = {
                'metrics_over_trainsteps_checkpoint':
                    {"target": 'pytorch_lightning.callbacks.ModelCheckpoint',
                     'params': {
                         "dirpath": os.path.join(ckptdir, 'trainstep_checkpoints'),
                         "filename": "{epoch:06}-{step:09}",
                         "verbose": True,
                         'save_top_k': -1,
                         'every_n_train_steps': 10000,
                         'save_weights_only': True
                     }
                     }
            }
            default_callbacks_cfg.update(default_metrics_over_trainsteps_ckpt_dict)

        callbacks_cfg = OmegaConf.merge(default_callbacks_cfg, callbacks_cfg)
        if 'ignore_keys_callback' in callbacks_cfg and hasattr(trainer_opt, 'resume_from_checkpoint'):
            callbacks_cfg.ignore_keys_callback.params['ckpt_path'] = trainer_opt.resume_from_checkpoint
        elif 'ignore_keys_callback' in callbacks_cfg:
            del callbacks_cfg['ignore_keys_callback']

        trainer_kwargs["callbacks"] = [instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg]

        trainer = Trainer.from_argparse_args(trainer_opt, **trainer_kwargs)
        trainer.logdir = logdir  ###

        # data
        data = instantiate_from_config(config.data)
        # NOTE according to https://pytorch-lightning.readthedocs.io/en/latest/datamodules.html
        # calling these ourselves should not be necessary but it is.
        # lightning still takes care of proper multiprocessing though
        data.prepare_data()
        data.setup()
        print("#### Data #####")
        for k in data.datasets:
            print(f"{k}, {data.datasets[k].__class__.__name__}, {len(data.datasets[k])}")

        # configure learning rate
        bs, base_lr = config.data.params.batch_size, config.model.base_learning_rate
        if not cpu:
            ngpu = len(lightning_config.trainer.gpus.strip(",").split(','))
        else:
            ngpu = 1
        if 'accumulate_grad_batches' in lightning_config.trainer:
            accumulate_grad_batches = lightning_config.trainer.accumulate_grad_batches
        else:
            accumulate_grad_batches = 1
        print(f"accumulate_grad_batches = {accumulate_grad_batches}")
        lightning_config.trainer.accumulate_grad_batches = accumulate_grad_batches
        if opt.scale_lr:
            model.learning_rate = accumulate_grad_batches * ngpu * bs * base_lr
            print(
                "Setting learning rate to {:.2e} = {} (accumulate_grad_batches) * {} (num_gpus) * {} (batchsize) * {:.2e} (base_lr)".format(
                    model.learning_rate, accumulate_grad_batches, ngpu, bs, base_lr))
        else:
            model.learning_rate = base_lr
            print("++++ NOT USING LR SCALING ++++")
            print(f"Setting learning rate to {model.learning_rate:.2e}")


        # allow checkpointing via USR1
        def melk(*args, **kwargs):
            # run all checkpoint hooks
            if trainer.global_rank == 0:
                print("Summoning checkpoint.")
                ckpt_path = os.path.join(ckptdir, "last.ckpt")
                trainer.save_checkpoint(ckpt_path)


        def divein(*args, **kwargs):
            if trainer.global_rank == 0:
                import pudb;
                pudb.set_trace()


        import signal

        signal.signal(signal.SIGUSR1, melk)
        signal.signal(signal.SIGUSR2, divein)

        # run
        if opt.train:
            try:
                trainer.fit(model, data)
            except Exception:
                melk()
                raise
        if not opt.no_test and not trainer.interrupted:
            trainer.test(model, data)
    except Exception:
        if opt.debug and trainer.global_rank == 0:
            try:
                import pudb as debugger
            except ImportError:
                import pdb as debugger
            debugger.post_mortem()
        raise
    finally:
        # move newly created debug project to debug_runs
        if opt.debug and not opt.resume and trainer.global_rank == 0:
            dst, name = os.path.split(logdir)
            dst = os.path.join(dst, "debug_runs", name)
            os.makedirs(os.path.split(dst)[0], exist_ok=True)
            os.rename(logdir, dst)
        if trainer.global_rank == 0:
            print(trainer.profiler.summary())

imports/ldm/main.py
jellyheadandrew/CHORUS
from setuptools import setup, find_packages

setup(
    name='ldm4chorus',
    version='0.0.1',
    description='',
    packages=find_packages(),
    install_requires=[
        'torch',
        'numpy',
        'tqdm',
    ],
)
imports/ldm/setup.py
jellyheadandrew/CHORUS
checkpoint_config = dict(interval=50)
log_config = dict(interval=50, hooks=[dict(type='TextLoggerHook')])
log_level = 'INFO'
load_from = None
resume_from = None
dist_params = dict(backend='nccl')
workflow = [('train', 1)]
opencv_num_threads = 0
mp_start_method = 'fork'
dataset_info = dict(
    dataset_name='coco',
    paper_info=dict(
        author=
        'Lin, Tsung-Yi and Maire, Michael and Belongie, Serge and Hays, James and Perona, Pietro and Ramanan, Deva and Doll{\'a}r, Piotr and Zitnick, C Lawrence',
        title='Microsoft coco: Common objects in context',
        container='European conference on computer vision',
        year='2014',
        homepage='http://cocodataset.org/'),
    keypoint_info=dict({
        0:
        dict(name='nose', id=0, color=[51, 153, 255], type='upper', swap=''),
        1:
        dict(
            name='left_eye',
            id=1,
            color=[51, 153, 255],
            type='upper',
            swap='right_eye'),
        2:
        dict(
            name='right_eye',
            id=2,
            color=[51, 153, 255],
            type='upper',
            swap='left_eye'),
        3:
        dict(
            name='left_ear',
            id=3,
            color=[51, 153, 255],
            type='upper',
            swap='right_ear'),
        4:
        dict(
            name='right_ear',
            id=4,
            color=[51, 153, 255],
            type='upper',
            swap='left_ear'),
        5:
        dict(
            name='left_shoulder',
            id=5,
            color=[0, 255, 0],
            type='upper',
            swap='right_shoulder'),
        6:
        dict(
            name='right_shoulder',
            id=6,
            color=[255, 128, 0],
            type='upper',
            swap='left_shoulder'),
        7:
        dict(
            name='left_elbow',
            id=7,
            color=[0, 255, 0],
            type='upper',
            swap='right_elbow'),
        8:
        dict(
            name='right_elbow',
            id=8,
            color=[255, 128, 0],
            type='upper',
            swap='left_elbow'),
        9:
        dict(
            name='left_wrist',
            id=9,
            color=[0, 255, 0],
            type='upper',
            swap='right_wrist'),
        10:
        dict(
            name='right_wrist',
            id=10,
            color=[255, 128, 0],
            type='upper',
            swap='left_wrist'),
        11:
        dict(
            name='left_hip',
            id=11,
            color=[0, 255, 0],
            type='lower',
            swap='right_hip'),
        12:
        dict(
            name='right_hip',
            id=12,
            color=[255, 128, 0],
            type='lower',
            swap='left_hip'),
        13:
        dict(
            name='left_knee',
            id=13,
            color=[0, 255, 0],
            type='lower',
            swap='right_knee'),
        14:
        dict(
            name='right_knee',
            id=14,
            color=[255, 128, 0],
            type='lower',
            swap='left_knee'),
        15:
        dict(
            name='left_ankle',
            id=15,
            color=[0, 255, 0],
            type='lower',
            swap='right_ankle'),
        16:
        dict(
            name='right_ankle',
            id=16,
            color=[255, 128, 0],
            type='lower',
            swap='left_ankle')
    }),
    skeleton_info=dict({
        0:
        dict(link=('left_ankle', 'left_knee'), id=0, color=[0, 255, 0]),
        1:
        dict(link=('left_knee', 'left_hip'), id=1, color=[0, 255, 0]),
        2:
        dict(link=('right_ankle', 'right_knee'), id=2, color=[255, 128, 0]),
        3:
        dict(link=('right_knee', 'right_hip'), id=3, color=[255, 128, 0]),
        4:
        dict(link=('left_hip', 'right_hip'), id=4, color=[51, 153, 255]),
        5:
        dict(link=('left_shoulder', 'left_hip'), id=5, color=[51, 153, 255]),
        6:
        dict(link=('right_shoulder', 'right_hip'), id=6, color=[51, 153, 255]),
        7:
        dict(
            link=('left_shoulder', 'right_shoulder'),
            id=7,
            color=[51, 153, 255]),
        8:
        dict(link=('left_shoulder', 'left_elbow'), id=8, color=[0, 255, 0]),
        9:
        dict(
            link=('right_shoulder', 'right_elbow'), id=9, color=[255, 128, 0]),
        10:
        dict(link=('left_elbow', 'left_wrist'), id=10, color=[0, 255, 0]),
        11:
        dict(link=('right_elbow', 'right_wrist'), id=11, color=[255, 128, 0]),
        12:
        dict(link=('left_eye', 'right_eye'), id=12, color=[51, 153, 255]),
        13:
        dict(link=('nose', 'left_eye'), id=13, color=[51, 153, 255]),
        14:
        dict(link=('nose', 'right_eye'), id=14, color=[51, 153, 255]),
        15:
        dict(link=('left_eye', 'left_ear'), id=15, color=[51, 153, 255]),
        16:
        dict(link=('right_eye', 'right_ear'), id=16, color=[51, 153, 255]),
        17:
        dict(link=('left_ear', 'left_shoulder'), id=17, color=[51, 153, 255]),
        18:
        dict(
            link=('right_ear', 'right_shoulder'), id=18, color=[51, 153, 255])
    }),
    joint_weights=[
        1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.2, 1.2, 1.5, 1.5, 1.0, 1.0, 1.2,
        1.2, 1.5, 1.5
    ],
    sigmas=[
        0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072, 0.062,
        0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089
    ])
evaluation = dict(interval=50, metric='mAP', save_best='AP')
optimizer = dict(type='Adam', lr=0.0015)
optimizer_config = dict(grad_clip=None)
lr_config = dict(
    policy='step',
    warmup='linear',
    warmup_iters=500,
    warmup_ratio=0.001,
    step=[200, 260])
total_epochs = 300
channel_cfg = dict(
    dataset_joints=17,
    dataset_channel=[[
        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
    ]],
    inference_channel=[
        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
    ])
data_cfg = dict(
    image_size=512,
    base_size=256,
    base_sigma=2,
    heatmap_size=[128],
    num_joints=17,
    dataset_channel=[[
        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
    ]],
    inference_channel=[
        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
    ],
    num_scales=1,
    scale_aware_sigma=False)
model = dict(
    type='AssociativeEmbedding',
    pretrained=
    'https://download.openmmlab.com/mmpose/pretrain_models/hrnet_w32-36af842e.pth',
    backbone=dict(
        type='HRNet',
        in_channels=3,
        extra=dict(
            stage1=dict(
                num_modules=1,
                num_branches=1,
                block='BOTTLENECK',
                num_blocks=(4, ),
                num_channels=(64, )),
            stage2=dict(
                num_modules=1,
                num_branches=2,
                block='BASIC',
                num_blocks=(4, 4),
                num_channels=(32, 64)),
            stage3=dict(
                num_modules=4,
                num_branches=3,
                block='BASIC',
                num_blocks=(4, 4, 4),
                num_channels=(32, 64, 128)),
            stage4=dict(
                num_modules=3,
                num_branches=4,
                block='BASIC',
                num_blocks=(4, 4, 4, 4),
                num_channels=(32, 64, 128, 256)))),
    keypoint_head=dict(
        type='AESimpleHead',
        in_channels=32,
        num_joints=17,
        num_deconv_layers=0,
        tag_per_joint=True,
        with_ae_loss=[True],
        extra=dict(final_conv_kernel=1),
        loss_keypoint=dict(
            type='MultiLossFactory',
            num_joints=17,
            num_stages=1,
            ae_loss_type='exp',
            with_ae_loss=[True],
            push_loss_factor=[0.001],
            pull_loss_factor=[0.001],
            with_heatmaps_loss=[True],
            heatmaps_loss_factor=[1.0])),
    train_cfg=dict(),
    test_cfg=dict(
        num_joints=17,
        max_num_people=30,
        scale_factor=[1],
        with_heatmaps=[True],
        with_ae=[True],
        project2image=True,
        align_corners=False,
        nms_kernel=5,
        nms_padding=2,
        tag_per_joint=True,
        detection_threshold=0.1,
        tag_threshold=1,
        use_detection_val=True,
        ignore_too_much=False,
        adjust=True,
        refine=True,
        flip_test=True))
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='BottomUpRandomAffine',
        rot_factor=30,
        scale_factor=[0.75, 1.5],
        scale_type='short',
        trans_factor=40),
    dict(type='BottomUpRandomFlip', flip_prob=0.5),
    dict(type='ToTensor'),
    dict(
        type='NormalizeTensor',
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]),
    dict(type='BottomUpGenerateTarget', sigma=2, max_num_people=30),
    dict(
        type='Collect',
        keys=['img', 'joints', 'targets', 'masks'],
        meta_keys=[])
]
val_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='BottomUpGetImgSize', test_scale_factor=[1]),
    dict(
        type='BottomUpResizeAlign',
        transforms=[
            dict(type='ToTensor'),
            dict(
                type='NormalizeTensor',
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225])
        ]),
    dict(
        type='Collect',
        keys=['img'],
        meta_keys=[
            'image_file', 'aug_data', 'test_scale_factor', 'base_size',
            'center', 'scale', 'flip_index'
        ])
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='BottomUpGetImgSize', test_scale_factor=[1]),
    dict(
        type='BottomUpResizeAlign',
        transforms=[
            dict(type='ToTensor'),
            dict(
                type='NormalizeTensor',
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225])
        ]),
    dict(
        type='Collect',
        keys=['img'],
        meta_keys=[
            'image_file', 'aug_data', 'test_scale_factor', 'base_size',
            'center', 'scale', 'flip_index'
        ])
]
data_root = 'data/coco'
data = dict(
    workers_per_gpu=2,
    train_dataloader=dict(samples_per_gpu=24),
    val_dataloader=dict(samples_per_gpu=1),
    test_dataloader=dict(samples_per_gpu=1),
    train=dict(
        type='BottomUpCocoDataset',
        ann_file='data/coco/annotations/person_keypoints_train2017.json',
        img_prefix='data/coco/train2017/',
        data_cfg=dict(
            image_size=512,
            base_size=256,
            base_sigma=2,
            heatmap_size=[128],
            num_joints=17,
            dataset_channel=[[
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
            ]],
            inference_channel=[
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
            ],
            num_scales=1,
            scale_aware_sigma=False),
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(
                type='BottomUpRandomAffine',
                rot_factor=30,
                scale_factor=[0.75, 1.5],
                scale_type='short',
                trans_factor=40),
            dict(type='BottomUpRandomFlip', flip_prob=0.5),
            dict(type='ToTensor'),
            dict(
                type='NormalizeTensor',
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]),
            dict(type='BottomUpGenerateTarget', sigma=2, max_num_people=30),
            dict(
                type='Collect',
                keys=['img', 'joints', 'targets', 'masks'],
                meta_keys=[])
        ],
        dataset_info=dict(
            dataset_name='coco',
            paper_info=dict(
                author=
                'Lin, Tsung-Yi and Maire, Michael and Belongie, Serge and Hays, James and Perona, Pietro and Ramanan, Deva and Doll{\'a}r, Piotr and Zitnick, C Lawrence',
                title='Microsoft coco: Common objects in context',
                container='European conference on computer vision',
                year='2014',
                homepage='http://cocodataset.org/'),
            keypoint_info=dict({
                0:
                dict(
                    name='nose',
                    id=0,
                    color=[51, 153, 255],
                    type='upper',
                    swap=''),
                1:
                dict(
                    name='left_eye',
                    id=1,
                    color=[51, 153, 255],
                    type='upper',
                    swap='right_eye'),
                2:
                dict(
                    name='right_eye',
                    id=2,
                    color=[51, 153, 255],
                    type='upper',
                    swap='left_eye'),
                3:
                dict(
                    name='left_ear',
                    id=3,
                    color=[51, 153, 255],
                    type='upper',
                    swap='right_ear'),
                4:
                dict(
                    name='right_ear',
                    id=4,
                    color=[51, 153, 255],
                    type='upper',
                    swap='left_ear'),
                5:
                dict(
                    name='left_shoulder',
                    id=5,
                    color=[0, 255, 0],
                    type='upper',
                    swap='right_shoulder'),
                6:
                dict(
                    name='right_shoulder',
                    id=6,
                    color=[255, 128, 0],
                    type='upper',
                    swap='left_shoulder'),
                7:
                dict(
                    name='left_elbow',
                    id=7,
                    color=[0, 255, 0],
                    type='upper',
                    swap='right_elbow'),
                8:
                dict(
                    name='right_elbow',
                    id=8,
                    color=[255, 128, 0],
                    type='upper',
                    swap='left_elbow'),
                9:
                dict(
                    name='left_wrist',
                    id=9,
                    color=[0, 255, 0],
                    type='upper',
                    swap='right_wrist'),
                10:
                dict(
                    name='right_wrist',
                    id=10,
                    color=[255, 128, 0],
                    type='upper',
                    swap='left_wrist'),
                11:
                dict(
                    name='left_hip',
                    id=11,
                    color=[0, 255, 0],
                    type='lower',
                    swap='right_hip'),
                12:
                dict(
                    name='right_hip',
                    id=12,
                    color=[255, 128, 0],
                    type='lower',
                    swap='left_hip'),
                13:
                dict(
                    name='left_knee',
                    id=13,
                    color=[0, 255, 0],
                    type='lower',
                    swap='right_knee'),
                14:
                dict(
                    name='right_knee',
                    id=14,
                    color=[255, 128, 0],
                    type='lower',
                    swap='left_knee'),
                15:
                dict(
                    name='left_ankle',
                    id=15,
                    color=[0, 255, 0],
                    type='lower',
                    swap='right_ankle'),
                16:
                dict(
                    name='right_ankle',
                    id=16,
                    color=[255, 128, 0],
                    type='lower',
                    swap='left_ankle')
            }),
            skeleton_info=dict({
                0:
                dict(
                    link=('left_ankle', 'left_knee'), id=0, color=[0, 255, 0]),
                1:
                dict(link=('left_knee', 'left_hip'), id=1, color=[0, 255, 0]),
                2:
                dict(
                    link=('right_ankle', 'right_knee'),
                    id=2,
                    color=[255, 128, 0]),
                3:
                dict(
                    link=('right_knee', 'right_hip'),
                    id=3,
                    color=[255, 128, 0]),
                4:
                dict(
                    link=('left_hip', 'right_hip'), id=4, color=[51, 153,
                                                                 255]),
                5:
                dict(
                    link=('left_shoulder', 'left_hip'),
                    id=5,
                    color=[51, 153, 255]),
                6:
                dict(
                    link=('right_shoulder', 'right_hip'),
                    id=6,
                    color=[51, 153, 255]),
                7:
                dict(
                    link=('left_shoulder', 'right_shoulder'),
                    id=7,
                    color=[51, 153, 255]),
                8:
                dict(
                    link=('left_shoulder', 'left_elbow'),
                    id=8,
                    color=[0, 255, 0]),
                9:
                dict(
                    link=('right_shoulder', 'right_elbow'),
                    id=9,
                    color=[255, 128, 0]),
                10:
                dict(
                    link=('left_elbow', 'left_wrist'),
                    id=10,
                    color=[0, 255, 0]),
                11:
                dict(
                    link=('right_elbow', 'right_wrist'),
                    id=11,
                    color=[255, 128, 0]),
                12:
                dict(
                    link=('left_eye', 'right_eye'),
                    id=12,
                    color=[51, 153, 255]),
                13:
                dict(link=('nose', 'left_eye'), id=13, color=[51, 153, 255]),
                14:
                dict(link=('nose', 'right_eye'), id=14, color=[51, 153, 255]),
                15:
                dict(
                    link=('left_eye', 'left_ear'), id=15, color=[51, 153,
                                                                 255]),
                16:
                dict(
                    link=('right_eye', 'right_ear'),
                    id=16,
                    color=[51, 153, 255]),
                17:
                dict(
                    link=('left_ear', 'left_shoulder'),
                    id=17,
                    color=[51, 153, 255]),
                18:
                dict(
                    link=('right_ear', 'right_shoulder'),
                    id=18,
                    color=[51, 153, 255])
            }),
            joint_weights=[
                1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.2, 1.2, 1.5, 1.5, 1.0,
                1.0, 1.2, 1.2, 1.5, 1.5
            ],
            sigmas=[
                0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072,
                0.062, 0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089
            ])),
    val=dict(
        type='BottomUpCocoDataset',
        ann_file='data/coco/annotations/person_keypoints_val2017.json',
        img_prefix='data/coco/val2017/',
        data_cfg=dict(
            image_size=512,
            base_size=256,
            base_sigma=2,
            heatmap_size=[128],
            num_joints=17,
            dataset_channel=[[
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
            ]],
            inference_channel=[
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
            ],
            num_scales=1,
            scale_aware_sigma=False),
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(type='BottomUpGetImgSize', test_scale_factor=[1]),
            dict(
                type='BottomUpResizeAlign',
                transforms=[
                    dict(type='ToTensor'),
                    dict(
                        type='NormalizeTensor',
                        mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
                ]),
            dict(
                type='Collect',
                keys=['img'],
                meta_keys=[
                    'image_file', 'aug_data', 'test_scale_factor', 'base_size',
                    'center', 'scale', 'flip_index'
                ])
        ],
        dataset_info=dict(
            dataset_name='coco',
            paper_info=dict(
                author=
                'Lin, Tsung-Yi and Maire, Michael and Belongie, Serge and Hays, James and Perona, Pietro and Ramanan, Deva and Doll{\'a}r, Piotr and Zitnick, C Lawrence',
                title='Microsoft coco: Common objects in context',
                container='European conference on computer vision',
                year='2014',
                homepage='http://cocodataset.org/'),
            keypoint_info=dict({
                0:
                dict(
                    name='nose',
                    id=0,
                    color=[51, 153, 255],
                    type='upper',
                    swap=''),
                1:
                dict(
                    name='left_eye',
                    id=1,
                    color=[51, 153, 255],
                    type='upper',
                    swap='right_eye'),
                2:
                dict(
                    name='right_eye',
                    id=2,
                    color=[51, 153, 255],
                    type='upper',
                    swap='left_eye'),
                3:
                dict(
                    name='left_ear',
                    id=3,
                    color=[51, 153, 255],
                    type='upper',
                    swap='right_ear'),
                4:
                dict(
                    name='right_ear',
                    id=4,
                    color=[51, 153, 255],
                    type='upper',
                    swap='left_ear'),
                5:
                dict(
                    name='left_shoulder',
                    id=5,
                    color=[0, 255, 0],
                    type='upper',
                    swap='right_shoulder'),
                6:
                dict(
                    name='right_shoulder',
                    id=6,
                    color=[255, 128, 0],
                    type='upper',
                    swap='left_shoulder'),
                7:
                dict(
                    name='left_elbow',
                    id=7,
                    color=[0, 255, 0],
                    type='upper',
                    swap='right_elbow'),
                8:
                dict(
                    name='right_elbow',
                    id=8,
                    color=[255, 128, 0],
                    type='upper',
                    swap='left_elbow'),
                9:
                dict(
                    name='left_wrist',
                    id=9,
                    color=[0, 255, 0],
                    type='upper',
                    swap='right_wrist'),
                10:
                dict(
                    name='right_wrist',
                    id=10,
                    color=[255, 128, 0],
                    type='upper',
                    swap='left_wrist'),
                11:
                dict(
                    name='left_hip',
                    id=11,
                    color=[0, 255, 0],
                    type='lower',
                    swap='right_hip'),
                12:
                dict(
                    name='right_hip',
                    id=12,
                    color=[255, 128, 0],
                    type='lower',
                    swap='left_hip'),
                13:
                dict(
                    name='left_knee',
                    id=13,
                    color=[0, 255, 0],
                    type='lower',
                    swap='right_knee'),
                14:
                dict(
                    name='right_knee',
                    id=14,
                    color=[255, 128, 0],
                    type='lower',
                    swap='left_knee'),
                15:
                dict(
                    name='left_ankle',
                    id=15,
                    color=[0, 255, 0],
                    type='lower',
                    swap='right_ankle'),
                16:
                dict(
                    name='right_ankle',
                    id=16,
                    color=[255, 128, 0],
                    type='lower',
                    swap='left_ankle')
            }),
            skeleton_info=dict({
                0:
                dict(
                    link=('left_ankle', 'left_knee'), id=0, color=[0, 255, 0]),
                1:
                dict(link=('left_knee', 'left_hip'), id=1, color=[0, 255, 0]),
                2:
                dict(
                    link=('right_ankle', 'right_knee'),
                    id=2,
                    color=[255, 128, 0]),
                3:
                dict(
                    link=('right_knee', 'right_hip'),
                    id=3,
                    color=[255, 128, 0]),
                4:
                dict(
                    link=('left_hip', 'right_hip'), id=4, color=[51, 153,
                                                                 255]),
                5:
                dict(
                    link=('left_shoulder', 'left_hip'),
                    id=5,
                    color=[51, 153, 255]),
                6:
                dict(
                    link=('right_shoulder', 'right_hip'),
                    id=6,
                    color=[51, 153, 255]),
                7:
                dict(
                    link=('left_shoulder', 'right_shoulder'),
                    id=7,
                    color=[51, 153, 255]),
                8:
                dict(
                    link=('left_shoulder', 'left_elbow'),
                    id=8,
                    color=[0, 255, 0]),
                9:
                dict(
                    link=('right_shoulder', 'right_elbow'),
                    id=9,
                    color=[255, 128, 0]),
                10:
                dict(
                    link=('left_elbow', 'left_wrist'),
                    id=10,
                    color=[0, 255, 0]),
                11:
                dict(
                    link=('right_elbow', 'right_wrist'),
                    id=11,
                    color=[255, 128, 0]),
                12:
                dict(
                    link=('left_eye', 'right_eye'),
                    id=12,
                    color=[51, 153, 255]),
                13:
                dict(link=('nose', 'left_eye'), id=13, color=[51, 153, 255]),
                14:
                dict(link=('nose', 'right_eye'), id=14, color=[51, 153, 255]),
                15:
                dict(
                    link=('left_eye', 'left_ear'), id=15, color=[51, 153,
                                                                 255]),
                16:
                dict(
                    link=('right_eye', 'right_ear'),
                    id=16,
                    color=[51, 153, 255]),
                17:
                dict(
                    link=('left_ear', 'left_shoulder'),
                    id=17,
                    color=[51, 153, 255]),
                18:
                dict(
                    link=('right_ear', 'right_shoulder'),
                    id=18,
                    color=[51, 153, 255])
            }),
            joint_weights=[
                1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.2, 1.2, 1.5, 1.5, 1.0,
                1.0, 1.2, 1.2, 1.5, 1.5
            ],
            sigmas=[
                0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072,
                0.062, 0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089
            ])),
    test=dict(
        type='BottomUpCocoDataset',
        ann_file='data/coco/annotations/person_keypoints_val2017.json',
        img_prefix='data/coco/val2017/',
        data_cfg=dict(
            image_size=512,
            base_size=256,
            base_sigma=2,
            heatmap_size=[128],
            num_joints=17,
            dataset_channel=[[
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
            ]],
            inference_channel=[
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
            ],
            num_scales=1,
            scale_aware_sigma=False),
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(type='BottomUpGetImgSize', test_scale_factor=[1]),
            dict(
                type='BottomUpResizeAlign',
                transforms=[
                    dict(type='ToTensor'),
                    dict(
                        type='NormalizeTensor',
                        mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
                ]),
            dict(
                type='Collect',
                keys=['img'],
                meta_keys=[
                    'image_file', 'aug_data', 'test_scale_factor', 'base_size',
                    'center', 'scale', 'flip_index'
                ])
        ],
        dataset_info=dict(
            dataset_name='coco',
            paper_info=dict(
                author=
                'Lin, Tsung-Yi and Maire, Michael and Belongie, Serge and Hays, James and Perona, Pietro and Ramanan, Deva and Doll{\'a}r, Piotr and Zitnick, C Lawrence',
                title='Microsoft coco: Common objects in context',
                container='European conference on computer vision',
                year='2014',
                homepage='http://cocodataset.org/'),
            keypoint_info=dict({
                0:
                dict(
                    name='nose',
                    id=0,
                    color=[51, 153, 255],
                    type='upper',
                    swap=''),
                1:
                dict(
                    name='left_eye',
                    id=1,
                    color=[51, 153, 255],
                    type='upper',
                    swap='right_eye'),
                2:
                dict(
                    name='right_eye',
                    id=2,
                    color=[51, 153, 255],
                    type='upper',
                    swap='left_eye'),
                3:
                dict(
                    name='left_ear',
                    id=3,
                    color=[51, 153, 255],
                    type='upper',
                    swap='right_ear'),
                4:
                dict(
                    name='right_ear',
                    id=4,
                    color=[51, 153, 255],
                    type='upper',
                    swap='left_ear'),
                5:
                dict(
                    name='left_shoulder',
                    id=5,
                    color=[0, 255, 0],
                    type='upper',
                    swap='right_shoulder'),
                6:
                dict(
                    name='right_shoulder',
                    id=6,
                    color=[255, 128, 0],
                    type='upper',
                    swap='left_shoulder'),
                7:
                dict(
                    name='left_elbow',
                    id=7,
                    color=[0, 255, 0],
                    type='upper',
                    swap='right_elbow'),
                8:
                dict(
                    name='right_elbow',
                    id=8,
                    color=[255, 128, 0],
                    type='upper',
                    swap='left_elbow'),
                9:
                dict(
                    name='left_wrist',
                    id=9,
                    color=[0, 255, 0],
                    type='upper',
                    swap='right_wrist'),
                10:
                dict(
                    name='right_wrist',
                    id=10,
                    color=[255, 128, 0],
                    type='upper',
                    swap='left_wrist'),
                11:
                dict(
                    name='left_hip',
                    id=11,
                    color=[0, 255, 0],
                    type='lower',
                    swap='right_hip'),
                12:
                dict(
                    name='right_hip',
                    id=12,
                    color=[255, 128, 0],
                    type='lower',
                    swap='left_hip'),
                13:
                dict(
                    name='left_knee',
                    id=13,
                    color=[0, 255, 0],
                    type='lower',
                    swap='right_knee'),
                14:
                dict(
                    name='right_knee',
                    id=14,
                    color=[255, 128, 0],
                    type='lower',
                    swap='left_knee'),
                15:
                dict(
                    name='left_ankle',
                    id=15,
                    color=[0, 255, 0],
                    type='lower',
                    swap='right_ankle'),
                16:
                dict(
                    name='right_ankle',
                    id=16,
                    color=[255, 128, 0],
                    type='lower',
                    swap='left_ankle')
            }),
            skeleton_info=dict({
                0:
                dict(
                    link=('left_ankle', 'left_knee'), id=0, color=[0, 255, 0]),
                1:
                dict(link=('left_knee', 'left_hip'), id=1, color=[0, 255, 0]),
                2:
                dict(
                    link=('right_ankle', 'right_knee'),
                    id=2,
                    color=[255, 128, 0]),
                3:
                dict(
                    link=('right_knee', 'right_hip'),
                    id=3,
                    color=[255, 128, 0]),
                4:
                dict(
                    link=('left_hip', 'right_hip'), id=4, color=[51, 153,
                                                                 255]),
                5:
                dict(
                    link=('left_shoulder', 'left_hip'),
                    id=5,
                    color=[51, 153, 255]),
                6:
                dict(
                    link=('right_shoulder', 'right_hip'),
                    id=6,
                    color=[51, 153, 255]),
                7:
                dict(
                    link=('left_shoulder', 'right_shoulder'),
                    id=7,
                    color=[51, 153, 255]),
                8:
                dict(
                    link=('left_shoulder', 'left_elbow'),
                    id=8,
                    color=[0, 255, 0]),
                9:
                dict(
                    link=('right_shoulder', 'right_elbow'),
                    id=9,
                    color=[255, 128, 0]),
                10:
                dict(
                    link=('left_elbow', 'left_wrist'),
                    id=10,
                    color=[0, 255, 0]),
                11:
                dict(
                    link=('right_elbow', 'right_wrist'),
                    id=11,
                    color=[255, 128, 0]),
                12:
                dict(
                    link=('left_eye', 'right_eye'),
                    id=12,
                    color=[51, 153, 255]),
                13:
                dict(link=('nose', 'left_eye'), id=13, color=[51, 153, 255]),
                14:
                dict(link=('nose', 'right_eye'), id=14, color=[51, 153, 255]),
                15:
                dict(
                    link=('left_eye', 'left_ear'), id=15, color=[51, 153,
                                                                 255]),
                16:
                dict(
                    link=('right_eye', 'right_ear'),
                    id=16,
                    color=[51, 153, 255]),
                17:
                dict(
                    link=('left_ear', 'left_shoulder'),
                    id=17,
                    color=[51, 153, 255]),
                18:
                dict(
                    link=('right_ear', 'right_shoulder'),
                    id=18,
                    color=[51, 153, 255])
            }),
            joint_weights=[
                1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.2, 1.2, 1.5, 1.5, 1.0,
                1.0, 1.2, 1.2, 1.5, 1.5
            ],
            sigmas=[
                0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072,
                0.062, 0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089
            ])))

imports/mmpose/associative_embedding_hrnet_w32_coco_512x512.py
jellyheadandrew/CHORUS
# Copyright (c) OpenMMLab. All rights reserved.
import copy
import os
import os.path as osp
import warnings
from argparse import ArgumentParser

import cv2
import mmcv
import numpy as np

from mmpose.apis import (collect_multi_frames, extract_pose_sequence,
                         get_track_id, inference_pose_lifter_model,
                         inference_top_down_pose_model, init_pose_model,
                         process_mmdet_results, vis_3d_pose_result)
from mmpose.core import Smoother
from mmpose.datasets import DatasetInfo
from mmpose.models import PoseLifter, TopDown

try:
    from mmdet.apis import inference_detector, init_detector

    has_mmdet = True
except (ImportError, ModuleNotFoundError):
    has_mmdet = False


def convert_keypoint_definition(keypoints, pose_det_dataset,
                                pose_lift_dataset):
    """Convert pose det dataset keypoints definition to pose lifter dataset
    keypoints definition.

    Args:
        keypoints (ndarray[K, 2 or 3]): 2D keypoints to be transformed.
        pose_det_dataset, (str): Name of the dataset for 2D pose detector.
        pose_lift_dataset (str): Name of the dataset for pose lifter model.
    """
    coco_style_datasets = [
        'TopDownCocoDataset', 'TopDownPoseTrack18Dataset',
        'TopDownPoseTrack18VideoDataset'
    ]
    if pose_det_dataset == 'TopDownH36MDataset' and \
            pose_lift_dataset == 'Body3DH36MDataset':
        return keypoints
    elif pose_det_dataset in coco_style_datasets and \
            pose_lift_dataset == 'Body3DH36MDataset':
        keypoints_new = np.zeros((17, keypoints.shape[1]))
        # pelvis is in the middle of l_hip and r_hip
        keypoints_new[0] = (keypoints[11] + keypoints[12]) / 2
        # thorax is in the middle of l_shoulder and r_shoulder
        keypoints_new[8] = (keypoints[5] + keypoints[6]) / 2
        # in COCO, head is in the middle of l_eye and r_eye
        # in PoseTrack18, head is in the middle of head_bottom and head_top
        keypoints_new[10] = (keypoints[1] + keypoints[2]) / 2
        # spine is in the middle of thorax and pelvis
        keypoints_new[7] = (keypoints_new[0] + keypoints_new[8]) / 2
        # rearrange other keypoints
        keypoints_new[[1, 2, 3, 4, 5, 6, 9, 11, 12, 13, 14, 15, 16]] = \
            keypoints[[12, 14, 16, 11, 13, 15, 0, 5, 7, 9, 6, 8, 10]]
        return keypoints_new
    else:
        raise NotImplementedError


def main():
    parser = ArgumentParser()
    parser.add_argument('det_config', help='Config file for detection')
    parser.add_argument('det_checkpoint', help='Checkpoint file for detection')
    parser.add_argument(
        'pose_detector_config',
        type=str,
        default=None,
        help='Config file for the 1st stage 2D pose detector')
    parser.add_argument(
        'pose_detector_checkpoint',
        type=str,
        default=None,
        help='Checkpoint file for the 1st stage 2D pose detector')
    parser.add_argument(
        'pose_lifter_config',
        help='Config file for the 2nd stage pose lifter model')
    parser.add_argument(
        'pose_lifter_checkpoint',
        help='Checkpoint file for the 2nd stage pose lifter model')
    parser.add_argument(
        '--video-path', type=str, default='', help='Video path')
    parser.add_argument(
        '--rebase-keypoint-height',
        action='store_true',
        help='Rebase the predicted 3D pose so its lowest keypoint has a '
        'height of 0 (landing on the ground). This is useful for '
        'visualization when the model do not predict the global position '
        'of the 3D pose.')
    parser.add_argument(
        '--norm-pose-2d',
        action='store_true',
        help='Scale the bbox (along with the 2D pose) to the average bbox '
        'scale of the dataset, and move the bbox (along with the 2D pose) to '
        'the average bbox center of the dataset. This is useful when bbox '
        'is small, especially in multi-person scenarios.')
    parser.add_argument(
        '--num-instances',
        type=int,
        default=-1,
        help='The number of 3D poses to be visualized in every frame. If '
        'less than 0, it will be set to the number of pose results in the '
        'first frame.')
    parser.add_argument(
        '--show',
        action='store_true',
        default=False,
        help='whether to show visualizations.')
    parser.add_argument(
        '--out-video-root',
        type=str,
        default='vis_results',
        help='Root of the output video file. '
        'Default not saving the visualization video.')
    parser.add_argument(
        '--device', default='cuda:0', help='Device for inference')
    parser.add_argument(
        '--det-cat-id',
        type=int,
        default=1,
        help='Category id for bounding box detection model')
    parser.add_argument(
        '--bbox-thr',
        type=float,
        default=0.9,
        help='Bounding box score threshold')
    parser.add_argument('--kpt-thr', type=float, default=0.3)
    parser.add_argument(
        '--use-oks-tracking', action='store_true', help='Using OKS tracking')
    parser.add_argument(
        '--tracking-thr', type=float, default=0.3, help='Tracking threshold')
    parser.add_argument(
        '--radius',
        type=int,
        default=8,
        help='Keypoint radius for visualization')
    parser.add_argument(
        '--thickness',
        type=int,
        default=2,
        help='Link thickness for visualization')
    parser.add_argument(
        '--smooth',
        action='store_true',
        help='Apply a temporal filter to smooth the pose estimation results. '
        'See also --smooth-filter-cfg.')
    parser.add_argument(
        '--smooth-filter-cfg',
        type=str,
        default='configs/_base_/filters/one_euro.py',
        help='Config file of the filter to smooth the pose estimation '
        'results. See also --smooth.')
    parser.add_argument(
        '--use-multi-frames',
        action='store_true',
        default=False,
        help='whether to use multi frames for inference in the 2D pose'
        'detection stage. Default: False.')
    parser.add_argument(
        '--online',
        action='store_true',
        default=False,
        help='inference mode. If set to True, can not use future frame'
        'information when using multi frames for inference in the 2D pose'
        'detection stage. Default: False.')

    assert has_mmdet, 'Please install mmdet to run the demo.'

    args = parser.parse_args()
    assert args.show or (args.out_video_root != '')
    assert args.det_config is not None
    assert args.det_checkpoint is not None

    video = mmcv.VideoReader(args.video_path)
    assert video.opened, f'Failed to load video file {args.video_path}'

    # First stage: 2D pose detection
    print('Stage 1: 2D pose detection.')

    print('Initializing model...')
    person_det_model = init_detector(
        args.det_config, args.det_checkpoint, device=args.device.lower())

    pose_det_model = init_pose_model(
        args.pose_detector_config,
        args.pose_detector_checkpoint,
        device=args.device.lower())

    assert isinstance(pose_det_model, TopDown), 'Only "TopDown"' \
        'model is supported for the 1st stage (2D pose detection)'

    # frame index offsets for inference, used in multi-frame inference setting
    if args.use_multi_frames:
        assert 'frame_indices_test' in pose_det_model.cfg.data.test.data_cfg
        indices = pose_det_model.cfg.data.test.data_cfg['frame_indices_test']

    pose_det_dataset = pose_det_model.cfg.data['test']['type']
    # get datasetinfo
    dataset_info = pose_det_model.cfg.data['test'].get('dataset_info', None)
    if dataset_info is None:
        warnings.warn(
            'Please set `dataset_info` in the config.'
            'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
            DeprecationWarning)
    else:
        dataset_info = DatasetInfo(dataset_info)

    pose_det_results_list = []
    next_id = 0
    pose_det_results = []

    # whether to return heatmap, optional
    return_heatmap = False

    # return the output of some desired layers,
    # e.g. use ('backbone', ) to return backbone feature
    output_layer_names = None

    print('Running 2D pose detection inference...')
    for frame_id, cur_frame in enumerate(mmcv.track_iter_progress(video)):
        pose_det_results_last = pose_det_results

        # test a single image, the resulting box is (x1, y1, x2, y2)
        mmdet_results = inference_detector(person_det_model, cur_frame)

        # keep the person class bounding boxes.
        person_det_results = process_mmdet_results(mmdet_results,
                                                   args.det_cat_id)

        if args.use_multi_frames:
            frames = collect_multi_frames(video, frame_id, indices,
                                          args.online)

        # make person results for current image
        pose_det_results, _ = inference_top_down_pose_model(
            pose_det_model,
            frames if args.use_multi_frames else cur_frame,
            person_det_results,
            bbox_thr=args.bbox_thr,
            format='xyxy',
            dataset=pose_det_dataset,
            dataset_info=dataset_info,
            return_heatmap=return_heatmap,
            outputs=output_layer_names)

        # get track id for each person instance
        pose_det_results, next_id = get_track_id(
            pose_det_results,
            pose_det_results_last,
            next_id,
            use_oks=args.use_oks_tracking,
            tracking_thr=args.tracking_thr)

        pose_det_results_list.append(copy.deepcopy(pose_det_results))

    # Second stage: Pose lifting
    print('Stage 2: 2D-to-3D pose lifting.')

    print('Initializing model...')
    pose_lift_model = init_pose_model(
        args.pose_lifter_config,
        args.pose_lifter_checkpoint,
        device=args.device.lower())

    assert isinstance(pose_lift_model, PoseLifter), \
        'Only "PoseLifter" model is supported for the 2nd stage ' \
        '(2D-to-3D lifting)'
    pose_lift_dataset = pose_lift_model.cfg.data['test']['type']

    if args.out_video_root == '':
        save_out_video = False
    else:
        os.makedirs(args.out_video_root, exist_ok=True)
        save_out_video = True

    if save_out_video:
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        fps = video.fps
        writer = None

    # convert keypoint definition
    for pose_det_results in pose_det_results_list:
        for res in pose_det_results:
            keypoints = res['keypoints']
            res['keypoints'] = convert_keypoint_definition(
                keypoints, pose_det_dataset, pose_lift_dataset)

    # load temporal padding config from model.data_cfg
    if hasattr(pose_lift_model.cfg, 'test_data_cfg'):
        data_cfg = pose_lift_model.cfg.test_data_cfg
    else:
        data_cfg = pose_lift_model.cfg.data_cfg

    # build pose smoother for temporal refinement
    if args.smooth:
        smoother = Smoother(
            filter_cfg=args.smooth_filter_cfg,
            keypoint_key='keypoints_3d',
            keypoint_dim=3)
    else:
        smoother = None

    num_instances = args.num_instances
    print('Running 2D-to-3D pose lifting inference...')
    for i, pose_det_results in enumerate(
            mmcv.track_iter_progress(pose_det_results_list)):
        # extract and pad input pose2d sequence
        pose_results_2d = extract_pose_sequence(
            pose_det_results_list,
            frame_idx=i,
            causal=data_cfg.causal,
            seq_len=data_cfg.seq_len,
            step=data_cfg.seq_frame_interval)
        # 2D-to-3D pose lifting
        pose_lift_results = inference_pose_lifter_model(
            pose_lift_model,
            pose_results_2d=pose_results_2d,
            dataset=pose_lift_dataset,
            with_track_id=True,
            image_size=video.resolution,
            norm_pose_2d=args.norm_pose_2d)

        # Pose processing
        pose_lift_results_vis = []
        for idx, res in enumerate(pose_lift_results):
            keypoints_3d = res['keypoints_3d']
            # exchange y,z-axis, and then reverse the direction of x,z-axis
            keypoints_3d = keypoints_3d[..., [0, 2, 1]]
            keypoints_3d[..., 0] = -keypoints_3d[..., 0]
            keypoints_3d[..., 2] = -keypoints_3d[..., 2]
            # rebase height (z-axis)
            if args.rebase_keypoint_height:
                keypoints_3d[..., 2] -= np.min(
                    keypoints_3d[..., 2], axis=-1, keepdims=True)
            res['keypoints_3d'] = keypoints_3d
            # add title
            det_res = pose_det_results[idx]
            instance_id = det_res['track_id']
            res['title'] = f'Prediction ({instance_id})'
            # only visualize the target frame
            res['keypoints'] = det_res['keypoints']
            res['bbox'] = det_res['bbox']
            res['track_id'] = instance_id
            pose_lift_results_vis.append(res)

        # Smoothing
        if smoother:
            pose_lift_results = smoother.smooth(pose_lift_results)

        # Visualization
        if num_instances < 0:
            num_instances = len(pose_lift_results_vis)
        img_vis = vis_3d_pose_result(
            pose_lift_model,
            result=pose_lift_results_vis,
            img=video[i],
            out_file=None,
            radius=args.radius,
            thickness=args.thickness,
            num_instances=num_instances)

        if save_out_video:
            if writer is None:
                writer = cv2.VideoWriter(
                    osp.join(args.out_video_root,
                             f'vis_{osp.basename(args.video_path)}'), fourcc,
                    fps, (img_vis.shape[1], img_vis.shape[0]))
            writer.write(img_vis)

    if save_out_video:
        writer.release()


if __name__ == '__main__':
    main()

imports/mmpose/demo/body3d_two_stage_video_demo.py
jellyheadandrew/CHORUS
# Copyright (c) OpenMMLab. All rights reserved.
import os
import warnings
from argparse import ArgumentParser

import cv2
import mmcv

from mmpose.apis import (get_track_id, inference_bottom_up_pose_model,
                         init_pose_model, vis_pose_tracking_result)
from mmpose.core import Smoother
from mmpose.datasets import DatasetInfo


def main():
    """Visualize the demo images."""
    parser = ArgumentParser()
    parser.add_argument('pose_config', help='Config file for pose')
    parser.add_argument('pose_checkpoint', help='Checkpoint file for pose')
    parser.add_argument('--video-path', type=str, help='Video path')
    parser.add_argument(
        '--show',
        action='store_true',
        default=False,
        help='whether to show visualizations.')
    parser.add_argument(
        '--out-video-root',
        default='',
        help='Root of the output video file. '
        'Default not saving the visualization video.')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--kpt-thr', type=float, default=0.5, help='Keypoint score threshold')
    parser.add_argument(
        '--pose-nms-thr',
        type=float,
        default=0.9,
        help='OKS threshold for pose NMS')
    parser.add_argument(
        '--use-oks-tracking', action='store_true', help='Using OKS tracking')
    parser.add_argument(
        '--tracking-thr', type=float, default=0.3, help='Tracking threshold')
    parser.add_argument(
        '--euro',
        action='store_true',
        help='(Deprecated, please use --smooth and --smooth-filter-cfg) '
        'Using One_Euro_Filter for smoothing.')
    parser.add_argument(
        '--smooth',
        action='store_true',
        help='Apply a temporal filter to smooth the pose estimation results. '
        'See also --smooth-filter-cfg.')
    parser.add_argument(
        '--smooth-filter-cfg',
        type=str,
        default='configs/_base_/filters/one_euro.py',
        help='Config file of the filter to smooth the pose estimation '
        'results. See also --smooth.')
    parser.add_argument(
        '--radius',
        type=int,
        default=4,
        help='Keypoint radius for visualization')
    parser.add_argument(
        '--thickness',
        type=int,
        default=1,
        help='Link thickness for visualization')

    args = parser.parse_args()

    assert args.show or (args.out_video_root != '')

    # build the pose model from a config file and a checkpoint file
    pose_model = init_pose_model(
        args.pose_config, args.pose_checkpoint, device=args.device.lower())

    dataset = pose_model.cfg.data['test']['type']
    dataset_info = pose_model.cfg.data['test'].get('dataset_info', None)
    if dataset_info is None:
        warnings.warn(
            'Please set `dataset_info` in the config.'
            'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
            DeprecationWarning)
        assert (dataset == 'BottomUpCocoDataset')
    else:
        dataset_info = DatasetInfo(dataset_info)

    video = mmcv.VideoReader(args.video_path)
    assert video.opened, f'Faild to load video file {args.video_path}'

    if args.out_video_root == '':
        save_out_video = False
    else:
        os.makedirs(args.out_video_root, exist_ok=True)
        save_out_video = True

    if save_out_video:
        fps = video.fps
        size = (video.width, video.height)
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        videoWriter = cv2.VideoWriter(
            os.path.join(args.out_video_root,
                         f'vis_{os.path.basename(args.video_path)}'), fourcc,
            fps, size)

    # optional
    return_heatmap = False

    # build pose smoother for temporal refinement
    if args.euro:
        warnings.warn(
            'Argument --euro will be deprecated in the future. '
            'Please use --smooth to enable temporal smoothing, and '
            '--smooth-filter-cfg to set the filter config.',
            DeprecationWarning)
        smoother = Smoother(
            filter_cfg='configs/_base_/filters/one_euro.py', keypoint_dim=2)
    elif args.smooth:
        smoother = Smoother(filter_cfg=args.smooth_filter_cfg, keypoint_dim=2)
    else:
        smoother = None

    # e.g. use ('backbone', ) to return backbone feature
    output_layer_names = None
    next_id = 0
    pose_results = []
    for cur_frame in mmcv.track_iter_progress(video):
        pose_results_last = pose_results

        pose_results, _ = inference_bottom_up_pose_model(
            pose_model,
            cur_frame,
            dataset=dataset,
            dataset_info=dataset_info,
            pose_nms_thr=args.pose_nms_thr,
            return_heatmap=return_heatmap,
            outputs=output_layer_names)

        # get track id for each person instance
        pose_results, next_id = get_track_id(
            pose_results,
            pose_results_last,
            next_id,
            use_oks=args.use_oks_tracking,
            tracking_thr=args.tracking_thr,
            use_one_euro=args.euro,
            fps=fps)

        # post-process the pose results with smoother
        if smoother:
            pose_results = smoother.smooth(pose_results)

        # show the results
        vis_frame = vis_pose_tracking_result(
            pose_model,
            cur_frame,
            pose_results,
            radius=args.radius,
            thickness=args.thickness,
            dataset=dataset,
            dataset_info=dataset_info,
            kpt_score_thr=args.kpt_thr,
            show=False)

        if args.show:
            cv2.imshow('Image', vis_frame)

        if save_out_video:
            videoWriter.write(vis_frame)

        if args.show and cv2.waitKey(1) & 0xFF == ord('q'):
            break

    if save_out_video:
        videoWriter.release()
    if args.show:
        cv2.destroyAllWindows()


if __name__ == '__main__':
    main()

imports/mmpose/demo/bottom_up_pose_tracking_demo.py
jellyheadandrew/CHORUS
import os
import os.path as osp
import platform
import shutil
import sys
import warnings
from setuptools import find_packages, setup


def readme():
    with open('README.md', encoding='utf-8') as f:
        content = f.read()
    return content


version_file = 'mmpose/version.py'


def get_version():
    with open(version_file, 'r') as f:
        exec(compile(f.read(), version_file, 'exec'))
    import sys

    # return short version for sdist
    if 'sdist' in sys.argv or 'bdist_wheel' in sys.argv:
        return locals()['short_version']
    else:
        return locals()['__version__']


def parse_requirements(fname='requirements.txt', with_version=True):
    """Parse the package dependencies listed in a requirements file but strips
    specific versioning information.

    Args:
        fname (str): path to requirements file
        with_version (bool, default=False): if True include version specs

    Returns:
        List[str]: list of requirements items

    CommandLine:
        python -c "import setup; print(setup.parse_requirements())"
    """
    import re
    import sys
    from os.path import exists
    require_fpath = fname

    def parse_line(line):
        """Parse information from a line in a requirements text file."""
        if line.startswith('-r '):
            # Allow specifying requirements in other files
            target = line.split(' ')[1]
            for info in parse_require_file(target):
                yield info
        else:
            info = {'line': line}
            if line.startswith('-e '):
                info['package'] = line.split('#egg=')[1]
            elif '@git+' in line:
                info['package'] = line
            else:
                # Remove versioning from the package
                pat = '(' + '|'.join(['>=', '==', '>']) + ')'
                parts = re.split(pat, line, maxsplit=1)
                parts = [p.strip() for p in parts]

                info['package'] = parts[0]
                if len(parts) > 1:
                    op, rest = parts[1:]
                    if ';' in rest:
                        # Handle platform specific dependencies
                        # http://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-platform-specific-dependencies
                        version, platform_deps = map(str.strip,
                                                     rest.split(';'))
                        info['platform_deps'] = platform_deps
                    else:
                        version = rest  # NOQA
                    info['version'] = (op, version)
            yield info

    def parse_require_file(fpath):
        with open(fpath, 'r') as f:
            for line in f.readlines():
                line = line.strip()
                if line and not line.startswith('#'):
                    for info in parse_line(line):
                        yield info

    def gen_packages_items():
        if exists(require_fpath):
            for info in parse_require_file(require_fpath):
                parts = [info['package']]
                if with_version and 'version' in info:
                    parts.extend(info['version'])
                if not sys.version.startswith('3.4'):
                    # apparently package_deps are broken in 3.4
                    platform_deps = info.get('platform_deps')
                    if platform_deps is not None:
                        parts.append(';' + platform_deps)
                item = ''.join(parts)
                yield item

    packages = list(gen_packages_items())
    return packages


def add_mim_extension():
    """Add extra files that are required to support MIM into the package.

    These files will be added by creating a symlink to the originals if the
    package is installed in `editable` mode (e.g. pip install -e .), or by
    copying from the originals otherwise.
    """

    # parse installment mode
    if 'develop' in sys.argv:
        # installed by `pip install -e .`
        if platform.system() == 'Windows':
            mode = 'copy'
        else:
            mode = 'symlink'
    elif 'sdist' in sys.argv or 'bdist_wheel' in sys.argv:
        # installed by `pip install .`
        # or create source distribution by `python setup.py sdist`
        mode = 'copy'
    else:
        return

    filenames = ['tools', 'configs', 'demo', 'model-index.yml']
    repo_path = osp.dirname(__file__)
    mim_path = osp.join(repo_path, 'mmpose', '.mim')
    os.makedirs(mim_path, exist_ok=True)

    for filename in filenames:
        if osp.exists(filename):
            src_path = osp.join(repo_path, filename)
            tar_path = osp.join(mim_path, filename)

            if osp.isfile(tar_path) or osp.islink(tar_path):
                os.remove(tar_path)
            elif osp.isdir(tar_path):
                shutil.rmtree(tar_path)

            if mode == 'symlink':
                src_relpath = osp.relpath(src_path, osp.dirname(tar_path))
                os.symlink(src_relpath, tar_path)
            elif mode == 'copy':
                if osp.isfile(src_path):
                    shutil.copyfile(src_path, tar_path)
                elif osp.isdir(src_path):
                    shutil.copytree(src_path, tar_path)
                else:
                    warnings.warn(f'Cannot copy file {src_path}.')
            else:
                raise ValueError(f'Invalid mode {mode}')


if __name__ == '__main__':
    add_mim_extension()
    setup(
        name='mmpose',
        version=get_version(),
        description='OpenMMLab Pose Estimation Toolbox and Benchmark.',
        author='MMPose Contributors',
        author_email='openmmlab@gmail.com',
        keywords='computer vision, pose estimation',
        long_description=readme(),
        long_description_content_type='text/markdown',
        packages=find_packages(exclude=('configs', 'tools', 'demo')),
        include_package_data=True,
        package_data={'mmpose.ops': ['*/*.so']},
        classifiers=[
            'Development Status :: 4 - Beta',
            'License :: OSI Approved :: Apache Software License',
            'Operating System :: OS Independent',
            'Programming Language :: Python :: 3',
            'Programming Language :: Python :: 3.5',
            'Programming Language :: Python :: 3.6',
            'Programming Language :: Python :: 3.7',
            'Programming Language :: Python :: 3.8',
            'Programming Language :: Python :: 3.9',
        ],
        url='https://github.com/open-mmlab/mmpose',
        license='Apache License 2.0',
        install_requires=parse_requirements('requirements/runtime.txt'),
        extras_require={
            'tests': parse_requirements('requirements/tests.txt'),
            'build': parse_requirements('requirements/build.txt'),
            'runtime': parse_requirements('requirements/runtime.txt'),
            'mim': parse_requirements('requirements/mminstall.txt'),
        },
        zip_safe=False)

imports/mmpose/setup.py
jellyheadandrew/CHORUS
#!/usr/bin/env bash
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import time

import torch
from mmcv import Config
from mmcv.cnn import fuse_conv_bn
from mmcv.parallel import MMDataParallel
from mmcv.runner.fp16_utils import wrap_fp16_model

from mmpose.datasets import build_dataloader, build_dataset
from mmpose.models import build_posenet


def parse_args():
    parser = argparse.ArgumentParser(
        description='MMPose benchmark a recognizer')
    parser.add_argument('config', help='test config file path')
    parser.add_argument(
        '--log-interval', default=10, help='interval of logging')
    parser.add_argument(
        '--fuse-conv-bn',
        action='store_true',
        help='Whether to fuse conv and bn, this will slightly increase'
        'the inference speed')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)
    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # build the dataloader
    dataset = build_dataset(cfg.data.val)
    data_loader = build_dataloader(
        dataset,
        samples_per_gpu=1,
        workers_per_gpu=cfg.data.workers_per_gpu,
        dist=False,
        shuffle=False)

    # build the model and load checkpoint
    model = build_posenet(cfg.model)
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        wrap_fp16_model(model)
    if args.fuse_conv_bn:
        model = fuse_conv_bn(model)
    model = MMDataParallel(model, device_ids=[0])

    # the first several iterations may be very slow so skip them
    num_warmup = 5
    pure_inf_time = 0

    # benchmark with total batch and take the average
    for i, data in enumerate(data_loader):

        torch.cuda.synchronize()
        start_time = time.perf_counter()
        with torch.no_grad():
            model(return_loss=False, **data)

        torch.cuda.synchronize()
        elapsed = time.perf_counter() - start_time

        if i >= num_warmup:
            pure_inf_time += elapsed
            if (i + 1) % args.log_interval == 0:
                its = (i + 1 - num_warmup) / pure_inf_time
                print(f'Done item [{i + 1:<3}],  {its:.2f} items / s')
    print(f'Overall average: {its:.2f} items / s')
    print(f'Total time: {pure_inf_time:.2f} s')


if __name__ == '__main__':
    main()

imports/mmpose/tools/analysis/benchmark_inference.py
jellyheadandrew/CHORUS
#!/usr/bin/env bash
# Copyright (c) OpenMMLab. All rights reserved.
"""This file is for benchmark data loading process. It can also be used to
refresh the memcached cache. The command line to run this file is:

$ python -m cProfile -o program.prof tools/analysis/benchmark_processing.py
configs/task/method/[config filename]

Note: When debugging, the `workers_per_gpu` in the config should be set to 0
during benchmark.

It use cProfile to record cpu running time and output to program.prof
To visualize cProfile output program.prof, use Snakeviz and run:
$ snakeviz program.prof
"""
import argparse

import mmcv
from mmcv import Config

from mmpose import __version__
from mmpose.datasets import build_dataloader, build_dataset
from mmpose.utils import get_root_logger


def main():
    parser = argparse.ArgumentParser(description='Benchmark data loading')
    parser.add_argument('config', help='train config file path')
    args = parser.parse_args()
    cfg = Config.fromfile(args.config)

    # init logger before other steps
    logger = get_root_logger()
    logger.info(f'MMPose Version: {__version__}')
    logger.info(f'Config: {cfg.text}')

    dataset = build_dataset(cfg.data.train)
    data_loader = build_dataloader(
        dataset,
        samples_per_gpu=1,
        workers_per_gpu=cfg.data.workers_per_gpu,
        dist=False,
        shuffle=False)

    # Start progress bar after first 5 batches
    prog_bar = mmcv.ProgressBar(
        len(dataset) - 5 * cfg.data.samples_per_gpu, start=False)
    for i, data in enumerate(data_loader):
        if i == 5:
            prog_bar.start()
        for _ in data['img']:
            if i < 5:
                continue
            prog_bar.update()


if __name__ == '__main__':
    main()

imports/mmpose/tools/analysis/benchmark_processing.py
jellyheadandrew/CHORUS
import os
from glob import glob
import argparse
import json
from copy import deepcopy
from tqdm import tqdm

import math
import random

from constants.ldm import DEFAULT_LDM_MODEL_KEY
from constants.filtering import DEFAULT_FILTER_SETTING_NUM
from constants.aggr_settings.quant import QUANT_AGGR_SETTINGS
from constants.aggr_settings.qual import QUAL_AGGR_SETTINGS
from constants.metadata import DEFAULT_TOTAL_AGGREGATION_NAME, DEFAULT_FILTERING_DIR, DEFAULT_PERSPECTIE_CAMERA_DIR, DEFAULT_AGGREGATION_DIR, DEFAULT_SEED

from utils.prepare_prompts import get_unique_categories, get_cps_for_given_category, prepare_cps_from_dirs
from utils.settings import load_filter_settings_combination, aggr_save_name
from utils.reproducibility import seed_everything
from utils.misc import load_category_id, extract_nonvector
from utils.aggregation.core import prepare_aggr_core
from utils.aggregation.load_input import load_aggr_input
from utils.aggregation.aggregator import Aggregator




def aggregate(
        cps,
        ldm_model_key,
        save_dir,
        cam_world_key,
        total_aggr_name,
        aggr_settings,
        filter_settings,
        visualize,
        skip_single_prompt,
        skip_done,
        verbose,
    ):
    # prepare aggregation core for running aggregation
    print("Preparing Aggregator Core...")

    # 'aggregation-results' save name (including 'total-aggregation')
    aggr_result_save_name = aggr_save_name(
        filter_setting_comb=filter_settings['filter_setting_comb'], 
        aggr_setting_name=aggr_settings['NAME']
    )

    # load aggregation settings and prepare 'aggregation_core' that contains all variables required, initialized
    aggr_core = prepare_aggr_core(verbose=verbose, **aggr_settings)
    
    # all unique 'augprompt:ID's per prompt, per category
    augid_per_cp = {
        category: {
            cp: sorted(list(set(
                [
                    retained_file_pth.split("/")[-1].split(".")[0] for retained_file_pth in 
                    sorted(list(glob(f"{filter_settings['filter_dir']}/{ldm_model_key}/{cp[0]}/settings:*/{cp[1]}/*.png")))
                    if int(retained_file_pth.split("/")[-3].replace("settings:","")) in filter_settings['filter_setting_comb']
                ]
            )))
            for cp in get_cps_for_given_category(cps, category)
        }
        for category in get_unique_categories(cps)
    }
    
    # if the ratio-of-data to aggregate is given, we random sample from 'augid_per_cp'
    if aggr_settings['DATA_RATIO'] < 1.:
        for category in augid_per_cp.keys(): 
            for cp in augid_per_cp[category]:
                # augprompt:id pairs
                augids = deepcopy(augid_per_cp[category][cp])
                # number of pairs to retain
                k = math.ceil(len(augids) * aggr_settings['DATA_RATIO'])
                # random-sample amount of data
                augid_per_cp[category][cp] = random.sample(augids, k=k)

    # clear screen if verbose
    print("Preparation Complete!")
    if not verbose: os.system("clear")

    ## iterate for all categories
    for category in get_unique_categories(cps):
        # (total-aggregation) result-save directory
        tot_result_save_dir = f"./{save_dir}/{ldm_model_key}/{category}/{total_aggr_name}"
        os.makedirs(tot_result_save_dir, exist_ok=True)
        # (total-aggregation) result-save paths
        tot_video_save_pth = f"{tot_result_save_dir}/{aggr_result_save_name}.mp4"
        tot_aggr_save_pth = f"{tot_result_save_dir}/{aggr_result_save_name}.pickle"
        tot_metadata_save_pth = f"{tot_result_save_dir}/{aggr_result_save_name}.json"
        
        # (total-aggregation) aggregator
        tot_aggregator = Aggregator(core=deepcopy(aggr_core))

        ## register inputs for aggregation (all prompt) 
        for cp in augid_per_cp[category].keys():
            # prompt
            prompt = cp[1]

            # result-save directory
            result_save_dir = f"./{save_dir}/{ldm_model_key}/{category}/{prompt}"
            os.makedirs(result_save_dir, exist_ok=True)
            # result-save paths
            video_save_pth = f"{result_save_dir}/{aggr_result_save_name}.mp4"
            aggr_save_pth = f"{result_save_dir}/{aggr_result_save_name}.pickle"
            metadata_save_pth = f"{result_save_dir}/{aggr_result_save_name}.json"

            # skip 'aggr' if already done (not visualize/skip-done/skip-single-prompt & aggregation/metadata exists)
            if os.path.exists(aggr_save_pth) and os.path.exists(metadata_save_pth) and os.path.exists(tot_aggr_save_pth) and os.path.exists(tot_metadata_save_pth) and (not visualize) and skip_done:
                if verbose:
                    print("\tContinueing since (aggregation) ALREADY DONE!!", tot_aggr_save_pth)
                continue
            # skip 'aggr' if already done (visualize/skip-done/skip-single-prompt & aggregation/video/metadata exists)
            if os.path.exists(aggr_save_pth) and os.path.exists(video_save_pth) and os.path.exists(metadata_save_pth) and os.path.exists(tot_aggr_save_pth) and os.path.exists(tot_video_save_pth) and os.path.exists(tot_metadata_save_pth) and visualize and skip_done:
                if verbose:
                    print("\tContinueing since (aggregation) ALREADY DONE!!", tot_aggr_save_pth)
                continue

            ## aggregate for 'single prompt'
            # declare 'aggregator'
            aggregator = Aggregator(core=deepcopy(aggr_core))

            ## register inputs for aggregation (per prompt)
            desc = f"Aggregating for [Category '{category}' / Prompt '{prompt}' / Aggregation Setting '{aggr_settings['NAME']}' / Filter Setting '{filter_settings['filter_setting_comb']}']"
            for augid in tqdm(augid_per_cp[category][cp], desc=desc):
                # augprompt
                augprompt = augid.split(":")[0]
                # image id
                _id = augid.split(":")[1]
                    
                # input file paths
                fps = dict(
                    cam_smpl_pth = f"{filter_settings['cam_dir']}/{ldm_model_key}/{category}/{prompt}/{augprompt}:{_id}.pickle",
                    human_smpl_pth = f"{filter_settings['human_dir']}/{ldm_model_key}/{category}/{prompt}/{augprompt}/{_id}.pickle"
                )
            
                # category id
                category_id = load_category_id(category)

                # aggregation-inputs
                aggr_input_list = load_aggr_input(
                    cam_world_key=cam_world_key,
                    category_id=category_id,
                    fps=fps,
                    zerobeta=True
                )

                # iterate for all 'aggr_input'
                for _aggr_input in aggr_input_list:
                    # '_aggr_input' (cam_R, cam_t, cam_K, mask, pose, betas) -> 'aggr_input' (cam_R, cam_t, cam_K, mask, pose, betas)
                    aggr_input = aggregator.preprocess_input(
                        cam_world_key=cam_world_key,
                        **_aggr_input
                    )
                    # register inputs to 'tot_aggregator'
                    tot_aggregator.register_cam_pose(
                        cam_world_key=cam_world_key,
                        add_human_for_vis=aggr_core.GLOBAL.add_human_for_vis, # we do not visualize posed-human with zero-human
                        **aggr_input
                    )
                    # register inputs to 'aggregator'
                    aggregator.register_cam_pose(
                        cam_world_key=cam_world_key,
                        add_human_for_vis=aggr_core.GLOBAL.add_human_for_vis, # we do not visualize posed-human with zero-human
                        **aggr_input
                    )

                    # symmetric augmentation (if on)
                    if aggr_core.GLOBAL.symmetric_augmentation:
                        # apply symmetric augmentation
                        aggr_input = aggregator.convert_left_to_right(
                            cam_world_key=cam_world_key,
                            **aggr_input
                        )
                        # register augmented inputs to 'tot_aggergator'
                        tot_aggregator.register_cam_pose(
                            cam_world_key=cam_world_key,
                            add_human_for_vis=aggr_core.GLOBAL.add_human_for_vis, # we do not visualize posed-human with zero-human
                            **aggr_input
                        )
                        # register augmented inputs to 'aggregator'
                        aggregator.register_cam_pose(
                            cam_world_key=cam_world_key,
                            add_human_for_vis=aggr_core.GLOBAL.add_human_for_vis, # we do not visualize posed-human with zero-human
                            **aggr_input
                        )

            ## aggregate all registered inputs (for single prompt)
            # aggregate
            aggregator.aggregate_all(
                camera_sampling=aggr_core[cam_world_key].camera_sampling,
                verbose=verbose
            )
            # save results
            if not skip_single_prompt:
                # export 'aggregator' results
                aggregator.export(aggr_save_pth)
                # save video
                if visualize: aggregator.export_video(save_pth=video_save_pth, prompt=prompt)
                # save metadata
                metadata = extract_nonvector({'filter_settings': filter_settings, 'aggr_core': aggr_core})
                with open(metadata_save_pth, 'w') as f: json.dump(metadata, f, indent=1)


        ## skip 'total-aggr' if already done (not visualize/skip-done/skip-single-prompt & aggregation/metadata exists)
        if os.path.exists(tot_aggr_save_pth) and os.path.exists(tot_metadata_save_pth) and (not visualize) and skip_done:
            if verbose:
                print("\tContinueing since (total-aggregation) ALREADY DONE!!", tot_aggr_save_pth)
            continue
        # skip 'total-aggr' if already done (visualize/skip-done/skip-single-prompt & aggregation/video/metadata exists)
        if os.path.exists(tot_aggr_save_pth) and os.path.exists(tot_video_save_pth) and os.path.exists(tot_metadata_save_pth) and visualize and skip_done:
            if verbose:
                print("\tContinueing since (total-aggregation) ALREADY DONE!!", tot_aggr_save_pth)
            continue


        ## aggregate all registered inputs (for all prompts)
        # aggregate
        tot_aggregator.aggregate_all(
            camera_sampling=aggr_core[cam_world_key].camera_sampling,
            verbose=verbose,
        )
        # save results
        tot_aggregator.export(tot_aggr_save_pth)
        # save video
        if visualize: tot_aggregator.export_video(save_pth=tot_video_save_pth, prompt="all prompts")
        # save metadata
        metadata = extract_nonvector({'filter_settings': filter_settings, 'aggr_core': aggr_core})
        with open(tot_metadata_save_pth, 'w') as f: json.dump(metadata, f, indent=1)


def run_aggregation(cps, args):
    ## aggregate for all 'args.aggr_setting_names'
    for aggr_setting_name in args.aggr_setting_names:
        # load aggregation setting
        aggr_settings = QUANT_AGGR_SETTINGS[aggr_setting_name] if args.eval_mode == "quant" else QUAL_AGGR_SETTINGS[aggr_setting_name]
        assert aggr_settings['NAME'] == aggr_setting_name, f"Mismatch in name for '{aggr_setting_name}'"

        # retrieve filter-settings (for using images)
        for filter_setting_comb in args.filter_setting_nums:
            # print setting informations
            if args.verbose: print(f"Aggregating 'aggr_settings':{aggr_setting_name} / 'filter_setting_comb': {filter_setting_comb}")

            # filter_setting_comb as list of filter_setting_num
            filter_setting_comb = [int(i) for i in filter_setting_comb.split(",")]

            # load 'filter_settings' from multiple 'settings:00x.txt' file (while checking consistency)
            filter_settings = load_filter_settings_combination(
                args=args, 
                filter_setting_comb=filter_setting_comb,
                args_to_ignore = [
                    'categories',
                    'prompts',
                    'ldm_model_key',
                    'total_aggr_name',
                    'save_dir',
                    'aggr_setting_names',
                    'filter_setting_nums',
                    'visualize',
                    'skip_single_prompt',
                    'skip_done',
                    'verbose',
                ]
            )

            # aggregate images and save 3d results
            aggregate(
                cps=cps,
                ldm_model_key=args.ldm_model_key,
                save_dir=args.save_dir,
                cam_world_key=args.cam_world_key,
                total_aggr_name=args.total_aggr_name,
                aggr_settings=aggr_settings,
                filter_settings=filter_settings,
                visualize=args.visualize,
                skip_single_prompt=args.skip_single_prompt,
                skip_done=args.skip_done,
                verbose=args.verbose,
            )
            

if __name__ == '__main__':
    """
    [Aggregating Images from Multiple Filter-Setting]
    NOTE To enable combinations:
        Write like   "...   --filter_settings '7,8' '9'"
            => Aggregates Results from Filter Settings [7,8] 
            => Aggregates Results from Filter Settings [9]
            => [7,8] and [9] are seperate experiments
        Write like   "...   --filter_settings '7' '8' '9'"
            => Aggregates Results from Filter Settings [7]
            => Aggregates Results from Filter Settings [8]
            => Aggregates Results from Filter Settings [9]
            => [7] and [8] and [9] are seperate experiments


    [Aggregation Settings]
    NOTE Provide Aggregation Setting Names
        Write like "... --aggr_settings 'quant:full'"
            => 'eval_mode' must be set as 'quant'
            => Aggregation settings are set as 'quant:full'
            => Results are Saved Under 'quant:full' directory
            
        Write like "... --aggr_settings 'qual:002'"
            => 'eval_mode' must be set as 'qual'
            => Aggregation settings are set as 'qual:002'
            => Results are Saved Under 'qual:002' directory
    """
    ## Arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--categories", nargs="+")
    parser.add_argument("--prompts", nargs="+")
    parser.add_argument("--ldm_model_key", type=str, default=DEFAULT_LDM_MODEL_KEY)
    
    parser.add_argument("--filter_dir", type=str, default=DEFAULT_FILTERING_DIR)
    parser.add_argument("--cam_dir", type=str, default=DEFAULT_PERSPECTIE_CAMERA_DIR)
    parser.add_argument("--save_dir", type=str, default=DEFAULT_AGGREGATION_DIR)
    parser.add_argument("--total_aggr_name", type=str, default=DEFAULT_TOTAL_AGGREGATION_NAME)

    parser.add_argument("--eval_mode", nargs="?", choices=["qual", "quant"])
    parser.add_argument("--cam_world_key", default='SMPL')
    
    parser.add_argument("--aggr_setting_names", nargs="+", type=str)
    parser.add_argument("--filter_setting_nums", nargs="+", type=str, default=[str(DEFAULT_FILTER_SETTING_NUM)]) # you can write as "--filter_settings '7,8' '9'"

    parser.add_argument("--visualize", action="store_true")
    parser.add_argument("--skip_single_prompt", action="store_true")
    parser.add_argument("--skip_done", action="store_true")
    parser.add_argument("--verbose", action="store_true")

    parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
    args = parser.parse_args()

    # load sorted (category / prompt) pairs
    cps = prepare_cps_from_dirs(args, image_dir=args.filter_dir, use_filter_setting=True)

    # seed for reproducible generation
    seed_everything(args.seed)

    # arguments
    if args.verbose:
        print("===Arguments===")
        print(args)

    # pose-canonicalized aggregation
    run_aggregation(cps=cps, args=args)
scripts/aggregation/aggregate.py
jellyheadandrew/CHORUS
import os
from glob import glob
import shutil

import argparse
from tqdm import tqdm

import random
import numpy as np
import pickle5 as pickle
import matplotlib.pyplot as plt

from constants.ldm import DEFAULT_LDM_MODEL_KEY
from constants.filtering import KEYPOINT_FILTERS, DEFAULT_KEYPOINT_FILTER_NUM, DEFAULT_FILTER_MULTI_PERSON_OBJECT_METHODS, DEFAULT_BBOX_OVERLAP_THRES, DEFAULT_KEYPOINT_CONFIDENCE_THRES
from constants.metadata import DEFAULT_IMAGE_DIR, DEFAULT_SEGMENTATION_DIR, DEFAULT_HUMAN_DIR, DEFAULT_KEYPOINT_DIR, DEFAULT_FILTERING_DIR, DEFAULT_SEED

from utils.settings import create_filter_settings, find_filter_setting_num
from utils.prepare_prompts import prepare_cps_from_dirs
from utils.filtering import bbox_iou
from utils.reproducibility import seed_everything
from utils.misc import load_category_id




def apply_filters(create_settings_only, cps, **kwargs):

    # retrieve settings
    settings_str, settings = create_filter_settings(**kwargs)
    
    # used for counting filtered samples in category
    curr_category = None

    # find setting number
    setting_num = find_filter_setting_num(settings_str, settings)

    # if create settings only
    if create_settings_only:
        if settings['verbose']:
            print("\tRunning 'create_settings_only'...")
        return

    # clear screen if not verbose
    if not settings['verbose']: os.system("clear")

    # run filtering for all (category, prompt) pairs
    for category, prompt in sorted(cps):
        # number of samples retained (for given category)
        if curr_category != category:
            num_retained_in_category = 0
            curr_category = category
    
        # all image path
        img_pths = sorted(list(glob(f"./{settings['image_dir']}/{settings['ldm_model_key']}/{category}/{prompt}/*/*.png")))
        num_retained_in_prompt = 0 # counter for useful samples (not filtered)

        # iterate through image path
        for img_pth in tqdm(img_pths, desc=f"Filtering Dataset for [Category '{category}' / Prompt '{prompt}' / Setting: {setting_num:03}]"):
            # print
            if settings['verbose']:
                print("Processing: ", img_pth)

            # augprompt
            augprompt = img_pth.split("/")[-2]
            # image id
            _id, _ext = img_pth.split("/")[-1].split(".")
            # result-save directory
            result_save_dir = f"./{settings['save_dir']}/{settings['ldm_model_key']}/{category}/settings:{setting_num:03}/{prompt}"
            os.makedirs(result_save_dir, exist_ok=True)
            os.makedirs(f"{result_save_dir}/FILTERED", exist_ok=True)
            # result-save paths
            retained_save_pth =f"{result_save_dir}/{augprompt}:{_id}.png"
            filtered_save_pth=f"{result_save_dir}/FILTERED/{augprompt}:{_id}"

            # skip if already done
            if os.path.exists(retained_save_pth) and settings['skip_done']:
                num_retained_in_prompt += 1
                num_retained_in_category += 1
                if settings['verbose']:
                    print(f"\tContinueing since ALREADY DONE (retained)!!", retained_save_pth) 
                continue
            if os.path.exists(filtered_save_pth) and settings['skip_done']:
                if settings['verbose']:
                    print(f"\tContinueing since ALREADY DONE (filtered)!!", filtered_save_pth) 
                continue

            ## loading human information & filtering
            # load human params
            human_pth = f"./{settings['human_dir']}/{settings['ldm_model_key']}/{category}/{prompt}/{augprompt}/{_id}.pickle"
            with open(human_pth, "rb") as handle: human_parameters = pickle.load(handle)

            # if there is no human, filtering is applied
            if human_parameters == "NO HUMANS":
                with open(filtered_save_pth, "w") as wf: wf.write("NO HUMAN")
                if settings['verbose']: print(f"\tFiltering {filtered_save_pth} since there doesn't exist any humans...")
                continue

            # if there is more than 1 human, filtering is applied (if skip_multi_person_image)
            num_humans = len(human_parameters['cams'])
            if num_humans >= 2 and settings['skip_multi_person_image']:
                with open(filtered_save_pth, "w") as wf: wf.write("MORE THAN 2 HUMAN")
                if settings['verbose']: print(f"\tFiltering {filtered_save_pth} since there exists more than 2 humans for single camera...")
                continue
        
            ## loading object information & filtering
            # category id
            category_id = load_category_id(category)
            
            #  load object mask & bbox (xyxy)
            instances = human_parameters['auxiliary']['instances_including_object']
            is_class = instances.pred_classes == category_id
            object_masks = instances.pred_masks[is_class]
            object_masks = object_masks.detach().cpu().numpy() # [N,img_size,img_size]
            object_bboxes_xyxy = instances[is_class].pred_boxes.tensor.detach().cpu().numpy() # [N,4]
            assert object_masks.shape[0] == object_bboxes_xyxy.shape[0]

            # if there are no objects in image, filtering is applied
            if object_masks.shape[0] == 0:
                with open(filtered_save_pth, "w") as wf:
                    wf.write("NO TARGET OBJECT")
                if settings['verbose']:
                    print(f"\tFiltering {filtered_save_pth} since there are no corresponding object mask (for given target category)")
                continue
            
            # if there are multiple objects in image, filtering is applied
            elif object_masks.shape[0] >= 2 and settings['skip_multi_object_image']:
                with open(filtered_save_pth, "w") as wf:
                    wf.write("MORE THAN 2 TARGET OBJECTS")
                if settings['verbose']:
                    print(f"\tFiltering {filtered_save_pth} since there are more than 2 object masks (for given target category)")
                continue

            ## checking object-bbox & human-bbox overlap and filtering
            # find overlapping bbox (criteria: intersection over smaller bbox area)
            overlap_bbox = []
            
            # iterate for all humans
            to_filter = False
            for human_idx in range(num_humans):
                human_bbox_xyxy = human_parameters['body_bbox_list_xyxy'][human_idx]
                for target_obj_idx in range(object_bboxes_xyxy.shape[0]):
                    target_obj_bbox_xyxy = object_bboxes_xyxy[target_obj_idx]
                    iou, human_bbox_area, obj_bbox_area, intersection_area = bbox_iou(human_bbox_xyxy, target_obj_bbox_xyxy, return_areas_too=True)
                    if intersection_area / obj_bbox_area > settings['bbox_overlap_thres']:
                        overlap_bbox.append(target_obj_idx)
            
                # if there are no restrictions for number of object-bboxes that overlap with human-bbox, use all object-bboxes
                if len(settings['allowed_bbox_num_per_human']) == 0:
                    object_mask = np.any(object_masks, axis=0)

                # if number of object-bboxes that overlap with human-bbox is not allowed, filtering is applied
                elif len(overlap_bbox) not in settings['allowed_bbox_num_per_human']:
                    to_filter = True

                # if number of object-bboxes that overlap with human-bbox is allowed, merge the overlapping object-bboxes (not all)
                else:
                    object_mask = np.any(object_masks[overlap_bbox], axis=0)

            # apply filtering
            if to_filter:
                with open(filtered_save_pth, "w") as wf:
                    wf.write("MORE THAN 2 VALID BBOX PER PERSON")
                if settings['verbose']:
                    print(f"\tFiltering {filtered_save_pth} since number of object-bbox that overlaps with human-bbox: {len(overlap_bbox)} is not allowed in: {settings['allowed_bbox_num_per_human']}")
                continue

            ## checking human 2D keypoints and filtering
            # retrieve human parameters (for checking consistency of results)
            mocap_output_list = human_parameters["mocap_output_list"]
                    
            ## checking coco-keypoints and filtering
            # load coco-keypoint prediction results (if we do filter)
            if len(settings['keypoints_that_must_exist']) > 0 or settings['keypoints_that_must_not_exist'] > 0:
                keypoint_coco_pth = f"./{settings['keypoint_dir']}/{settings['ldm_model_key']}/{category}/{prompt}/{augprompt}/{_id}.pickle"
                with open(keypoint_coco_pth, "rb") as handle:
                    keypoints_per_human = pickle.load(handle)
                assert len(keypoints_per_human) == len(mocap_output_list) # check the consistency via number of humans

            # iterate for all humans
            to_filter = False
            for human_idx in range(num_humans):
                # if all of elements in the "subset of coco-keypoints that 'one of them must exist'" has lower confidence than threshold, filtering is applied
                for query_idx_tuple in settings['keypoints_that_must_exist']:
                    # note that at least one joint in 'query_idx_tuple' must have higher confidence than threshold
                    all_must_exist_kp_under_thres = True
                    for query_idx in query_idx_tuple:
                        x_img, y_img, confidence = keypoints_per_human[human_idx]['keypoints'][query_idx][:].tolist()
                        if confidence >= settings['keypoint_confidence_threshold']:
                            all_must_exist_kp_under_thres = False
                            break
                    if all_must_exist_kp_under_thres:
                        to_filter = True
                        break

                # if all of elements in the "subset of coco-keypoints that 'one of them must-not exist'" has higher confidence than threshold, filtering is applied
                for query_idx_tuple in settings['keypoints_that_must_not_exist']:
                    # note that at least one joint in 'query_idx_tuple' must have higher confidence than threshold
                    all_mustnot_exist_kp_over_thres = True
                    for query_idx in query_idx_tuple:
                        x_img, y_img, confidence = keypoints_per_human[human_idx]['keypoints'][query_idx][:].tolist()
                        if confidence < settings['keypoint_confidence_threshold']:
                            all_mustnot_exist_kp_over_thres = False
                            break
                    if all_mustnot_exist_kp_over_thres:
                        to_filter = True
                        break

            # apply filtering
            if to_filter:
                with open(filtered_save_pth, "w") as wf:
                    wf.write("FILTERED VIA COCO KEYPOINTS")
                if settings['verbose']:
                    print(f"\tFiltering {filtered_save_pth} via 2D COCO-keypoint!!!")
                continue
                                
            ## after all these steps of filtering, sample is retained. add to the number for tracking
            num_retained_in_prompt += 1
            num_retained_in_category += 1

            ## save as 'png' format if not filtered
            with open(retained_save_pth, "w") as wf:
                wf.write(f"useful_in_category:{num_retained_in_category:06}/useful_in_prompt:{num_retained_in_prompt:06}")


def filtering(cps, args):
    # prepare keypoint filters
    if args.keypoint_filter_nums is None:
        if args.verbose:
            print("'keypoint_filter_nums' not specified: Running for all keypoint filters!!!")
        args.keypoint_filter_nums = sorted(list(KEYPOINT_FILTERS.keys()))
    else:
        pass

    ## creating settings first
    keypoint_filter_nums = sorted(list(KEYPOINT_FILTERS.keys()))
    for kpfilter_idx in keypoint_filter_nums:
        # only save settings
        apply_filters(
            create_settings_only=True,
            cps=None,
            ldm_model_key=args.ldm_model_key,
            image_dir=args.image_dir,
            seg_dir=args.seg_dir,
            human_dir=args.human_dir,
            save_dir=args.save_dir,
            keypoint_dir=args.keypoint_dir,
            use_template=args.use_template,
            keypoints_which = KEYPOINT_FILTERS[kpfilter_idx],
            filter_multi_person_object_methods = args.filter_multi_person_object_methods,
            bbox_overlap_thres=args.bbox_overlap_thres,
            keypoint_confidence_threshold=args.keypoint_confidence_threshold,
            skip_done=args.skip_done,
            verbose=args.verbose,
        )

    ## do filtering
    keypoint_filter_nums = args.keypoint_filter_nums
    for kpfilter_idx in keypoint_filter_nums:
        # Run filtering
        apply_filters(
            create_settings_only=False,
            cps=cps,
            ldm_model_key=args.ldm_model_key,
            image_dir=args.image_dir,
            seg_dir=args.seg_dir,
            human_dir=args.human_dir,
            save_dir=args.save_dir,
            keypoint_dir=args.keypoint_dir,
            use_template=args.use_template,
            keypoints_which = KEYPOINT_FILTERS[kpfilter_idx],
            filter_multi_person_object_methods = args.filter_multi_person_object_methods,
            bbox_overlap_thres=args.bbox_overlap_thres,
            keypoint_confidence_threshold=args.keypoint_confidence_threshold,
            skip_done=args.skip_done,
            verbose=args.verbose,
        )
            

if __name__ == "__main__":
    ## Arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--categories", nargs="+")
    parser.add_argument("--ldm_model_key", type=str, default=DEFAULT_LDM_MODEL_KEY)

    parser.add_argument("--image_dir", type=str, default=DEFAULT_IMAGE_DIR)
    parser.add_argument("--seg_dir", type=str, default=DEFAULT_SEGMENTATION_DIR)
    parser.add_argument("--human_dir", type=str, default=DEFAULT_HUMAN_DIR)
    parser.add_argument("--keypoint_dir", type=str, default=DEFAULT_KEYPOINT_DIR)
    parser.add_argument("--save_dir", type=str, default=DEFAULT_FILTERING_DIR)
    parser.add_argument("--use_template", action="store_true") # --> overrides 'seg_dir', 'human_dir', 'keypoint_dir'

    parser.add_argument("--keypoint_filter_nums", nargs="+", type=int, default=[DEFAULT_KEYPOINT_FILTER_NUM], choices=sorted(list(KEYPOINT_FILTERS.keys())))
    parser.add_argument("--filter_multi_person_object_methods", default=DEFAULT_FILTER_MULTI_PERSON_OBJECT_METHODS) # multi-person/object filtering
    parser.add_argument("--bbox_overlap_thres", type=float, default=DEFAULT_BBOX_OVERLAP_THRES) # bbox-overlap filtering
    parser.add_argument("--keypoint_confidence_threshold", type=float, default=DEFAULT_KEYPOINT_CONFIDENCE_THRES) # keypoint filtering confidence threshold

    parser.add_argument("--skip_done", action="store_true")
    parser.add_argument("--verbose", action="store_true")

    parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
    args = parser.parse_args()


    # load sorted (category / prompt) pairs
    cps = prepare_cps_from_dirs(args, image_dir=args.image_dir, use_filter_setting=False)

    # seed for reproducible generation
    seed_everything(args.seed)

    # arguments
    if args.verbose:
        print("===Arguments===")
        print(args)

    # filtering
    filtering(cps=cps, args=args)
scripts/aggregation/filter_dataset.py
jellyheadandrew/CHORUS
import os
from glob import glob
import argparse
from tqdm import tqdm

import pickle5 as pickle
from copy import deepcopy
from easydict import EasyDict

import torch
import cv2

import imports.frankmocap.mocap_utils.demo_utils as demo_utils

from constants.ldm import DEFAULT_LDM_MODEL_KEY
from constants.frankmocap import DEFAULT_EXCONF_FRANKMOCAP, DEFAULT_MINOVERLAP_FRANKMOCAP
from constants.metadata import DEFAULT_IMAGE_DIR, DEFAULT_SEGMENTATION_DIR, DEFAULT_HUMAN_DIR, DEFAULT_SEED

from utils.prepare_prompts import prepare_cps_from_dirs
from utils.prepare_frankmocap import prepare_frankmocap_regressor, prepare_frankmocap_visualizer
from utils.postprocess import process_bbox_mask, process_mocap_predictions, process_remove_none, process_remove_overlap, process_segmentation, bbox_xy_to_wh
from utils.reproducibility import seed_everything



def extract_smpl(hparams, img_pth, seg_pth, bodymocap, visualizer):
    # load image
    image_bgr=cv2.imread(img_pth)

    # load segmentation results
    with open(seg_pth, "rb") as handle:
        instances_orig=pickle.load(handle)
    instances_orig=instances_orig.to("cuda")


    # post-process all objects (excluding humans) before post-processing for humans 
    if hparams.remocc:
        # remove excessive occlusions for same categories
        instances = process_segmentation(
            deepcopy(instances_orig),
            minoverlap=hparams.minoverlap, 
            exconf=hparams.exconf,
            verbose=hparams.verbose,
        )
        # print progress
        if hparams.verbose:
            print("[Log] Removed excessive occlusions from image segmentation results.")
    else:
        instances = deepcopy(instances_orig)

    # find human predictions (2D)
    is_person = (instances.pred_classes == 0)
    human_indices = torch.tensor(list(range(len(instances.pred_classes))), device=is_person.device)[is_person]
    bboxes_person = (instances[is_person].pred_boxes.tensor.cpu().numpy()) # bbox type: xyxy
    masks_person = instances[is_person].pred_masks
    confidence_person = instances[is_person].scores.cpu().numpy() # ndarray of shape [num_human,]

    # if no humans were detected during segmentation
    if bboxes_person.shape[0] == 0:
        if hparams.verbose:
            print("No human detected --> ",img_pth)

        # save image with "no human" message
        image_bgr = cv2.putText(
            image_bgr, 
            text="NO HUMAN DETECTED", 
            org = (image_bgr.shape[0]-150, image_bgr.shape[0]-20),
            fontFace=cv2.FONT_HERSHEY_COMPLEX_SMALL,
            fontScale=0.5,
            color=(0,0,255))
        return None, image_bgr

    # if humans were detected during segmentation, post-process bbox adaptable for human-prediction model
    body_bbox_list = bbox_xy_to_wh(bboxes_person)  # numpy ndarray, shape: [num_human, 4], xyxy -> xywh
    confidence_list = confidence_person.tolist()

    # calculate the overlapping bboxes and remove if too much is overlapping
    if hparams.remocc:
        keepidx=process_remove_overlap(
            bbox_list=body_bbox_list.tolist(), 
            confidence_list=confidence_list, 
            minoverlap=hparams.minoverlap, 
            exconf=hparams.exconf
        )
    else:
        keepidx=list(range(len(body_bbox_list)))
    
    # if there are no human remaining, return image only
    if len(keepidx) == 0:
        if hparams.verbose:
            print("After removing overlapping, no human exists.")
        return None, image_bgr


    # regress body pose
    mocap_output_list = bodymocap.regress(image_bgr, body_bbox_list[keepidx])
    
    # if erroneous mocap-output exists, remove (cases include extremely small bboxes)
    if None in mocap_output_list:
        if hparams.verbose:
            print(f"Erroneous human prediction exists for {img_pth}\n--> Removing Erroneous")
    mocap_output_list, final_keepidx = process_remove_none(mocap_output_list, keepidx)
    
    # if there are no humans remaining, return image only
    if len(mocap_output_list) == 0:
        if hparams.verbose:
            print("After removing erroneous, no human exists.")
        return None, image_bgr

    # post-process remaining human bboxes & masks
    final_bboxes_person, final_masks_person = process_bbox_mask(bboxes_person, masks_person, final_keepidx)

    # extract SMPL mesh for rendering (vertices, faces in image space)
    pred_mesh_list = demo_utils.extract_mesh_from_output(mocap_output_list)

    # render SMPL mesh on image
    if visualizer is not None:
        res_img = visualizer.visualize(
            image_bgr,
            pred_mesh_list = pred_mesh_list, 
            body_bbox_list = body_bbox_list[final_keepidx])
    else:
        res_img = None

    # process mocap predictions to save (NOTE: order-changing happens inside this function)
    human_parameters = process_mocap_predictions( # boxes in human_parameters['bboxes'] --> Shape: (N,4), xyxy format
        mocap_predictions=mocap_output_list,
        bboxes=final_bboxes_person, # xyxy
        masks=final_masks_person,
        image_size=max(image_bgr.shape)
    )
    for key in human_parameters:
        human_parameters[key]=human_parameters[key].to("cpu")

    # direct outputs
    human_parameters['mocap_output_list'] = mocap_output_list
    human_parameters['body_bbox_list'] = body_bbox_list[final_keepidx] # xywh format, (N,4)
    human_parameters['body_bbox_list_xyxy'] = bboxes_person[final_keepidx] # xyxy format, (N,4)
    human_parameters['confidence_list'] = confidence_person[final_keepidx] # (N,)

    # original SMPL outputs
    human_parameters['rotmat'] = torch.stack([torch.tensor(pred_output['pred_rotmat'][0]) for pred_output in mocap_output_list], dim=0)
    human_parameters['betas'] = torch.stack([torch.tensor(pred_output['pred_betas'][0]) for pred_output in mocap_output_list], dim=0)
    
    # SMPL joints (not permutated)
    human_parameters['joint_vertices_smpl'] = torch.stack([torch.tensor(pred_output['pred_joints_smpl'][0]) for pred_output in mocap_output_list], dim=0)

    # re-order (left-bbox to right-bbox)
    indices = human_parameters.pop('indices')
    for key in ['rotmat', 'betas', 'joint_vertices_smpl']:
        human_parameters[key] = human_parameters[key][indices]
        
    # meshes & bboxes
    human_parameters['pred_mesh_list'] = pred_mesh_list
    human_parameters['body_bbox_list'] = body_bbox_list[final_keepidx]

    # indices of removed humans in original instances
    removed_human_idx=list(range(len(human_indices)))
    for survived_idx in final_keepidx:
        removed_human_idx.remove(survived_idx)

    # auxiliary information
    human_parameters['auxiliary'] = dict()
    human_parameters['auxiliary']['instances_including_object_original'] = instances_orig.to("cpu") # original instances
    human_parameters['auxiliary']['instances_including_object'] = instances.to("cpu") # set of instances where all humans and post-processed objects exist
    human_parameters['auxiliary']['total_human_indices'] = human_indices # human indices used to query 'instances including object'
    human_parameters['auxiliary']['removed_human_indices']=human_indices[removed_human_idx].detach().cpu().numpy() # indices of removed humans in original instances
    human_parameters['auxiliary']['not_removed_human_indices']=human_indices[final_keepidx].detach().cpu().numpy() # indices of remaining humans in original instances
    
    return human_parameters, res_img


def human_prediction(
        cps,
        ldm_model_key,
        image_dir,
        seg_dir,
        save_dir,
        exconf,
        minoverlap,
        remocc,
        use_visualizer,
        skip_done,
        verbose,
    ):
    
    # hyperparameters
    hparams = EasyDict(dict(exconf=exconf,minoverlap=minoverlap,remocc=remocc,verbose=verbose))

    # prepare bodymocap model
    bodymocap = prepare_frankmocap_regressor()
    
    # prepare visualizer to render human    
    visualizer = prepare_frankmocap_visualizer() if use_visualizer else None

    # clear screen if not verbose
    if not verbose: os.system("clear")
    
    # run human prediction for all (category, prompt) pairs
    for category, prompt in cps:
        # all image path
        img_pths = sorted(list(glob(f"./{image_dir}/{ldm_model_key}/{category}/{prompt}/*/*.png")))

        # run for all images
        for img_pth in tqdm(img_pths, desc=f"Predicting Humans for [Category '{category}' / Prompt '{prompt}']"):
            # print
            if verbose:
                print("Processing: ", img_pth)
            
            # augprompt
            augprompt = img_pth.split("/")[-2]
            # image id
            _id, _ext = img_pth.split("/")[-1].split(".")            
            # result-save-directory
            result_save_dir = f"./{save_dir}/{ldm_model_key}/{category}/{prompt}/{augprompt}"
            os.makedirs(result_save_dir, exist_ok=True)
            # result-save paths
            img_save_pth = os.path.join(result_save_dir, f"{_id}.png")
            human_save_pth = os.path.join(result_save_dir, f"{_id}.pickle")

            # skip if already done
            if use_visualizer and os.path.exists(img_save_pth) and os.path.exists(human_save_pth) and skip_done:
                if verbose:
                    print("\tContinueing since ALREADY DONE!!", human_save_pth)
                continue
            if not use_visualizer and os.path.exists(human_save_pth) and skip_done:
                if verbose:
                    print("\tContinueing since ALREADY DONE!!", human_save_pth)
                continue

            # load segmentation results
            seg_pth = img_pth.replace(".png", ".pickle").replace(image_dir, seg_dir)
            if not os.path.exists(seg_pth):
                if verbose:
                    print("\tContinueing because there is no segmentation result: ", seg_pth)
                continue

            # extract human (as SMPL) with post-processing / (optional) render SMPL human on image
            human_params, rendered_img = extract_smpl(
                hparams=hparams,
                img_pth=img_pth, 
                seg_pth=seg_pth, 
                bodymocap=bodymocap,
                visualizer=visualizer
            )
            # error cases
            assert not (human_params is None and rendered_img is None), "error case during human prediction"
            
            # if no human was detected
            if human_params is None:
                cv2.imwrite(img_save_pth, rendered_img)
                with open(human_save_pth, "wb") as handle:
                    pickle.dump("NO HUMANS", handle, protocol=pickle.HIGHEST_PROTOCOL)
                    
            # if human was detected
            else:
                # save human-rendered image if it exists 
                if use_visualizer:
                    imglength=rendered_img.shape[0]
                    cv2.imwrite(img_save_pth, rendered_img[:,imglength:])

                # save SMPL results
                with open(human_save_pth, "wb") as handle:
                    pickle.dump(human_params, handle, protocol=pickle.HIGHEST_PROTOCOL)


if __name__ == "__main__":
    ## Arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--categories", nargs="+")
    parser.add_argument("--ldm_model_key", type=str, default=DEFAULT_LDM_MODEL_KEY)

    parser.add_argument("--image_dir", type=str, default=DEFAULT_IMAGE_DIR)
    parser.add_argument("--seg_dir", type=str, default=DEFAULT_SEGMENTATION_DIR)
    parser.add_argument("--save_dir", type=str, default=DEFAULT_HUMAN_DIR)

    parser.add_argument("--exconf", type=float, default=DEFAULT_EXCONF_FRANKMOCAP) # exception confidence (maximum)
    parser.add_argument("--minoverlap", type=float, default=DEFAULT_MINOVERLAP_FRANKMOCAP) # minimum overlap ratio
    parser.add_argument("--disable_remocc", action="store_true")
    parser.add_argument("--use_visualizer", action="store_true")

    parser.add_argument("--skip_done", action="store_true")
    parser.add_argument("--verbose", action="store_true")
    
    parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
    args = parser.parse_args()



    # load sorted (category / prompt) pairs
    cps = prepare_cps_from_dirs(args, image_dir=args.image_dir, use_filter_setting=False)

    # seed for reproducible generation
    seed_everything(args.seed)

    # arguments
    if args.verbose:
        print("===Arguments===")
        print(args)

    # run human prediction
    human_prediction(
        cps=cps,
        ldm_model_key=args.ldm_model_key,
        image_dir=args.image_dir,
        seg_dir=args.seg_dir,
        save_dir=args.save_dir,
        exconf=args.exconf,
        minoverlap=args.minoverlap,
        remocc=not args.disable_remocc,
        use_visualizer=args.use_visualizer,
        skip_done=args.skip_done,
        verbose=args.verbose,
    )
scripts/aggregation/predict_humans.py
jellyheadandrew/CHORUS
import os
from glob import glob
import argparse
import pickle5 as pickle
from tqdm import tqdm

import numpy as np
import cv2

from mmpose.apis import inference_top_down_pose_model, init_pose_model, vis_pose_result, process_mmdet_results
from mmdet.apis import inference_detector, init_detector

from constants.ldm import DEFAULT_LDM_MODEL_KEY
from constants.keypoint_detection import MMPOSE_DETECTOR_SETTINGS, DEFAULT_KEYPOINT_THRESHOLD, DEFAULT_KEYPOINT_BBOX_THRESHOLD, DEFAULT_MMPOSE_DETECTOR_KEY
from constants.metadata import DEFAULT_IMAGE_DIR, DEFAULT_HUMAN_DIR, DEFAULT_KEYPOINT_DIR, DEFAULT_SEED

from utils.prepare_prompts import prepare_cps_from_dirs
from utils.reproducibility import seed_everything


def keypoint_prediction(
        cps,
        ldm_model_key,
        thres, # only for visualization
        skip_done,
        verbose,
        image_dir,
        human_dir, 
        save_dir,
        bbox_threshold,
        pose_config,
        pose_checkpoint,
    ):
    # load pretrained pose model
    pose_model = init_pose_model(pose_config, pose_checkpoint, device = "cuda")
    
    # clear screen if not verbose
    if not verbose: os.system("clear")
    
    # run keypoint prediction for all (category, prompt) pairs
    for category, prompt in cps:
        # all image path
        img_pths = sorted(list(glob(f"./{image_dir}/{ldm_model_key}/{category}/{prompt}/*/*.png")))
        for img_pth in tqdm(img_pths, desc=f"Predicting Keypoint Filters for [Category '{category}' / Prompt '{prompt}']"):
            # print
            if verbose:
                print("Processing: ", img_pth)

            # augprompt
            augprompt = img_pth.split("/")[-2]
            # image id
            _id, _ext = img_pth.split("/")[-1].split(".")
            # result-save directory
            result_save_dir = f"./{save_dir}/{ldm_model_key}/{category}/{prompt}/{augprompt}"
            os.makedirs(result_save_dir, exist_ok=True)
            # result-save paths
            img_save_pth = os.path.join(result_save_dir, f"{_id}.png")
            keypoint_save_pth = os.path.join(result_save_dir, f"{_id}.pickle")
    
            # skip if already done
            if os.path.exists(img_save_pth) and os.path.exists(keypoint_save_pth) and skip_done:
                if verbose:
                    print("\tContinueing since ALREADY DONE!!", keypoint_save_pth)
                continue
        
            # load human params
            human_pth = f"./{human_dir}/{ldm_model_key}/{category}/{prompt}/{augprompt}/{_id}.pickle"
            with open(human_pth, 'rb') as handle:
                human_parameters = pickle.load(handle)

            # if there is no human, return empty list
            if human_parameters == "NO HUMANS":
                if verbose:
                    print(f"\tNo humans for --> {human_pth}")
                person_results = []
            # if there is human, return bbox-list & confidence-list 
            else:
                # bbox-list / confidence-list
                body_bbox_list = human_parameters['body_bbox_list_xyxy'] # (shape: Nx4, xyxy)
                confidence_list = human_parameters['confidence_list'] # (shape: N)

                # conversion to mmdet format (Kx3, last column is confidence)
                bbox_conf_arr = np.concatenate([body_bbox_list, np.expand_dims(confidence_list, axis=-1)], axis=-1)
                person_results = [None] * len(bbox_conf_arr)
                # stack to list for all people
                for i in range(len(bbox_conf_arr)):
                    person_results[i] = {'bbox': bbox_conf_arr[i]}
            

            # inference pose
            pose_results, returned_outputs = inference_top_down_pose_model(
                pose_model,
                img_pth,
                person_results,
                bbox_thr=bbox_threshold, # 0.0
                format='xyxy',
                dataset=pose_model.cfg.data.test.type
            )

            # render keypoint prediction to image
            vis_result = vis_pose_result(
                pose_model,
                img_pth,
                pose_results,
                kpt_score_thr=thres,
                dataset=pose_model.cfg.data.test.type,
                show=False
            )

            # save rendered image
            cv2.imwrite(img_save_pth, vis_result)
            # save keypoint prediction (as numpy file)
            with open(keypoint_save_pth, "wb") as handle:
                pickle.dump(pose_results, handle, protocol=pickle.HIGHEST_PROTOCOL)


if __name__ == "__main__":
    ## Arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--categories", nargs="+")
    parser.add_argument("--ldm_model_key", type=str, default=DEFAULT_LDM_MODEL_KEY)

    parser.add_argument("--image_dir", type=str, default=DEFAULT_IMAGE_DIR)
    parser.add_argument("--human_dir", type=str, default=DEFAULT_HUMAN_DIR)
    parser.add_argument("--save_dir", type=str, default=DEFAULT_KEYPOINT_DIR)
        
    parser.add_argument("--detector_setting",
        nargs="?",
        choices=sorted(list(MMPOSE_DETECTOR_SETTINGS.keys())),
        default=DEFAULT_MMPOSE_DETECTOR_KEY
    )
    parser.add_argument("--thres", type=float, default=DEFAULT_KEYPOINT_THRESHOLD) # keypoint confidence threshold for visualization
    parser.add_argument("--bbox_threshold", type=float, default=DEFAULT_KEYPOINT_BBOX_THRESHOLD) # bbox where [confidence>bbox_threshold] used only
    
    parser.add_argument("--skip_done", action="store_true")
    parser.add_argument("--verbose", action="store_true")

    parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
    args = parser.parse_args()

    # load sorted (category / prompt) pairs
    cps = prepare_cps_from_dirs(args, image_dir=args.image_dir, use_filter_setting=False)

    # seed for reproducible generation
    seed_everything(args.seed)

    # arguments
    if args.verbose:
        print("===Arguments===")
        print(args)
    
    # predict 2d keypoints
    keypoint_prediction(
        cps=cps,
        ldm_model_key=args.ldm_model_key,
        thres=args.thres, 
        skip_done=args.skip_done,
        verbose=args.verbose,
        image_dir=args.image_dir,
        human_dir=args.human_dir, 
        save_dir=args.save_dir,
        bbox_threshold=args.bbox_threshold,
        pose_config=MMPOSE_DETECTOR_SETTINGS[args.detector_setting][0],
        pose_checkpoint=MMPOSE_DETECTOR_SETTINGS[args.detector_setting][1],
    )
scripts/aggregation/predict_keypoints.py
jellyheadandrew/CHORUS
import os
from glob import glob
import argparse

import pickle5 as pickle
from tqdm import tqdm
from easydict import EasyDict

import math
import numpy as np
import torch
import torch.nn as nn

import cv2
import matplotlib.pyplot as plt

from constants.ldm import DEFAULT_LDM_MODEL_KEY
from constants.filtering import DEFAULT_FILTER_SETTING_NUM
from constants.perspective import DEFAULT_FOV, DEFAULT_STOP_THRES, DEFAULT_LR, DEFAULT_EPOCHS
from constants.metadata import DEFAULT_FILTERING_DIR, DEFAULT_PERSPECTIE_CAMERA_DIR, DEFAULT_SEED

from utils.prepare_prompts import prepare_cps_from_dirs
from utils.reproducibility import seed_everything
from utils.settings import load_filter_settings
from utils.perspective_conversion import apply_wp_to_p


def convert_to_perspective(cps, args):
    # clear screen if not verbose
    if not args.verbose: os.system("clear")

    ## iterate for all setting numbers
    for filter_setting_num in args.filter_setting_nums:
        # load filter_settings to arguments
        filter_settings = load_filter_settings(args, filter_setting_num)
    
        # run 'weak perspective->perspective' for all (category, prompt) pairs
        for category, prompt in cps:
            # all filtered images path
            filtered_img_pths = sorted(list(glob(f"./{filter_settings['filter_dir']}/{filter_settings['ldm_model_key']}/{category}/settings:{filter_setting_num:03}/{prompt}/*.png")))

            # all filtered image path
            for filtered_img_pth in tqdm(filtered_img_pths, desc=f"Converting WP->P for [Category '{category}' / Prompt '{prompt}' / Setting '{filter_setting_num:03}']"):
                # print
                if filter_settings['verbose']:
                    print("Processing: ", filtered_img_pth)

                # augprompt
                augprompt = filtered_img_pth.split("/")[-1].split(":")[0]
                # image id
                _id, _ext = filtered_img_pth.split("/")[-1].split(":")[1].split(".")
                # result-save directory
                result_save_dir = f"./{filter_settings['save_dir']}/{filter_settings['ldm_model_key']}/{category}/{prompt}"
                os.makedirs(result_save_dir, exist_ok=True)
                # result-save paths
                cam_save_pth = f"{result_save_dir}/{augprompt}:{_id}.pickle"

                # image pth / frankmocap pth
                img_pth = f"./{filter_settings['image_dir']}/{filter_settings['ldm_model_key']}/{category}/{prompt}/{augprompt}/{_id}.png"
                human_pth = f"./{filter_settings['human_dir']}/{filter_settings['ldm_model_key']}/{category}/{prompt}/{augprompt}/{_id}.pickle"
                
                # skip if already done
                if os.path.exists(cam_save_pth) and filter_settings['skip_done']:
                    if filter_settings['verbose']:
                        print("\tContinueing since ALREADY DONE!!", cam_save_pth)
                    continue

                ## convert 'weak perspective->perspective'
                persp_cam_list, human_parameters = apply_wp_to_p(
                    filter_settings=filter_settings, 
                    img_pth=img_pth, 
                    human_pth=human_pth,
                    mocap_output_list=None
                )
                
                # save perspective camera
                with open(cam_save_pth, "wb") as handle:
                    if args.save_with_human_params:
                        tosave = {
                            'human_parameters': human_parameters,
                            'persp_cam_list': persp_cam_list
                        }
                        pickle.dump(tosave, handle, protocol=pickle.HIGHEST_PROTOCOL)
                    else:
                        tosave = {
                            'persp_cam_list': persp_cam_list
                        }
                        pickle.dump(tosave, handle, protocol=pickle.HIGHEST_PROTOCOL)


if __name__ == "__main__":
    ## Arguments
    parser = argparse.ArgumentParser()
    # important
    parser.add_argument("--categories", nargs="+")
    parser.add_argument("--prompts", nargs="+")
    parser.add_argument("--ldm_model_key", type=str, default=DEFAULT_LDM_MODEL_KEY)
    parser.add_argument("--filter_setting_nums", nargs="+", type=int, default=[DEFAULT_FILTER_SETTING_NUM])

    parser.add_argument("--filter_dir", type=str, default=DEFAULT_FILTERING_DIR)
    parser.add_argument("--save_dir", type=str, default=DEFAULT_PERSPECTIE_CAMERA_DIR)

    parser.add_argument("--use_openpose_joints", action="store_true")
    parser.add_argument("--ignore_joints_outside", action="store_true")
    parser.add_argument("--freeze_foc", action="store_true")
    parser.add_argument("--fovy", type=float, nargs="?", default=DEFAULT_FOV)
    parser.add_argument("--lr", type=float, default=DEFAULT_LR)
    parser.add_argument("--thres", type=float, default=DEFAULT_STOP_THRES)
    parser.add_argument("--epochs", type=int, default=DEFAULT_EPOCHS)
    parser.add_argument("--save_with_human_params", action="store_true")

    parser.add_argument("--skip_done", action="store_true")
    parser.add_argument("--verbose", action="store_true")
    
    parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
    args = parser.parse_args()

    # load sorted (category / prompt) pairs
    cps = prepare_cps_from_dirs(args, image_dir=args.filter_dir, use_filter_setting=True)

    # seed for reproducible generation
    seed_everything(args.seed)
    
    # arguments
    if args.verbose:
        print("===Arguments===")
        print(args)
    
    # weak perspective -> perspective
    convert_to_perspective(cps=cps, args=args)
scripts/aggregation/predict_perspective.py
jellyheadandrew/CHORUS
import os
from glob import glob
import pickle5 as pickle
import argparse
from easydict import EasyDict
from tqdm import tqdm

import torch
import cv2

from detectron2 import model_zoo
from detectron2.utils.logger import setup_logger
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer, ColorMode
from detectron2.data import MetadataCatalog
from detectron2.structures.instances import Instances
from detectron2.projects import point_rend

from constants.datasets import LVIS_METADATA_PTH
from constants.ldm import DEFAULT_LDM_MODEL_KEY
from constants.segmentation import COCO_SEG_CONFIG_PTH, COCO_SEG_WEIGHTS_PTH, LVIS_SEG_CONFIG_PTH, LVIS_SEG_WEIGHTS_PTH, DEFAULT_SEGMENTATION_THRESHOLD
from constants.metadata import DEFAULT_IMAGE_DIR, DEFAULT_SEGMENTATION_DIR, DEFAULT_SEED

from utils.prepare_prompts import prepare_cps_from_dirs
from utils.reproducibility import seed_everything




def segmentation_coco(
        cps, 
        ldm_model_key,
        image_dir, 
        save_dir,
        thres, 
        skip_done, 
        verbose,
    ):

    # make arguments as easydict
    hparams = EasyDict(dict(threshold=thres))

    # setup coco metadata
    setup_logger()
    coco_metadata = MetadataCatalog.get("coco_2017_val")

    # get segmentation model (coco)
    cfg = get_cfg()
    point_rend.add_pointrend_config(cfg) # --> Add PointRend-specific config
    cfg.merge_from_file(COCO_SEG_CONFIG_PTH)
    cfg.MODEL.WEIGHTS = COCO_SEG_WEIGHTS_PTH
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = hparams.threshold
    cfg.MODEL.DEVICE = "cuda"

    # get segmentation model
    SegmentationModel=DefaultPredictor(cfg)
    
    # clear screen if not verbose
    if not verbose: os.system("clear")
    
    # run segmentation for images
    for category, prompt in cps:
        # all image paths
        img_pths = sorted(list(glob(f"./{image_dir}/{ldm_model_key}/{category}/{prompt}/*/*.png")))
        
        # iterate through image path
        for img_pth in tqdm(img_pths, desc=f"Predicting Segmentations for [Category '{category}' / Prompt '{prompt}']"):
            if verbose: print("Processing: ", img_pth)
            
            # augprompt
            augprompt = img_pth.split("/")[-2]
            # Image ID
            _id, _ext = img_pth.split("/")[-1].split(".")
            # Create result-save directory
            result_save_dir = f"./{save_dir}/{ldm_model_key}/{category}/{prompt}/{augprompt}"
            os.makedirs(result_save_dir, exist_ok=True)

            # save path (segmentation results, visualized image)
            img_save_pth = os.path.join(result_save_dir, f"{_id}.{_ext}")
            seg_save_pth = os.path.join(result_save_dir, f"{_id}.pickle")
            
            # skip if exists
            if os.path.exists(img_save_pth) and os.path.exists(seg_save_pth) and skip_done:
                if verbose: print("Already processed: ", img_save_pth)
                continue
    
            # read image
            im=cv2.imread(img_pth)

            # run segmentation
            outputs=SegmentationModel(im)
            instances=outputs["instances"]

            # visualize segmentation on image
            v = Visualizer(im[:,:,::-1], coco_metadata, scale=0.5, instance_mode=ColorMode.IMAGE_BW)
            rend_result=v.draw_instance_predictions(instances.to("cpu")).get_image()
            
            # save visualized image
            cv2.imwrite(img_save_pth, rend_result[:,:,::-1])
            
            # save segmentation information
            with open(seg_save_pth, "wb") as handle:
                pickle.dump(instances.to("cpu"), handle, protocol=pickle.HIGHEST_PROTOCOL)






def segmentation_lvis(
        cps,
        ldm_model_key,
        image_dir, 
        save_dir,
        thres, 
        skip_done, 
        verbose,
    ):

    # make arguments as easydict
    hparams = EasyDict(dict(threshold=thres))

    # setup coco metadata
    setup_logger()
    coco_metadata = MetadataCatalog.get("coco_2017_val")

    # setup lvis metadata
    with open(LVIS_METADATA_PTH, "rb") as handle:
        lvis_metadata = pickle.load(handle)

    # get segmentation model (lvis)
    cfg = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file(LVIS_SEG_CONFIG_PTH))
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(LVIS_SEG_WEIGHTS_PTH)
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = hparams.threshold
    cfg.MODEL.DEVICE = "cuda"

    # get segmentation model
    SegmentationModel=DefaultPredictor(cfg)
    
    # get segmentation model (coco) to segment human only
    cfg_coco = get_cfg()
    point_rend.add_pointrend_config(cfg_coco) # --> Add PointRend-specific config
    cfg_coco.merge_from_file(COCO_SEG_CONFIG_PTH)
    cfg_coco.MODEL.WEIGHTS = COCO_SEG_WEIGHTS_PTH
    cfg_coco.MODEL.ROI_HEADS.SCORE_THRESH_TEST = hparams.threshold
    cfg_coco.MODEL.DEVICE = "cuda"

    # get segmentation model (only for human)
    HumanSegmentationModel=DefaultPredictor(cfg_coco)

    # clear screen if not verbose
    if not verbose: os.system("clear")

    # run segmentation for images
    for category, prompt in cps:
        # all image paths
        img_pths = sorted(list(glob(f"./{image_dir}/{ldm_model_key}/{category}/{prompt}/*/*.png")))
        
        # iterate through image path
        for img_pth in tqdm(img_pths, desc=f"Category: {category} / Prompt: {prompt}"):
            if verbose: print("Processing: ", img_pth)
            # augprompt
            augprompt = img_pth.split("/")[-2]
            # Image ID
            _id, _ext = img_pth.split("/")[-1].split(".")
            # Create result-save directory
            result_save_dir = f"./{save_dir}/{ldm_model_key}/{category}/{prompt}/{augprompt}"
            os.makedirs(result_save_dir, exist_ok=True)

            # save path (segmentation results, visualized image)
            img_save_pth = os.path.join(result_save_dir, f"{_id}.{_ext}")
            seg_save_pth = os.path.join(result_save_dir, f"{_id}.pickle")
            
            # skip if exists
            if os.path.exists(img_save_pth) and os.path.exists(seg_save_pth) and skip_done:
                if verbose: print("Already processed: ", img_save_pth)
                continue
    
            # read image
            im=cv2.imread(img_pth)

            # run segmentation
            instances_lvis=SegmentationModel(im)['instances'] # for object
            instances_coco=HumanSegmentationModel(im)['instances'] 
            instances_coco=instances_coco[instances_coco.pred_classes == 0] # for human
            instances = Instances.cat([instances_lvis, instances_coco]).to('cpu')
            
            # indices for object (lvis)
            is_lvis = torch.tensor([True] * len(instances_lvis) + [False] * len(instances_coco))
            instances.is_lvis = is_lvis

            # visualize segmentation on image
            v = Visualizer(im[:,:,::-1], lvis_metadata, scale=0.5, instance_mode=ColorMode.IMAGE_BW)
            rend_result=v.draw_instance_predictions(instances.to("cpu")).get_image()
            
            # save visualized image
            cv2.imwrite(img_save_pth, rend_result[:,:,::-1])
            
            # save segmentation information
            with open(seg_save_pth, "wb") as handle:
                pickle.dump(instances, handle, protocol=pickle.HIGHEST_PROTOCOL)


def segmentation(mode, **kwargs):
    assert mode in ["coco", "lvis"], f"Segmentation Mode: {mode} --> Not implemented..."
    # coco categories
    if mode == "coco":
        segmentation_coco(**kwargs)
    # lvis categories
    if mode == "lvis":
        segmentation_lvis(**kwargs)


if __name__ == "__main__":
    ## Arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--categories", nargs="+")
    parser.add_argument("--ldm_model_key", type=str, default=DEFAULT_LDM_MODEL_KEY)

    parser.add_argument("--image_dir", type=str, default=DEFAULT_IMAGE_DIR)
    parser.add_argument("--save_dir", type=str, default=DEFAULT_SEGMENTATION_DIR)

    parser.add_argument("--mode", type=str, choices=["coco", "lvis"])
    parser.add_argument("--threshold", type=float, default=DEFAULT_SEGMENTATION_THRESHOLD)

    parser.add_argument("--skip_done", action="store_true")
    parser.add_argument("--verbose", action="store_true")

    parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
    args = parser.parse_args()

    # load sorted (category / prompt) pairs
    cps = prepare_cps_from_dirs(args, image_dir=args.image_dir, use_filter_setting=False)

    # seed for reproducible generation
    seed_everything(args.seed)

    # print information
    if args.verbose:
        print("===Arguments===")
        print(args)

    # run segmentation
    segmentation(
        mode=args.mode,
        cps=cps, 
        ldm_model_key=args.ldm_model_key,
        image_dir=args.image_dir, 
        save_dir=args.save_dir,
        thres=args.threshold, 
        skip_done=args.skip_done, 
        verbose=args.verbose,
    )
scripts/aggregation/predict_segmentations.py
jellyheadandrew/CHORUS
import os
from glob import glob
from copy import deepcopy
import argparse
import pickle5 as pickle
import json
from tqdm import tqdm

import torch
import numpy as np

import cv2
import matplotlib.pyplot as plt

from constants.ldm import DEFAULT_LDM_MODEL_KEY
from constants.filtering import DEFAULT_FILTER_SETTING_NUM
from constants.aggr_settings.quant import QUANT_AGGR_SETTINGS
from constants.datasets import EXTENDED_COCO_EFT_VAL_SAVE_PTH, COCO_VAL_IMAGE_DIR
from constants.evaluation import DEFAULT_DOWNSCALE_LENGTH, DEFAULT_NUM_ENFORCE_MULTIVIEW, DEFAULT_PRECISON_THRESHOLDS
from constants.metadata import DEFAULT_TOTAL_AGGREGATION_NAME, DEFAULT_SEED, DEFAULT_AGGREGATION_DIR, DEFAULT_QUANT_EVAL_DIR

from utils.reproducibility import seed_everything
from utils.prepare_prompts import prepare_cps_from_dirs, get_unique_categories
from utils.misc import load_category_id
from utils.settings import aggr_save_name
from utils.aggregation.aggregator import Aggregator
from utils.aggregation.core import prepare_aggr_core
from utils.evaluation import compute_pap, compute_mean_pap



def evaluate_pap(
        cps,
        ldm_model_key,
        aggr_dir,
        save_dir,
        aggr_settings,
        filter_setting_comb,
        eft_dataset_pth,
        num_multiview_thres,
        specified_image_ids,
        downscale_length,
        thresholds,
        skip_done,
        verbose
    ):
    # print progress
    print("Loading COCO-EFT Dataset for Quantitative Evaluation...")
    
    ## load extended COCO-EFT data
    # open file
    with open(eft_dataset_pth, "rb") as handle: test_dataset = pickle.load(handle)
    # data-to-use
    all_test_data = test_dataset['eft_data']
    
    # use statistics & only use categories with number-of-images over 'args.num_multiview_thres'
    stats = test_dataset['stats']

    ## load metadata of aggregation-results
    # 'aggregation-results' save name
    aggr_result_save_name =  aggr_save_name(
        filter_setting_comb=filter_setting_comb, 
        aggr_setting_name=aggr_settings['NAME']
    )

    ## load aggregation settings and prepare 'aggregation_core' --> overwritten with pretrained-results later
    print("Preparing Aggregator Core...")
    aggr_core = prepare_aggr_core(verbose=verbose, **aggr_settings)

    # clear screen if verbose
    print("Preparation Complete!")
    if not verbose: os.system("clear")

    ## run evaluation for all (category, prompt) pairs
    for category, prompt in sorted(cps):
        # if number-of-images are insufficient for 'multiview evaluation', skip
        if category not in stats.keys():
            if verbose: print(f"not using category '{category}' since not in test-dataset")
            continue
        if stats[category] < num_multiview_thres:
            if verbose: print(f"not using category '{category}' since the number of images in test-dataset under {num_multiview_thres}")
            continue

        ## declare aggregator
        aggregator = Aggregator(core=deepcopy(aggr_core))
        
        ## load aggregation metadata (filter-settings, aggr-core)
        # metadata path
        metadata_pth = f"{aggr_dir}/{ldm_model_key}/{category}/{prompt}/{aggr_result_save_name}.json"
        # load metadata
        with open(metadata_pth, 'r') as rf: metadata = json.load(rf)
        # retrieve filter-settings
        filter_settings = metadata['filter_settings']

        ## load aggregation results
        # aggregation path
        aggr_pth = f"{aggr_dir}/{ldm_model_key}/{category}/{prompt}/{aggr_result_save_name}.pickle"
        # load aggregation to 'aggregator'
        aggregator.load(load_pth=aggr_pth)

        ## preapre for evaluation
        # category id
        category_id = load_category_id(category)

        # placeholder for saving quantitative-evaluation results
        evaluate_results = {
            'aggr_pth': aggr_pth,
            'category': category,
            'results': dict()
        }

        ## prepare save-path
        # save-directory and path
        quant_eval_save_dir = f"./{save_dir}/{ldm_model_key}/{category}"
        os.makedirs(quant_eval_save_dir, exist_ok=True)
        quant_eval_save_pth = f"{quant_eval_save_dir}/{prompt}:{aggr_result_save_name}.json"
        # skip if already exists
        if os.path.exists(quant_eval_save_pth) and skip_done:
            if verbose:
                print("\tContinueing since (evaluation) ALREADY DONE!!", quant_eval_save_pth)
            continue

        ## iterate for all test-eft-data
        num_done = 0
        pbar = tqdm(all_test_data, desc=f"Quantitative Evaluation for [Category '{category}' / Settings '{aggr_result_save_name}' / Num-Eval-Done ({num_done}/{stats[category]})]")
        for test_data in pbar:
            # if 'test_data' contains object category, test it
            if category in test_data['target_category_list']:
                test_data = deepcopy(test_data)
            else:
                continue
            
            # prepare input
            # placeholder for saving inputs
            _aggr_input = {}

            # load image
            img_name = test_data['imageName']
            img_id = img_name.replace(".jpg", "").replace("COCO_val2014_", "")
            img_pth = f'{COCO_VAL_IMAGE_DIR}/{img_name}'
            img = cv2.cvtColor(cv2.imread(img_pth), code=cv2.COLOR_BGR2RGB)
            
            # skip if 'specified_image_ids' did not come out
            if specified_image_ids is not None:
                if img_id not in specified_image_ids:
                    if verbose:
                        print(f"Continue '{img_id}' since 'specified_image_ids': {specified_image_ids} -> didn't come out")
                    continue 

            # load perspective camera
            persp_cam = test_data['persp_cam_list'][0]
            _aggr_input['cam_R'] = persp_cam['R']
            _aggr_input['cam_t'] = persp_cam['t']
            _aggr_input['cam_K'] = persp_cam['K']

            # load object mask
            _aggr_input['mask'] = test_data['mask'][category] # HxW
            H, W = _aggr_input['mask'].shape
            assert len(_aggr_input['mask'].shape) == 2

            # pose
            _aggr_input['pose'] = test_data['pred_aa'] # 1x72
            assert _aggr_input['pose'].shape == (1,72)

            # betas
            _aggr_input['betas'] = torch.zeros_like(test_data['pred_betas']) # 1x10
            assert _aggr_input['betas'].shape == (1,10)
            
            # preprocess
            aggr_input = aggregator.preprocess_input(cam_world_key='SMPL', **_aggr_input)
            
            # if evaluating, we must render the mask. 
            rendered_masks_per_thresh, target_mask_size = aggregator.render_discretized_distribution(
                cam_world_key='SMPL',
                cam_R = aggr_input['cam_R'],
                cam_t = aggr_input['cam_t'],
                cam_K = aggr_input['cam_K'],
                pose = aggr_input['pose'],
                betas = aggr_input['betas'],
                mask_shape = aggr_input['mask'].shape,
                downscale_length = downscale_length,
                use_normalized_K = False,
                thresholds = thresholds,
                verbose = verbose
            )
            
            # downsample target mask
            target_H, target_W = target_mask_size
            target_mask = cv2.resize(aggr_input['mask'], dsize=(target_W, target_H), interpolation=cv2.INTER_AREA)

            # compute PAP values
            PAP, PAP_strict, PAP_hoa, PAP_strict_hoa, target_area, target_area_hoa = compute_pap(
                pred=rendered_masks_per_thresh,
                target_mask=target_mask,
                verbose=verbose
            )

            # evaluate results
            evaluate_results['results'][img_name] = {
                'PAP': PAP,
                'PAP_strict': PAP_strict,
                'PAP_hoa': PAP_hoa,
                'PAP_strict_hoa': PAP_strict_hoa,
                'target_area': target_area,
                'target_area_hoa': target_area_hoa
            }

            # set description
            num_done += 1
            desc = f"Quantitative Evaluation for [Category '{category}' / Settings '{aggr_result_save_name}' / Num-Eval-Done ({num_done}/{stats[category]})]"
            pbar.set_description(desc)

        # check number-of-data used for testing
        assert num_done == stats[category]

        ## calculate 'mean-projective-average-precision'
        evaluate_results = compute_mean_pap(evaluate_results)

        ## save the evaluation results
        with open(quant_eval_save_pth, "w") as wf: json.dump(evaluate_results, wf, indent=1)


def run_quant_eval(cps, args):

    ## iterate for all 'args.aggr_setting_names'
    for aggr_setting_name in args.aggr_setting_names:
        # load aggregation setting
        aggr_settings = QUANT_AGGR_SETTINGS[aggr_setting_name]
        assert aggr_settings['NAME'] == aggr_setting_name, f"Mismatch in name for '{aggr_setting_name}'"

        # iterate for filter-settings-combinations
        for filter_setting_comb in args.filter_setting_nums:
            # print setting informations
            if args.verbose: print(f"Quant-Evaluating 'aggr_settings':{aggr_setting_name} / 'filter_setting_comb': {filter_setting_comb}")

            # filter_setting_comb as list of filter_setting_num
            filter_setting_comb = [int(i) for i in filter_setting_comb.split(",")]

            ## evaluate aggregation result with 'projective average precision' metrics and save results
            evaluate_pap(
                cps=cps,
                ldm_model_key=args.ldm_model_key,
                aggr_dir=args.aggr_dir,
                save_dir=args.save_dir,
                aggr_settings=aggr_settings,
                filter_setting_comb=filter_setting_comb,
                eft_dataset_pth=args.eft_dataset_pth,
                num_multiview_thres=args.num_multiview_thres,
                downscale_length=args.downscale_length,
                thresholds=args.thresholds,
                specified_image_ids=args.specified_image_ids,
                skip_done=args.skip_done,
                verbose=args.verbose
            )


if __name__ == "__main__":
    ## Arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--categories", nargs="+")
    parser.add_argument("--prompts", nargs="+", default=[DEFAULT_TOTAL_AGGREGATION_NAME])
    parser.add_argument("--ldm_model_key", type=str, default=DEFAULT_LDM_MODEL_KEY)

    parser.add_argument("--aggr_dir", type=str, default=DEFAULT_AGGREGATION_DIR)
    parser.add_argument("--save_dir", type=str, default=DEFAULT_QUANT_EVAL_DIR)

    parser.add_argument("--aggr_setting_names", nargs="+", type=str)
    parser.add_argument("--filter_setting_nums", nargs="+", type=str, default=[str(DEFAULT_FILTER_SETTING_NUM)]) # you can write as "--filter_settings '7,8' '9'"

    parser.add_argument("--eft_dataset_pth", type=str, default=EXTENDED_COCO_EFT_VAL_SAVE_PTH)
    parser.add_argument("--num_multiview_thres", type=int, default=DEFAULT_NUM_ENFORCE_MULTIVIEW)

    parser.add_argument("--downscale_length", type=int, default=DEFAULT_DOWNSCALE_LENGTH)
    parser.add_argument("--thresholds", nargs="+", type=float, default=DEFAULT_PRECISON_THRESHOLDS)

    parser.add_argument("--specified_image_ids", nargs="+", type=str)
    parser.add_argument("--skip_done", action="store_true")
    parser.add_argument("--verbose", action="store_true")

    parser.add_argument("--seed", default=DEFAULT_SEED)
    args = parser.parse_args()

    # load sorted (category / prompt) pairs
    cps = prepare_cps_from_dirs(args, image_dir=args.aggr_dir, use_filter_setting=False)

    # seed for reproducible generation
    seed_everything(args.seed)

    # arguments
    if args.verbose:
        print("===Arguments===")
        print(args)

    # evaluate PAP for 'trained-results'
    run_quant_eval(cps=cps, args=args)
scripts/evaluation/evaluate_pap.py
jellyheadandrew/CHORUS
import os
import argparse
from copy import deepcopy
from easydict import EasyDict
import pickle5 as pickle
from tqdm import tqdm
import json

import numpy as np
import torch

import cv2
import matplotlib.pyplot as plt

from pycocotools.coco import COCO

from imports.frankmocap.bodymocap.models import SMPL

from constants.frankmocap import BODY_MOCAP_SMPL_PTH
from constants.prepare_eft import DEFAULT_EFT_MIN_BBOX_OVERLAP, DEFAULT_EFT_REPROJECTION_ERROR_THRESH, DEFAULT_EFT_FOV, DEFAULT_EFT_LR, DEFAULT_EFT_JOINT_LOSS_THRES, DEFAULT_EFT_EPOCHS
from constants.datasets import COCO_VAL_ANNOTATION_PTH, COCO_VAL_IMAGE_DIR, COCO_EFT_VAL_PTH, COCO_CLASS_NAME2ID, EXTENDED_COCO_EFT_VAL_SAVE_PTH
from constants.metadata import DEFAULT_SEED

from utils.reproducibility import seed_everything
from utils.filtering import intersection_over_smaller_bbox
from utils.transformations import rotmat_to_axis_angle_batch
from utils.frankmocap_conversion import convert_smpl_to_img
from utils.perspective_conversion import apply_wp_to_p


def extend_coco_eft_dataset(args):
    ## prepare smpl model & coco-constants
    smpl = SMPL(BODY_MOCAP_SMPL_PTH, batch_size=1, create_transl=False)
    coco=COCO(COCO_VAL_ANNOTATION_PTH)
    person_category_id = 1

    ## open COCO-EFT (val) dataset
    # load eft-dataset
    with open(COCO_EFT_VAL_PTH, "r") as rf: eft_data = json.load(rf)

    # print EFT-data version
    if args.verbose:
        print("EFT data: ver {}".format(eft_data['ver']))

    # original eft data
    orig_eft_data = eft_data['data']

    # default categories: all coco-categories
    if args.categories is None: args.categories = list(COCO_CLASS_NAME2ID.keys())

    # placeholder to save 'extended data'
    eft_data_save = dict()
    stats_per_category = dict()

    # clear screen if not verbose
    if not args.verbose: os.system("clear")

    ## iterate for all 'categories'
    for category in args.categories:
        # skip 'person' category
        if category == 'person':
            continue

        # get category_id (note: not contiguous-id)
        category_id = coco.getCatIds(catNms=[category])[0]

        # count number-of-processed data
        num_filtered = 0
        desc = f"Extending COCO-EFT Dataset for [Category: '{category}'] --> Num-Filtered: ({num_filtered})"
        
        ## iterate for all data in 'eft-dataset' (note: each data corresponds to 'smpl-body annotation')
        pbar = tqdm(enumerate(sorted(orig_eft_data, key = lambda x: x['imageName'])), desc=desc)
        for eft_idx, orig_data in pbar:
            # if this data is already processed (and saved) for any previous 'category', we append info from current 'category'
            if eft_idx in eft_data_save.keys():
                # print existence
                if args.verbose: print(f"eft_idx:{eft_idx} already exists in 'eft_data_save'")
                # copy data
                data = deepcopy(eft_data_save[eft_idx])
                is_overridden = True
                if args.verbose: print(f"while {data['target_category_list']} exists, we are now Updating for {category}")
            else:
                is_overridden = False
                data = deepcopy(orig_data)

            ## ids
            # load image
            img_name = data['imageName']
            img_pth = f'{COCO_VAL_IMAGE_DIR}/{img_name}'
            img_id = int(img_name.replace("COCO_val2014_", "").replace(".jpg", ""))
            img = cv2.cvtColor(cv2.imread(img_pth), code=cv2.COLOR_BGR2RGB)

            # coco-annotation-ids for category 'person'
            person_annIds = coco.getAnnIds(imgIds=img_id, catIds=[person_category_id], iscrowd=None)
            
            # coco-annotation-ids for category
            object_annIds = coco.getAnnIds(imgIds=img_id, catIds=[category_id], iscrowd=None)

            ## filter 'multi-person' / 'multi-object' images
            # skip if multiple objects exist in the image
            if len(object_annIds) != 1:
                if args.verbose:
                    print(f"Continue {category} for this image since multiple objects or no object")
                continue

            # skip if multiple people exist in the image
            if len(person_annIds) != 1:
                if args.verbose:
                    print("Continue since multiple person or no person")
                continue

            ## load object-bboxes & object-masks
            # load coco-annotations from 'coco-annotation-ids'
            person_anns = coco.loadAnns(person_annIds)
            object_anns = coco.loadAnns(object_annIds)

            # declare placeholder for saving 'bbox', 'mask' and precomputed 'bbox-bbox_overlap'
            if is_overridden:
                assert 'bbox' in data.keys()
                assert category not in data['bbox'].keys()
            else:
                assert 'bbox' not in data.keys()
                data['bbox'] = {}
            if is_overridden:
                assert 'mask' in data.keys()
                assert category not in data['mask'].keys()
            else:
                assert 'mask' not in data.keys()
                data['mask'] = {}
            if is_overridden:
                assert 'bbox_overlap_per_category' in data.keys()
                assert category not in data['bbox_overlap_per_category']
            else:
                assert 'bbox_overlap_per_category' not in data.keys()
                data['bbox_overlap_per_category'] = {}

            # mask and bbox for category 'person'
            person_mask = coco.annToMask(person_anns[0])
            [x,y,w,h] = person_anns[0]['bbox']
            person_bbox_xyxy = [x,y,x+w,y+h]
                        
            # mask and bbox for category 'object'
            object_mask = coco.annToMask(object_anns[0])
            [x,y,w,h] = object_anns[0]['bbox']
            object_bbox_xyxy = [x,y,x+w,y+h]

            ## filter by 'bbox-bbox_overlap'
            # compute 'bbox-overlap', and save (even if filtered)
            bbox_overlap = intersection_over_smaller_bbox(person_bbox_xyxy, object_bbox_xyxy)
            data['bbox_overlap_per_category'][category] = bbox_overlap

            # filter images with small 'bbox-overlap'
            if bbox_overlap < args.min_bbox_overlap:
                if args.verbose:
                    print("Continue since bbox bbox_overlap is too small")
                continue

            ## fill in 'target-category', 'mask' and 'bbox'
            # fill in 'target category list'
            if is_overridden:
                assert 'target_category_list' in data.keys()
                assert category not in data.keys()
                data['target_category_list'].append(category)
            else:
                assert 'target_category_list' not in data.keys()
                data['target_category_list'] = [category]
            
            # fill in 'mask' and 'bbox' of person/object
            if is_overridden:
                assert 'bbox' in data.keys()
                assert 'mask' in data.keys()
                assert 'person' in data['bbox'].keys()
                assert 'person' in data['mask'].keys()
                assert category not in data['bbox'].keys()
                assert category not in data['mask'].keys()
                data['bbox'][category] = object_bbox_xyxy
                data['mask'][category] = (object_mask > 0).astype(np.uint8)                
            else:
                data['bbox']['person'] = person_bbox_xyxy
                data['mask']['person'] = (person_mask > 0).astype(np.uint8)
                data['bbox'][category] = object_bbox_xyxy
                data['mask'][category] = (object_mask > 0).astype(np.uint8)
                

            ## compute & fill in 'perspective camera'
            if is_overridden:
                assert 'pred_aa' in data.keys()
                assert 'pred_betas' in data.keys()
                assert 'persp_cam_list' in data.keys()
                assert 'smpl_joints_2d' in data.keys()
                persp_cam_smpl_list = data['persp_cam_list']
                pred_aa = data['pred_aa']
                pred_betas = data['pred_betas']
                smpl_joints_2d = data['smpl_joints_2d']
            else:
                assert 'pred_aa' not in data.keys()
                assert 'pred_betas' not in data.keys()
                assert 'persp_cam_list' not in data.keys()
                assert 'smpl_joints_2d' not in data.keys()

                ## prepare data required
                # bbox scale & center
                bbox_scale = data['bbox_scale']
                bbox_center = data['bbox_center']
                # weak-perspective camera
                pred_camera = np.array(data['parm_cam'])
                cam_param_scale = pred_camera[0]
                cam_param_trans = pred_camera[1:]
                # smpl pose & betas -> set betas as zero
                pred_betas = np.reshape(np.array( data['parm_shape'], dtype=np.float32),(1,10))            #(1,10)
                pred_betas = torch.from_numpy(pred_betas)
                pred_pose_rotmat = np.reshape(np.array( data['parm_pose'], dtype=np.float32), (1,24,3,3))  #(1,24,3,3)
                pred_pose_rotmat = torch.from_numpy(pred_pose_rotmat)
                
                axis, theta = rotmat_to_axis_angle_batch(pred_pose_rotmat[0]) # (24,3), (24,)
                pred_aa = axis * theta[:,None]
                pred_aa = pred_aa.reshape(1,72)

                ## retrieve smpl-mesh
                # infer smpl mesh
                smpl_output = smpl( 
                    betas=pred_betas, 
                    body_pose=pred_aa[:,3:],
                    global_orient=pred_aa[:,:3], 
                    pose2rot=True
                )
                smpl_output_zerorot_zerobeta = smpl( 
                    betas=torch.zeros_like(pred_betas), 
                    body_pose=pred_aa[:,3:],
                    global_orient=torch.zeros_like(pred_aa[:,:3]), 
                    pose2rot=True
                )
                
                # smpl vertices & joints in '3d-smpl-world'
                smpl_joints_3d = smpl_output.joints.detach().cpu().numpy()[0]

                # conversion of joints in '3d-smpl-world' to '2d-image-world'
                smpl_joints_2d = convert_smpl_to_img(
                    smpl_3d=smpl_joints_3d, 
                    img=img, 
                    bbox_center=bbox_center, 
                    bbox_scale=bbox_scale, 
                    cam_param_scale=cam_param_scale, 
                    cam_param_trans=cam_param_trans
                )

                # create a 'mocap_output_list' to use as input for 'weak-perspective to perspective' conversion
                mocap_output_list = [dict(
                    pred_joints_img = smpl_joints_2d,
                    pred_camera=pred_camera,
                    pred_joints_smpl_zerorot_zerobeta = smpl_output_zerorot_zerobeta.joints, # 1x49x3
                    pred_body_pose = pred_aa, # 1 x 72
                )]
                
                # prepare filter_settings
                filter_settings = EasyDict(
                    verbose=args.verbose,
                    use_openpose_joints=args.use_openpose_joints,
                    ignore_joints_outside=args.ignore_joints_outside,
                    freeze_foc=args.freeze_foc,
                    fovy=args.fovy,
                    lr=args.lr,
                    epochs=args.epochs,
                    thres=args.joint_loss_thres,
                    save_with_human_params=False,
                )

                # convert 'weak-perspective' to 'perspective'
                persp_cam_smpl_list, _ = apply_wp_to_p(
                    filter_settings=filter_settings,
                    mocap_output_list=mocap_output_list, 
                    img_size=img.shape[:2]
                )

                if persp_cam_smpl_list[0]['min_loss'] > args.reprojection_error_thresh:
                    if args.verbose:
                        print(f"continue eft-idx: {eft_idx} -> since high perspective reprojection error")
                    continue

                # add to data
                data['persp_cam_list'] = persp_cam_smpl_list
                data['pred_aa'] = pred_aa
                data['pred_betas'] = pred_betas
                data['smpl_joints_2d'] = smpl_joints_2d
            
            # visualize
            if args.visualize:
                plt.figure(figsize=(20,20))
                plt.suptitle(f"category:{category}/currently_filterd_and_saved:{num_filtered+1}/eft_idx:{eft_idx:09}")
                
                # show person mask
                plt.subplot(3,3,1); plt.imshow(data['mask']['person'])
                [x,y,x_,y_] = data['bbox']['person']
                # show person bbox & joints
                img_person = cv2.rectangle(img.copy(), (int(x), int(y)), (int(x_), int(y_)), (255,0,0), 5)
                plt.subplot(3,3,2); plt.imshow(img_person)
                # show joints
                for joint_idx in range(len(smpl_joints_2d)):
                    x, y = smpl_joints_2d[joint_idx]
                    plt.plot(x, y, marker="o", markersize=5, markeredgecolor="blue", markerfacecolor = "blue")

                # show object mask
                plt.subplot(3,3,3); plt.imshow(data['mask'][category])
                [x,y,x_,y_] = data['bbox'][category]
                # show object bbox & joints
                img_object = cv2.rectangle(img.copy(), (int(x), int(y)), (int(x_), int(y_)), (255,0,0), 5)
                plt.subplot(3,3,4); plt.imshow(img_object)
                
                human_idx = 0
                assert len(mocap_output_list) == len(persp_cam_smpl_list)
                # prepare results
                persp_cam_smpl = persp_cam_smpl_list[human_idx] # perspective camera
                R = torch.tensor(persp_cam_smpl['R']).to("cuda") # extrinsic: 4x3x3 or 3x3
                t = torch.tensor(persp_cam_smpl['t']).to("cuda") # extrinsic: 4x3 or 3,  
                K = torch.tensor(persp_cam_smpl['K']).to("cuda") # intrinsic & focal_length: 4x3x3 or 3x3
                if R.ndim != 3:
                    assert R.shape == (3,3)
                    assert t.shape == (3,)
                    assert K.shape == (3,3)
                    R = torch.stack([R]*4, dim=0)
                    t = torch.stack([t]*4, dim=0)
                    K = torch.stack([K]*4, dim=0)

                    global_zerorot_joint = persp_cam_smpl_list[human_idx]['joints_used_SMPL']
                    global_zerorot_joint_batch = np.stack([global_zerorot_joint]*4, axis=0)
                else:        
                    global_zerorot_joint_batch = persp_cam_smpl_list[human_idx]['joints_used_SMPL_batch'] # Joints in SMPL space
                focal_length = persp_cam_smpl['focal_length'] # intrinsic & focal_length: 4, or scalar
                # min_losses = persp_cam_smpl['min_losses'] # joint reprojection loss: 4, or scalar

                homo=torch.matmul(K,(torch.matmul(R, torch.tensor(global_zerorot_joint_batch).to("cuda").permute(0,2,1))\
                                    + t[:,:,None])).permute(0,2,1) # 4xJx3
                projected_joint_img_smpl = homo[:,:,:2].clone()
                projected_joint_img_smpl[:,:,0]/=homo[:,:,2] # Homogenize
                projected_joint_img_smpl[:,:,1]/=homo[:,:,2] # Homogenize
                projected_joint_img_smpl = projected_joint_img_smpl.detach().cpu().numpy() # 4xJx2
                # to write on subplot subtitles
                if 'idx2name' not in persp_cam_smpl_list[human_idx]:
                    smpl_idx2name = ['ONLY LearnedK-Zerobeta Done']
                else:
                    smpl_idx2name = persp_cam_smpl_list[human_idx]['idx2name']

                for i in range(len(smpl_idx2name)):
                    plt.subplot(3,3,i+5); plt.imshow(img)
                    plt.gca().set_title(smpl_idx2name[i])
                    # GT joints
                    for idx in range(len(smpl_joints_2d)):
                        x = smpl_joints_2d[idx][0]; y=smpl_joints_2d[idx][1]
                        plt.plot(x, y, marker="o", markersize=5, markeredgecolor="green", markerfacecolor="green")
                    # Projected joints
                    for idx in range(len(projected_joint_img_smpl[i])):
                        x = projected_joint_img_smpl[i][idx][0]; y = projected_joint_img_smpl[i][idx][1]
                        plt.plot(x, y, marker="o", markersize=5, markeredgecolor="orange", markerfacecolor="orange")
                plt.show()
                plt.close()

            ## update to 'eft_data_save'
            # assert
            if is_overridden: assert eft_idx in eft_data_save.keys()
            else: assert eft_idx not in eft_data_save.keys()
            
            # save & update 'num_filtered'
            eft_data_save[eft_idx] = data
            num_filtered += 1
            
            # update description
            desc = f"Extending COCO-EFT Dataset for [Category: '{category}'] --> Num-Filtered: ({num_filtered})"
            pbar.set_description(desc)

        if args.verbose:
            print(f"Num-filtered for {category} when min bbox iou is {args.min_bbox_overlap} --> {num_filtered}")       
        
        ## add statistics for given object category
        stats_per_category[category] = num_filtered

    # placeholder for 'results-to-save' extended-coco-eft dataset
    to_save = {
        'eft_data': list(eft_data_save.values()), 
        'stats': stats_per_category,
        'min_bbox_overlap': args.min_bbox_overlap,
        'args': args,
    }

    # print statistics
    if args.verbose: print(to_save['stats'])

    # save results
    with open(args.save_pth, "wb") as handle: pickle.dump(to_save, handle, protocol=pickle.HIGHEST_PROTOCOL)


if __name__== "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--categories", nargs="+")

    parser.add_argument("--original_dataset_pth", type=str, default=COCO_EFT_VAL_PTH)
    parser.add_argument("--save_pth", default=EXTENDED_COCO_EFT_VAL_SAVE_PTH)

    parser.add_argument("--min_bbox_overlap", type=float, default=DEFAULT_EFT_MIN_BBOX_OVERLAP)
    parser.add_argument("--reprojection_error_thresh", type=float, default=DEFAULT_EFT_REPROJECTION_ERROR_THRESH)
    parser.add_argument("--use_openpose_joints", action="store_true")
    parser.add_argument("--ignore_joints_outside", action="store_true")
    parser.add_argument("--freeze_foc", action="store_true")
    parser.add_argument("--fovy", type=float, nargs="?", default=DEFAULT_EFT_FOV)
    parser.add_argument("--lr", type=float, default=DEFAULT_EFT_LR)
    parser.add_argument("--joint_loss_thres", type=float, default=DEFAULT_EFT_JOINT_LOSS_THRES)
    parser.add_argument("--epochs", type=int, default=DEFAULT_EFT_EPOCHS)

    parser.add_argument("--visualize", action="store_true")
    parser.add_argument("--verbose", action='store_true')

    parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
    args = parser.parse_args()

    # seed for reproducible generation
    seed_everything(args.seed)

    # extend coco-eft dataset (validation-set) to use as test-set
    extend_coco_eft_dataset(args=args)
scripts/evaluation/extend_eft.py
jellyheadandrew/CHORUS
import os
from glob import glob
import argparse
import json
from easydict import EasyDict
from tqdm import tqdm

from constants.ldm import DEFAULT_LDM_MODEL_KEY
from constants.metadata import DEFAULT_QUANT_EVAL_DIR


def report_pap(args):
    ## placeholder for saving 'mean' evaluation results
    eval_total_results = dict()

    # evaluation-result paths
    eval_pths = sorted(list(glob(f"./{args.quant_eval_dir}/{args.ldm_model_key}/*/*.json")))
    # iterate for all quantitative-evaluation results
    for eval_pth in tqdm(eval_pths, desc="Reporting mPAP metrics..."):
        # open quantitative-evaluation results
        with open(eval_pth, "r") as rf:
            data = json.load(rf)
        category = data["category"]
        # skip if not in 'args.categories'
        if args.categories is not None:
            if category not in args.categories: continue
        
        # model_name
        model_name = data["aggr_pth"].split("/")[-1].replace(".pickle", "")

        # retrieve results        
        results = data["results"]
        mPAP = results["mPAP"]
        mPAP_strict = results["mPAP_strict"]
        mPAP_hoa = results["mPAP_hoa"]
        mPAP_strict_hoa = results["mPAP_strict_hoa"]

        # if the case of computing 'mean' mPAP, add to placeholder
        if args.categories is None and model_name not in eval_total_results.keys():
            eval_total_results[model_name] = dict()
            eval_total_results[model_name]['category_count'] = 1
            eval_total_results[model_name]['mPAP_mean'] = mPAP
            eval_total_results[model_name]['mPAP_strict_mean'] = mPAP_strict
            eval_total_results[model_name]['mPAP_hoa_mean'] = mPAP_hoa
            eval_total_results[model_name]['mPAP_strict_hoa_mean'] = mPAP_strict_hoa

        elif args.categories is None and model_name in eval_total_results.keys():
            eval_total_results[model_name]['category_count'] += 1
            eval_total_results[model_name]['mPAP_mean'] += mPAP
            eval_total_results[model_name]['mPAP_strict_mean'] += mPAP_strict
            eval_total_results[model_name]['mPAP_hoa_mean'] += mPAP_hoa
            eval_total_results[model_name]['mPAP_strict_hoa_mean'] += mPAP_strict_hoa

    # in the 'mean' mPAP evaluation case
    if args.categories is None:
        for model_name in eval_total_results.keys():
            total_count = eval_total_results[model_name]['category_count']
            mPAP_avg = eval_total_results[model_name]['mPAP_mean'] / total_count
            mPAP_strict_avg = eval_total_results[model_name]['mPAP_strict_mean'] / total_count
            mPAP_hoa_avg = eval_total_results[model_name]['mPAP_hoa_mean'] / total_count
            mPAP_strict_hoa_avg = eval_total_results[model_name]['mPAP_strict_hoa_mean'] / total_count

            print(f"\n==Aggregation Settings: {model_name}==")
            print(f"Average mPAP: {mPAP_avg * 100:.2f}")
            print(f"Average mPAP_strict: {mPAP_strict_avg * 100:.2f}")
            print(f"Average mPAP_hoa: {mPAP_hoa_avg * 100:.2f}")
            print(f"Average mPAP_strict_hoa: {mPAP_strict_hoa_avg * 100:.2f}")
            print(f"Categories used: {total_count}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--categories", nargs="+", type=str) # set as 'None' for default
    parser.add_argument("--ldm_model_key", type=str, default=DEFAULT_LDM_MODEL_KEY)
    parser.add_argument("--quant_eval_dir", default=DEFAULT_QUANT_EVAL_DIR)
    args = parser.parse_args()
    
    report_pap(args=args)
scripts/evaluation/report_pap.py
jellyheadandrew/CHORUS
import os
import argparse
from omegaconf import OmegaConf

import torch
from torch import autocast
import numpy as np

from PIL import Image

from tqdm import tqdm
from einops import rearrange

from imports.ldm.ldm.util import load_model_from_config
from imports.ldm.ldm.models.diffusion.ddim import DDIMSampler
from imports.ldm.ldm.models.diffusion.plms import PLMSSampler

from constants.ldm import DEFAULT_BATCH_SIZE, DEFAULT_LDM_MODEL_KEY, AUGMENTATION_PROMPT_PTH, MODELS, GetLDMOptions
from constants.metadata import DEFAULT_PROMPT_DIR, DEFAULT_IMAGE_DIR, DEFAULT_SEED

from utils.prepare_prompts import prepare_cps, prepare_augdict
from utils.reproducibility import seed_everything




def gen_images(
        cps, 
        augdict, 
        ldm_model_key,
        batch_size, 
        num_batch_per_augprompt,
        augprompts_not_to_use,
        skip_done,
        verbose,
        save_dir,
    ):

    # Create result directory
    result_dir=f"./{save_dir}/{ldm_model_key}"
    os.makedirs(result_dir, exist_ok=True)

    # get options
    opt = GetLDMOptions(model_key=ldm_model_key, models=MODELS)   
    
    # load config & model
    config = OmegaConf.load(f"{opt.config}")
    model = load_model_from_config(config, f"{opt.ckpt}", verbose)
    
    # send model to device
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    model = model.to(device)
    
    # diffusion sampler
    if opt.plms:
        sampler = PLMSSampler(model)
    else:
        sampler = DDIMSampler(model)

    # clear screen if not verbose
    if not verbose: os.system('clear')

    # if opt.fixed_code, we use fixed initial noisy latent for generation
    start_code = None
    if opt.fixed_code:
        start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device)
    # set precision
    precision_scope = autocast

    # check if you can continue
    for augprompt_not_to_use in augprompts_not_to_use:
        for prompttype_prompt in augdict.keys():
            if augprompt_not_to_use in augdict[prompttype_prompt]:
                augdict[prompttype_prompt].remove(augprompt_not_to_use)
    
    # print augprompts we are using
    if verbose:
        print("\n\n\n\nAugprompts we are using: ")
        for prompttype_prompt in augdict:
            print(f"{prompttype_prompt} => {augdict[prompttype_prompt]}")
        print("\n\n\n\n")


    # generate images
    for category, prompt in cps:
        # declar saving directory
        final_savedir = f"{result_dir}/{category}/{prompt}"
        os.makedirs(final_savedir, exist_ok=True)

        # create a batch per prompt+augmentation
        num_batch = len(augdict[(category, prompt)]) * num_batch_per_augprompt

        # skip if exists
        choose2skip = True
        for augprompt in augdict[(category, prompt)]:
            # make save directory IF there's no directory
            os.makedirs(f"{final_savedir}/{augprompt}", exist_ok=True)
            # check to skip or not
            if verbose:
                print("augprompt: ",augprompt)
            if len(os.listdir(os.path.join(final_savedir, augprompt))) != num_batch_per_augprompt * batch_size:
                choose2skip = False
        if choose2skip and skip_done:
            if verbose:
                print("\n\tContinueing: ", prompt)
            continue

        
        # generate images
        base_count = 0
        with torch.no_grad():
            with precision_scope("cuda"):
                with model.ema_scope():
                    for n in tqdm(range(num_batch), desc=f"Generating Images for [Category '{category}' / Prompt '{prompt}']"):
                        # retrieve augmentation
                        augprompt = augdict[(category, prompt)][n // num_batch_per_augprompt]
                        
                        # make a batch of augmented prompts
                        if augprompt == "original":
                            datas = [[prompt] * batch_size]
                        else:
                            datas = [[prompt+augprompt] * batch_size]

                        # continue or not
                        if len(os.listdir(os.path.join(final_savedir, augprompt))) // batch_size > n % num_batch_per_augprompt and skip_done:
                            if verbose:
                                print(f"\tSkipping batch:{n % num_batch_per_augprompt +1} for prompt+augmentation:{datas[0][0]}")
                            base_count += batch_size
                            continue
                        else:
                            if verbose:
                                print(f"\n\nBatch: {n % num_batch_per_augprompt +1} for for prompt+augmentation:{datas[0][0]} --> Generating...\n\n")
                        
                        # generate data batch
                        if verbose: batch_iterable = tqdm(datas, desc="data")
                        else: batch_iterable = datas
                        for data in batch_iterable:
                            if verbose:
                                print(data)
                            uc = None
                            if opt.scale != 1.0:
                                uc = model.get_learned_conditioning(batch_size * [""])
                            if isinstance(data, tuple):
                                data = list(data)
                            c = model.get_learned_conditioning(data)
                            shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
                            samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
                                                            conditioning=c,
                                                            batch_size=batch_size,
                                                            shape=shape,
                                                            verbose=False,
                                                            unconditional_guidance_scale=opt.scale,
                                                            unconditional_conditioning=uc,
                                                            eta=opt.ddim_eta,
                                                            x_T=start_code,
                                                            )
                            x_samples_ddim = model.decode_first_stage(samples_ddim)
                            x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
                            # save images
                            for x_sample in x_samples_ddim:
                                x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
                                Image.fromarray(x_sample.astype(np.uint8)).save(
                                    os.path.join(final_savedir, augprompt, f"{base_count:06}.png"))
                                base_count += 1




if __name__ == "__main__":
    ## Arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--categories", type=str, nargs="+")
    parser.add_argument("--ldm_model_key", type=str, default=DEFAULT_LDM_MODEL_KEY)
    
    parser.add_argument("--disable_txt_prompts", action="store_true")
    parser.add_argument("--augprompts_path", type=str, default=AUGMENTATION_PROMPT_PTH)
    parser.add_argument("--augprompts_not_to_use", nargs="+")

    parser.add_argument("--prompt_dir", type=str, default=DEFAULT_PROMPT_DIR)
    parser.add_argument("--save_dir", type=str, default=DEFAULT_IMAGE_DIR)

    parser.add_argument("--batch_size", type=int, default=DEFAULT_BATCH_SIZE)
    parser.add_argument("--num_batch_per_augprompt", type=int)

    parser.add_argument("--skip_done", action="store_true")
    parser.add_argument("--verbose", action="store_true")

    parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
    args = parser.parse_args()
    

    # prepare prompts
    cps = prepare_cps(args.prompt_dir, args.categories, not args.disable_txt_prompts)

    # prepare prompt augmentations
    augdict = prepare_augdict(cps, args.augprompts_path, args.verbose)

    # seed for reproducible generation
    seed_everything(args.seed)

    # prepare augmentations to skip
    if args.augprompts_not_to_use is None:
        args.augprompts_not_to_use = []

    # print information
    if args.verbose:
        print(f"Running for {len(cps)} number of (category, prompt) pairs\n\n")
        print("\n====================================\n")
        print("\nAugprompts:")
        for prompt in augdict.keys():
            print(f"{prompt} => ")
            for augprompt in augdict[prompt]:
                print("\t",augprompt)
        print("\n====================================\n")

        print(f"You have specified augprompts not to use: {args.augprompts_not_to_use}")
        # arguments
        print("===Arguments===")
        print(args)


    # generate images
    gen_images(
        cps=cps, 
        augdict=augdict, 
        ldm_model_key=args.ldm_model_key,
        batch_size=args.batch_size, 
        num_batch_per_augprompt=args.num_batch_per_augprompt,
        augprompts_not_to_use=args.augprompts_not_to_use,
        skip_done=args.skip_done,
        verbose=args.verbose,
        save_dir=args.save_dir,
    )
scripts/generation/generate_images.py
jellyheadandrew/CHORUS
import os
import openai
import argparse
import random
from tqdm import tqdm
import time

from constants.metadata import DEFAULT_PROMPT_DIR
from constants.prompt_generation import DEFAULT_TEMPLATE, DEFAULT_TEMPERATURE, DEFAULT_AVAILABLE_M, DEFAULT_SWITCHING_PART

from dotenv import load_dotenv
load_dotenv()

def set_api_key():
    # Export API Key Environment Variables and save under OpenAI Instance
    openai.organization = os.getenv("OPENAI_API_ORG_KEY")
    openai.api_key = os.getenv("OPENAI_API_KEY")



def generate_prompts(
        categories, 
        save_dir, 
        num_prompts,
        verbose
    ):
    # print log
    if verbose:
        print(f"\nRunning for {len(categories)} number of categories...\n"
              f"Saving prompts under \"{save_dir}\"...\n\n")
    
    # Iterate through all categories
    for category in categories:
        # create saving directory
        os.makedirs(save_dir, exist_ok=True)

        # if m is None, randomly choose between 3~20
        if num_prompts is None:
            m = random.choice(DEFAULT_AVAILABLE_M)
        else:
            m = num_prompts
            
        # Assert m to be positive integer
        assert (m >= 1 and type(m) == int) or m is None, "Only positive integer allowed."
        
        # choose which to switch
        switch = random.choice(DEFAULT_SWITCHING_PART)
        # Instantiated Template
        instantiated_template = DEFAULT_TEMPLATE.replace("[M]", f"{m}").replace("[CATEGORY]", category).replace("[SWITCHING-PART]", switch)
        if verbose:
            print(f"INSTANTIATED-DEFAULT_TEMPLATE: {instantiated_template}\n\n")


        # HOI prompts to fill in
        HOI_prompts = []

        pbar = tqdm(total=m, desc=f"Generating Prompts for: [Category '{category}']")
        while len(HOI_prompts) != m:
            # reset if iterating again
            HOI_prompts = []


            # Chat-Completion
            # NOTE: OpenAI API does not support random-seeding. 
            # Reference: https://community.openai.com/t/is-there-a-way-to-set-a-a-random-seed-for-responses-with-temperature-0/4164 
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=[{
                    "role": "user",
                    "content": instantiated_template}],
                temperature=DEFAULT_TEMPERATURE,
            )

            # iterate through all available lines
            for idx, HOI_prompt in enumerate(response['choices'][0].message['content'].split("\n")):
                # Remove any remaining whitespace or \n at the start-end of the sentence
                HOI_prompt = HOI_prompt.strip()

                # If there is number, remove the number
                if f"{idx+1}." in HOI_prompt[:3]:
                    # Remove the number and remaining whitespaces
                    HOI_prompt = HOI_prompt.replace(f"{idx+1}.", "").strip()
                    
                # If HOI prompt is a blank, continue
                if HOI_prompt == "":
                    continue
                
                # If the sentence ends with ".", remove it
                if HOI_prompt[-1] == ".":
                    # Remove the "." and remaining whitespaces
                    HOI_prompt = HOI_prompt[:-1].strip()

                # Add to HOI prompts
                HOI_prompts.append(HOI_prompt)
                pbar.update(1)

            if len(HOI_prompts) != m:
                pbar = tqdm(total=m, desc=f"Restart Category: {category}")
                continue
            else:
                if verbose:
                    print(f"=========CATEGORY: {category}=========")
                    print(response['choices'][0].message['content'])
                break

        # Save to txt file
        with open(os.path.join(save_dir, f"{category}.txt"), "w") as wf:
            wf.write("\n".join(HOI_prompts))
        
    



if __name__=="__main__":
    # Arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--categories", nargs="+", type=str)
    parser.add_argument("--save_dir", type=str, default=DEFAULT_PROMPT_DIR)
    parser.add_argument("--num_prompts", type=int, default=None) # If None, m is randomly chosen between 3~20
    parser.add_argument("--verbose", action="store_true")
    args = parser.parse_args()

    # Set OpenAI API Key
    set_api_key()

    # generate HOI prompts
    generate_prompts(
        categories=args.categories, 
        save_dir=args.save_dir, 
        num_prompts=args.num_prompts,
        verbose=args.verbose
    )
scripts/generation/generate_prompts.py
jellyheadandrew/CHORUS
from setuptools import setup, find_packages

setup(
    name='chorus',
    version='0.0.1',
    description='',
    packages=find_packages(),
    install_requires=[],
)
setup.py
jellyheadandrew/CHORUS
import os
from glob import glob
import shutil

import pickle5 as pickle
from tqdm import tqdm
from copy import deepcopy
from easydict import EasyDict

import math
import numpy as np
import torch

import cv2
import matplotlib.pyplot as plt

import open3d as o3d
import mayavi.mlab as mlab
from tvtk.util import ctf

from smplx.lbs import blend_shapes, vertices2joints

from constants.mayavi_cams import VIDEO_CAMS

from utils.transformations import get_azimuth, batch_rodrigues, batch_rigid_transform, perspective_projection_grid
from utils.misc import to_np_torch_recursive
from utils.aggregation.load_lbs_weight import precompute_lbs_weights

from visualization.visualize_video import generate_video_from_imgs
from visualization.visualize_open3d import return_cube, return_coordinate_frame, return_triangle_mesh, return_camera_frustum



class Aggregator:
    def __init__(self, core: EasyDict):
        super().__init__()

        # core
        self.core = deepcopy(core)

        # whether to use torch or np & device
        self.use_torch = core.GLOBAL.use_torch
        self.device = core.GLOBAL.device

        ## cache for registering inputs for aggregation
        self.cache = EasyDict({_cam_world_key: EasyDict() for _cam_world_key in self.core.keys()}) 
        self.num_processed = EasyDict({_cam_world_key: 0 for _cam_world_key in self.core.keys()})


        ## full_aggregations: aggregating 'all' images (clustering disabled)
        # declare for each worlds
        self.full_aggregations = EasyDict({_cam_world_key: EasyDict() for _cam_world_key in self.core.keys()})
        del self.full_aggregations['GLOBAL']
        
        # add required variables for each world
        for _cam_world_key in self.full_aggregations.keys():
            # voting grid - mask occupancy
            self.full_aggregations[_cam_world_key].visual_hull = np.zeros(
                [
                    self.core[_cam_world_key].N_x, 
                    self.core[_cam_world_key].N_y, 
                    self.core[_cam_world_key].N_z,
                ], dtype=float
            )
            # voting grid - image occupancy
            self.full_aggregations[_cam_world_key].full_image_visual_hull = np.zeros(
                [
                    self.core[_cam_world_key].N_x,
                    self.core[_cam_world_key].N_y,
                    self.core[_cam_world_key].N_z,
                ], dtype=float
            )
            # camera-pose (R,t,K) history
            self.full_aggregations[_cam_world_key].aggr_input_history = {}
            # geometries for visualization (mayavi)
            self.full_aggregations[_cam_world_key].geometries = EasyDict(
                    grid=dict( # create grid using lines
                        length_x=self.core[_cam_world_key].length_x,
                        length_y=self.core[_cam_world_key].length_y,
                        length_z=self.core[_cam_world_key].length_z,
                        center=self.core[_cam_world_key].center,
                        color=self.core[_cam_world_key].grid_color,
                    ),
                    coordinate_frame=dict(size=1.0), # create coordinate frame
                    template = dict( # 3. create template mesh in canonical space
                        verts=self.core[_cam_world_key].v_template_holistic,
                        faces=self.core[_cam_world_key].v_template_faces_holistic,
                        color=self.core[_cam_world_key].v_template_color
                    ),
                )
            # geometries for visualization (open3d)
            self.full_aggregations[_cam_world_key].o3d_geometries = EasyDict(
                    grid=return_cube( # create grid using lines
                        length_x=self.core[_cam_world_key].length_x,
                        length_y=self.core[_cam_world_key].length_y,
                        length_z=self.core[_cam_world_key].length_z,
                        center=self.core[_cam_world_key].center,
                        color=self.core[_cam_world_key].grid_color,
                    ),
                    coordinate_frame=return_coordinate_frame(size=1.0), # create coordinate frame
                    template = return_triangle_mesh( # create template mesh in canonical space
                        verts=self.core[_cam_world_key].v_template_holistic,
                        faces=self.core[_cam_world_key].v_template_faces_holistic,
                        color=self.core[_cam_world_key].v_template_color
                    ),
                )
        

        ## aggregations: aggregating images (clustering enabled) in the specified world
        # declare for each worlds
        self.aggregations = EasyDict({_cam_world_key: EasyDict() for _cam_world_key in self.core.keys()})
        del self.aggregations['GLOBAL']

        # iterate for worlds
        for _cam_world_key in self.aggregations.keys():
            # dictionary of body_part_name: List[tag1, tag2, ...] OR None
            aggr_world_name_per_body_parts = self.core[_cam_world_key].aggr_world_name_per_body_parts 

            # each elements of dictionary would be world name
            self.aggregations[_cam_world_key].aggr_world_name_per_body_parts = {
                body_part_name: 
                    list(
                        set(
                            aggr_world_name_per_body_parts[body_part_name] + [_cam_world_key]
                        )
                    ) if aggr_world_name_per_body_parts[body_part_name] is not None 
                    else [_cam_world_key]
                for body_part_name in self.core[_cam_world_key].body_part_names
            }

            # each elements of dictionary would be voting grid (mask occupancy)
            self.aggregations[_cam_world_key].visual_hull_per_body_parts = {
                body_part_name: {
                    _aggr_world_key: np.zeros( # shape: NxNxN, dtype: float
                        [
                            self.core[_aggr_world_key].N_x, 
                            self.core[_aggr_world_key].N_y, 
                            self.core[_aggr_world_key].N_z,
                        ], dtype=float
                    )
                    for _aggr_world_key in self.aggregations[_cam_world_key].aggr_world_name_per_body_parts[body_part_name]
                } for body_part_name in self.core[_cam_world_key].body_part_names
            }

            # each elements of dictionary would be voting grid (image occupancy)
            self.aggregations[_cam_world_key].full_image_visual_hull_per_body_parts = {
                body_part_name: { 
                    _aggr_world_key: np.zeros( # shape: NxNxN, dtype: float
                        [ 
                            self.core[_aggr_world_key].N_x, 
                            self.core[_aggr_world_key].N_y, 
                            self.core[_aggr_world_key].N_z
                        ], dtype=float
                    )
                    for _aggr_world_key in self.aggregations[_cam_world_key].aggr_world_name_per_body_parts[body_part_name]
                } for body_part_name in self.core[_cam_world_key].body_part_names
            }
            
            # camera-pose (R,t,K) history
            self.aggregations[_cam_world_key].aggr_input_history_per_body_parts = {
                body_part_name: {
                    _aggr_world_key: EasyDict()
                    for _aggr_world_key in self.aggregations[_cam_world_key].aggr_world_name_per_body_parts[body_part_name]
                } for body_part_name in self.core[_cam_world_key].body_part_names
            }
            
            # geometries for visualization (mayavi)
            self.aggregations[_cam_world_key].geometries_per_body_parts = {
                body_part_name: {
                    _aggr_world_key: EasyDict(
                        grid=dict( # create Grid using lines
                            length_x=self.core[_aggr_world_key].length_x,
                            length_y=self.core[_aggr_world_key].length_y,
                            length_z=self.core[_aggr_world_key].length_z,
                            center=self.core[_aggr_world_key].center,
                            color=self.core[_aggr_world_key].grid_color,
                            ),
                        coordinate_frame=dict(size=1.0), # create coordinate frame
                        template = dict( # create Template in Canonical Frame
                            verts=self.core[_aggr_world_key].v_template_holistic,
                            faces=self.core[_aggr_world_key].v_template_faces_holistic,
                            color=self.core[_aggr_world_key].v_template_color
                        ),
                    ) for _aggr_world_key in self.aggregations[_cam_world_key].aggr_world_name_per_body_parts[body_part_name]
                } for body_part_name in self.core[_cam_world_key].body_part_names
            }
            # geometries for visualization (open3d)
            self.aggregations[_cam_world_key].o3d_geometries_per_body_parts = {
                body_part_name: {
                    _aggr_world_key: EasyDict(
                        grid=return_cube( # create Grid using lines
                            length_x=self.core[_aggr_world_key].length_x,
                            length_y=self.core[_aggr_world_key].length_y,
                            length_z=self.core[_aggr_world_key].length_z,
                            center=self.core[_aggr_world_key].center,
                            color=self.core[_aggr_world_key].grid_color,
                            ),
                        coordinate_frame=return_coordinate_frame(size=1.0), # create coordinate frame
                        template = return_triangle_mesh( # create Template in Canonical Frame
                            verts=self.core[_aggr_world_key].v_template_holistic,
                            faces=self.core[_aggr_world_key].v_template_faces_holistic,
                            color=self.core[_aggr_world_key].v_template_color
                        ),
                    ) for _aggr_world_key in self.aggregations[_cam_world_key].aggr_world_name_per_body_parts[body_part_name]
                } for body_part_name in self.core[_cam_world_key].body_part_names
            }
    
        # send all tensors to torch.Tensor or np.ndarray
        self.core = to_np_torch_recursive(self.core, use_torch=self.use_torch, device=self.device)
        self.full_aggregations = to_np_torch_recursive(self.full_aggregations, use_torch=self.use_torch, device=self.device)
        self.aggregations = to_np_torch_recursive(self.aggregations, use_torch=self.use_torch, device=self.device)
        


    ## using default settings from 'core'. hierarchy: GLOBAL > world name
    def query_from_core(self, name=None, _cam_world_key=None):
        if name in self.core.GLOBAL.keys():
            return self.core.GLOBAL[name]
        elif name in self.core[_cam_world_key].keys():
            return self.core[_cam_world_key][name]
        else:
            assert False, f"'{name}' does not exist in 'core.GLOBAL' & 'core.{_cam_world_key}'"


    ## left-to-right conversion functions
    def convert_left_smpl_to_right(self,cam_R, cam_t, cam_K, mask, pose, betas):
        """
            1. negate the y-rotation & z-rotation of 'left-(body part) angle'
            2. assign this negated 'left-(body part) angle' to 'right-(body part) angle'
        """
        # convert R (r_1, r_2, r_3 are column-vectors of R^-1)
        cam_R = cam_R.clone()
        cam_R[0,1] *= -1 # r_1's y
        cam_R[0,2] *= -1 # r_2's z
        cam_R[1,0] *= -1 # r_2's x
        cam_R[2,0] *= -1 # r_3's x
        # convert x-element of 't' (negation)
        cam_t = cam_t.clone() 
        cam_t[0] *= -1
        # convert width-dimension of 'mask'
        mask = mask.cpu().numpy()
        mask = mask[:,::-1].copy()
        mask = torch.tensor(mask, device=self.device)
        ## convert y-rotation & z-rotation of pose parameters
        assert pose.shape == (1,72), "smpl pose must be 48 in dimensions"
        pose = pose.clone()
        # negate all y-rotation & z-rotation
        pose[:,1::3] *= -1
        pose[:,2::3] *= -1
        # map 'left-(body part) angle' to 'right-(body part) angle' (reference: mmhuman3d.core.conventions.keypoints_mapping.KEYPOINTS_FACTORY)
        mapped_pose = torch.zeros_like(pose)
        mapped_pose[:,0:3] = pose[:,0:3]                # 1. pelvis <-> pelvis
        mapped_pose[:,3:6] = pose[:,6:9]                # 2. left-hip <-> right-hip
        mapped_pose[:,6:9] = pose[:,3:6]                # 3. left-hip <-> right-hip
        mapped_pose[:,9:12] = pose[:,9:12]              # 4. spine1 <-> spine1
        mapped_pose[:,12:15] = pose[:,15:18]            # 5. left-knee <-> right-knee
        mapped_pose[:,15:18] = pose[:,12:15]            # 6. left-knee <-> right-knee
        mapped_pose[:,18:21] = pose[:,18:21]            # 7. spine2 <-> spine2
        mapped_pose[:,21:24] = pose[:,24:27]            # 8. left-ankle <-> right-ankle
        mapped_pose[:,24:27] = pose[:,21:24]            # 9. left-ankle <-> right-ankle
        mapped_pose[:,27:30] = pose[:,27:30]            # 10. spine3 <-> spine3
        mapped_pose[:,30:33] = pose[:,33:36]            # 11. left-foot <-> right-foot
        mapped_pose[:,33:36] = pose[:,30:33]            # 12. left-foot <-> right-foot
        mapped_pose[:,36:39] = pose[:,36:39]            # 13. neck <-> neck
        mapped_pose[:,39:42] = pose[:,42:45]            # 14. left-collar <-> right-collar
        mapped_pose[:,42:45] = pose[:,39:42]            # 15. left-collar <-> right-collar
        mapped_pose[:,45:48] = pose[:,45:48]            # 16. head <-> head
        mapped_pose[:,48:51] = pose[:,51:54]            # 17. left-shoulder <-> right-shoulder
        mapped_pose[:,51:54] = pose[:,48:51]            # 18. left-shoulder <-> right-shoulder
        mapped_pose[:,54:57] = pose[:,57:60]            # 19. left-elbow <-> right-elbow
        mapped_pose[:,57:60] = pose[:,54:57]            # 20. left-elbow <-> right-elbow
        mapped_pose[:,60:63] = pose[:,63:66]            # 21. left-wrist <-> right-wrist
        mapped_pose[:,63:66] = pose[:,60:63]            # 22. left-wrist <-> right-wrist
        mapped_pose[:,66:69] = pose[:,69:72]            # 23. left-hand <-> right-hand
        mapped_pose[:,69:72] = pose[:,66:69]            # 24. left-hand <-> right-hand
        
        return cam_R, cam_t, cam_K, mask, mapped_pose, betas

    ## left-to-right conversion functions
    def convert_left_to_right(self, cam_world_key, cam_R, cam_t, cam_K, mask, pose, betas):
        converter = getattr(self, f"convert_left_{cam_world_key.lower()}_to_right")
        cam_R,cam_t,cam_K,mask,pose,betas = converter(cam_R,cam_t,cam_K,mask,pose,betas)

        # return as keyword-arguments
        return {
            'cam_R': cam_R,
            'cam_t': cam_t,
            'cam_K': cam_K,
            'mask': mask,
            'pose': pose,
            'betas': betas
        }
    

    ## preprocessing input
    def preprocess_input(self, cam_world_key, cam_R, cam_t, cam_K, mask, pose, betas):
        pose_prep_func = getattr(self, f"prepare_{cam_world_key.lower()}_pose_for_input")
        full_pose = pose_prep_func(pose)
                
        # betas as zeros
        betas = torch.zeros_like(betas)
        
        return {
            'cam_R': cam_R,
            'cam_t': cam_t,
            'cam_K': cam_K,
            'mask': mask,
            'pose': full_pose,
            'betas': betas
        }


    ## preparing smpl pose for input
    def prepare_smpl_pose_for_input(self, pose):
        # remove first 3 dimensions (global rotation)
        assert pose.shape[-1] in [72, 69] and (pose.shape[0] == 1 or pose.ndim == 1)
        pose = pose.clone().squeeze()
        if pose.shape == (72,):
            pose = pose[3:]
        if pose.shape == (69,):
            pose = pose[:]
        
        # batchify (1 x 69)
        pose = pose.unsqueeze(0)

        # global orient as 0
        global_orient = torch.zeros([1,3]).to(pose.dtype).to(pose.device)
        
        full_pose = torch.cat([global_orient, pose], dim=-1) # 1x72
        
        return full_pose
        

    ## register single image (cam, body pose) for 'later-aggregation'
    def register_cam_pose(
        self,
        cam_world_key, # world that conducts aggregation
        pose,
        betas,
        cam_R,
        cam_t,
        cam_K,
        mask,
        add_human_for_vis,
        auxiliary_geometry=EasyDict(),
        frustum_size=None,
        use_normalized_K=None,
        cam_color=None,
        tag_auxiliary:bool=None,
        auxiliary_exist_ok:bool=None, # raise error if auxiliary member of same key already exists in cache
        max_cam_save_num:int=None,
    ):
        ## 'key' for registering to cache
        cache_key = f'cache:{len(self.cache[cam_world_key].keys()):05}'

        ## add input-information to cache
        self.cache[cam_world_key][cache_key] = {
            'cam_world_key': cam_world_key, 
            'pose': pose, 
            'betas': betas,
            'cam_R': cam_R, 
            'cam_t': cam_t, 
            'cam_K': cam_K, 
            'mask': mask,
            'add_human_for_vis': add_human_for_vis,
            'auxiliary_geometry': auxiliary_geometry,
            'frustum_size': self.query_from_core('frustum_size', cam_world_key) if frustum_size is None else frustum_size,
            'use_normalized_K': self.query_from_core('use_normalized_K', cam_world_key) if use_normalized_K is None else use_normalized_K,
            'cam_color': self.query_from_core('cam_color', cam_world_key) if cam_color is None else cam_color, 
            'tag_auxiliary': self.query_from_core('tag_auxiliary', cam_world_key) if tag_auxiliary is None else tag_auxiliary,
            'auxiliary_exist_ok': self.query_from_core('auxiliary_exist_ok', cam_world_key) if auxiliary_exist_ok is None else auxiliary_exist_ok,
            'max_cam_save_num': self.query_from_core('max_cam_save_num', cam_world_key) if max_cam_save_num is None else max_cam_save_num
        }


    ## camera-weighting methods
    def return_camera_weights(self, method):
        ## dict to save 'camera-weights'
        camera_weights_per_world = dict()

        ## method 1: total-uniform
        if method == "total-uniform":
            for _cam_world_key in self.cache.keys():
                # declare dictionary to save camera-weight for cameras in world
                camera_weights_per_world[_cam_world_key] = dict()
                # iterate for 'registered inputs' in cache
                for cache_key in self.cache[_cam_world_key].keys():
                    camera_weights_per_world[_cam_world_key][cache_key] = 1.

        ## method 2: divide by 3 azimuth-regions and uniform 'within region'
        elif method == "azimuth-uniform":
            for _cam_world_key in self.cache.keys():
                # declare dictionary to save camera-weight for cameras in world
                camera_weights_per_world[_cam_world_key] = dict()
                # 60 degrees delimited azimuth space
                azim_range_cams = {0: [], 1: [], 2: []}
                # iterate for 'registered inputs' in cache
                for cache_key in self.cache[_cam_world_key].keys():
                    # find camera's location of origin
                    cam_R = self.cache[_cam_world_key][cache_key]['cam_R']
                    cam_t = self.aggregations[_cam_world_key][cache_key]['cam_t']
                    
                    # x, z position of camera
                    location = -cam_R.T @ cam_t
                    x = location[0]; z = location[2]

                    # check azimuth & register cam to 'azimuth-region'
                    sqrt_3 = math.sqrt(3)
                    if z>= sqrt_3 * x and z>=-sqrt_3 * x:
                        azim_range_cams[0].append(cache_key)
                    elif z<sqrt_3 * x and z<-sqrt_3 * x:
                        azim_range_cams[0].append(cache_key)
                    elif z >= 0 and z < sqrt_3 * x:
                        azim_range_cams[1].append(cache_key)
                    elif z < 0 and z >= sqrt_3 * x:
                        azim_range_cams[1].append(cache_key)
                    elif z >= 0 and z<-sqrt_3 * x:
                        azim_range_cams[2].append(cache_key)
                    elif z < 0 and z >= -sqrt_3 * x:
                        azim_range_cams[2].append(cache_key)
                    else:
                        assert False

                # calculate camera-weights for '_cam_world_key'
                weight = dict()
                weight[0] = 1. / len(azim_range_cams[0]) if len(azim_range_cams[0]) > 0 else 1.
                weight[1] = 1. / len(azim_range_cams[1]) if len(azim_range_cams[1]) > 0 else 1.
                weight[2] = 1. / len(azim_range_cams[2]) if len(azim_range_cams[2]) > 0 else 1.

                # update 'camera-weight' per world'
                for azim_region_key in azim_range_cams.keys():
                    for cache_key in azim_range_cams[azim_region_key]:
                        assert cache_key not in camera_weights_per_world[_cam_world_key].keys(), "Duplicate keys."
                        camera_weights_per_world[_cam_world_key][cache_key] = weight[azim_region_key]
                        
        elif method == "azimuth-uniform-fine":
            N = 12
            d_theta = 360. / N
            for _cam_world_key in self.cache.keys():
                # declare dictionary to save camera-weight for cameras in world
                camera_weights_per_world[_cam_world_key] = dict()
                # azimuth regions
                azim_range_cams = {i: [] for i in range(N)}
                # iterate for 'registered inputs' in cache
                for cache_key in self.cache[_cam_world_key].keys():
                    # find camera's location of origin
                    cam_R = self.cache[_cam_world_key][cache_key]['cam_R']
                    cam_t = self.cache[_cam_world_key][cache_key]['cam_t']
                    # x, z position of camera
                    location = -cam_R.T @ cam_t
                    x = location[0]; z = location[2]

                    # check azimuth & register cam to 'azimuth-region'
                    azimuth = get_azimuth(x,z)
                    i = int(azimuth / d_theta) % N
                    azim_range_cams[i].append(cache_key)

                # calculate camera-weights for '_cam_world_key'
                weight = dict()
                for i in range(N):
                    weight[i] = 1. / len(azim_range_cams[i]) if len(azim_range_cams[i]) > 0 else 1.

                # update 'camera-weight' per world'
                for azim_region_key in azim_range_cams.keys():
                    for cache_key in azim_range_cams[azim_region_key]:
                        assert cache_key not in camera_weights_per_world[_cam_world_key].keys(), "Duplicate keys."
                        camera_weights_per_world[_cam_world_key][cache_key] = weight[azim_region_key]

        else:
            assert False, "Invalid camera-sampling method."

        return camera_weights_per_world
    

    ## aggregate all images (cam/body pose) in 'cache'
    def aggregate_all(self, camera_sampling="all", verbose=False):
        # 'camera-weighting' method
        camera_weights_per_world = self.return_camera_weights(method=camera_sampling) # camera-weights per world
        # print message
        for _cam_world_key in self.cache.keys():
            if verbose:
                print(f"==== aggregating {len(self.cache[_cam_world_key].keys())} registered inputs in cache ====")
            
            # iterate through registered inputs in cache
            for cache_key, aggr_input in self.cache[_cam_world_key].items():
                # prepare 'aggr_input' as torch/numpy
                aggr_input = to_np_torch_recursive(aggr_input, use_torch=self.use_torch, device=self.device)
                # run visual hull
                self.aggregate_single_input(
                    camera_weight=camera_weights_per_world[_cam_world_key][cache_key],
                    verbose=verbose,    
                    **aggr_input
                )
                # print which registered-input is aggregated 
                if verbose:
                    print(f"\taggregation for {cache_key} done...")

            # re-declare 'clean' cache after aggregation
            self.cache[_cam_world_key] = EasyDict()

            # print success for aggregation
            if verbose:
                print(f"==== aggregation for '{_cam_world_key}' succesful ====")


    ## apply projection and remove if outside image
    def retain_projection_inside_img(self, projected_grid, img_size):
        # unpack 'img_size'
        H, W = img_size

        # round-down to integer-values (to obtain pixel-values) 
        projected_grid = projected_grid.type(torch.long) # Shape: [2,*grid_size]
        
        # remove indices if 'projected-points' are outside the image
        above_zero = torch.logical_and(projected_grid[0] >= 0, projected_grid[1] >= 0)
        below_img_size = torch.logical_and(projected_grid[0] < W, projected_grid[1] < H)
        inside_img = torch.logical_and(above_zero, below_img_size) # Shape: [*grid_size]
        projected_grid = projected_grid[:,inside_img] # Shape: [2, *grid_size]

        return projected_grid, inside_img


    ## save mask occupancy & image occupancy per 'voxel in grid'
    def uplift_mask_image_occupancy(self, projected_grid, mask, img_size):
        # remove indices if 'projected-points' are outside the image
        projected_grid, inside_img = self.retain_projection_inside_img(projected_grid, img_size=img_size)

        # check mask-occupancy of 'projected-points'
        mask_occupancy_grid = torch.zeros_like(inside_img, dtype=torch.float32)
        mask_occupancy_grid[inside_img] = mask[projected_grid[1], projected_grid[0]].type(torch.float32) # Shape: [N_x, N_y, N_z] / dtype: long
        # check image-occupancy of 'projected-points'
        image_occupancy_grid = torch.zeros_like(inside_img, dtype=torch.float32)
        image_occupancy_grid[inside_img] = torch.ones_like(mask)[projected_grid[1], projected_grid[0]].type(torch.float32) # Shape: [N_x, N_y, N_z] / dtype: long

        return mask_occupancy_grid, image_occupancy_grid, inside_img


    ## aggregate single input given
    def aggregate_single_input(
        self,
        camera_weight,
        verbose, 
        cam_world_key,
        cam_R,
        cam_t,
        cam_K,
        mask,
        pose,
        betas,
        add_human_for_vis,
        frustum_size,
        use_normalized_K,
        cam_color,
        auxiliary_geometry,
        tag_auxiliary,
        auxiliary_exist_ok,
        max_cam_save_num,
    ):
        # mask shape
        H, W = mask.shape
        
        # camera-intrinsics
        if use_normalized_K:
            focal_length = cam_K[0,0]
        else:
            focal_length = cam_K[0,0] / W

        # remove batch-dimension for 'pose/betas'
        pose = pose.squeeze(); betas = betas.squeeze()

        ## we don't need to re-compute aggregation result (mask-occupancy / image-occupancy) for every body-part
        mask_occupancy_precomputed = {
            (_cam_world_key, _aggr_world_key): None 
            for _aggr_world_key in self.aggregations.keys()
            for _cam_world_key in self.aggregations.keys()
        }
        image_occupancy_precomputed = {
            (_cam_world_key, _aggr_world_key): None 
            for _aggr_world_key in self.aggregations.keys()
            for _cam_world_key in self.aggregations.keys()
        }
        deformed_human_mesh_precomputed = {
            (_cam_world_key, _aggr_world_key): None
            for _aggr_world_key in self.aggregations.keys()
            for _cam_world_key in self.aggregations.keys()
        }

        ## run aggregation for all inputs [_cam_world_key -> body_part_name -> _aggr_world_key]
        for _cam_world_key in self.aggregations.keys():
            # update number of 'cameras done'
            cam_index = self.num_processed[_cam_world_key]
            self.num_processed[_cam_world_key] += 1
                        
            ## iterate for 'body parts'
            for body_part_name in self.core[_cam_world_key].body_part_names:
                # run visual hull for all [world-cluster-_aggr_world_key] grids
                for _aggr_world_key in self.aggregations[_cam_world_key].aggr_world_name_per_body_parts[body_part_name]:
                    # skip if not related for aggregation
                    if _aggr_world_key != cam_world_key and _cam_world_key != cam_world_key:
                        continue
                    
                    # use precomputed visual hull results (if exists)
                    if mask_occupancy_precomputed[(_cam_world_key, _aggr_world_key)] is not None:
                        # load visual hull and accum visual hull
                        mask_occupancy_grid = mask_occupancy_precomputed[(_cam_world_key, _aggr_world_key)]
                        image_occupancy_grid = image_occupancy_precomputed[(_cam_world_key, _aggr_world_key)]
                        deformed_human_mesh = deepcopy(deformed_human_mesh_precomputed[(_cam_world_key, _aggr_world_key)])

                    # precompute visual hull results (if not exists)
                    else:    
                        # print progress
                        if verbose:
                            print(f"precomputing 'visual-hull' & 'full-image-visual-hull' for '_cam_world_key': {_cam_world_key} / '_aggr_world_key': {_aggr_world_key} (given 'cam_world_key': {cam_world_key})")
                         
                        ## translations before/after 'skinning-deformation'
                        translations = dict(
                            smpl2cam_trans = (-1.) * self.core[_cam_world_key].canoncam2smpl_trans, # 1x3 or None
                            aggr2smpl_trans = self.core[_aggr_world_key].canoncam2smpl_trans, # 1x3 or None
                            deformed2world_trans = None,
                        )

                        ## canonical-grid for 'aggregation' (aggregation world)
                        aggr_canon_grid = self.core[_aggr_world_key].canon_grid.clone()
                        N_x, N_y, N_z = aggr_canon_grid.shape[1:]
                        
                        ## situation 1: '(canonical) aggregation world' -> '(pose-deformed) aggregation world' (same as camera-world)
                        if _aggr_world_key == cam_world_key:
                            # translate grid in 'aggregation world' before deformation (no translation)
                            pre_translation = torch.zeros([3], dtype=torch.float32).to(self.device)
                            canon_grid_pre_transform = aggr_canon_grid + pre_translation[:,None,None,None] # 3 x N x N x N
                            # forward deformation at 'aggregation world': [N_all_voxels x 3] X [(3*J)] => [N_all_voxels x 3] X [V x 3, F x 3]
                            deformed_canon_grid, deformed_human_mesh, post_translation = self.skinning_deformation(
                                cam_world_key,
                                canon_grid_pre_transform.reshape(3,-1).T, 
                                pose,
                                betas,
                                return_human=add_human_for_vis,
                            )
                            # translate grid in '(deformed) aggregation world'
                            cam_world_grid = deformed_canon_grid + post_translation
                            

                        # situation 2: 'aggregation world' -> '(canonical) camera world' -> '(pose-deformed) camera world'
                        elif _cam_world_key == cam_world_key:
                            # translate grid in 'aggregation world' to 'camera world' before deformation
                            pre_translation = translations['aggr2smpl_trans'] + translations['smpl2cam_trans']
                            canon_grid_pre_transform = aggr_canon_grid + pre_translation[:,None,None,None] # [3 x N x N x N]
                            # forward deformation at 'camera world': [N_all_voxels x 3] X [(3*J)] => [N_all_voxels x 3] X [V x 3, F x 3]
                            deformed_canon_grid, deformed_human_mesh, post_translation = self.skinning_deformation( 
                                cam_world_key,
                                canon_grid_pre_transform.reshape(3,-1).T,
                                pose, 
                                betas,
                                return_human=add_human_for_vis,
                            )
                            # translate grid in '(deformed) camera world'
                            cam_world_grid = deformed_canon_grid + post_translation


                        ## project 'deformed-grid' (in camera-world) into camera
                        projected_grid = perspective_projection_grid(
                            grid=cam_world_grid, 
                            K=cam_K, 
                            R=cam_R, 
                            t=cam_t, 
                            grid_size=(N_x, N_y, N_z), 
                            use_normalized_K=use_normalized_K, 
                            img_size=(H,W)
                        ) # Shape: 2 x N_x x N_y x N_z
                        

                        ## check 'image-occupancy' and 'mask-occupancy'
                        mask_occupancy_grid, image_occupancy_grid, _ = self.uplift_mask_image_occupancy(projected_grid, mask, img_size=(H,W))

                        ## save to 'precompute'
                        mask_occupancy_precomputed[(_cam_world_key, _aggr_world_key)] = mask_occupancy_grid.clone()
                        image_occupancy_precomputed[(_cam_world_key, _aggr_world_key)] = image_occupancy_grid.clone()
                        deformed_human_mesh_precomputed[(_cam_world_key, _aggr_world_key)] = deepcopy(deformed_human_mesh)

                    ## check if mask-rays pass 'interaction region' (N_inter_points x 3, LongTensor)
                    # voxels within 'interaction-region'
                    interaction_regions = self.core[_cam_world_key].interaction_regions[body_part_name][_aggr_world_key]
                    
                    # does mask-rays pass 'interaction region (pass)'?
                    dont_aggregate_for_body_part = False
                    for indices in interaction_regions['pass']:
                        interaction_region_mask_occupancy = mask_occupancy_grid[indices[:,0], indices[:,1], indices[:,2]].type(torch.long) # N_interaction_points
                        if torch.sum(interaction_region_mask_occupancy) == 0:
                            dont_aggregate_for_body_part = True
                            break
                    
                    # does mask-rays avoid 'interaction region (block)'?
                    for indices in interaction_regions['block']:
                        interaction_region_mask_occupancy = mask_occupancy_grid[indices[:,0], indices[:,1], indices[:,2]].type(torch.long) # N_interaction_points
                        if torch.sum(interaction_region_mask_occupancy) > 0:
                            dont_aggregate_for_body_part = True
                            break
                    
                    # if 'interaction-region' does not lie in 'mask-occupancy', continue
                    if dont_aggregate_for_body_part:
                        if verbose:
                            print(f"\tcontinue body_part: '{body_part_name}' for aggregation")
                        continue


                    ## aggregate mask-occupancy & image-occupancy
                    self.aggregations[_cam_world_key].visual_hull_per_body_parts[body_part_name][_aggr_world_key] += mask_occupancy_grid * camera_weight
                    self.aggregations[_cam_world_key].full_image_visual_hull_per_body_parts[body_part_name][_aggr_world_key] += image_occupancy_grid * camera_weight

                    ## add 'inputs' to history
                    if len(self.aggregations[_cam_world_key].aggr_input_history_per_body_parts[body_part_name][_aggr_world_key]) < max_cam_save_num:
                        # key for saving to history
                        history_key = f'cam:{cam_index:05}'
                        
                        # print which 'aggr_input' you register to history
                        if verbose:
                            print(f"\t{body_part_name}: {len(self.aggregations[_cam_world_key].aggr_input_history_per_body_parts[body_part_name][_aggr_world_key])}")
                        
                        # add 'aggr_input' to history
                        self.aggregations[_cam_world_key].aggr_input_history_per_body_parts[body_part_name][_aggr_world_key][history_key] = {
                            'pose': pose.clone().squeeze().detach().cpu().numpy(), # 3*J
                            'betas': betas.clone().squeeze().detach().cpu().numpy(), # 10
                            'cam_R': cam_R.clone().squeeze().detach().cpu().numpy(), # 3x3
                            'cam_t': cam_t.clone().squeeze().detach().cpu().numpy(), # 3,
                            'cam_K': cam_K.clone().squeeze().detach().cpu().numpy(), # 3x3
                            'mask': mask.clone().squeeze().detach().cpu().numpy(), # HxW
                        }

                        # add camera frustum
                        self.aggregations[_cam_world_key].o3d_geometries_per_body_parts[body_part_name][_aggr_world_key][history_key] = return_camera_frustum(
                            frustum_size=frustum_size, 
                            focal_length=focal_length.detach().cpu().numpy(), 
                            cam_R=cam_R.detach().cpu().numpy(), 
                            cam_t=cam_t.detach().cpu().numpy(), 
                            color=cam_color
                        )
                        self.aggregations[_cam_world_key].geometries_per_body_parts[body_part_name][_aggr_world_key][f'cam:{cam_index:05}'] = {
                            'frustum_size': frustum_size, 
                            'focal_length': focal_length.detach().cpu().numpy(), 
                            'cam_R': cam_R.detach().cpu().numpy(), 
                            'cam_t': cam_t.detach().cpu().numpy(), 
                            'color': cam_color
                        }
                            
                        # add human for visualization
                        if add_human_for_vis:
                            # add human-mesh
                            self.add_auxiliary_geometry(
                                _cam_world_key=_cam_world_key, 
                                body_part_name=body_part_name, 
                                _aggr_world_key=_aggr_world_key,
                                name = f"posed:{cam_index:05}",
                                geometry = deformed_human_mesh,
                                exist_ok=auxiliary_exist_ok,
                                tag_auxiliary=tag_auxiliary,
                            )

                    else:
                        self.aggregations[_cam_world_key].aggr_input_history_per_body_parts[body_part_name][_aggr_world_key][f'cam:{cam_index:05}'] = None
                        self.aggregations[_cam_world_key].o3d_geometries_per_body_parts[body_part_name][_aggr_world_key][f'cam:{cam_index:05}'] = None
                        self.aggregations[_cam_world_key].geometries_per_body_parts[body_part_name][_aggr_world_key][f'cam:{cam_index:05}'] = None

                        # add human-mesh
                        if add_human_for_vis:
                            self.add_auxiliary_geometry(
                                _cam_world_key=_cam_world_key, 
                                body_part_name=body_part_name, 
                                _aggr_world_key=_aggr_world_key,
                                name = f"posed:{cam_index:05}",
                                geometry = None, 
                                exist_ok=auxiliary_exist_ok,
                                tag_auxiliary=tag_auxiliary,
                            )                            
                    
                    # add manual auxiliary geometry
                    for name, geometry in auxiliary_geometry.items():
                        self.add_auxiliary_geometry(
                            _cam_world_key=_cam_world_key, 
                            body_part_name=body_part_name, 
                            _aggr_world_key=_aggr_world_key,
                            name = name,
                            geometry = deepcopy(geometry),
                            exist_ok=auxiliary_exist_ok,
                            tag_auxiliary=tag_auxiliary,
                        )

        ## for 'full_aggregations', aggregate no matter what body-part
        _cam_world_key = cam_world_key
        # Load Accum Visual Hull
        mask_occupancy_grid = mask_occupancy_precomputed[(_cam_world_key, _cam_world_key)]
        image_occupancy_grid = image_occupancy_precomputed[(_cam_world_key, _cam_world_key)]
        deformed_human_mesh = deepcopy(deformed_human_mesh_precomputed[(_cam_world_key, _cam_world_key)])
        assert mask_occupancy_grid is not None
        assert image_occupancy_grid is not None

        self.full_aggregations[cam_world_key].visual_hull += mask_occupancy_grid * camera_weight
        self.full_aggregations[cam_world_key].full_image_visual_hull += image_occupancy_grid * camera_weight
        assert f'cam_pose:{cam_index:05}' not in self.full_aggregations[cam_world_key].aggr_input_history.keys(), "Duplicate keys."
        if len(self.full_aggregations[cam_world_key].aggr_input_history) < max_cam_save_num:
            self.full_aggregations[cam_world_key].aggr_input_history[f'cam:{cam_index:05}'] = {
                'pose': pose.clone().squeeze().detach().cpu().numpy(), # 3*J
                'betas': betas.clone().squeeze().detach().cpu().numpy(), # 10
                'cam_R': cam_R.clone().squeeze().detach().cpu().numpy(), # 3x3
                'cam_t': cam_t.clone().squeeze().detach().cpu().numpy(), # 3,
                'cam_K': cam_K.clone().squeeze().detach().cpu().numpy(), # 3x3
                'mask': mask.clone().squeeze().detach().cpu().numpy(), # HxW
            }
            self.full_aggregations[cam_world_key].o3d_geometries[f'cam:{cam_index:05}'] = \
                return_camera_frustum(
                    frustum_size=frustum_size, 
                    focal_length=focal_length.detach().cpu().numpy(), 
                    cam_R=cam_R.detach().cpu().numpy(), 
                    cam_t=cam_t.detach().cpu().numpy(), 
                    color=cam_color
                )
            self.full_aggregations[cam_world_key].geometries[f'cam:{cam_index:05}'] = \
                dict(
                    frustum_size=frustum_size, 
                    focal_length=focal_length.detach().cpu().numpy(), 
                    cam_R=cam_R.detach().cpu().numpy(), 
                    cam_t=cam_t.detach().cpu().numpy(), 
                    color=cam_color
                )
            # add human if needed
            if add_human_for_vis:
                assert deformed_human_mesh is not None
                self.add_auxiliary_geometry(
                    _cam_world_key=cam_world_key,
                    body_part_name=None, 
                    _aggr_world_key=None,
                    name = f"posed:{cam_index:05}",
                    geometry = deformed_human_mesh, 
                    exist_ok=auxiliary_exist_ok,
                    tag_auxiliary=tag_auxiliary,
                    )

            # add manual auxiliary geometry
            for name, geometry in auxiliary_geometry.items():
                self.add_auxiliary_geometry(
                    _cam_world_key=cam_world_key, 
                    body_part_name=None, 
                    _aggr_world_key=None,
                    name = name,
                    geometry = deepcopy(geometry),
                    exist_ok=auxiliary_exist_ok,
                    tag_auxiliary=tag_auxiliary,
                    )                            
        else:
            self.full_aggregations[cam_world_key].aggr_input_history[f'cam:{cam_index:05}'] = None
            self.full_aggregations[cam_world_key].o3d_geometries[f'cam:{cam_index:05}'] = None
            self.full_aggregations[cam_world_key].geometries[f'cam:{cam_index:05}'] = None
            # add human if needed
            if add_human_for_vis:
                assert deformed_human_mesh is not None
                self.add_auxiliary_geometry(
                    _cam_world_key=_cam_world_key, 
                    body_part_name=None, 
                    _aggr_world_key=None,
                    name = f"posed:{cam_index:05}",
                    geometry = None, 
                    exist_ok=auxiliary_exist_ok,
                    tag_auxiliary=tag_auxiliary,
                )


    ## skinning-deformation using lbs-weights
    def skinning_deformation(
        self, 
        cam_world_key,
        points, # to deform, (Nx3)
        pose, # (Jx3),
        betas, # (10,)
        pose2rot: bool = True, 
        return_human: bool = False,
        backward: bool = False,
        backward_lbs_weight_grid=None,
        ):
        # forward-skinning / backward-skinning
        if backward_lbs_weight_grid is None:
            lbs_weight_grid = self.core[cam_world_key].lbs_weight_grid.clone()[None] # 1 x N_x x N_y x N_z
        else:
            lbs_weight_grid = backward_lbs_weight_grid

        # [N_all_voxels x 3] X [72] => [3 x N_all_voxels] and [Deformed Template] and [Dict of Translations]
        deformation_func = getattr(self, f"skinning_deformation_{cam_world_key.lower()}")
        deformed_grid, deformed_human_mesh, J_transformed = deformation_func(
            points=points, 
            pose=pose, 
            betas=betas, 
            pose2rot=pose2rot, 
            return_human=return_human,
            lbs_weight_grid=lbs_weight_grid,
            backward=backward, 
        )
        deformed2world_trans = torch.zeros([3], dtype=points.dtype, device=points.device)
        return deformed_grid, deformed_human_mesh, deformed2world_trans


    ## lbs-skinning for smpl
    def skinning_deformation_smpl(
            self, 
            points, 
            pose, 
            betas, 
            pose2rot, 
            return_human,
            lbs_weight_grid,
            backward:bool = False,
        ):
        """ [N_all_voxels x 3] X [72] -> [3 x N_all_voxels] and [Current Human] """
        # batchify inputs
        betas = betas.unsqueeze(dim=0).to(self.device) # 1 x 10
        pose = pose.unsqueeze(dim=0).to(self.device) # 1 x 72
        points = points.unsqueeze(dim=0).to(self.device) # 1 x N_points x 3
        N_x, N_y, N_z, N_joint = self.core['SMPL'].lbs_weight_grid.shape
                        
        # apply 'blend-skinning' and get joints of 'zero-pose' template mesh
        v_shaped = self.core.GLOBAL.smpl_info['v_template_T'] + blend_shapes(betas, self.core.GLOBAL.smpl_info['shapedirs'])
        J_tpose_movedbybeta = vertices2joints(self.core.GLOBAL.smpl_info['J_regressor'], v_shaped)
        
        ## compute 'pose-blend' offsets
        # identity
        identity = torch.eye(3).to(self.device) # Shape: 3
        
        # if we use angle-axis for pose
        if pose2rot:
            # joint-pose as rotation matrices
            rot_mats = batch_rodrigues(pose.view(-1, 3)).view([1, -1, 3, 3]) # Shape: 1 x (J+1) x 3 x 3
            # pose-blend features
            pose_feature = (rot_mats[:, 1:, :, :] - identity).view([1, -1]) # Shape: 1 x 9J
            # pose-blend offsets
            pose_offsets = torch.matmul(pose_feature, self.core.GLOBAL.smpl_info['posedirs']).view(1, -1, 3) # Shape: (1 x 9J) x (9J x 3V) -> 1 x 3V -> 1 x V x 3
        
        # if we use rotation-matrix for pose
        else:
            # joint-pose as rotation matrices
            rot_mats = pose.view(1, -1, 3, 3) # Shape: 1 x (J+1) x 3 x 3
            # pose-blend features
            pose_feature = (rot_mats[:, 1:, :, :] - identity).view([1, -1]) # Shape: 1 x 9J
            # pose-blend offsets
            pose_offsets = torch.matmul(pose_feature, self.core.GLOBAL.smpl_info['posedirs']).view(1, -1, 3) # Shape: (1 x 9J) x (9J x 3V) -> 1 x 3V -> 1 x V x 3
                        
        ## add 'pose-blend' offsets to 'shape-blend' + 'template'
        v_posed = pose_offsets + v_shaped.to(self.device)
        
        ## get transformation matrices (for all joints)
        J_transformed, A_tpose2current = batch_rigid_transform(
            rot_mats=rot_mats,
            joints=J_tpose_movedbybeta, 
            parents=self.core.GLOBAL.smpl_info['parents'], 
            dtype=torch.float32
        )
        
        ## calculate 'A' (transformation matrices) is for 'user-defined-template space -> t-pose-template space'
        # if pre-computed, use the 'A' value
        if 'A_custom2tpose' in self.core.GLOBAL.smpl_info.keys():
            # 'A' (transformation matrices) for 't-pose-template space' -> 'user-defined-template space'
            A_custom2tpose = self.core.GLOBAL.smpl_info['A_custom2tpose'] # 1 x (J+1) x 4 x 4
            # 'A' (transformation matrices) for 't-pose-template space' -> 'camera space' 
            A_custom2current = A_tpose2current @ A_custom2tpose # 1 x (J+1) x 4 x 4
        
        # if not pre-computed, run pre-computation if 'user-defined zero_pose' is not entirely 0
        elif self.core.GLOBAL.smpl_info['zero_pose'].type(torch.bool).any():
            # if we use angle-axis for pose
            if pose2rot:            
                # joint-pose as rotation matrices
                rot_mats_custom2tpose = batch_rodrigues(self.core.GLOBAL.smpl_info["zero_pose"].view(-1, 3)).view([1, -1, 3, 3]) # Shape: 1 x (J+1) x 3 x 3 
            else:
                rot_mats_custom2tpose = pose.view(1, -1, 3, 3) # 1 x (J+1) x 3 x 3
            
            # joint-locations for 't-pose template'
            J_tpose = vertices2joints(self.core.GLOBAL.smpl_info['J_regressor'], self.core.GLOBAL.smpl_info['v_template_T'])

            # 'A' (transformation matrices) for 't-pose-template space' -> 'user-defined-template space'
            _, A_tpose2custom = batch_rigid_transform(
                rot_mats=rot_mats_custom2tpose, 
                joints=J_tpose, 
                parents=self.core.GLOBAL.smpl_info['parents'], 
                dtype=torch.float32
            )

            # 'A' (transformation matrices) for 'user-defined-template space' -> 't-pose-template space'
            A_custom2tpose = torch.inverse(A_tpose2custom) # 1 x (J+1) x 4 x 4
            # update precomputed 'A' to core
            self.core.GLOBAL.smpl_info["A_custom2tpose"] = A_custom2tpose
            # 'A' (transformation matrices) for 't-pose-template space' -> 'camera space' 
            A_custom2current = A_tpose2current @ A_custom2tpose # 1 x (J+1) x 4 x 4

        # if 'user-defined zero-pose' is entirely 0, use t-pose-template
        else:
            A_custom2current = A_tpose2current


        ## calculate 'forward' pose-deformation matrices (canon->deformed) for all voxels in grid
        t_canon2deform = torch.matmul(
            lbs_weight_grid.reshape(1, N_x * N_y * N_z, N_joint), 
            A_custom2current.view(1, A_custom2current.shape[1], 16)
        ).view(1, -1, 4, 4) # (1 x N_points x J) x (1 x J x 16) -> 1 x N_points x 16 -> 1 x N_points x 4 x 4

        ## calculate 'backward' pose-deformation matrices (deformed->canon) by computing inverse of 'forward' matrices
        t_final = torch.inverse(t_canon2deform) if backward else t_canon2deform

        ## apply 'canonical-space grid' to 'current-space grid' transformation
        # homogeneous coordinates of points (to transform)
        points_homo = torch.cat([points, torch.ones_like(points)[...,0:1]], dim=-1) # 1 x N_points x 4
        
        # apply transformation
        points_current = torch.matmul(t_final, points_homo.unsqueeze(dim=-1)) # (1 x N_points x 4 x 4) x (1 x N_points x 4 x 1) -> (1 x N_points x 4 x 1)
        
        # dehomogenization
        points_current = points_current[:,:,:3,0] # 1 x N_points x 3
        current_grid = points_current[0] # N_points x 3

        ## return human-mesh (smpl-mesh) in 'current-space'
        if not return_human:
            current_human_mesh = None
        else:
            # human mesh lbs-weights
            W = self.core.GLOBAL.smpl_info['lbs_weights'].unsqueeze(dim=0).expand([1, -1, -1]) # 1 x V x (1+J)
            
            # number of joints
            num_joints = self.core.GLOBAL.smpl_info['J_regressor'].shape[0] 
            
            # transformations for each-vertex on 'smpl-mesh'
            T = torch.matmul(
                W, 
                A_tpose2current.view(1, num_joints, 16)
            ).view(1, -1, 4, 4) # (1 x V x (J + 1)) x (1 x (J + 1) x 16) -> (1 x V x 16) -> (1 x V x 4 x 4)
            
            # homogeneous coordinates of each-vertex on 'smpl-mesh'
            v_posed_homo = torch.cat([v_posed, torch.ones_like(v_posed)[...,0:1]], dim=2) # 1 x N_vertices x 4

            # apply transformation
            v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))
            verts = v_homo[:, :, :3, 0]

            ## current human-mesh as 'geometry-dict'
            current_human_mesh_dict = {
                'vertices': verts.squeeze().detach().cpu().numpy().astype(np.float32), 
                'triangles': self.core.SMPL.v_template_faces_holistic.squeeze().detach().cpu().numpy().astype(np.long),
                'vertex_colors': self.core.SMPL.v_template_color.squeeze().detach().cpu().numpy(),
            }

            ## current human-mesh as 'open3d' 
            current_human_mesh = o3d.geometry.TriangleMesh()
            current_human_mesh.vertices = o3d.utility.Vector3dVector(current_human_mesh_dict['vertices'])
            current_human_mesh.triangles = o3d.utility.Vector3iVector(current_human_mesh_dict['triangles'])
            
            # color open3d objects 
            if current_human_mesh_dict['vertex_colors'].shape == (3,): 
                current_human_mesh.paint_uniform_color(current_human_mesh_dict['vertex_colors'])
            else:
                current_human_mesh.vertex_colors = o3d.utility.Vector3iVector(current_human_mesh_dict['vertex_colors'])
            current_human_mesh.compute_vertex_normals()

        return current_grid, current_human_mesh, J_transformed
    

    ## discretize & render binary mask per 'discretization-thresholds'
    def render_discretized_distribution(
            self,
            cam_world_key,
            cam_R,
            cam_t,
            cam_K,
            pose,
            betas,
            mask_shape,
            downscale_length,
            use_normalized_K,
            thresholds,
            verbose
        ):

        # prepare inputs
        cam_R = to_np_torch_recursive(cam_R, use_torch=self.use_torch, device=self.device)
        cam_t = to_np_torch_recursive(cam_t, use_torch=self.use_torch, device=self.device)
        cam_K = to_np_torch_recursive(cam_K, use_torch=self.use_torch, device=self.device)

        # (downscaled) target-mask shape
        H, W = mask_shape
        downscale_factor = downscale_length / min(H, W)
        target_H = int(H * downscale_factor) if H > W else downscale_length
        target_W = int(W * downscale_factor) if H < W else downscale_length

        ## 'downscale-factor' for match pixel-scale in cam_K 
        cam_K[0,0] *= downscale_factor
        cam_K[1,1] *= downscale_factor
        cam_K[0,2] *= downscale_factor
        cam_K[1,2] *= downscale_factor

        # retrieve shape of grid
        N_x, N_y, N_z = self.core[cam_world_key].indexgrid.shape[1:]

        ## deform 'learned-probability in canonical space' to 'pose-deformed space'
        deformed_probability_grids_total = self.deform_probability_field(
            cam_world_key=cam_world_key, 
            pose=pose,
            betas=betas,
            return_current_human=True,
            deformation_method="backward-deformation",
            only_return_full=True,
            verbose=verbose,
        )
        
        # deformed-human is same for all body-part clusters: variable to save precomputed human
        human_mask_precomputed = None

        # save rendered-masks per 'discretization thresholds'
        rendered_masks_per_thresh = EasyDict()

        # iterate for 'thresholds'
        for threshold in sorted(thresholds):
            ## render object-mask
            # placeholder to save 'rendered-masks'
            rendered_masks = EasyDict()
            
            # apply 'discretization' to probability in 3D
            threeDmask = deformed_probability_grids_total.full_aggregations[cam_world_key]['deformed_prob_grid'] >= threshold # Shape: NxNxN, bool
    
            # voxel-grid
            cam_world_grid = self.core[cam_world_key].canon_grid.clone() # voxel-grid in 'pose-deformed space'. Shape: 3 x N x N x N 

            # voxels above with 'probability above-threshold'
            cam_world_grid_overthresh = cam_world_grid[:,threeDmask].T # N_over, 3
            
            # project 'binary-grid' (in camera-world) into camera
            projected_grid = perspective_projection_grid(
                grid=cam_world_grid_overthresh,
                K=cam_K,
                R=cam_R,
                t=cam_t,
                grid_size=(len(cam_world_grid_overthresh),),
                use_normalized_K=use_normalized_K,
                img_size=(target_H,target_W)
            ) # Shape: 2 x N_over

            # remove indices if 'projected-points' are outside the image (note: dtype turns to torch.long)
            projected_grid, inside_img = self.retain_projection_inside_img(projected_grid, img_size=(target_H,target_W)) # 2 x N_over_insideimg

            # W,H -> H,W
            projected_grid = torch.flip(projected_grid, dims=[0]) # 2 x N_over_insideimg            

            ## create 'object-rendered-mask' by checking binary occupancies
            # generate as a sparse-tensor (regarding 'coalescence')
            pred_mask = torch.sparse_coo_tensor(
                indices=projected_grid, # 2 x N_over_insideimg
                values=torch.ones_like(projected_grid[0]),
                size=(target_H,target_W)
            ).to_dense()
            # convert as long dtype
            pred_mask = (pred_mask > 0).type(torch.long)
            
            ## fill in the results (object mask)
            rendered_masks['object_mask'] = pred_mask.clone().cpu().numpy() # still torch

            ## render human-mask
            # use precomputed if exists
            if human_mask_precomputed is not None:  
                human_mask = human_mask_precomputed.clone()

            # precompute human-mask if not exists
            else:
                # render 'human-mesh' vertices
                deformed_human = deformed_probability_grids_total.full_aggregations[cam_world_key]['deformed_human_mesh']
                deformed_human_verts = torch.tensor(deformed_human['verts'], dtype=torch.float32).to(self.device) # Shape: V x 3

                # project 'human-mesh' vertices (in camera-world) into camera
                projected_grid_human = perspective_projection_grid(
                    grid=deformed_human_verts,
                    K=cam_K,
                    R=cam_R,
                    t=cam_t,
                    grid_size=(len(deformed_human_verts),),
                    use_normalized_K=use_normalized_K,
                    img_size=(target_H,target_W)
                ) # Shape: 2 x V

                # remove indices if 'projected-points' are outside the image (note: dtype turns to torch.long)
                projected_grid_human, inside_img_human = self.retain_projection_inside_img(projected_grid_human, img_size=(target_H,target_W)) # 2 x V_over_insideimg

                # W,H -> H,W
                projected_grid_human = torch.flip(projected_grid_human, dims=[0]) # 2 x V_over_insideimg
                
                ## create 'human-rendered-mask' by checking binary occupancies
                # generate as a sparse-tensor (regarding 'coalescence')
                human_mask = torch.sparse_coo_tensor(
                    indices=projected_grid_human, # 2 x V_over_insideimg
                    values=torch.ones_like(projected_grid_human[0]),
                    size=(target_H,target_W)
                ).to_dense()
                # convert as long dtype
                human_mask = (human_mask > 0).type(torch.long)

                # save as precompute
                human_mask_precomputed = human_mask.clone()

            ## fill in the results (object mask)
            rendered_masks['human_mask'] = human_mask.clone().cpu().numpy()

            # save to rendered_masks_per_thersh
            rendered_masks_per_thresh[f"{threshold:.6f}"] = rendered_masks

        return rendered_masks_per_thresh, (target_H, target_W)


    ## deform probability field in 'canonical space' to 'pose-deformed space'
    def deform_probability_field(self,
            cam_world_key,
            pose,
            betas,
            return_current_human: bool = True,
            deformation_method = "backward-deformation", # default
            only_return_full=False, # only for evaluation
            eps=torch.tensor([1e-7], dtype=torch.float32),
            verbose=False,
        ):
        ## prepare inputs
        pose = to_np_torch_recursive(pose, use_torch=self.use_torch, device=self.device)
        betas = to_np_torch_recursive(betas, use_torch=self.use_torch, device=self.device)
        eps = to_np_torch_recursive(eps, use_torch=self.use_torch, device=self.device)
        pose = pose.squeeze()
        betas = betas.squeeze()
        
        ## placeholder to return computed values
        deformed_probability_grids = EasyDict(dict(full_aggregations=None, aggregations=None))

        ## we don't need to re-compute deformation result (deformed human-mesh / backward lbs-weights)
        deformed_human_mesh_precomputed = {
            (_cam_world_key, _aggr_world_key): None
            for _aggr_world_key in list(self.aggregations.keys()) + [None] # None used for full
            for _cam_world_key in self.aggregations.keys()
        }
        backward_lbs_weights_precomputed = {
            (_cam_world_key, _aggr_world_key): None
            for _aggr_world_key in list(self.aggregations.keys()) + [None] # None used for full
            for _cam_world_key in self.aggregations.keys()
        }

        ## run deformation for 'full_aggregations'
        # placeholder for saving results
        deformed_probability_grids['full_aggregations'] = EasyDict()
        
        # iterate for all 'camera-world' keys
        for _cam_world_key in self.full_aggregations.keys():
            # for 'full_aggregations', aggregation-world and camera-world are equivalent
            _aggr_world_key = _cam_world_key
                       
            # skip if '_cam_world_key' is not equal to 'cam_world_key' (since '_aggr_world_key == _cam_world_key' in this case)
            if _cam_world_key != cam_world_key:
                deformed_probability_grids['full_aggregations'][_cam_world_key] = None
                continue

            # placeholder for saving results
            deformed_probability_grids['full_aggregations'][_cam_world_key] = EasyDict()
                        
            ## 'deformed' human-mesh
            # use precomputed if exists
            if deformed_human_mesh_precomputed[(_cam_world_key, _aggr_world_key)] is not None:
                deformed_human_mesh = deformed_human_mesh_precomputed[(_cam_world_key, _aggr_world_key)].clone()
            
            # precompute 'deformed' human-mesh if not exists
            else:
                # translations before/after 'skinning-deformation'
                translations = dict(
                    smpl2cam_trans = (-1.) * self.core[_cam_world_key].canoncam2smpl_trans, # 1x3 or None
                    aggr2smpl_trans = self.core[_aggr_world_key].canoncam2smpl_trans, # 1x3 or None
                    deformed2world_trans = None,
                )

                # canonical-grid in 'aggregation world'
                aggr_canon_grid = self.core[_aggr_world_key].canon_grid.clone()
                N_x, N_y, N_z = aggr_canon_grid.shape[1:]
                
                ## situation 1: '(canonical) aggregation world' -> '(pose-deformed) aggregation world' (same as camera-world)
                if _aggr_world_key == cam_world_key:
                    # translate grid in 'aggregation world' before deformation (no translation)
                    pre_translation = torch.zeros([3], dtype=torch.float32).to(self.device)
                    canon_grid_pre_transform = aggr_canon_grid + pre_translation[:,None,None,None] # 3 x N x N x N
                    # forward deformation at 'aggregation world': [N_all_voxels x 3] X [(3*J)] => [N_all_voxels x 3] X [V x 3, F x 3]
                    deformed_canon_grid, deformed_human_mesh, post_translation = self.skinning_deformation(
                        cam_world_key,
                        canon_grid_pre_transform.reshape(3,-1).T, 
                        pose,
                        betas,
                        return_human=return_current_human,
                    )
                    
                ## situation 2: 'aggregation world' -> '(canonical) camera world' -> '(pose-deformed) camera world'
                elif _cam_world_key == cam_world_key:
                    # translate grid in 'aggregation world' to 'camera world' before deformation
                    pre_translation = translations['aggr2smpl_trans'] + translations['smpl2cam_trans']
                    canon_grid_pre_transform = aggr_canon_grid + pre_translation[:,None,None,None] # [3 x N x N x N]
                    # forward deformation at 'camera world': [N_all_voxels x 3] X [(3*J)] => [N_all_voxels x 3] X [V x 3, F x 3]
                    deformed_canon_grid, deformed_human_mesh, post_translation = self.skinning_deformation( 
                        cam_world_key,
                        canon_grid_pre_transform.reshape(3,-1).T,
                        pose, 
                        betas,
                        return_human=return_current_human,
                    )
            
                # save precomputed 'deformed' human-mesh
                deformed_human_mesh_precomputed[(_cam_world_key, _aggr_world_key)] = deformed_human_mesh

            # learned probability field in 'canonical world'. Shape: N x N x N
            canon_visual_hull = self.full_aggregations[_cam_world_key].visual_hull.to(self.device)
            canon_full_image_visual_hull = self.full_aggregations[_cam_world_key].full_image_visual_hull.to(self.device)
            canon_probability_field = canon_visual_hull / torch.where(canon_full_image_visual_hull > eps, canon_full_image_visual_hull, eps)

            ## 'backward' lbs weights
            # deformation-methods
            if deformation_method != "backward-deformation":
                assert deformation_method in ['linear', 'negative-log-mixing'], "Unavailable 'deformation-method'"
                assert False, f"Method '{deformation_method}' deprecated..."

            # use precomputed if exists
            if backward_lbs_weights_precomputed[(_cam_world_key, _aggr_world_key)] is not None:
                backward_lbs_weight_grid = backward_lbs_weights_precomputed[(_cam_world_key, _aggr_world_key)]
            
            # precompute 'backward' lbs-weights if not exists
            else:
                # use deformed human-mesh as template
                v_template_deformed = torch.tensor(np.asarray(deformed_human_mesh.vertices), dtype=torch.float32)
                
                # eqiuspaced-grid in camera world
                equispace_grid_cam_world = self.core[cam_world_key].canon_grid.clone() # 3 x N x N x N

                # compute 'backward' lbs-weights
                backward_lbs_weight_grid = precompute_lbs_weights(
                    canon_grid=equispace_grid_cam_world,
                    v_template=v_template_deformed,
                    lbs_weights=self.core[cam_world_key].lbs_weights.clone(),
                    lbs_precompute_settings=self.core[cam_world_key].lbs_precompute_settings,
                    verbose=verbose
                ).to(self.device).type(torch.float32)
        
                # save to precomputed 'backward' lbs weights
                backward_lbs_weights_precomputed[(_cam_world_key, _aggr_world_key)] = backward_lbs_weight_grid

            ## apply 'backward' deformation
            # align before 'backward' deformation
            grid_deformed_world = equispace_grid_cam_world.reshape(3,-1).T - post_translation
            
            # 'backward' skinning. [N_all_voxels x 3] X [(3*J)] => [N_all_voxels x 3] X [V x 3, F x 3]
            canon_grid_pre_transform, _, _ = self.skinning_deformation(
                cam_world_key,
                grid_deformed_world,
                pose, 
                betas,
                return_human=return_current_human,
                backward=True,
                backward_lbs_weight_grid=backward_lbs_weight_grid,
            )

            # align to 'aggregation world'
            aggr_canon_grid = canon_grid_pre_transform.T.reshape(3,N_x,N_y,N_z) - pre_translation[:,None,None,None] # Shape: 3 x N x N x N
            
            # retrieve indices that fall within 'canonical-space' grid when backward-deformation applied
            leftbottom = self.core[_aggr_world_key].center - 0.5*torch.tensor(
                [
                    self.core[_aggr_world_key].length_x,
                    self.core[_aggr_world_key].length_y,
                    self.core[_aggr_world_key].length_z
                ], 
                dtype=torch.float32
            ).to(self.device)
            aggr_canon_index_grid = aggr_canon_grid - leftbottom[:,None,None,None] # Shape: 3 x N x N x N
            aggr_canon_index_grid /= self.core[_aggr_world_key].voxel_size # Index-Scale. Shape: 3 x N x N x N
            closest_equivoxel_indices = aggr_canon_index_grid.type(torch.long)

            # find indices of voxels that are warped to 'inside' of (canonical) aggr-grid
            is_above_zero = (closest_equivoxel_indices >= 0).all(dim=0) # 3 x N x N x N
            is_x_below_max = closest_equivoxel_indices[0] < N_x
            is_y_below_max = closest_equivoxel_indices[1] < N_y
            is_z_below_max = closest_equivoxel_indices[2] < N_z
            is_below_max = torch.stack([is_x_below_max, is_y_below_max, is_z_below_max], dim=0).all(dim=0)
            is_inside_canongrid = torch.logical_and(is_above_zero, is_below_max) # NxNxN, bool
            
            # if voxels are 'backward' warped to outside of 'canonical' grid, fill probability with 0 value
            closest_equivoxel_indices = torch.where(is_inside_canongrid[None], closest_equivoxel_indices, 0)
            
            # if voxels are 'backward' warped to inside of 'canonical' grid, fill probability with queried values
            deformed_prob_grid = canon_probability_field[
                closest_equivoxel_indices[0,:,:,:],
                closest_equivoxel_indices[1,:,:,:],
                closest_equivoxel_indices[2,:,:,:],
            ]
            deformed_prob_grid = deformed_prob_grid * is_inside_canongrid.type(torch.float)

            # save to return
            deformed_probability_grids['full_aggregations'][_cam_world_key]['deformed_prob_grid'] = deformed_prob_grid.clone().squeeze().detach().cpu().numpy() # N x N x N
            deformed_probability_grids['full_aggregations'][_cam_world_key]['deformed_human_mesh'] = {
                "verts": np.asarray(deformed_human_mesh.vertices),
                "faces": np.asarray(deformed_human_mesh.triangles),
                "color": np.asarray(deformed_human_mesh.vertex_colors) if np.asarray(deformed_human_mesh.vertex_colors).shape[0] != 0 else None
            }

        # run Accumulation for All [_cam_world_key -> body_part_name -> _aggr_world_key]
        deformed_probability_grids['aggregations'] = EasyDict()
        for _cam_world_key in self.aggregations.keys():
            # if only return 'full_aggregations', then skip
            if only_return_full:
                continue

            # placeholder for saving results
            deformed_probability_grids['aggregations'][_cam_world_key] = EasyDict()

            # iterate for all 'body parts'
            for body_part_name in self.core[_cam_world_key].body_part_names:
                # placeholder for saving results
                deformed_probability_grids['aggregations'][_cam_world_key][body_part_name] = EasyDict()
                
                # iterate for all 'aggregation-world key' (for given body-part)
                for _aggr_world_key in self.aggregations[_cam_world_key].aggr_world_name_per_body_parts[body_part_name]:

                    # skip if '_cam_world_key'/'_aggr_world_key' is not equal to 'cam_world_key'
                    if _aggr_world_key != cam_world_key and _cam_world_key != cam_world_key:
                        deformed_probability_grids['aggregations'][_cam_world_key][body_part_name][_aggr_world_key] = None
                        continue
                    
                    # placeholder for saving results
                    deformed_probability_grids['aggregations'][_cam_world_key][body_part_name][_aggr_world_key] = EasyDict()
                    

                    ## 'deformed' human-mesh
                    # use precomputed if exists
                    if deformed_human_mesh_precomputed[(_cam_world_key, _aggr_world_key)] is not None:
                        deformed_human_mesh = deformed_human_mesh_precomputed[(_cam_world_key, _aggr_world_key)]

                    # precompute 'deformed' human-mesh if not exists
                    else:
                        # translations before/after 'skinning-deformation'
                        translations = dict(
                            smpl2cam_trans = (-1.) * self.core[_cam_world_key].canoncam2smpl_trans, # 1x3 or None
                            aggr2smpl_trans = self.core[_aggr_world_key].canoncam2smpl_trans, # 1x3 or None
                            deformed2world_trans = None,
                        )

                        # canonical-grid in 'aggregation world'
                        aggr_canon_grid = self.core[_aggr_world_key].canon_grid.clone()
                        N_x, N_y, N_z = aggr_canon_grid.shape[1:]
                        
                        ## situation 1: '(canonical) aggregation world' -> '(pose-deformed) aggregation world' (same as camera-world)
                        if _aggr_world_key == cam_world_key:
                            # translate grid in 'aggregation world' before deformation (no translation)
                            pre_translation = torch.zeros([3], dtype=torch.float32).to(self.device)
                            canon_grid_pre_transform = aggr_canon_grid + pre_translation[:,None,None,None] # 3 x N x N x N
                            # forward deformation at 'aggregation world': [N_all_voxels x 3] X [(3*J)] => [N_all_voxels x 3] X [V x 3, F x 3]
                            _, deformed_human_mesh, _ = self.skinning_deformation(
                                cam_world_key,
                                canon_grid_pre_transform.reshape(3,-1).T, 
                                pose,
                                betas,
                                return_human=return_current_human,
                            )                            

                        ## situation 2: 'aggregation world' -> '(canonical) camera world' -> '(pose-deformed) camera world'        
                        elif _cam_world_key == cam_world_key:
                            # translate grid in 'aggregation world' to 'camera world' before deformation
                            pre_translation = translations['aggr2smpl_trans'] + translations['smpl2cam_trans']
                            canon_grid_pre_transform = aggr_canon_grid + pre_translation[:,None,None,None] # [3 x N x N x N]
                            # forward deformation at 'camera world': [N_all_voxels x 3] X [(3*J)] => [N_all_voxels x 3] X [V x 3, F x 3]
                            deformed_canon_grid, deformed_human_mesh, post_translation = self.skinning_deformation(
                                cam_world_key,
                                canon_grid_pre_transform.reshape(3,-1).T,
                                pose, 
                                betas,
                                return_human=return_current_human,
                            )

                        # save precomputed 'deformed' human-mesh
                        deformed_human_mesh_precomputed[(_cam_world_key, _aggr_world_key)] = deformed_human_mesh

                    # Deform log-(1-probability) field in canonical space and interpolate
                    canon_visual_hull = self.aggregations[_cam_world_key].visual_hull_per_body_parts[body_part_name][_aggr_world_key].to(self.device)
                    canon_full_image_visual_hull = self.aggregations[_cam_world_key].full_image_visual_hull_per_body_parts[body_part_name][_aggr_world_key].to(self.device)
                    # Canonical probability field, NxNxN
                    canon_probability_field = canon_visual_hull / torch.where(canon_full_image_visual_hull > eps, canon_full_image_visual_hull, eps)
                    
                    ## 'backward' lbs weights
                    # deformation-methods
                    if deformation_method != "backward-deformation":
                        assert deformation_method in ['linear', 'negative-log-mixing'], "Unavailable 'deformation-method'"
                        assert False, f"Method '{deformation_method}' deprecated..."

                    # use precomputed if exists
                    if backward_lbs_weights_precomputed[(_cam_world_key, _aggr_world_key)] is not None:
                        backward_lbs_weight_grid = backward_lbs_weights_precomputed[(_cam_world_key, _aggr_world_key)]

                    # precompute 'backward' lbs-weights if not exists
                    else:
                        # use deformed human-mesh as template
                        v_template_deformed = torch.tensor(np.asarray(deformed_human_mesh.vertices), dtype=torch.float32)
                        
                        # eqiuspaced-grid in camera world
                        equispace_grid_cam_world = self.core[cam_world_key].canon_grid.clone() # 3 x N x N x N
                        
                        # compute 'backward' lbs-weights
                        backward_lbs_weight_grid = precompute_lbs_weights(
                            canon_grid=equispace_grid_cam_world,
                            v_template=v_template_deformed,
                            lbs_weights=self.core[cam_world_key].lbs_weights.clone(),
                            **self.core[cam_world_key].lbs_precompute_settings
                        ).to(self.device).type(torch.float32)

                        # save to precomputed 'backward' lbs weights
                        backward_lbs_weights_precomputed[(_cam_world_key, _aggr_world_key)] = backward_lbs_weight_grid
                            
                    ## apply 'backward' deformation
                    # align before 'backward' deformation
                    grid_deformed_world = equispace_grid_cam_world.reshape(3,-1).T - post_translation
                    
                    # 'backward' skinning. [N_all_voxels x 3] X [(3*J)] => [N_all_voxels x 3] X [V x 3, F x 3]
                    canon_grid_pre_transform, _, _ = self.skinning_deformation(
                        cam_world_key,
                        grid_deformed_world,
                        pose, 
                        betas,
                        return_human=return_current_human,
                        backward=True,
                        backward_lbs_weight_grid=backward_lbs_weight_grid,
                    )

                    # align to 'aggregation world'
                    aggr_canon_grid = canon_grid_pre_transform.T.reshape(3,N_x,N_y,N_z) - pre_translation[:,None,None,None] # [3 x N x N x N]
                    
                    # retrieve indices that fall within 'canonical-space' grid when backward-deformation applied
                    leftbottom = self.core[_aggr_world_key].center - 0.5*torch.tensor(
                        [
                            self.core[_aggr_world_key].length_x,
                            self.core[_aggr_world_key].length_y,
                            self.core[_aggr_world_key].length_z
                        ], dtype=torch.float32
                    ).to(self.device)
                    aggr_canon_index_grid = aggr_canon_grid - leftbottom[:,None,None,None] # Shape: 3 x N x N x N
                    aggr_canon_index_grid /= self.core[_aggr_world_key].voxel_size # Index-Scale. Shape: 3 x N x N x N
                    closest_equivoxel_indices = aggr_canon_index_grid.type(torch.long)

                    # find indices of voxels that are warped to 'inside' of (canonical) aggr-grid
                    is_above_zero = (closest_equivoxel_indices >= 0).all(dim=0) # 3 x N x N x N
                    is_x_below_max = closest_equivoxel_indices[0] < N_x
                    is_y_below_max = closest_equivoxel_indices[1] < N_y
                    is_z_below_max = closest_equivoxel_indices[2] < N_z
                    is_below_max = torch.stack([is_x_below_max, is_y_below_max, is_z_below_max], dim=0).all(dim=0)
                    is_inside_canongrid = torch.logical_and(is_above_zero, is_below_max) # NxNxN, bool

                    # if voxels are 'backward' warped to outside of 'canonical' grid, fill probability with 0 value
                    closest_equivoxel_indices = torch.where(
                        is_inside_canongrid[None], closest_equivoxel_indices, 0)
                    
                    # if voxels are 'backward' warped to inside of 'canonical' grid, fill probability with queried values
                    deformed_prob_grid = canon_probability_field[
                        closest_equivoxel_indices[0,:,:,:],
                        closest_equivoxel_indices[1,:,:,:],
                        closest_equivoxel_indices[2,:,:,:],
                    ]
                    deformed_prob_grid = deformed_prob_grid * is_inside_canongrid.type(torch.float)

                    ## get the probability
                    deformed_probability_grids['aggregations'][_cam_world_key][body_part_name][_aggr_world_key]['deformed_prob_grid'] = deformed_prob_grid.clone().squeeze().detach().cpu().numpy() # NxNxN
                    deformed_probability_grids['aggregations'][_cam_world_key][body_part_name][_aggr_world_key]['deformed_human_mesh'] = {
                        "verts": np.asarray(deformed_human_mesh.vertices),
                        "faces": np.asarray(deformed_human_mesh.triangles),
                        "color": np.asarray(deformed_human_mesh.vertex_colors) if np.asarray(deformed_human_mesh.vertex_colors).shape[0] != 0 else None
                    }

        # if only returning 'full_aggregations', delete the 'aggregations'
        if only_return_full: del deformed_probability_grids['aggregations']

        return deformed_probability_grids


    ## visualize as open3d geometries
    def visualize_current_o3d_geometries(self, verbose):
        # visualize 'full_aggregations'
        for _cam_world_key in self.full_aggregations.keys():
            # print progress
            if verbose:
                print(f"visualization for 'full_aggregations' in _cam_world_key: {_cam_world_key}\n")
            
            # visualization via open3d
            o3d.visualization.draw_geometries(list(self.full_aggregations[_cam_world_key].o3d_geometries.values()))
        
        # visualize 'aggregations'
        for _cam_world_key in self.aggregations.keys():
            if _cam_world_key == 'GLOBAL':
                continue

            # iterate for 'body parts'
            for body_part_name in self.core[_cam_world_key].body_part_names:
                if verbose:
                    print(f"visualization for 'aggregations': _cam_world_key:{_cam_world_key} / body part: {body_part_name} / _aggr_world_key: {self.aggregations[_cam_world_key].aggr_world_name_per_body_parts[body_part_name]}==========\n")
                
                # placeholder to save geometries
                geometries_cluster = []
                for _aggr_world_key in self.aggregations[_cam_world_key].aggr_world_name_per_body_parts[body_part_name]:
                    # update list of geometries
                    geometries_cluster += list(self.aggregations[_cam_world_key].o3d_geometries_per_body_parts[body_part_name][_aggr_world_key].values())
                    
                # visualization via open3d
                o3d.visualization.draw_geometries(geometries_cluster)


    ## auxiliary geometries for additional visualization
    def add_auxiliary_geometry(
            self, 
            _cam_world_key, 
            body_part_name,
            _aggr_world_key,
            name, 
            geometry, 
            exist_ok=False,
            tag_auxiliary=False
        ):
        assert not exist_ok and not tag_auxiliary, "Other settings not implemented..."
        
        ## add auxiliary geometry for 'aggregations'
        if body_part_name is not None:
            assert _aggr_world_key is not None, "'_aggr_world_key' must be provided when 'body_part_name' is given."

            # if duplicate key exists in set of geometries
            if name in self.aggregations[_cam_world_key].o3d_geometries_per_body_parts[body_part_name][_aggr_world_key].keys():
                # raise error 
                assert exist_ok, f"There already exists \"{name}\" in self.aggregations.{_cam_world_key}.{body_part_name}.{_aggr_world_key}"
                # remove existing key and save new 'auxiliary geometry'
                del self.aggregations[_cam_world_key].o3d_geometries_per_body_parts[body_part_name][_aggr_world_key][name]
            
            # add 'tags' in front of key for 'auxiliary geometry'
            if tag_auxiliary:
                name = f"{_cam_world_key}:{body_part_name}:{_aggr_world_key}:{name}"

            # save to set of geometries
            self.aggregations[_cam_world_key].o3d_geometries_per_body_parts[body_part_name][_aggr_world_key][name] = geometry
            self.aggregations[_cam_world_key].geometries_per_body_parts[body_part_name][_aggr_world_key][name] = { # check 'return_triangle_mesh' function
                "verts": np.asarray(geometry.vertices),
                "faces": np.asarray(geometry.triangles),
                "color": np.asarray(geometry.vertex_colors) if np.asarray(geometry.vertex_colors).shape[0] != 0 else None,
                "translation": None,
            }

        ## add auxiliary geometry for 'full_aggregations'
        else:
            assert _aggr_world_key is None, "'_aggr_world_key' must not be provided when 'body_part_name' is not given."
            
            # if duplicate key exists in set of geometries
            if name in self.full_aggregations[_cam_world_key].o3d_geometries.keys():
                # raise error
                assert exist_ok, f"There already exists \"{name}\" in self.full_aggregations['{_cam_world_key}']"
                # remove existing key and save new 'auxiliary geometry'
                del self.full_aggregations[_cam_world_key].o3d_geometries[name]

            # add 'tags' in front of key for 'auxiliary geometry'
            if tag_auxiliary:
                name = f"auxiliary:{_cam_world_key}:{body_part_name}:{_aggr_world_key}:{name}"

            # save to set of geometries
            self.full_aggregations[_cam_world_key].o3d_geometries[name] = geometry
            self.full_aggregations[_cam_world_key].geometries[name] = { # hint: check 'return_triangle_mesh' function
                "verts": np.asarray(geometry.vertices),
                "faces": np.asarray(geometry.triangles),
                "color": np.asarray(geometry.vertex_colors) if np.asarray(geometry.vertex_colors).shape[0] != 0 else None,
                "translation": None,
            }


    ## generate video for single 'aggregation result'
    def generate_video_single(
            self, 
            _cam_world_key, 
            body_part_name, # if None, 'full_aggregations' is visualized
            _aggr_world_key, # if None, 'full_aggregations' is visualized
            save_pth, 
            offscreen,  
            title,
            eps=1e-7,
            auxiliary_geometries=dict(), 
            normalize_max=False,
            tmp_image_cache = None,
            window_size = 500,
            create_video=False,
            title_size=0.1,
            title_color = (0,0,0),
            alpha = 0.8,
            opacity = 0.3, 
            mesh_representation = 'surface', # 'surface', 'wireframe', 'points', 'mesh', 'fancymesh'
            background_color = (1,1,1),
            outline_color = (0,0,0),
            scalar_field_vmin = 0.05,
            scalar_field_vmax = 1.0,
        ):

        # use mayavi.mlab's pipeline's visualizer as off-screen
        if offscreen: mlab.options.offscreen = offscreen
            
        # raise error if '_cam_world_key' not provided
        assert _cam_world_key is not None
                
        # declare mayavi figure
        fig = mlab.figure(size=(window_size, window_size), bgcolor=background_color)

        # voxel-grid settings
        length_x = self.core[_cam_world_key].length_x
        length_y = self.core[_cam_world_key].length_y
        length_z = self.core[_cam_world_key].length_z
        center = self.core[_cam_world_key].center.cpu().numpy()
        voxel_size = self.core[_cam_world_key].voxel_size
        voxel_resolution_x = int(round(length_x / voxel_size, 0))
        voxel_resolution_y = int(round(length_y / voxel_size, 0))
        voxel_resolution_z = int(round(length_z / voxel_size, 0))

        ## visualize probability grid from 'aggregations' as video
        if body_part_name is not None:
            assert _aggr_world_key is not None, "'_aggr_world_key' must be provided when 'body_part_name' is given."
            
            # calculate 'probability grid'
            visual_hull = self.aggregations[_cam_world_key].visual_hull_per_body_parts[body_part_name][_aggr_world_key].detach().cpu().numpy()
            full_image_visual_hull = self.aggregations[_cam_world_key].full_image_visual_hull_per_body_parts[body_part_name][_aggr_world_key].detach().cpu().numpy()
            prob_visual_hull = visual_hull / np.where(full_image_visual_hull > eps, full_image_visual_hull, eps)
            if normalize_max:
                prob_visual_hull = prob_visual_hull / max(prob_visual_hull.max(), eps)

            # geometries
            geometries = self.aggregations[_cam_world_key].geometries_per_body_parts[body_part_name][_aggr_world_key]

        ## visualize probability grid from 'full_aggregations' as video
        else:
            assert _aggr_world_key is None, "'_aggr_world_key' must not be provided when 'body_part_name' is not given."
        
            # calculate 'probability grid'
            visual_hull = self.full_aggregations[_cam_world_key].visual_hull.detach().cpu().numpy()
            full_image_visual_hull = self.full_aggregations[_cam_world_key].full_image_visual_hull.detach().cpu().numpy()
            prob_visual_hull = visual_hull / np.where(full_image_visual_hull > eps, full_image_visual_hull, eps)
            if normalize_max:
                prob_visual_hull = prob_visual_hull / max(prob_visual_hull.max(), eps)

            # geometries
            geometries = self.full_aggregations[_cam_world_key].geometries

        # create scalar-field volume source for visualization (automatically added to 'mayavi.mlab' pipeline when declared)
        mlab.pipeline.volume(mlab.pipeline.scalar_field(prob_visual_hull), vmin=scalar_field_vmin, vmax=scalar_field_vmax)
        mlab.outline(color=outline_color)
        mlab.orientation_axes()

        # visualize human
        for key in list(geometries.keys()):
            # skip if geometry is not human
            if key[:6] != "posed:" and key != "template":
                continue

            # retrieve 'vertices', 'faces(triangles)', 'vertex colors' of 'human' as 'np.ndarray' on 'cpu'
            posed_human = geometries[key]
            vertices = to_np_torch_recursive(deepcopy(posed_human['verts']), use_torch=False, device='cpu') # V x 3
            triangles = to_np_torch_recursive(deepcopy(posed_human['faces']), use_torch=False, device='cpu') # F x 3
            vertex_colors = to_np_torch_recursive(deepcopy(posed_human['color']), use_torch=False, device='cpu') # V x 3

            # broadcast 'vertex_colors' to match dimensions of 'vertices'
            if vertex_colors.shape == (3,):
                vertex_colors = np.broadcast_to(vertex_colors, (vertices.shape[0], 3))

            # move geometry to origin / normalize vertices / move geometry to 'normalized center'
            vertices -= center
            vertices /= voxel_size
            vertices += 0.5*np.array([[voxel_resolution_x,voxel_resolution_y,voxel_resolution_z]])
            
            # add human-mesh to 'mayavi.mlab' pipeline
            posed_human = mlab.triangular_mesh(
                vertices[:,0], 
                vertices[:,1], 
                vertices[:,2], 
                triangles, 
                representation=mesh_representation,
                opacity=opacity,
                scalars = np.arange(vertices.shape[0])
            )

            # add texture to 'posed human'
            posed_human.module_manager.scalar_lut_manager.lut.table = np.c_[
                (vertex_colors * 255), 
                alpha * 255 * np.ones(vertices.shape[0])
            ].astype(np.uint8)

            # update to 'mayavi.mlab' engine
            posed_human.mlab_source.update()
            posed_human.parent.update()

        # visualize auxiliary geometries (these are meshes)
        for key in list(auxiliary_geometries.keys()):
            # retrieve 'vertices', 'faces(triangles)', 'vertex colors' of 'auxiliary geometry' as 'np.ndarray' on 'cpu'
            vertices = to_np_torch_recursive(deepcopy(auxiliary_geometries[key]['verts']), use_torch=False, device='cpu') # V x 3
            triangles = to_np_torch_recursive(deepcopy(auxiliary_geometries[key]['faces']), use_torch=False, device='cpu') # F x 3
            vertex_colors = to_np_torch_recursive(deepcopy(auxiliary_geometries[key]['color']), use_torch=False, device='cpu') # V x 3

            # broadcast 'vertex colors' to match dimensions of 'vertices' 
            if vertex_colors.shape == (3,): vertex_colors = np.broadcast_to(vertex_colors, (vertices.shape[0], 3))

            # move geometry to origin / normalize vertices / move geometry to 'normalized center'
            vertices -= center
            vertices /= voxel_size
            vertices += 0.5*np.array([[voxel_resolution_x,voxel_resolution_y,voxel_resolution_z]]) 

            # add auxiliary geometry (mesh) to 'mayavi.mlab' pipeline
            auxiliary_mesh = mlab.triangular_mesh(
                vertices[:,0], 
                vertices[:,1], 
                vertices[:,2], 
                triangles, 
                representation=mesh_representation,
                opacity=opacity,
                scalars = np.arange(vertices.shape[0]) # used to assign colors based on 'scalars'
            )

            # add texture to auxiliary geometry
            auxiliary_mesh.module_manager.scalar_lut_manager.lut.table = np.c_[
                (vertex_colors * 255), 
                alpha * 255 * np.ones(vertices.shape[0])
            ].astype(np.uint8)

            # update to 'mayavi.mlab' engine
            auxiliary_mesh.mlab_source.update()
            auxiliary_mesh.parent.update()

        # set title
        mlab.title(title, size=title_size, color=title_color)

        # default 'image-cache'
        if tmp_image_cache is None: tmp_image_cache = save_pth.replace(".mp4", "")
        
        # if 'image-cache' already exists, remove the directory and 're-generate'
        if os.path.exists(tmp_image_cache): shutil.rmtree(tmp_image_cache)
        os.makedirs(tmp_image_cache)

        # render geometries from 'various views' and save to 'image-cache'
        for cam_idx in tqdm(range(len(VIDEO_CAMS))):
            # set camera for render
            camera = fig.scene.camera 
            camera.position = VIDEO_CAMS[cam_idx]['position']
            camera.focal_point = VIDEO_CAMS[cam_idx]['focal_point']
            camera.view_angle = VIDEO_CAMS[cam_idx]['view_angle']
            camera.view_up = VIDEO_CAMS[cam_idx]['view_up']
            camera.clipping_range = VIDEO_CAMS[cam_idx]['clipping_range']
            camera.compute_view_plane_normal()
            # save image
            mlab.savefig(f"{tmp_image_cache}/{cam_idx:03}.png", size=(window_size,window_size))

        # close 'mayavi.mlab' figure
        mlab.close(all=True)

        # create video
        if create_video:
            generate_video_from_imgs(tmp_image_cache)

    ## export current 'aggregation results' as video
    def export_video(
            self,
            save_pth,
            prompt,
            delete_cam = False,
            tmp_image_cache = None,
            window_size_per_cluster = 500,
            title_size = 0.1,
            delete_existing_image_cache = True,
            auxiliary_geometries = dict(),
            normalize_max = False,
            offscreen = True,
            verbose = True,
        ):
        # use mayavi.mlab's pipeline's visualizer as off-screen
        if offscreen: mlab.options.offscreen = offscreen

        # default 'image-cache'
        if tmp_image_cache is None: tmp_image_cache = save_pth.replace(".mp4", "")
        
        # if 'image-cache' already exists, remove the directory and 're-generate'
        if delete_existing_image_cache and os.path.exists(tmp_image_cache): shutil.rmtree(tmp_image_cache)
        os.makedirs(tmp_image_cache, exist_ok=True)

        ## generate images
        # generate images for 'full_aggregations'
        directory_num = 0
        for _cam_world_key in self.full_aggregations.keys(): # iterate for '_cam_world_key'
            # number of images aggregated
            num_aggregated = len(self.full_aggregations[_cam_world_key].aggr_input_history.keys())
            # title template
            title = \
                f"prompt:{prompt}\n" \
                f"num_aggregated:{num_aggregated}\n" \
                f"camera-world: {_cam_world_key}\n" \
                f"body-part: none ('full_aggregations')\n" \
                f"aggregation-world: same as camera world ('full_aggregations')"
            # print title
            if verbose:
                print(title)
            
            ## generate images for single 'aggregation result' in 'full_aggregations'
            self.generate_video_single(
                _cam_world_key=_cam_world_key,
                body_part_name=None,
                _aggr_world_key=None,
                save_pth=save_pth, 
                offscreen=offscreen,
                auxiliary_geometries=auxiliary_geometries, 
                normalize_max=normalize_max, 
                tmp_image_cache=f"{tmp_image_cache}/{directory_num:04}",
                window_size = window_size_per_cluster,
                create_video=False,
                title=title,
                title_size=title_size,
            )

            # increase 'directory number'
            directory_num += 1

        # generate images for 'aggregations'
        for _cam_world_key in self.aggregations.keys(): # iterate for '_cam_world_key'
            for body_part_name in self.core[_cam_world_key].body_part_names: # iterate for 'body part'
                for _aggr_world_key in self.aggregations[_cam_world_key].aggr_input_history_per_body_parts[body_part_name].keys(): # iterate for '_aggr_world_key'
                    # number of images aggregated
                    num_aggregated = len(self.aggregations[_cam_world_key].aggr_input_history_per_body_parts[body_part_name][_aggr_world_key].keys())
                    # title tempalte
                    title = \
                        f"prompt:{prompt}\n" \
                        f"num_aggregated:{num_aggregated}\n" \
                        f"camera-world: {_cam_world_key}\n" \
                        f"body-part:{body_part_name}\n" \
                        f"aggregation-world:{_aggr_world_key}"
                    # print title
                    if verbose:
                        print(title)

                    ## generate images for single 'aggregation result' in 'aggregations'
                    self.generate_video_single(
                        _cam_world_key=_cam_world_key,
                        body_part_name=body_part_name,
                        _aggr_world_key=_aggr_world_key,
                        save_pth=save_pth, 
                        offscreen=offscreen,
                        auxiliary_geometries=auxiliary_geometries, 
                        normalize_max=normalize_max, 
                        tmp_image_cache=f"{tmp_image_cache}/{directory_num:04}",
                        window_size = window_size_per_cluster,
                        create_video=False,
                        title=title,
                        title_size=title_size,
                    )

                    # increase 'directory number'
                    directory_num += 1
            

        ## merge all images 'per full_aggregations/aggregations', 'per body-parts' to single image
        # all image paths
        all_image_pths = glob(f"{tmp_image_cache}/*/*.png")
        all_body_part_indices = sorted(list(set([img_pth.split("/")[-2] for img_pth in all_image_pths])))
        all_image_ids = sorted(list(set([img_pth.split("/")[-1].split(".")[0] for img_pth in all_image_pths])))

        # number-of-images, width, height in 'merged frame'
        num_imgs_per_frame = len(all_body_part_indices)
        width = int(math.ceil(math.sqrt(num_imgs_per_frame)))
        height = int(math.ceil(num_imgs_per_frame / width))

        # merge
        for _id in all_image_ids:
            # placeholder for saving merged-images
            total_img = np.zeros([window_size_per_cluster * height, window_size_per_cluster * width, 3], dtype=np.uint8)
            
            # iterate for body-parts
            for body_part_index in all_body_part_indices:
                # open 'single-aggregation' result
                img_pth = f"{tmp_image_cache}/{body_part_index}/{_id}.png"
                image = cv2.imread(img_pth) # Shape: h, w, 3
                
                # fill in the 'single-aggregation' result into 'merged-image'
                body_part_index = int(body_part_index)
                h = body_part_index // width
                w = body_part_index % width
                total_img[
                    window_size_per_cluster*h:window_size_per_cluster*h+window_size_per_cluster, 
                    window_size_per_cluster*w:window_size_per_cluster*w+window_size_per_cluster
                ] = image
                
            # save merged-image
            cv2.imwrite(filename=f"{tmp_image_cache}/{_id}.png",img=total_img)

        # clean sub-directories (each contains 'single-aggregation' results)
        for body_part_index in all_body_part_indices:
            if os.path.exists(f"{tmp_image_cache}/{body_part_index}"): shutil.rmtree(f"{tmp_image_cache}/{body_part_index}")
            
        ## generate mp4 video file
        generate_video_from_imgs(tmp_image_cache)

    ## export aggregation-results
    def export(self, save_pth=None, save_core=True):
        ## export 'aggregations'
        # copy of 'self.aggregations'
        aggregations = deepcopy(self.aggregations)
        
        # remove open3d objects
        for _cam_world_key in aggregations.keys():
            aggregations[_cam_world_key].o3d_geometries_per_body_parts = None

        ## export 'full_aggregations'
        # copy of 'self.full_aggregations'
        full_aggregations = deepcopy(self.full_aggregations)

        # remove open3d objects
        for _cam_world_key in full_aggregations.keys():
            full_aggregations[_cam_world_key].o3d_geometries = None
        
        # save as 'np.ndarrays'
        aggregations = to_np_torch_recursive(aggregations, use_torch=False, device="cpu")
        full_aggregations = to_np_torch_recursive(full_aggregations, use_torch=False, device="cpu")
        
        ## export core
        core = self.export_core(to_cpu=True, to_numpy=True) if save_core else None 


        ## export as dictionary
        export = EasyDict({
            "aggregations": aggregations,
            "full_aggregations": full_aggregations,
            "core": core,
            "use_torch": self.use_torch,
            "device": self.device,
        })
            
        ## return if wanted
        if save_pth is None: return export
        
        ## save if not
        with open(save_pth, "wb") as handle: pickle.dump(export, handle, protocol=pickle.HIGHEST_PROTOCOL)


    ## export 'core'
    def export_core(self, to_cpu=True, to_numpy=True):
        # copy of 'self.core'
        core = deepcopy(self.core)

        # convert to 'np.ndarray' or 'torch.Tensor' / set device
        device = "cpu" if to_cpu else self.device
        core = to_np_torch_recursive(core, use_torch = not to_numpy, device = device)

        return core


    ## load pretrained
    def load(self, load_file=None, load_pth=None, create_o3d_geometries=False):
        # directly load from opened 'load_file'
        if load_file is not None:
            assert load_pth is None, "If 'load_file' is provided, 'load_pth' must not be provided."
        
        # open the 'save_pth' and load file
        else:
            assert load_pth is not None, "If 'load_file' is not provided, 'load_pth' must be provided."
            with open(load_pth, "rb") as handle: load_file = pickle.load(handle)
        
        # load pretrained
        self.aggregations = load_file['aggregations']
        self.full_aggregations = load_file['full_aggregations']
        self.use_torch = load_file['use_torch']
        self.device = load_file['device']

        # create open3d-geometries from loaded pretrained
        if create_o3d_geometries: self._create_o3d_geometries()

        # convert to 'np.ndarray' or 'torch.Tensor'
        self.aggregations = to_np_torch_recursive(self.aggregations, use_torch=self.use_torch, device=self.device)
        self.full_aggregations = to_np_torch_recursive(self.full_aggregations, use_torch=self.use_torch, device=self.device)

        # load 'core' if exist
        if load_file['core'] is not None:
            # load
            self.core = load_file['core']            
            # convert to 'np.ndarray' or 'torch.Tensor'
            self.core = to_np_torch_recursive(self.core, use_torch=self.use_torch, device=self.device)


    ## create open3d-geometries from 'meta-information' of geometries
    def _create_o3d_geometries(self):
        ## create open3d-geometries for 'full_aggregations'
        # iterate for 'camera world'
        for _cam_world_key in self.full_aggregations.keys():
            # placeholder for saving results
            if self.full_aggregations[_cam_world_key].o3d_geometries is None: self.full_aggregations[_cam_world_key].o3d_geometries = EasyDict()
            
            # iterate for geometries in 'full_aggregations'
            for name in self.full_aggregations[_cam_world_key].geometries.keys():
                # 'meta-information' of geometry
                geo_args = self.full_aggregations[_cam_world_key].geometries[name]

                # generator function for open3d objects
                if name == 'grid': o3d_generator = return_cube
                elif name == 'coordinate_frame': o3d_generator = return_coordinate_frame
                elif name == 'template': o3d_generator = return_triangle_mesh
                elif name[:6] == 'posed:': o3d_generator = return_triangle_mesh
                elif name[:4] == 'cam:': o3d_generator = return_camera_frustum
                else: assert False, f"geometry type for '{name}' not implemented"

                # create open3d geometries
                self.full_aggregations[_cam_world_key].o3d_geometries[name] = o3d_generator(**geo_args)
                
        ## create open3d-geometries for 'aggregations'
        # iterate for 'camera world'
        for _cam_world_key in self.aggregations.keys():
            # placeholder for saving results
            if self.aggregations[_cam_world_key].o3d_geometries_per_body_parts is None: self.aggregations[_cam_world_key].o3d_geometries_per_body_parts = EasyDict()
            
            # iterate for 'body part'
            for body_part_name in self.aggregations[_cam_world_key].geometries_per_body_parts.keys():
                # placeholder for saving results
                self.aggregations[_cam_world_key].o3d_geometries_per_body_parts[body_part_name] = EasyDict()
                
                # itearte for 'aggregation world'
                for _aggr_world_key in self.aggregations[_cam_world_key].geometries_per_body_parts[body_part_name].keys():
                    # placeholder for saving results
                    self.aggregations[_cam_world_key].o3d_geometries_per_body_parts[body_part_name][_aggr_world_key] = EasyDict()
                    
                    # iterate for geometries in 'aggregations'
                    for name in self.aggregations[_cam_world_key].geometries_per_body_parts[body_part_name][_aggr_world_key].keys():
                        # 'meta-information' of geometry
                        geo_args = self.aggregations[_cam_world_key].geometries_per_body_parts[body_part_name][_aggr_world_key][name]
                        
                        # generator function for open3d objects
                        if name == 'grid': o3d_generator = return_cube
                        elif name == 'coordinate_frame': o3d_generator = return_coordinate_frame
                        elif name == 'template': o3d_generator = return_triangle_mesh
                        elif name[:6] == 'posed:': o3d_generator = return_triangle_mesh
                        elif name[:4] == 'cam:': o3d_generator = return_camera_frustum
                        else: assert False, f"geometry type for '{name}' not implemented"

                        # create open3d geometries
                        self.aggregations[_cam_world_key].o3d_geometries_per_body_parts[body_part_name][_aggr_world_key][name] = o3d_generator(**geo_args)
utils/aggregation/aggregator.py
jellyheadandrew/CHORUS
import numpy as np
from easydict import EasyDict
from copy import deepcopy

from constants.frankmocap import BODY_MOCAP_SMPL_PTH
from constants.aggr_settings.default_aggr_core import DEFAULT_AGGR_SETTINGS

from utils.aggregation.load_smpl_info import load_smpl_info
from utils.aggregation.load_voxelgrid import load_voxelgrid
from utils.aggregation.load_body_parts import load_body_parts
from utils.aggregation.load_lbs_weight import load_lbs_weights
from utils.misc import to_np_torch_recursive


def prepare_aggr_core(verbose, **aggr_settings):
    ## replace default settings with provided 'aggr_settings' if exists 
    core_settings = deepcopy(DEFAULT_AGGR_SETTINGS)
    if 'USE_SMPL' in aggr_settings.keys(): core_settings['USE_SMPL'] = aggr_settings['USE_SMPL']
    if 'USE_STAR_POSE_TEMPLATE' in aggr_settings.keys(): core_settings['GLOBAL']['USE_STAR_POSE_TEMPLATE'] = aggr_settings['USE_STAR_POSE_TEMPLATE']
    if 'ZEROBETA' in aggr_settings.keys(): core_settings['GLOBAL']['ZEROBETA'] = aggr_settings['ZEROBETA']
    if 'LEARNEDK' in aggr_settings.keys(): core_settings['GLOBAL']['LEARNEDK'] = aggr_settings['LEARNEDK']
    if 'SYMMETRIC_AUGMENTATION' in aggr_settings.keys(): core_settings['GLOBAL']['SYMMETRIC_AUGMENTATION'] = aggr_settings['SYMMETRIC_AUGMENTATION']
    if 'GRID_RESOLUTION' in aggr_settings.keys(): core_settings['GRID_RESOLUTION'] = aggr_settings['GRID_RESOLUTION']
    if 'GRID_LENGTH_SMPL' in aggr_settings.keys(): core_settings['SMPL']['GRID_LENGTH'] = aggr_settings['GRID_LENGTH_SMPL']
    if 'BODY_PART_DEFINE_METHOD_SMPL' in aggr_settings.keys(): core_settings['SMPL']['BODY_PART_DEFINE_METHOD'] = aggr_settings['BODY_PART_DEFINE_METHOD_SMPL']
    if 'CAMERA_SAMPLING_SMPL' in aggr_settings.keys(): core_settings['SMPL']['CAMERA_SAMPLING'] = aggr_settings['CAMERA_SAMPLING_SMPL']
    if 'LBS_MERGE_METHOD_SMPL' in aggr_settings.keys(): core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_MERGE_METHOD'] = aggr_settings['LBS_MERGE_METHOD_SMPL']
    if 'LBS_SMOOTHING_METHOD_SMPL' in aggr_settings.keys(): core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_SMOOTHING_METHOD'] = aggr_settings['LBS_SMOOTHING_METHOD_SMPL']
    if 'LBS_NUM_NEIGHBORS_SMPL' in aggr_settings.keys(): core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_NUM_NEIGHBORS'] = aggr_settings['LBS_NUM_NEIGHBORS_SMPL']
    if 'LBS_NUM_SMOOTH_TIMES_SMPL' in aggr_settings.keys(): core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_NUM_SMOOTH_TIMES'] = aggr_settings['LBS_NUM_SMOOTH_TIMES_SMPL']
    if 'LBS_Y_SCALE_SMPL' in aggr_settings.keys(): core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_Y_SCALE'] = aggr_settings['LBS_Y_SCALE_SMPL']
    if 'LBS_DIST_ORIGIN_THRESH_SMPL' in aggr_settings.keys(): core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_DIST_ORIGIN_THRESH'] = aggr_settings['LBS_DIST_ORIGIN_THRESH_SMPL']
    if 'LBS_IDENTITY_CENTER_SMPL' in aggr_settings.keys(): core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_IDENTITY_CENTER'] = aggr_settings['LBS_IDENTITY_CENTER_SMPL']
    if 'LBS_IDENTITY_DIST_THRESHOLD_SMPL' in aggr_settings.keys(): core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_IDENTITY_DIST_THRESHOLD'] = aggr_settings['LBS_IDENTITY_DIST_THRESHOLD_SMPL']

    ## for demonstration (in colab notebook)
    if 'demo' in aggr_settings:
        if aggr_settings['demo']:
            core_settings['GLOBAL']['DEVICE'] = 'cuda'
            core_settings['GLOBAL']['INTERACTION_REGIONS_PRECOMPUTE_SETTINGS']['PART_INTERACTION_REGION_PRECOMPUTE_DEVICE'] = 'cuda'
            core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['PRECOMPUTE_DEVICE'] = 'cuda'

    ## prepare aggregation core
    return _aggr_core(verbose=verbose, core_settings=core_settings)


def _aggr_core(verbose, core_settings):
    # aggr_core
    aggr_core = EasyDict()
    smpl_info = load_smpl_info(
        model_path=BODY_MOCAP_SMPL_PTH,
        use_star_pose_template=core_settings['GLOBAL']['USE_STAR_POSE_TEMPLATE']
    )
    
    # GLOBAL core components
    aggr_core.GLOBAL = EasyDict(dict(
        use_star_pose_template = core_settings['GLOBAL']['USE_STAR_POSE_TEMPLATE'], 
        smpl_info = smpl_info,
        visualize_pose_num = core_settings['GLOBAL']['VISUALIZE_POSE_NUM'], # how many current spaces will we visualize per cluster?
        auxiliary_exist_ok = core_settings['GLOBAL']['AUXILIARY_EXIST_OK'],
        tag_auxiliary = core_settings['GLOBAL']['TAG_AUXILIARY'],
        add_human_for_vis = core_settings['GLOBAL']['ADD_HUMAN_FOR_VIS'],
        use_torch = core_settings['GLOBAL']['USE_TORCH'], # use torch or not
        device = core_settings['GLOBAL']['DEVICE'],
        interaction_region_precompute_settings = { # for precomputing interaction region
            'part_interaction_region_precompute_device': core_settings['GLOBAL']['INTERACTION_REGIONS_PRECOMPUTE_SETTINGS']['PART_INTERACTION_REGION_PRECOMPUTE_DEVICE'],
            'part_interaction_region_precompute_chunk_size': core_settings['GLOBAL']['INTERACTION_REGIONS_PRECOMPUTE_SETTINGS']['PART_INTERACTION_REGION_PRECOMPUTE_CHUNK_SIZE'],
        },
        zerobeta = core_settings['GLOBAL']['ZEROBETA'], # which inputs to use
        learned_k = core_settings['GLOBAL']['LEARNED_K'],
        symmetric_augmentation=core_settings['GLOBAL']['SYMMETRIC_AUGMENTATION'], # whether to use symmetric augmentation or not
        max_cam_save_num = core_settings['GLOBAL']['MAX_CAM_SAVE_NUM'],
    ))
    
    # SMPL core components (critical for all grids that learn in SMPL-canonical space)
    aggr_core.SMPL = EasyDict(dict(
        voxel_size = core_settings['SMPL']['GRID_LENGTH']/core_settings['GRID_RESOLUTION'],
        voxel_resolution = core_settings['GRID_RESOLUTION'],
        center = core_settings['SMPL']['CENTER'],
        o3d_voxel_color_filled = core_settings['SMPL']['O3D_VOXEL_COLOR_FILLED'],
        o3d_voxel_color_unfilled = core_settings['SMPL']['O3D_VOXEL_COLOR_UNFILLED'],
        o3d_voxel_sphere_resolution = core_settings['SMPL']['O3D_VOXEL_SPHERE_RESOLUTION'],
        grid_color = core_settings['SMPL']['GRID_COLOR'],
        frustum_size = core_settings['SMPL']['FRUSTUM_SIZE'],
        camera_sampling = core_settings['SMPL']['CAMERA_SAMPLING'], # camera sampling/weighting methods
        use_normalized_K = core_settings['SMPL']['USE_NORMALIZED_K'],
        manual_selection = core_settings['SMPL']['MANUAL_SELECTION'],
        cam_color = core_settings['SMPL']['CAM_COLOR'],
        canoncam2smpl_trans = core_settings['SMPL']['CANONCAM2SMPL_TRANS'], # for smpl, canonical world is smpl-template world itself
        v_template_holistic = smpl_info['v_template'].clone().squeeze().cpu().numpy(), # template mesh vertices of smpl (1x6980x3)
        v_template_faces_holistic = smpl_info['faces'].clone().cpu().numpy(),
        v_template_local = smpl_info['v_template'].clone().squeeze().cpu().numpy(), # template mesh vertices of SMPL (1x6980x3)
        v_template_faces_local = smpl_info['faces'].clone().cpu().numpy(),
        v_template_color = core_settings['SMPL']['V_TEMPLATE_COLOR'], # silver
        body_part_define_method = core_settings['SMPL']['BODY_PART_DEFINE_METHOD'], # clustering method
        body_segmentation_json_pth = core_settings['SMPL']['BODY_SEGMENTATION_JSON_PTH'], # meshcapade segmentation pth
        lbs_weights=smpl_info['lbs_weights'],
        lbs_precompute_settings = dict( # lbs-precompute settings
            precompute_device = core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['PRECOMPUTE_DEVICE'],
            chunk_size = core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['CHUNK_SIZE'],
            lbs_num_neighbors = core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_NUM_NEIGHBORS'],
            lbs_merge_method = core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_MERGE_METHOD'],
            lbs_smoothing_method = core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_SMOOTHING_METHOD'], 
            lbs_num_smoothing_times = core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_NUM_SMOOTH_TIMES'],
            lbs_y_scale = core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_Y_SCALE'],
            lbs_dist_origin_thresh = core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_DIST_ORIGIN_THRESH'],
            lbs_identity_center = core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_IDENTITY_CENTER'],
            lbs_identity_dist_threshold = core_settings['SMPL']['LBS_PRECOMPUTE_SETTINGS']['LBS_IDENTITY_DIST_THRESHOLD'],
        ),
        length_x = None, # voxel_size * N
        length_y = None, # voxel_size * N
        length_z = None, # voxel_size * N
        N_x = None, # N
        N_y = None, # N
        N_z = None, # N
        start_point = None, # -length_x/2, -length_y/2, -length_z/2
        voxel_radius = None,
        indexgrid = None, # 3xNxNxN, torch.int64
        canon_grid = None, # 3xNxNxN, torch.float64
        body_part_names = None,
        body_parts = None,
        epsilon_per_body_part = None, 
        aggr_world_name_per_body_parts = None,
        inter_regions = None,
        lbs_weight_grid = None,
    ))

    # disable core components if you want to
    if not core_settings['USE_SMPL']: aggr_core['SMPL'] = None; del aggr_core['SMPL']; assert 'SMPL' not in aggr_core.keys()

    # fill in missing core-components
    aggr_core = load_voxelgrid(aggr_core) # load voxelgrid
    aggr_core = to_np_torch_recursive(
        aggr_core,
        use_torch=aggr_core.GLOBAL.use_torch,
        device=aggr_core.GLOBAL.device
    )
    aggr_core = load_body_parts(aggr_core, verbose=verbose) # load body parts for semantic clustering
    aggr_core = load_lbs_weights(aggr_core, verbose=verbose) # load precomputed lbs weights
    aggr_core = to_np_torch_recursive(
        aggr_core,
        use_torch=aggr_core.GLOBAL.use_torch,
        device=aggr_core.GLOBAL.device
    )

    return aggr_core
utils/aggregation/core.py
jellyheadandrew/CHORUS
import json
from easydict import EasyDict

import numpy as np
import torch

from constants.aggr_settings.semantic_clustering import EPSILON_DEFAULT, BODY_PART_CORRESPONDENCE



def load_body_parts(aggr_core, verbose=False):
    # print status
    if verbose:
        print("Creating interaction regions for clustering!!!")
    
    ## iterate for all tags in 'aggr_core'
    for tag in aggr_core.keys():
        # no need to set body-parts for 'GLOBAL'
        if tag == 'GLOBAL':
            continue
        
        # prepare data
        v_template_base = aggr_core[tag].v_template_holistic
        body_segmentation_json_pth = aggr_core[tag].body_segmentation_json_pth
        body_part_define_method = aggr_core[tag].body_part_define_method

        # get body-part mesh-vertex indices for each tag
        if tag == "SMPL":
            body_parts, epsilon_per_body_part, aggr_world_name_per_body_parts = load_segmentation_smpl(
                body_segmentation_json_pth = body_segmentation_json_pth,
                body_part_define_method = body_part_define_method,
                v_template=v_template_base,
                verbose=verbose,
            )
            
        # update the 'aggr_core' body_parts, epsilon_per_body_part, aggr_world_name_per_body_parts
        aggr_core[tag].body_parts = body_parts
        aggr_core[tag].epsilon_per_body_part = epsilon_per_body_part # corresponds to epsilon (interaction-dist-threshold) per body-part, per tag
        aggr_core[tag].aggr_world_name_per_body_parts = aggr_world_name_per_body_parts # tags per body-parts | used for which tag to compute interaction-region

        # check if duplicate body_part names exist
        body_part_names = sorted(list(aggr_core[tag].body_parts.keys()))
        assert sorted(list(set(body_part_names))) == body_part_names, f"No duplicate body-part name allowed. Check the aggregation settings"
        aggr_core[tag].body_part_names = body_part_names # originally 'all_cluster_keys'

        # print body-part-names
        if verbose:
            print(f"Body part names: {body_part_names}")

        ## iterate through body parts
        interaction_regions = EasyDict()
        for body_part_name in body_parts.keys():
            # generate interaction region
            assert body_part_name not in interaction_regions, "Duplicate part_name in 'inter_regions'"
            interaction_regions[body_part_name] = EasyDict()

            ## iterate through all tags
            for _tag in epsilon_per_body_part[body_part_name].keys():
                # first, translate the '_tag' canon-grid to 'tag' canon-grid
                aligned_canon_grid = \
                    aggr_core[_tag].canon_grid + \
                    aggr_core[_tag].canoncam2smpl_trans[:,None,None,None] - \
                    aggr_core[tag].canoncam2smpl_trans[:,None,None,None]

                # second, compute the interaction-region based on inter_dist_thresh ('epsilon')
                interaction_region_voxel_indices = compute_interaction_region(
                    canon_grid = aligned_canon_grid,
                    v_template = aggr_core[tag].v_template_holistic,
                    body_part_vert_indices = body_parts[body_part_name],
                    inter_dist_thresh = epsilon_per_body_part[body_part_name][tag],
                    verbose=verbose,
                    **aggr_core.GLOBAL.interaction_region_precompute_settings
                )

                # third, update the 'interaction_regions' with voxel-indices
                interaction_regions[body_part_name][tag] = interaction_region_voxel_indices

        ## update 'aggr_core' with 'interaction_regions'
        aggr_core[tag].interaction_regions = interaction_regions        
        
    # print ending status of loading body parts
    if verbose:
        print("Loading Body Parts Finished!!!")
    
    return aggr_core        



def compute_interaction_region(
    canon_grid, # 3xNxNxN
    v_template, # Vx3
    body_part_vert_indices, # List of indices
    inter_dist_thresh, # Float value
    precompute_device="cpu",
    chunk_size=50000,
    verbose=False,
    **kwargs
    ):
    # dict to save 'interaction_regions' to pass or block (each values are list of 3D voxel indices, close to mesh)
    interaction_regions = {}
    
    # 'pass' or 'block' images if 3D rays pass through interaction_region
    for pass_or_block in ["pass", "block"]:         
        ## iterate for 'pass' and 'block' 
        interaction_region = []
        for _body_part_vert_indices in body_part_vert_indices[pass_or_block]:
            # shape of canonical voxel-grid: (3, N_x, N_y, N_z)
            N_x, N_y, N_z = canon_grid.shape[1:]

            # N_body_part_vertices x N_all_voxels x 3
            displacements = \
                canon_grid.to(precompute_device).reshape(3,-1).T[None,:,:] \
                - v_template.to(precompute_device)[_body_part_vert_indices][:,None,:]
            
            ## precomputation kills cpu if implemented naively: 'precompute per chunks'
            # placeholder to save distance between 'canon-voxel' and 'body_part vertex'
            distances = torch.empty(displacements.shape[:2], dtype=torch.float32) # N_body_part_vertices x N_all_voxels
            
            # iterate for all 'chunks'
            displacements_done = 0
            while displacements_done < displacements.shape[1]:
                # print progress
                if verbose:
                    print("\n===== precomputing interaction regions ... =====")
                    print("[interaction-region] precomputing: ", displacements_done)
                
                # which chunk to precompute
                start_idx = displacements_done
                end_idx = min(displacements_done + chunk_size, displacements.shape[1])

                # compute 'distance between 'body-part vertices' and 'canon-grid voxels'
                distances[:,start_idx:end_idx] = displacements[:,start_idx:end_idx].square().sum(dim=-1).sqrt()

                # update 'displacements_done' to move to next chunk
                displacements_done = end_idx
                
            # remove used 'displacements' variable & print 'end' for precomputation
            del displacements
            if verbose:
                print("===== precomputing interaction regions: done! =====\n")
            

            ## reshape
            distances = distances.reshape(distances.shape[0],N_x,N_y,N_z) # shape: N_body_part_vertices, N_x, N_y, N_z

            # boolean array that checks whether voxel is inside 'interaction_region' or not
            interaction_region_grid = torch.where(distances.min(dim=0)[0] < inter_dist_thresh, 1, 0).type(torch.float32) # shape: N_x x N_y x N_z
            
            # N_inter_points x 3, long
            interaction_region_indices = torch.nonzero(interaction_region_grid)

            # update interaction_region
            interaction_region.append(interaction_region_indices) 

        # update interaction_regions
        interaction_regions[pass_or_block] = interaction_region

    return interaction_regions # dict{list[N_inter_points x 3]} / dtype: long



def load_segmentation_smpl(
    body_segmentation_json_pth,
    body_part_define_method,
    verbose=False,
    **kwargs
    ):
    # get body_parts indices
    with open(body_segmentation_json_pth, "r") as rf:
        body_segmentation_loaded = json.load(rf)

    # we must return these
    body_parts = dict()
    epsilon_per_body_part = dict()
    aggr_world_name_per_body_parts = dict()

    # print body-part-define method
    if verbose:
        print("body_part_define_method: ", body_part_define_method)

    ## merge 'body_segmentation's to obtain 'body_part' mesh vertex indices 
    for body_part_name, data in BODY_PART_CORRESPONDENCE[body_part_define_method].items():
        # unpack data
        correspondence_map = data['correspondence_map']
        interaction_threshold_per_world = data['interaction_threshold_per_world']

        ## there are two types of body parts
        # 'pass' -> at least one vertex must lie within 2D occupancy
        # 'block' -> any vertex must not lie within 2D occupancy
        body_parts[body_part_name] = {'pass': [], 'block': []}
        
        # merge to get body part (pass)
        for correspondence in correspondence_map['pass']:
            # this is a list of list. Each list contains vertex indices (empty list means no body-part clustering)
            body_parts[body_part_name]['pass'].append(list())
            # for distinct list of indices, each list must contain at least one point that lies within 2D occupancy
            for body_seg_key in correspondence:
                body_parts[body_part_name]['pass'][-1] += body_segmentation_loaded[body_seg_key]
            body_parts[body_part_name]['pass'][-1] = list(set(body_parts[body_part_name]['pass'][-1]))

        # merge to get body part (block)
        for correspondence in correspondence_map['block']:
            # this is a list of list. Each list contains vertex indices (empty list means no body-part clustering)
            body_parts[body_part_name]['block'].append(list())
            # for distinct list of indices, each list must contain all points that does not lie within 2D occupancy
            for body_seg_key in correspondence:
                body_parts[body_part_name]['block'][-1] += body_segmentation_loaded[body_seg_key]
            body_parts[body_part_name]['block'][-1] = list(set(body_parts[body_part_name]['block'][-1]))
            
        # save epsilon (interaction-distance threshold) per body-part
        epsilon_per_body_part[body_part_name] = interaction_threshold_per_world

        # world-tag per body parts
        aggr_world_name_per_body_parts[body_part_name] = list(interaction_threshold_per_world.keys())

    return fill_in_missing_parts(body_parts, epsilon_per_body_part, aggr_world_name_per_body_parts)



def fill_in_missing_parts(body_parts, epsilon_per_body_part, aggr_world_name_per_body_parts):
    # Check for invalid keys in 'epsilon_per_body_part'
    for k in epsilon_per_body_part.keys():
        if k not in body_parts.keys():
            assert False, "Invalid key in 'epsilon_per_body_part'"

    # Check for invalid keys in 'aggr_world_name_per_body_parts
    for k in aggr_world_name_per_body_parts.keys():
        if k not in body_parts.keys():
            assert False, "Invalid key in 'aggr_world_name_per_body_parts'"

    # If keys for 'epsilon_per_body_part' is differnt from 'body_parts', fill default values
    for k in body_parts.keys():
        if k not in epsilon_per_body_part.keys():
            epsilon_per_body_part[k] = {'SMPL': EPSILON_DEFAULT}

    # If keys for 'aggr_world_name_per_body_parts' is different from 'body_parts', fill with default values
    for k in body_parts.keys():
        if k not in aggr_world_name_per_body_parts.keys():
            aggr_world_name_per_body_parts[k] = ['SMPL']
        else:
            # If tag is not added to aggr_world_name_per_body_parts, add it
            if 'SMPL' not in aggr_world_name_per_body_parts[k]:
                aggr_world_name_per_body_parts[k].append('SMPL')

    # Wrap the lists with integer values as a list for changed format of inter-region computation
    for k in body_parts.keys():
        if type(body_parts[k]) == list:
            assert len(body_parts[k]) >= 1
            if type(body_parts[k][0]) == int:
                body_parts[k] = [body_parts[k]]
            else:
                for seg_part in body_parts[k]:
                    assert type(seg_part) == list
        elif type(body_parts[k]) == dict:
            assert "pass" in body_parts[k].keys()
            assert "block" in body_parts[k].keys()
            assert len(body_parts[k]["pass"]) >= 1
            # body_parts[k] = 

    # If not specified, then Wrap the "list of lists" to a dictionary of "pass" and "block"
    for k in body_parts.keys():
        if type(body_parts[k]) == dict:
            assert "pass" in body_parts[k]
            assert "block" in body_parts[k]
            assert type(body_parts[k]["pass"]) == list 
            if len(body_parts[k]["pass"]) > 0: assert type(body_parts[k]["pass"][0]) == list
            assert type(body_parts[k]["block"]) == list
            if len(body_parts[k]["block"]) > 0: assert type(body_parts[k]["block"][0]) == list
        elif type(body_parts[k]) == list and type(body_parts[k][0]) == list:
            body_parts[k] = {"pass": body_parts[k], "block": []}
        else:
            assert False, f"body_parts here looks like this for key: {k} -> {body_parts[k]}"

    # final check for invalid
    assert tuple(sorted(list(body_parts.keys()))) == tuple(sorted(list(epsilon_per_body_part.keys())))
    assert tuple(sorted(list(body_parts.keys()))) == tuple(sorted(list(aggr_world_name_per_body_parts.keys())))
    for k in body_parts.keys():
        assert tuple(sorted(list(aggr_world_name_per_body_parts[k]))) == tuple(sorted(list(epsilon_per_body_part[k].keys()))), "The tags under 'epsilon_per_body_part' and 'aggr_world_name_per_body_parts' must be same."

    return body_parts, epsilon_per_body_part, aggr_world_name_per_body_parts
utils/aggregation/load_body_parts.py
jellyheadandrew/CHORUS
import pickle5 as pickle

import torch

def load_aggr_input(
        cam_world_key,
        category_id,
        fps,
        zerobeta=True,
    ):
    # keys for fp
    cam_world_key = cam_world_key.lower()
    cam_key = f'cam_{cam_world_key.lower()}_pth'
    human_key = f'human_{cam_world_key.lower()}_pth'
    seg_key = 'human_smpl_pth' # post-processed 'seg' info is saved under here!

    # load cam
    with open(fps[cam_key], "rb") as handle:
        cam_data = pickle.load(handle)
        persp_cam_list = cam_data['persp_cam_list']

    # load smpl parameters
    with open(fps[human_key], "rb") as handle:
        human_parameters = pickle.load(handle)
        mocap_output_list = human_parameters['mocap_output_list']

    # load post-processed segmentations
    with open(fps[seg_key], "rb") as handle:
        human_parameters = pickle.load(handle)
        instances = human_parameters['auxiliary']['instances_including_object']
    
    # prepare cam_R, cam_t, cam_K, thetas, betas
    aggr_input_list = []
    for i in range(len(persp_cam_list)):
        # prepare segmentation
        is_class = instances.pred_classes == category_id
        mask = torch.any(instances.pred_masks[is_class], dim=0)

        # prepare cam
        cam_R = torch.tensor(persp_cam_list[i]['R'])
        cam_t = torch.tensor(persp_cam_list[i]['t'])
        cam_K = torch.tensor(persp_cam_list[i]['K'])

        # pose (thetas)
        pose = torch.tensor(mocap_output_list[i]['pred_body_pose'], dtype=torch.float32)
        betas = torch.zeros([10], dtype=torch.float32) if zerobeta else torch.tensor(mocap_output_list[i]['betas'], dtype=torch.float32)
        
        # add to 'input' list
        aggr_input_list.append(
            {
                'cam_R': cam_R,
                'cam_t': cam_t,
                'cam_K': cam_K,
                'mask': mask,
                'pose': pose,
                'betas': betas,
            }
        )

    return aggr_input_list
utils/aggregation/load_input.py
jellyheadandrew/CHORUS
import numpy as np
import torch


def load_lbs_weights(aggr_core, verbose=False):
    ## iterate for all tags in 'aggr_core'
    for base_tag in aggr_core.keys():
        # no need to set lbs-weights for 'GLOBAL'
        if base_tag == 'GLOBAL':
            continue

        # prepare data
        canon_grid = aggr_core[base_tag].canon_grid
        v_template = aggr_core[base_tag].v_template_holistic
        lbs_weights = aggr_core[base_tag].lbs_weights
        lbs_precompute_settings = aggr_core[base_tag].lbs_precompute_settings
        
        # compute lbs weight for all voxels in grid
        lbs_weight_grid = precompute_lbs_weights(
            canon_grid=canon_grid,
            v_template=v_template,
            lbs_weights=lbs_weights,
            lbs_precompute_settings=lbs_precompute_settings,
            verbose=verbose
        )

        # update to 'aggr_core' with precomputed 'lbs_weight_grid'
        aggr_core[base_tag].lbs_weight_grid = lbs_weight_grid
    
    return aggr_core



def precompute_lbs_weights(
    canon_grid, # 3xNxNxN
    v_template, # Vx3
    lbs_weights, # VxJ
    lbs_precompute_settings,
    verbose=False,
    ):
 
    ## prepare data
    # calculate the distance from grid to all smpl mesh vertices
    N_x, N_y, N_z = canon_grid.shape[1:]; batch_size = 1 

    # ravel from the shape of (3xNxNxN)
    canon_grid = canon_grid.to(lbs_precompute_settings['precompute_device'])
    grid_raveled = canon_grid.reshape(3, -1).T

    # smpl mesh template (Vx3)
    v_template = v_template.to(lbs_precompute_settings['precompute_device'])

    # lbs-weights per vertex (VxJ)
    lbs_weights = lbs_weights.to(lbs_precompute_settings['precompute_device'])


    ## find nearest neighbor vertices
    # calculate distance from all points to all vertices in SMPL
    displacements = grid_raveled[None,:,:] - v_template.squeeze()[:,None,:] # 6980 x N_all_voxels, 3
    
    ## precomputation kills cpu if implemented naively: 'precompute per chunks'
    # placeholder to save distance between 'canon-voxel' and 'body_part vertex'
    distances = torch.empty(displacements.shape[:2], dtype=torch.float32)
    
    # iterate for all 'chunks'
    displacements_done = 0
    while displacements_done < displacements.shape[1]:
        # print progress
        if verbose:
            print("\n===== precomputing lbs-weights for all voxels ... =====")
            print("[lbs-weights] precomputing: ", displacements_done)
    
        # which chunk to precompute
        start_idx = displacements_done
        end_idx = min(displacements_done + lbs_precompute_settings['chunk_size'], displacements.shape[1])

        # compute 'distance between 'body-part vertices' and 'canon-grid voxels'
        distances[:,start_idx:end_idx] = displacements[:,start_idx:end_idx].square().sum(dim=-1).sqrt()

        # update 'displacements done' to move to next chunk
        displacements_done = end_idx
        
    # remove used 'dsplacements' variable & print 'end' for precomputation
    del displacements
    if verbose:
        print("===== precomputing lbs weights: done! =====\n")
    
    ## get top-K closest neighbors (output: N_neighbors x N_all_voxels)
    # top-K neighbors per voxel
    topk_neighbor_distances, topk_neighbor_indices = torch.topk(
        distances, 
        k=lbs_precompute_settings['lbs_num_neighbors'], 
        dim=0, 
        largest = False, 
        sorted=True
    )

    # top-1 neighbor per voxel
    top1_neighbor_distances = topk_neighbor_distances[0] # N_all_voxels


    ## query the nearest vertices skinning weights (N_neighbors x N_all_voxels x J)
    # lbs_weights of 'K' neighbors
    lbs_weights_of_neighbors = lbs_weights[topk_neighbor_indices]
    
    # shape
    N_all_voxels, N_joint = lbs_weights_of_neighbors.shape[1:]
    
    ## calculate skinning weights
    # method 1: identity (no canonicalization/warping)
    if lbs_precompute_settings['lbs_merge_method'] == "identity":
        lbs_weights_per_voxel = torch.zeros(
            [batch_size, N_all_voxels, N_joint], 
            device=lbs_precompute_settings['precompute_device']
        )
        lbs_weights_per_voxel[...,0] = 1 # batch x N_all_voxels x J
        force_identity = None

    # method 2: nearest-neighbor naive mean
    elif lbs_precompute_settings['lbs_merge_method'] == "nn_mean":
        lbs_weights_per_voxel = lbs_weights_of_neighbors.mean(dim=0).expand([batch_size, -1, -1]) # batch x N_all_voxels x J
        force_identity = None

    # method 3: nearest neighbor inverse-distance weighted mean
    elif lbs_precompute_settings['lbs_merge_method'] == "nn_inv_distance_mean":
        exactly_smpl_vertice = (top1_neighbor_distances == 0)[None,:] # 1 x N_all_voxels
        to_replace = torch.zeros_like(topk_neighbor_distances); to_replace[0, :] = 1 # N_neighbors x N_all_voxels
        inverse_neighbor_distances = torch.where(exactly_smpl_vertice, torch.zeros_like(topk_neighbor_distances), 1./topk_neighbor_distances)  # N_neighbors x N_all_voxels
        neighbor_weights = inverse_neighbor_distances / inverse_neighbor_distances.sum(dim=0, keepdim=True).clamp(min=1e-10) # Shape: N_neighbors x N_all_voxels
        assert not torch.any(torch.isnan(neighbor_weights)), "NaN problems"
        neighbor_weights += exactly_smpl_vertice.float() * to_replace # N_neighbors x N_all_voxels
        lbs_weights_per_voxel = (lbs_weights_of_neighbors * neighbor_weights[...,None]).sum(dim=0).expand([batch_size,-1,-1]) # batch x N_all_voxels x J
        force_identity = None

    # method 4: nearest neighbor inverse-distance weighted mean + ratio identity mixing (deweighting)
    elif lbs_precompute_settings['lbs_merge_method'] == "nn_inv_distance_mean+identity_mixing_ratio":
        # inverse distance weighted sum for nearest neighbors
        exactly_smpl_vertice = (topk_neighbor_distances[0] == 0)[None,:] # 1 x N_all_voxels
        to_replace = torch.zeros_like(topk_neighbor_distances); to_replace[0, :] = 1 # N_neighbors x N_all_voxels
        inverse_neighbor_distances = torch.where(exactly_smpl_vertice, torch.zeros_like(topk_neighbor_distances), 1./topk_neighbor_distances)  # N_neighbors x N_all_voxels
        neighbor_weights = inverse_neighbor_distances / inverse_neighbor_distances.sum(dim=0, keepdim=True).clamp(min=1e-10) # Shape: N_neighbors x N_all_voxels
        assert not torch.any(torch.isnan(neighbor_weights)), "NaN problems"
        neighbor_weights += exactly_smpl_vertice.float() * to_replace # N_neighbors x N_all_voxels
        lbs_weights_per_voxel = (lbs_weights_of_neighbors * neighbor_weights[...,None]).sum(dim=0).expand([batch_size,-1,-1]) # batch x N_all_voxels x J
        # identity skinning weights
        identity_skin_weights = torch.zeros([batch_size, N_all_voxels, N_joint], device=lbs_precompute_settings['precompute_device'])
        identity_skin_weights[...,0] = 1 # batch x N_all_voxels x J
        # soft mixing weights for identity and nearest neighbor skinning weights
        dists_from_origin = (grid_raveled - lbs_precompute_settings["lbs_identity_center"][None].to(lbs_precompute_settings['precompute_device'])).square().sum(dim=-1).sqrt() # N_all_voxels
        force_identity = (dists_from_origin > lbs_precompute_settings["lbs_identity_dist_threshold"]).float() # N_all_voxels, 0. or 1.
        ratios = top1_neighbor_distances / dists_from_origin # N_all_voxels. NOTE Closer to 1, bigger the influence of identity
        weights_for_identity = zero2one_to_zero2inf(ratios, method='arctanh', y_scale=lbs_precompute_settings["lbs_y_scale"]) # N_all_voxels
        weights_for_identity = 1 - torch.exp(-weights_for_identity)
        weights_for_identity = weights_for_identity * (1 - force_identity) + torch.ones_like(weights_for_identity) * force_identity
        # soft mixing
        weights_for_identity = weights_for_identity[None,:,None]
        lbs_weights_per_voxel = identity_skin_weights * weights_for_identity + lbs_weights_per_voxel * (1 - weights_for_identity)

    # method 5: nearest neighbor inverse-distance weighted mean + linear identity mixing (deweighting)
    elif lbs_precompute_settings['lbs_merge_method'] == "nn_inv_distance_mean+identity_mixing_linear":
        # inverse distance weighted sum for nearest neighbors
        exactly_smpl_vertice = (topk_neighbor_distances[0] == 0)[None,:] # 1 x N_all_voxels
        to_replace = torch.zeros_like(topk_neighbor_distances); to_replace[0, :] = 1 # N_neighbors x N_all_voxels
        inverse_neighbor_distances = torch.where(exactly_smpl_vertice, torch.zeros_like(topk_neighbor_distances), 1./topk_neighbor_distances)  # N_neighbors x N_all_voxels
        neighbor_weights = inverse_neighbor_distances / inverse_neighbor_distances.sum(dim=0, keepdim=True).clamp(min=1e-10) # Shape: N_neighbors x N_all_voxels
        assert not torch.any(torch.isnan(neighbor_weights)), "NaN problems"
        neighbor_weights += exactly_smpl_vertice.float() * to_replace # N_neighbors x N_all_voxels
        lbs_weights_per_voxel = (lbs_weights_of_neighbors * neighbor_weights[...,None]).sum(dim=0).expand([batch_size,-1,-1]) # batch x N_all_voxels x J
        # identity skinning weights
        identity_skin_weights = torch.zeros([batch_size, N_all_voxels, N_joint], device=lbs_precompute_settings['precompute_device'])
        identity_skin_weights[...,0] = 1 # batch x N_all_voxels x J
        # soft mixing weights for identity and nearest neighbor skinning weights
        dists_from_origin = (grid_raveled - lbs_precompute_settings["lbs_identity_center"][None].to(lbs_precompute_settings['precompute_device'])).square().sum(dim=-1).sqrt() # N_all_voxels
        force_identity = (dists_from_origin > lbs_precompute_settings["lbs_identity_dist_threshold"]).float() # N_all_voxels, 0. or 1.
        identity_distances = lbs_precompute_settings["lbs_dist_origin_thresh"] * torch.ones_like(top1_neighbor_distances) # N_all_voxels
        weights_for_identity = top1_neighbor_distances / (identity_distances + top1_neighbor_distances) # N_all_voxels. NOTE top1_neighbor_distances small means weights for identity small
        weights_for_identity = weights_for_identity * (1 - force_identity) + torch.ones_like(weights_for_identity) * force_identity
        # soft mixing
        weights_for_identity = weights_for_identity[None,:,None] # 1 x N_all_voxels x 1
        lbs_weights_per_voxel = identity_skin_weights * weights_for_identity + lbs_weights_per_voxel * (1 - weights_for_identity) # batch x N_all_voxels x J

    # method 6: nearest neighbor inverse-distance weighted mean + linear & ratio identity mixing (deweighting) --> selected method
    elif lbs_precompute_settings['lbs_merge_method'] == "nn_inv_distance_mean+identity_mixing_linear_ratio_with_constant":
        # inverse distance weighted sum for nearest neighbors
        exactly_smpl_vertice = (topk_neighbor_distances[0] == 0)[None,:] # 1 x N_all_voxels
        to_replace = torch.zeros_like(topk_neighbor_distances); to_replace[0, :] = 1 # N_neighbors x N_all_voxels
        inverse_neighbor_distances = torch.where(exactly_smpl_vertice, torch.zeros_like(topk_neighbor_distances), 1./topk_neighbor_distances)  # N_neighbors x N_all_voxels
        neighbor_weights = inverse_neighbor_distances / inverse_neighbor_distances.sum(dim=0, keepdim=True).clamp(min=1e-10) # Shape: N_neighbors x N_all_voxels
        assert not torch.any(torch.isnan(neighbor_weights)), "NaN problems"
        neighbor_weights += exactly_smpl_vertice.float() * to_replace # N_neighbors x N_all_voxels
        lbs_weights_per_voxel = (lbs_weights_of_neighbors.to(lbs_precompute_settings['precompute_device']) * neighbor_weights[...,None].to(lbs_precompute_settings['precompute_device'])).sum(dim=0).expand([batch_size,-1,-1]) # batch x N_all_voxels x J
        # identity skinning weights
        identity_skin_weights = torch.zeros([batch_size, N_all_voxels, N_joint], device=lbs_precompute_settings['precompute_device'])
        identity_skin_weights[...,0] = 1 # batch x N_all_voxels x J
        # soft mixing weights for identity, NOT via origin. --> set as 0.75
        force_identity = (top1_neighbor_distances > lbs_precompute_settings["lbs_identity_dist_threshold"]).float() # N_all_voxels, 0. or 1.
        force_identity = force_identity.to(lbs_precompute_settings['precompute_device'])
        # identity distance ratio        --> lbs_dist_origin_thresh: set as 0.5
        identity_distances = lbs_precompute_settings["lbs_dist_origin_thresh"] * torch.ones_like(top1_neighbor_distances) # N_all_voxels
        dist_from_body_ratios = top1_neighbor_distances / identity_distances
        # weights for identity
        weights_for_identity = zero2one_to_zero2inf(dist_from_body_ratios, method='arctanh', y_scale=lbs_precompute_settings["lbs_y_scale"]) # N_all_voxels
        weights_for_identity = 1 - torch.exp(-weights_for_identity)
        weights_for_identity = weights_for_identity.to(lbs_precompute_settings['precompute_device'])
        weights_for_identity = weights_for_identity * (1 - force_identity) + torch.ones_like(weights_for_identity) * force_identity
        # soft mixing
        weights_for_identity = weights_for_identity[None,:,None] # 1 x N_all_voxels x 1
        lbs_weights_per_voxel = identity_skin_weights * weights_for_identity + lbs_weights_per_voxel * (1 - weights_for_identity) # batch x N_all_voxels x J

    # method not implemented
    else:
        raise NotImplementedError

    ## reshape to 'grid' format & apply laplace smoothing
    # reshape skinning weights to grid format
    lbs_weights_per_voxel = lbs_weights_per_voxel / lbs_weights_per_voxel.sum(dim=-1, keepdim=True) # Normalize
    lbs_weights_per_voxel = lbs_weights_per_voxel.reshape(-1,N_x,N_y,N_z,N_joint) # NOTE batch x N_x x N_y x N_z x J
    # laplace smoothing for skinning weight
    if lbs_precompute_settings['lbs_smoothing_method'] == "laplace-smoothing":
        lbs_weights_per_voxel = laplace_smooth_weights(
            lbs_weights=lbs_weights_per_voxel, 
            smooth_times=lbs_precompute_settings['lbs_num_smoothing_times']
        )
    elif lbs_precompute_settings['lbs_smoothing_method'] == "laplace-smoothing+identity_thresh_preserve":
        lbs_weights_per_voxel = laplace_smooth_weights(
            lbs_weights=lbs_weights_per_voxel, 
            smooth_times=lbs_precompute_settings['lbs_num_smoothing_times'], 
            force_identity=force_identity
        )
    
    return lbs_weights_per_voxel.squeeze() # N_x x N_y x N_z x J



# [0,1] -> [0,inf] converter
def zero2one_to_zero2inf(ratio: torch.Tensor, method='tan', **kwargs):
	if method == 'tan':
		ratio = (ratio * np.pi / 2).clamp(min=0.0, max=np.pi / 2)
		output = kwargs['y_scale'] * torch.tan(ratio)
	elif method == 'arctanh':
		ratio = ratio.clamp(min=0.0, max=1.0)
		output = kwargs['y_scale'] * torch.arctanh(ratio)
	else:
		assert False, "Not implemented."
	return output



# laplace-smoothing skinning weights in grid
def laplace_smooth_weights(
        lbs_weights: torch.Tensor,
        smooth_times=3, 
        force_identity: torch.Tensor = None
    ):
    ## apply force identity for 'far' points
	if force_identity is not None:
		batch, N_x, N_y, N_z, N_joint = lbs_weights.shape
		force_identity = force_identity[None,:,None] # 1 x N_all_voxels x 1
		identity_weights = torch.zeros_like(lbs_weights).reshape(batch, N_x * N_y * N_z, N_joint); identity_weights[...,0] = 1 # batch x N_all_voxels x J
		
    ## apply laplace smoothing
	for _ in range(smooth_times):
		# laplace smoothing
		mean=(lbs_weights[:,2:,1:-1,1:-1,:]+lbs_weights[:,:-2,1:-1,1:-1,:]+\
			lbs_weights[:,1:-1,2:,1:-1,:]+lbs_weights[:,1:-1,:-2,1:-1,:]+\
			lbs_weights[:,1:-1,1:-1,2:,:]+lbs_weights[:,1:-1,1:-1,:-2,:])/6.0
		lbs_weights[:,1:-1,1:-1,1:-1,:]=(lbs_weights[:,1:-1,1:-1,1:-1,:]-mean)*0.7+mean
		
		# for outside identity_dist_threshold, update with identity 1
		if force_identity is not None:
			lbs_weights = lbs_weights.reshape(batch, N_x * N_y * N_z, N_joint) # batch x N_all_voxels x N_joint
			lbs_weights = lbs_weights * (1-force_identity) + identity_weights * force_identity  
			lbs_weights = lbs_weights.reshape(batch, N_x, N_y, N_z, N_joint)
		# normalize
		sums=lbs_weights.sum(dim=-1,keepdim=True)
		lbs_weights=lbs_weights/sums
                
	return lbs_weights
utils/aggregation/load_lbs_weight.py
jellyheadandrew/CHORUS


import pickle5 as pickle

import numpy as np
import torch

from smplx.utils import Struct
from smplx.lbs import lbs

from utils.misc import to_np


def load_smpl_info(
        model_path = None,
        gender = "neutral",
        age = "adult",
        use_star_pose_template = True,
    ):
    # batch size
    batch_size = 1

    # translation of bones
    create_transl = False # by default
    transl = None # by default
    apply_trans = False # by default

    # data-struct containing all information for smpl
    with open(model_path, "rb") as smpl_file:
        data_struct = Struct(**pickle.load(smpl_file, encoding='latin1'))

    # shapedirs
    shapedirs = data_struct.shapedirs; num_betas = 10
    shapedirs = torch.tensor(to_np(shapedirs[:,:,:num_betas]), dtype=torch.float32)

    # J_regressor
    J_regressor = torch.tensor(to_np(data_struct.J_regressor), dtype=torch.float32)

    # Pose blend shape basis: V x 3 x 207 (hence, 6890 x 3 x 207) --> Coefficients of PCA
    num_pose_basis = data_struct.posedirs.shape[-1]
    posedirs = np.reshape(data_struct.posedirs, [-1, num_pose_basis]).T # 207 x 20670
    posedirs = torch.tensor(to_np(posedirs), dtype=torch.float32)

    # indices of parents for each joints
    parents = torch.tensor(to_np(data_struct.kintree_table[0])).long()
    parents[0] = -1

    # linear blend skinning weights
    lbs_weights = torch.tensor(to_np(data_struct.weights), dtype=torch.float32)

    # auxiliary
    joint_mapper = None
    faces = data_struct.f
    faces = torch.tensor(to_np(faces, dtype=np.int64), dtype=torch.long)
    pose2rot = True

    # template pose (t-pose)
    v_template_T = data_struct.v_template
    v_template_T = torch.tensor(to_np(v_template_T), dtype=torch.float32) # B x V x 3

    # template pose (star-pose)
    if use_star_pose_template:
        zero_pose = torch.zeros([1,72], dtype=torch.float32)
        zero_pose[:,5] = np.pi / 6 # left hip z-axis rotation 30 degrees clockwise
        zero_pose[:,8] = -np.pi / 6 # right hip z-axis rotation 30 degrees counterclockwise
        zero_betas = torch.zeros([1,10], dtype=torch.float32)
        v_template, joints_v_template = lbs(
            betas=zero_betas,
            pose=zero_pose,
            v_template=v_template_T,
            shapedirs=shapedirs,
            posedirs=posedirs,
            J_regressor=J_regressor,
            parents=parents,
            lbs_weights=lbs_weights,
            pose2rot=pose2rot
        )
    else:
        zero_pose = torch.zeros([1,72], dtype=torch.float32)
        zero_betas = torch.zeros([1,10], dtype=torch.float32)
        v_template, joints_v_template = lbs(
            betas=zero_betas,
            pose=zero_pose,
            v_template=v_template_T,
            shapedirs=shapedirs,
            posedirs=posedirs,
            J_regressor=J_regressor,
            parents=parents,
            lbs_weights=lbs_weights,
            pose2rot=pose2rot
        )
    
    # add t-pose joints, too
    _check, joints_v_template_T = lbs(
        betas=torch.zeros_like(zero_betas),
        pose=torch.zeros_like(zero_pose),
        v_template=v_template_T,
        shapedirs=shapedirs,
        posedirs=posedirs,
        J_regressor=J_regressor,
        parents=parents,
        lbs_weights=lbs_weights,
        pose2rot=pose2rot
    )
    assert torch.allclose(v_template_T.unsqueeze(0), _check)
    
    # return smpl constants
    smpl_info = {
        "v_template_T": v_template_T.unsqueeze(dim=0), # 1 x 6980 x 3
        "v_template": v_template, # 1 x 6980 x 3
        "joints_v_template": joints_v_template, # 1 x 24 x 3
        "joints_v_template_T": joints_v_template_T, # 1 x 24 x 3
        "zero_pose": zero_pose, # 1 x 72
        "zero_betas": zero_betas, # 1 x 10
        "shapedirs": shapedirs,
        "posedirs": posedirs,
        "J_regressor": J_regressor,
        "parents": parents,
        "lbs_weights": lbs_weights,
        "pose2rot": pose2rot,
        "faces": faces,
        "SMPL2SMPLX_TRANS": np.array([0.012,-0.184,-0.0185]),
        "SMPLX2SMPL_TRANS": -np.array([0.012,-0.184,-0.0185]),
    }

    return smpl_info
utils/aggregation/load_smpl_info.py
jellyheadandrew/CHORUS
import numpy as np

from utils.misc import get_3d_indexgrid_ijk


def load_voxelgrid(aggr_core):
    for tag in aggr_core.keys():
        # no need to set voxel-grid for 'GLOBAL'
        if tag == 'GLOBAL':
            continue

        # voxel properties
        voxel_size = aggr_core[tag].voxel_size
        voxel_resolution = aggr_core[tag].voxel_resolution
        voxel_radius = voxel_size / 8
        length_x = length_y = length_z = voxel_size * voxel_resolution
        N_x = N_y = N_z = voxel_resolution
        center = aggr_core[tag].center
        center = center if type(center) == np.ndarray else np.array(center)
        start_point = center - np.array([length_x/2, length_y/2, length_z/2])

        # index-grid. By querying with [:,i,j,k], you can get the index [i,j,k].
        indexgrid = get_3d_indexgrid_ijk(N_x, N_y, N_z)

        # canon-space grid. By querying with [:,i,j,k], you can get the canon-world-coordinate values of that grid index.
        canon_grid = start_point.reshape(3,1,1,1) + voxel_size * indexgrid.astype(np.float32) + voxel_size / 2

        # load these information to 'aggr_core'
        aggr_core[tag].length_x = length_x
        aggr_core[tag].length_y = length_y
        aggr_core[tag].length_z = length_z
        aggr_core[tag].N_x = N_x
        aggr_core[tag].N_y = N_y
        aggr_core[tag].N_z = N_z
        aggr_core[tag].start_point = start_point
        aggr_core[tag].voxel_radius = voxel_radius
        aggr_core[tag].indexgrid = indexgrid
        aggr_core[tag].canon_grid = canon_grid

    return aggr_core
utils/aggregation/load_voxelgrid.py
jellyheadandrew/CHORUS
import numpy as np

## Computes mean PAP for all images (within category)
def compute_mean_pap(evaluate_results):
    # average values
    mPAP = 0.
    mPAP_strict = 0.
    mPAP_hoa = 0.
    mPAP_strict_hoa = 0.
    PAP_count = 0
    PAP_count_hoa = 0
    for img_name in evaluate_results['results'].keys():
        image_results = evaluate_results['results'][img_name]
        if image_results['target_area'] == 0:
            # since the target does not exist, we do not account as evaluation
            continue
        elif image_results['target_area_hoa'] == 0:
            # just skip the hoa accumulation
            mPAP += image_results['PAP']
            mPAP_strict += image_results['PAP_strict']
            PAP_count += 1
        else:
            # accumulate for both                    
            mPAP += image_results['PAP']
            mPAP_strict += image_results['PAP_strict']
            PAP_count += 1
            # human occlusion aware
            mPAP_hoa += image_results['PAP_hoa']
            mPAP_strict_hoa += image_results['PAP_strict_hoa']
            PAP_count_hoa += 1
    
    ## add average to dict
    # 'mean-projective-average-precision'
    mPAP /= max(1,PAP_count)
    mPAP_strict /= max(1,PAP_count)
    # 'mean-projective-average-precision' (human occlusion aware)
    mPAP_hoa /= max(1,PAP_count_hoa)
    mPAP_strict_hoa /= max(1,PAP_count_hoa)
    # add to dict
    evaluate_results['results']['mPAP'] = mPAP
    evaluate_results['results']['mPAP_strict'] = mPAP_strict
    # hoa
    evaluate_results['results']['mPAP_hoa'] = mPAP_hoa
    evaluate_results['results']['mPAP_strict_hoa'] = mPAP_strict_hoa
    
    return evaluate_results


def compute_pap(pred, target_mask, verbose):
    """
        pred: A dictionary where
            -key: threshold
            -value: dict of
                - 'human_mask': Human mask rendered
                - 'object_mask': Object mask rendered at 'threshold'
                        
        target: A target object mask    
    """
    # precision-recall list
    precision_recall_list, precision_recall_hoa_list = return_precision_recall_list(
        pred=pred, 
        target_mask=target_mask,
        verbose=verbose
    )
    
    PAP, PAP_strict, target_area = projective_average_precision(precision_recall_list)
    PAP_hoa, PAP_strict_hoa, target_area_hoa = projective_average_precision(precision_recall_hoa_list)
    assert target_area >= target_area_hoa

    return PAP, PAP_strict, PAP_hoa, PAP_strict_hoa, target_area, target_area_hoa


def return_precision_recall_list(pred, target_mask, verbose):
    ## placeholder to save pap values
    precision_recall_list = []
    precision_recall_hoa_list = []

    # iterate for thresholds
    for _threshold in pred.keys():
        # threshold
        threshold = float(_threshold)

        # object mask and human mask
        pred_mask = pred[_threshold]['object_mask']
        human_mask = pred[_threshold]['human_mask']

        # human-occlusion aware masks
        pred_mask_hoa = np.logical_and(pred_mask, 1-human_mask)
        target_mask_hoa = np.logical_and(target_mask, 1-human_mask)

        # pixel-wise precision & recall
        target_area = float(target_mask.sum())
        pred_area = float(pred_mask.sum())
        intersection = np.logical_and(target_mask, pred_mask)
        precision = float(intersection.sum() / max(1, pred_mask.sum()))
        recall = float(intersection.sum() / max(1,target_mask.sum()))
        if pred_area > 0 and verbose: print(f"======\nThreshold:{threshold}\nPrecision:{precision}\nRecall:{recall}\nPrediction Area:{pred_area}")

        # human-occlusion aware precision & recall
        target_area_hoa = float(target_mask_hoa.sum())
        pred_area_hoa = float(pred_mask_hoa.sum())
        intersection_hoa = np.logical_and(target_mask_hoa, pred_mask_hoa)
        precision_hoa = float(intersection_hoa.sum() / max(1,pred_mask_hoa.sum()))
        recall_hoa = float(intersection_hoa.sum() / max(1,target_mask_hoa.sum()))
        if pred_area_hoa > 0 and verbose: print(f"======\nThreshold:{threshold}\nPrecision (human-occlusion-aware):{precision_hoa}\nRecall (human-occlusion-aware):{recall_hoa}\nPrediction Area (human-occlusion-aware):{pred_area_hoa}")


        # save-to-list
        precision_recall_list.append(
            (threshold, precision, recall, pred_area, target_area)
        )
        precision_recall_hoa_list.append(
            (threshold, precision_hoa, recall_hoa, pred_area_hoa, target_area_hoa)
        )

    return precision_recall_list, precision_recall_hoa_list
    

def projective_average_precision(precision_recall_list: list):
    PAP = 0. # original version -> do not accumulate when precision starts to become 0.
    PAP_strict = 0. # strict -> set precision as 0 if nothing is predicted
    # what we have is a list of 3-tuple (theshold, precision, recall)
    precision_to_accumulate = precision_recall_list[0][1] # precision
    recall_before = precision_recall_list[0][2]
    target_area = precision_recall_list[0][-1]
    for threshold, precision, recall, prediction_area, target_area_ in precision_recall_list[1:]:                        
        assert target_area == target_area_
        
        # get the width of the bin in precision-recall graph
        assert recall <= recall_before
        width = recall_before - recall
        
        # accumulate PAP
        if prediction_area == 0:
            PAP += width * precision_to_accumulate
            PAP_strict += 0.
        else:
            PAP += width * precision_to_accumulate
            PAP_strict += width * precision_to_accumulate

        # update "precision_to_accumulate"
        if precision > precision_to_accumulate:
            precision_to_accumulate = precision
        else:
            # keep
            pass

        # update "recall before"
        recall_before = recall

    return PAP, PAP_strict, target_area
utils/evaluation.py
jellyheadandrew/CHORUS
""" Modified 'def get_iou()' from https://github.com/charlesq34/votenet-1/blob/master/utils/box_util.py """
def bbox_iou(bb1, bb2, return_areas_too=False):
    """
    Calculate the Intersection over Union (IoU) of two bounding boxes.

    Parameters
    ----------
    bb1 : ndarray
        x1, y1, x2, y2
        The (x1, y1) position is at the top left corner,
        the (x2, y2) position is at the bottom right corner
    bb2 : ndarray
        x1, y1, x2, y2
        The (x, y) position is at the top left corner,
        the (x2, y2) position is at the bottom right corner

    Returns
    -------
    float
        in [0, 1]
    """
    assert bb1[0] < bb1[2]
    assert bb1[1] < bb1[3]
    assert bb2[0] < bb2[2]
    assert bb2[1] < bb2[3]

    # determine the coordinates of the intersection rectangle
    x_left = max(bb1[0], bb2[0])
    y_top = max(bb1[1], bb2[1])
    x_right = min(bb1[2], bb2[2])
    y_bottom = min(bb1[3], bb2[3])



    # compute the area of both AABBs
    bb1_area = (bb1[2] - bb1[0]) * (bb1[3] - bb1[1])
    bb2_area = (bb2[2] - bb2[0]) * (bb2[3] - bb2[1])


    if x_right < x_left or y_bottom < y_top:
        if return_areas_too:
            return 0.0, bb1_area, bb2_area, 0.0
        else:
            return 0.0
        
    # The intersection of two axis-aligned bounding boxes is always an
    # axis-aligned bounding box
    intersection_area = (x_right - x_left) * (y_bottom - y_top)


    # 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)
    assert iou >= 0.0
    assert iou <= 1.0
    if return_areas_too:
        return iou, bb1_area, bb2_area, intersection_area
    else:
        return iou



def intersection_over_smaller_bbox(bb1, bb2):
    """
    Calculate the 'Intersection over Smaller Boudning-Box' of two bounding boxes.

    Parameters
    ----------
    bb1 : ndarray
        x1, y1, x2, y2
        The (x1, y1) position is at the top left corner,
        the (x2, y2) position is at the bottom right corner
    bb2 : ndarray
        x1, y1, x2, y2
        The (x, y) position is at the top left corner,
        the (x2, y2) position is at the bottom right corner

    Returns
    -------
    float
        in [0, 1]
    """
    # compute areas of 'bb1', 'bb2', 'intersection'
    _, bb1_area, bb2_area, intersection_area = bbox_iou(bb1, bb2, return_areas_too=True)
    
    # compute 'intersection-over-small-bounding-box'
    iosbb = intersection_area / min(bb1_area, bb2_area)

    return iosbb
utils/filtering.py
jellyheadandrew/CHORUS
import numpy as np
import cv2

from constants.frankmocap import HMR_INPUT_SIZE


""" reference: https://github.com/facebookresearch/frankmocap """
## convert from 'bounding-box' world to 'original-image' world
def convert_bbox_to_oriIm(data3D, boxScale_o2n, bboxTopLeft, imgSizeW, imgSizeH):
    imgSize = np.array([imgSizeW,imgSizeH])
    data3D /= boxScale_o2n
    data3D[:,:2] += bboxTopLeft - imgSize * 0.5 + (HMR_INPUT_SIZE * 0.5) / boxScale_o2n
    return data3D


""" reference: https://github.com/facebookresearch/frankmocap """
## convert 'weak-perspective' 3d-joints to 2d-joints
def convert_3djoint_to_2djoint(smpl_joints_3d_vis, imgshape):

    smpl_joints_2d_vis = smpl_joints_3d_vis[:,:2]       # 3D is in camera comaera coordinate with origin on the image center
    smpl_joints_2d_vis[:,0] += imgshape[1]*0.5          # offset to move the origin on the top left
    smpl_joints_2d_vis[:,1] += imgshape[0]*0.5          # offset to move the origin on the top left

    return smpl_joints_2d_vis


""" reference: https://github.com/facebookresearch/frankmocap """
## convert 'smpl' world to 'bounding-box' world
def convert_smpl_to_bbox(data3D, scale, trans, bAppTransFirst=False):
    

    if bAppTransFirst:              # for hand model
        data3D[:,0:2] += trans      # apply translation
        data3D *= scale             # apply scaling
    else:
        data3D *= scale             # apply scaling
        data3D[:,0:2] += trans      # apply translation
    
    data3D *= (HMR_INPUT_SIZE * 0.5)         # 112 is originated from hmr's input size (224,224)

    return data3D


""" reference: https://github.com/facebookresearch/frankmocap """
## generate transformation-matrix
def get_transform(center, scale, res, rot=0):
    h = 200 * scale     # h becomes the original bbox max(height, min). 
    t = np.zeros((3, 3))
    t[0, 0] = float(res[1]) / h         # this becomes a scaling factor to rescale original bbox -> res size (default: 224x224)
    t[1, 1] = float(res[0]) / h
    t[0, 2] = res[1] * (-float(center[0]) / h + .5)
    t[1, 2] = res[0] * (-float(center[1]) / h + .5)
    t[2, 2] = 1
    if not rot == 0:
        rot = -rot # To match direction of rotation from cropping
        rot_mat = np.zeros((3,3))
        rot_rad = rot * np.pi / 180
        sn,cs = np.sin(rot_rad), np.cos(rot_rad)
        rot_mat[0,:2] = [cs, -sn]
        rot_mat[1,:2] = [sn, cs]
        rot_mat[2,2] = 1
        # Need to rotate around center
        t_mat = np.eye(3)
        t_mat[0,2] = -res[1]/2
        t_mat[1,2] = -res[0]/2
        t_inv = t_mat.copy()
        t_inv[:2,2] *= -1
        t = np.dot(t_inv,np.dot(rot_mat,np.dot(t_mat,t)))
    return t


""" reference: https://github.com/facebookresearch/frankmocap """
## transform pixel-location to different frame
def transform(pt, center, scale, res, invert=0, rot=0):
    # generate transformation-matrix
    t = get_transform(center, scale, res, rot=rot)
    
    # to invert-or-not
    if invert:
        t = np.linalg.inv(t)

    new_pt = np.array([pt[0]-1, pt[1]-1, 1.]).T
    new_pt = np.dot(t, new_pt)
    return new_pt[:2].astype(int)+1


""" reference: https://github.com/facebookresearch/frankmocap """
## crop image with given bounding-box
def crop_bboxInfo(img, center, scale, res =(HMR_INPUT_SIZE,HMR_INPUT_SIZE)):
    # upper left point
    ul = np.array(transform([1, 1], center, scale, res, invert=1))-1
    # bottom right point
    br = np.array(transform([res[0]+1,
                             res[1]+1], center, scale, res, invert=1))-1


    ## new shape (bbox)
    new_shape = [br[1] - ul[1], br[0] - ul[0]]
    if len(img.shape) > 2:
        new_shape += [img.shape[2]]
    if new_shape[0] <1  or new_shape[1] <1:
        return None, None, None
    new_img = np.zeros(new_shape, dtype=np.uint8)

    if new_img.shape[0] ==0:
        return None, None, None

    #Compute bbox for legacy format
    bboxScale_o2n = res[0] / new_img.shape[0]             #224 / 531 (legacy: 531)

    # range to fill new array
    new_x = max(0, -ul[0]), min(br[0], len(img[0])) - ul[0]
    new_y = max(0, -ul[1]), min(br[1], len(img)) - ul[1]
    # range to sample from original image
    old_x = max(0, ul[0]), min(len(img[0]), br[0])
    old_y = max(0, ul[1]), min(len(img), br[1])

    if new_y[0] <0 or new_y[1]<0 or new_x[0] <0 or new_x[1]<0 :
        return None, None, None

    new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[0]:old_y[1],
                                                        old_x[0]:old_x[1]]

    # bbox-top-left in original = (old_x[0], old_y[0] )
    bboxTopLeft_inOriginal = (ul[0], ul[1] )

    if new_img.shape[0] <20 or new_img.shape[1]<20:
        return None, None, None

    new_img = cv2.resize(new_img, res)

    return new_img, bboxScale_o2n, np.array(bboxTopLeft_inOriginal)


def convert_smpl_to_img(smpl_3d, img, bbox_center, bbox_scale, cam_param_scale, cam_param_trans):
    #Crop image using cropping information
    _, boxScale_o2n, bboxTopLeft = crop_bboxInfo(img, bbox_center, bbox_scale, (HMR_INPUT_SIZE, HMR_INPUT_SIZE) )
    # from bbox space -> img space
    smpl_3d_vis = smpl_3d.copy()
    smpl_3d_vis = convert_smpl_to_bbox(smpl_3d_vis, cam_param_scale, cam_param_trans)
    smpl_3d_vis = convert_bbox_to_oriIm(smpl_3d_vis, boxScale_o2n, bboxTopLeft, img.shape[1], img.shape[0])
    smpl_2d = convert_3djoint_to_2djoint(smpl_3d_vis, img.shape) # 49x2

    return smpl_2d
utils/frankmocap_conversion.py
jellyheadandrew/CHORUS
from copy import deepcopy
from typing import Union
from easydict import EasyDict

import numpy as np
import torch

from constants.datasets import COCO_CLASS_NAME2ID, LVIS_CLASS_NAME2ID, CATEGORY_EXCEPTIONS


def to_np(array, dtype=np.float32):
    if 'scipy.sparse' in str(type(array)):
        array = array.todense()
    return np.array(array, dtype=dtype)


def to_np_torch_recursive(
    X: Union[dict, EasyDict, np.ndarray, torch.Tensor], 
    use_torch=True, 
    device="cuda",
    np_float_type = np.float32,
    np_int_type = np.long,
    torch_float_type = torch.float32,
    torch_int_type = torch.int64,
    ):

    ## recursive approach
    # for dictionaries, run array-to-tensor recursively
    if type(X) == dict or type(X) == EasyDict:
        for key in X.keys():
            if type(X[key]) in [dict, EasyDict, np.ndarray, torch.Tensor]:
                X[key] = to_np_torch_recursive(X[key], use_torch, device)
    # for np.ndarrays, send to torch.Tensor
    elif type(X) == np.ndarray:
        if use_torch:
            X = torch.tensor(X, device=device)
    # for torch.Tensor, set the device only
    elif type(X) == torch.Tensor:
        if use_torch:
            X = X.to(device)
        else:
            X = X.detach().cpu().numpy()

    ## dtype conversion
    if type(X) == torch.Tensor:
        if X.dtype in [torch.float64, torch.float32, torch.float16, torch.bfloat16]:
            X = X.type(torch_float_type)
        elif X.dtype in [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64]:
            X = X.type(torch_int_type)
        else:
            pass
    elif type(X) == np.ndarray:
        if X.dtype in [np.float32, np.float16, np.float64]:
            X = X.astype(np_float_type)
        elif X.dtype in [np.int64, np.int32, np.int16]:
            X = X.astype(np_int_type)
        else:
            pass

    return X


def extract_nonvector(X):
    ## recursive approach
    # for dictionaries, extract non-vector (non np.ndarray / non torch.Tensor) recursively
    new_X = deepcopy(X)
    
    # if X is dict
    if type(X) == dict or type(X) == EasyDict:
        for key in X.keys():
            if type(X[key]) in [dict, EasyDict, list]:
                new_X[key] = extract_nonvector(X[key])    
            elif type(X[key]) in [np.ndarray, torch.Tensor]:
                new_X[key] = "REMOVED (np.ndarray/torch.Tensor)"
            else:
                pass

    # if X is list
    if type(X) == list:
        for i, element in enumerate(X):
            if type(element) in [dict, EasyDict, list]:
                new_X[i] = extract_nonvector(element)
            elif type(element) in [np.ndarray, torch.Tensor]:
                new_X[i] = "REMOVED (np.ndarray/torch.Tensor)"
            else:
                pass

    return new_X


def get_3d_indexgrid_ijk(N_x, N_y, N_z, raveled=False):
    # Create indice grid
    indices =np.mgrid[0:N_x,0:N_y,0:N_z]
    if raveled:
        indices = np.stack([indices[0].ravel()+1,
                            indices[1].ravel()+1,
                            indices[2].ravel()+1], axis=-1)
        # -> Shape: [(N_x+1) * (N_y+1) * (N_z+1), 3]
        # -> NOTE: Order would be 
        #            0 0 0
        #            0 0 1
        #            0 0 2
        #            ...
        #            0 1 0
        #            0 1 1
        #            ...
        #            1 0 0
        #            1 0 1
        #            ...
    return indices


def load_category_id(category):
    # load category: coco
    if category in list(COCO_CLASS_NAME2ID.keys()):
        category_id = COCO_CLASS_NAME2ID[category]
    # load category: lvis
    elif category in list(LVIS_CLASS_NAME2ID.keys()):
        category_id = LVIS_CLASS_NAME2ID[category]
    # load category: exceptions (e.g., 'surfboard-demo' -> 'surfboard' id)
    elif category in CATEGORY_EXCEPTIONS.keys():
        category_id = CATEGORY_EXCEPTIONS[category]
    else:
        assert False, f"method not implemented for [category: {category}] yet"

    return category_id    
utils/misc.py
jellyheadandrew/CHORUS
import pickle5 as pickle
from tqdm import tqdm

import numpy as np
import torch
import torch.nn as nn

import cv2

from utils.transformations import batch_rodrigues, get_focal_length, convert_Kparams_Kmat


def apply_wp_to_p(
        filter_settings, 
        img_pth=None, 
        human_pth=None, 
        mocap_output_list=None, 
        img_size=None,
        device='cuda'
    ):
    # image size
    if img_size is None: H, W = cv2.imread(img_pth).shape[:2]
    else: H, W = img_size

    # load mocap-output-list
    if mocap_output_list is None:
        with open(human_pth, "rb") as handle: 
            human_parameters = pickle.load(handle)
            mocap_output_list = [] if human_parameters == "NO HUMANS" else human_parameters["mocap_output_list"]

    # number of humans
    num_humans = len(mocap_output_list)
    if filter_settings['verbose']:
        print(f"\tNumber of predicted humans in this case: {num_humans}")
    
    # a list to save element (min_loss, R, t, K) 
    persp_cam_list = []

    # iterate for all humans
    for human_idx in range(num_humans):
        ## load input joints (3d-smpl) & target joints (2d-projection) for optimization
        # smpl joints rendered on image
        pred_joints_img=mocap_output_list[human_idx]['pred_joints_img'][:,:2] # -> Shape: (49,2), xy format

        # smpl joints are set as globally zero-rotation (aligned at center, same front direction)
        pred_joints_smpl=mocap_output_list[human_idx]['pred_joints_smpl_zerorot_zerobeta'] # (1,49,3), direct output of smpl model
        pred_joints_smpl = pred_joints_smpl[0]

        # if using openpose joints, use all 24+25 joints for target. if not, use 24 smpl joints for target.
        if not filter_settings['use_openpose_joints']:
            pred_joints_img = pred_joints_img[25:]
            pred_joints_smpl = pred_joints_smpl[25:]
            
        # if ignoring joints outside the image range [0,H] X [0,W], filter the target joints.
        if filter_settings['ignore_joints_outside']:
            joints_on_image_below = np.logical_and(pred_joints_img[:,0] < W, pred_joints_img[:,1] < H)
            joints_on_image_above = np.logical_and(pred_joints_img[:,0] > 0, pred_joints_img[:,1] > 0)
            joints_on_image = np.logical_and(joints_on_image_above, joints_on_image_below)
            pred_joints_img = pred_joints_img[joints_on_image]
            pred_joints_smpl = pred_joints_smpl[joints_on_image]
            
        # number of joints used for optimization
        num_target_joints = pred_joints_smpl.shape[0]

        # convert joints vector to tensor
        pred_joints_smpl_t = torch.tensor(pred_joints_smpl, device=device)
        pred_joints_img_t = torch.tensor(pred_joints_img, device=device)

        ## load focal length
        focal_length_init_t=torch.tensor(get_focal_length(1.0, filter_settings['fovy']))
        focal_length = nn.Parameter(focal_length_init_t, requires_grad=True)
        
        ## load translation
        # orthographic camera (scale, tx, ty) in smpl scale
        cams=mocap_output_list[human_idx]['pred_camera']

        # set initial translation
        t_init = torch.tensor([cams[1], cams[2], 2.0 * focal_length / cams[0]])
        t_param = nn.Parameter(t_init, requires_grad=True)
        
        # initialize rotation
        R_init = mocap_output_list[human_idx]['pred_body_pose'][0,:3] # (3,), global rotation in axis-angle (unnormalized)
        R_param = nn.Parameter(torch.Tensor(R_init), requires_grad=True)

        ## declare optimizer
        if filter_settings['freeze_foc']:
            optimizer = torch.optim.Adam([R_param, t_param], lr=filter_settings['lr'])
        else:
            optimizer = torch.optim.Adam([R_param, t_param, focal_length], lr=filter_settings['lr'])  

        ## optimize perspective camera
        min_loss = 1e10
        if filter_settings['verbose']: pbar = tqdm(range(filter_settings['epochs']))
        else: pbar = range(filter_settings['epochs'])
        
        for epoch in pbar:
            # clear optimizer grad
            optimizer.zero_grad()
            
            # prepare R, t, K
            R_w2c_rotmat = batch_rodrigues(R_param.unsqueeze(0)).squeeze().float().to(device)
            t_w2c = t_param.float().to(device)
            K_param = focal_length.float().to(device) * torch.Tensor([1,1,0,0]).to(device) + torch.Tensor([0,0,0.5,0.5 * H / W]).to(device)
            K_param = K_param * torch.Tensor([W]).to(device) # pixel-scale K 
            K = convert_Kparams_Kmat(K_param)

            # project smpl joints to image
            tmp=torch.matmul(K,(torch.matmul(R_w2c_rotmat, pred_joints_smpl_t.T)+t_w2c.expand(pred_joints_smpl_t.shape).T)).T
            projected_joint_img = tmp[:,:2].clone()
            projected_joint_img[:,0]/=tmp[:,2]
            projected_joint_img[:,1]/=tmp[:,2]

            # calculate l2-error between projected joints and target joints
            loss = torch.linalg.norm(pred_joints_img_t - projected_joint_img)
            loss = loss.divide(num_target_joints)

            # backward pass
            loss.backward()

            # whether early termination occurs
            if loss.item() < filter_settings['thres']:
                R_best = R_param.clone().detach()
                t_best = t_param.clone().detach()
                K_best = K_param.clone().detach()
                min_loss = loss.item()
                if filter_settings['verbose']:
                    print(f'Early termination with min_loss: {min_loss} and epoch: {epoch}')
                break

            # update the parameters
            optimizer.step()

            # update minimum loss
            if loss.item() < min_loss:
                    R_best = R_param.clone().detach()
                    t_best = t_param.clone().detach()
                    K_best = K_param.clone().detach()
                    min_loss = loss.item()

        # print when max epoch reached
        if filter_settings['verbose']:
            print(f'Reached max epoch: {epoch} with min_loss: {min_loss}')

        # (R, t, K) as convenient format (loss, SO(3) matrix, 3-dim vec, 3x3 matrix)
        R_best2save = batch_rodrigues(R_best.unsqueeze(0)).squeeze().float().to(device)
        t_best2save = t_best.float().to(device)
        K_best2save = convert_Kparams_Kmat(K_best)

        # save the results to list
        persp_cam_list.append(dict(
            min_loss=min_loss,
            R=R_best2save.detach().cpu().numpy(),
            t=t_best2save.detach().cpu().numpy(),
            K=K_best2save.detach().cpu().numpy(),
            joints_used_GT_image=pred_joints_img_t.detach().cpu().numpy(),
            joints_used_optimized_image=projected_joint_img.detach().cpu().numpy(),
            joints_used_SMPL=pred_joints_smpl_t.detach().cpu().numpy(),
            focal_length=focal_length.detach().cpu().numpy(),
            img_size=(H,W),
            filter_settings=filter_settings,                    
        ))
    
    # raise error if returned list has number-of-human mismatch
    assert len(persp_cam_list) == num_humans, f"There must be {num_humans} number of outputs, same as number of human in image."

    return persp_cam_list, (human_parameters if filter_settings['save_with_human_params'] else None)
utils/perspective_conversion.py
jellyheadandrew/CHORUS
import os
from glob import glob

import numpy as np
import torch

from detectron2.structures.boxes import BoxMode
from detectron2.structures.masks import BitMasks

from constants.ldm import MODELS
from constants.datasets import COCO_CLASS_ID2NAME, LVIS_CLASS_ID2NAME



def bbox_xy_to_wh(bbox):
    if not isinstance(bbox, (tuple, list)):
        original_shape = bbox.shape
        bbox = bbox.reshape((-1, 4))
    else:
        original_shape = None
    bbox = BoxMode.convert(
        box=bbox, from_mode=BoxMode.XYXY_ABS, to_mode=BoxMode.XYWH_ABS
    )
    if original_shape is not None:
        return bbox.reshape(original_shape)
    return bbox



def bbox_wh_to_xy(bbox):
    if not isinstance(bbox, (tuple, list)):
        original_shape = bbox.shape
        bbox = bbox.reshape((-1, 4))
    else:
        original_shape = None
    bbox = BoxMode.convert(
        box=bbox, from_mode=BoxMode.XYWH_ABS, to_mode=BoxMode.XYXY_ABS
    )
    if original_shape is not None:
        return bbox.reshape(original_shape)
    return bbox



def make_bbox_square(bbox, bbox_expansion=0.0):
    bbox = np.array(bbox)
    original_shape = bbox.shape
    bbox = bbox.reshape(-1, 4)
    center = np.stack(
        (bbox[:, 0] + bbox[:, 2] / 2, bbox[:, 1] + bbox[:, 3] / 2), axis=1
    )
    b = np.expand_dims(np.maximum(bbox[:, 2], bbox[:, 3]), 1)
    b *= 1 + bbox_expansion
    square_bboxes = np.hstack((center - b / 2, b, b))
    return square_bboxes.reshape(original_shape)



def local_to_global_cam(bboxes, cams, L):
    square_bboxes = make_bbox_square(bbox_xy_to_wh(bboxes))
    global_cams = []
    for cam, bbox in zip(cams, square_bboxes):
        x, y, b, _ = bbox
        X = np.stack((x, y))
        # Bbox space [0, b]
        s_crop = b * cam[0] / 2
        t_crop = cam[1:] + 1 / cam[0]

        # Global image space [0, 1]
        s_og = s_crop / L
        t_og = t_crop + X / s_crop

        # Normalized global space [-1, 1]
        s = s_og * 2
        t = t_og - 0.5 / s_og
        global_cams.append(np.concatenate((np.array([s]), t)))
    return np.stack(global_cams)



def process_bbox_mask(bboxes_person, masks_person, final_keepidx):
    return bboxes_person[final_keepidx], masks_person[final_keepidx]



def process_remove_overlap(bbox_list, confidence_list, minoverlap=0.8, exconf=0.98):
    # list to keep remaining outputs
    keepidx=list(range(len(bbox_list)))

    # iterate through bboxes
    sorted_bbox_list = sorted(zip(list(range(len(bbox_list))), bbox_list, confidence_list), key=lambda tup: tup[-1])
    for original_idx, bbox, confidence in sorted_bbox_list:
        # if confidence is over threshold, keep it no matter what
        if confidence >= exconf:
            pass
        
        # else, check the bbox
        else:
            # area of confidence bbox
            confidence_area = bbox[-2] * bbox[-1]

            # compare with other bboxes
            for idx, comparison_bbox in enumerate(bbox_list):
                # if already removed, or is the same, continue
                if idx not in keepidx:
                    continue
                elif idx == original_idx:
                    continue
                # else, compare
                else:
                    # area of confidence bbox for comparison
                    comp_confidence_area = comparison_bbox[-2] * comparison_bbox[-1]
                
                    # area of intersection
                    start_xs=[bbox[0], comparison_bbox[0]]
                    end_xs=[bbox[0]+bbox[2], comparison_bbox[0]+comparison_bbox[2]]
                    start_ys=[bbox[1], comparison_bbox[1]]
                    end_ys=[bbox[1]+bbox[3], comparison_bbox[1]+comparison_bbox[3]]
                    intersection_w = max([min(end_xs) - max(start_xs), 0.])
                    intersection_h = max([min(end_ys) - max(start_ys), 0.])
                    inter = intersection_w * intersection_h
                    
                    # if [intersection / smaller bounding box]
                    if  inter / confidence_area >= minoverlap or inter / comp_confidence_area >= minoverlap:
                        keepidx.remove(original_idx)
                        break
    
    return keepidx
                    


# for objects only
def process_segmentation(instances, minoverlap=0.8, exconf=0.98, verbose=False): # Use minoverlap=0.8, exconf=0.98!!!
    # print progress
    if verbose:
        print("[Log] Post-processing started for segmentation instances...")
    
    # to not affect person data, we pre-extract the information
    is_person = instances.pred_classes == 0
    survived_instance_idxs=torch.tensor(list(range(len(instances.pred_classes))), device=is_person.device)[is_person]
    
    # for all classes
    for class_id in list(set(instances.pred_classes.cpu().numpy().tolist())):
        # skip if class_id denotes human
        if class_id == 0:
            continue
        
        # if class_id denotes object: lvis
        if hasattr(instances, "is_lvis"):
            # if contains "is_lvis", class must be a LVIS category 
            is_class = instances.pred_classes == class_id
            assert instances.is_lvis[is_class].any(), "Class must be a LVIS category if instances contain 'is_lvis' attribute"
            # retrieve class name
            class_name = LVIS_CLASS_ID2NAME[class_id]
        
        # if class_id denotes object: coco
        else:
            class_name = COCO_CLASS_ID2NAME[class_id]

        # indices of object in instances
        is_class=instances.pred_classes==class_id
        index_in_instances = torch.tensor(list(range(len(instances.pred_classes))), device=is_class.device)[is_class]
        
        # bboxes, masks, confidences
        bboxes_class = instances[is_class].pred_boxes.tensor.cpu().numpy() # xyxy
        masks_class = instances[is_class].pred_masks
        confidence_class = instances[is_class].scores.cpu().numpy() # ndarray of shape [B,]
        
        # process bboxes & confidences
        class_bbox_list = bbox_xy_to_wh(bboxes_class) 
        confidence_list = confidence_class.tolist()

        # remove excessive overlaps
        keepidx = process_remove_overlap(class_bbox_list.tolist(), confidence_list, minoverlap=minoverlap, exconf=exconf)
        survived_instance_idxs=torch.cat([survived_instance_idxs,index_in_instances[keepidx]])
        
        # notify if removal occured for specific class
        if len(keepidx) < len(instances[is_class].pred_classes):
            if verbose:
                print(f"[Log] Some of instances of CLASS: \"{class_name}\" have been removed due to excessive occlusion.")

    # sort remaining instances
    survived_instance_idxs = torch.sort(survived_instance_idxs).values.detach().cpu().numpy()
    
    # print progress
    if len(survived_instance_idxs) < len(instances.pred_classes):
        if verbose:
            print("[Log] There has been removal of some instances. Postprocessing complete.")
    elif len(survived_instance_idxs) == len(instances.pred_classes):
        if verbose:
            print("[Log] No removed instance. Postprocessing complete.")
    else:
        assert False

    return instances[survived_instance_idxs]



def process_remove_none(mocap_output_lst, keepidx):
    # list to save remaining outputs
    mocap_output_lst2return=[]
    final_keepidx=[]

    # iterate through all human outputs
    for pred_output, idx in zip(mocap_output_lst,keepidx):
        # remove if none
        if pred_output is not None:
            mocap_output_lst2return.append(pred_output)
            final_keepidx.append(idx)
    
    return mocap_output_lst2return, final_keepidx



def process_mocap_predictions(
    mocap_predictions, bboxes, image_size, masks=None
):
    verts = np.stack([p["pred_vertices_smpl"] for p in mocap_predictions])
    # All faces are the same, so just need one copy.
    faces = np.expand_dims(mocap_predictions[0]["faces"].astype(np.int32), 0)
    max_dim = np.max(bbox_xy_to_wh(bboxes)[:, 2:], axis=1)
    local_cams = []
    for b, pred in zip(max_dim, mocap_predictions): # mocap predictions: same as pred_output_list
        local_cam = pred["pred_camera"].copy()
        scale_o2n = pred["bbox_scale_ratio"] * b / 224
        local_cam[0] /= scale_o2n
        local_cam[1:] /= local_cam[:1]
        local_cams.append(local_cam)
    local_cams = np.stack(local_cams)
    global_cams = local_to_global_cam(bboxes, local_cams, image_size)
    inds = np.argsort(bboxes[:, 0])  # Sort from left to right
    person_parameters = {
        "bboxes": bboxes[inds].astype(np.float32),
        "cams": global_cams[inds].astype(np.float32),
        "faces": faces,
        "local_cams": local_cams[inds].astype(np.float32),
        "verts": verts[inds].astype(np.float32),
        "indices": inds # --> indices for making the order from left bbox to right bbox
    }
    for k, v in person_parameters.items():
        person_parameters[k] = torch.from_numpy(v).cuda()
    if masks is not None:
        full_boxes = torch.tensor([[0, 0, image_size, image_size]] * len(bboxes))
        full_boxes = full_boxes.float().cuda()
        masks = BitMasks(masks).crop_and_resize(boxes=full_boxes, mask_size=image_size)
        person_parameters["masks"] = masks[inds].cuda()
    return person_parameters
utils/postprocess.py
jellyheadandrew/CHORUS
# Copyright (c) Facebook, Inc. and its affiliates.
"""
Wrapper for Human Pose Estimator using BodyMocap.
See: https://github.com/facebookresearch/frankmocap
"""
import torch

from imports.frankmocap.bodymocap.body_mocap_api import BodyMocap
from imports.frankmocap.renderer.visualizer import Visualizer

from constants.frankmocap import BODY_MOCAP_REGRESSOR_CKPT, BODY_MOCAP_SMPL_PTH, VISUALIZER_MODE


# mocap for body
def prepare_frankmocap_regressor():
    assert torch.cuda.is_available(), "Current version only supports GPU"
    
    # Set mocap regressor
    bodymocap = BodyMocap(
        regressor_checkpoint=BODY_MOCAP_REGRESSOR_CKPT, 
        smpl_dir=None, 
        smplModelPath=BODY_MOCAP_SMPL_PTH
    )
    return bodymocap


def prepare_frankmocap_visualizer():
    assert torch.cuda.is_available(), "Current version only supports GPU"

    # Prepare Pytorch3d Visualizer
    visualizer = Visualizer(VISUALIZER_MODE)
    return visualizer

utils/prepare_frankmocap.py
jellyheadandrew/CHORUS
import os
from glob import glob


def prepare_cps(prompt_dir, categories, use_txt_prompts: bool):
    if categories is None and use_txt_prompts:
        categories = [
            prompt_pth.split("/")[-1].replace(".txt", "")
            for prompt_pth in sorted(list(glob(f"{prompt_dir}/*.txt")))
        ]
    elif categories is None and not use_txt_prompts:
        categories = [
            prompt_dir.split("/")[-1]
            for prompt_dir in sorted(list(glob(f"{prompt_dir}/*")))
        ]
    else:
        pass

    # final cps to return
    final_cps = []

    # load category-prompt pairs
    for category in categories:
        if not use_txt_prompts:
            cps = [tuple(prompt_pth.split("/")[-2:]) 
                    for prompt_pth 
                    in list(glob(f"{prompt_dir}/{category}/*"))]
        else:
            txt_pth = f"{prompt_dir}/{category}.txt"
            with open(txt_pth, "r") as f:
                cps = [
                        tuple(
                            [os.path.splitext(os.path.basename(txt_pth))[0], 
                            line.strip()]
                        ) for line in f.read().split("\n")
                    ]

        # remove empty string
        cps = [tuple([category, prompt]) 
                    for category, prompt in cps
                    if len(prompt) > 0
                ]
        # remove prompt starting with #
        cps = [tuple([category, prompt]) 
                    for category, prompt in cps
                    if prompt[0] != "#"
                ]
        # remove . in end
        cps = [tuple([category, prompt[:-1]]) if prompt[-1] == "." 
                    else tuple([category, prompt]) 
                    for category, prompt in cps]
        # remove duplicates
        cps = sorted(list(set(cps)))
        # append to final_cps
        final_cps += cps
 
    return final_cps


def prepare_augdict(cps, augmentation_path, verbose):
    with open(augmentation_path, "r") as f:
        # read augmentations
        augprompts = [line.strip() for line in f.read().split("\n") if line.strip() != ""]
        augprompts = [line.strip() for line in augprompts if line.strip()[0] != "#"]
        # header starts with "AUGPROMPT_COMMON": use augmentations for all prompts
        if augprompts[0] == "AUGPROMPT_COMMON":
            if verbose:
                print("Using Common Augprompts over All Prompts...")
            augprompts = augprompts[1:]
            tmp = dict()
            for cp in cps:
                tmp[cp] = augprompts
            augdict = tmp
        # header starts with "AUGPROMPT_PER_PROMPT": augmentations specified for each prompt     
        elif augprompts[0] == "AUGPROMPT_PER_PROMPT":
            if verbose:
                print("Using Prompt-specific Augprompts...")
            augprompts = augprompts[1:]
            tmp = dict()
            current_category = ""
            current_prompt = ""
            for line in augprompts:
                if line[:12] == "category:":
                    current_category = line.replace("category:", "").strip()
                elif line[:7] == "PROMPT:":
                    current_prompt = line.replace("PROMPT:", "").strip()
                    assert (current_category, current_prompt) not in tmp.keys(), f"Duplicate prompt in {augmentation_path} file"
                    assert (current_category, current_prompt) in cps, f"In {augmentation_path}, there are invalid prompts."
                    tmp[(current_category, current_prompt)] = []
                elif current_category == "" or current_prompt == "":
                    continue
                else:
                    tmp[(current_category, current_prompt)].append(line.strip())
            augdict = tmp

        # invalid header => raise error    
        else:
            assert False, "Augprompt-header must be specified."
    
    # sort augdict
    augdict = dict(sorted(augdict.items()))
    assert tuple(sorted(list(augdict.keys()))) == tuple(sorted(cps)), f"Mismatch in augprompt file: '{augmentation_path}' & cps we are to use: {cps}"

    return augdict


def prepare_cps_from_dirs(args, image_dir, use_filter_setting=False):
    # Categories. If None, type them in.
    if args.categories is None:
        categories = ["*"] # all
    else:
        assert type(args.categories) == list
        categories = args.categories

    # Parse Categories.
    if type(categories) == str:
        categories = [category.strip() for category in categories.split(",")]
    else:
        assert type(categories) == list

    # To Save cps.
    cps = []

    # If using setting (usually After Filtering)
    if use_filter_setting:
        for category in categories:
            cps += [
                # modified for "settings" input
                tuple([prompt_pth.split("/")[-3],prompt_pth.split("/")[-1]]) 
                for prompt_pth in list(glob(f"{image_dir}/*/{category}/settings:*/*"))
            ]
    # If not using setting (usually Before Filtering)
    else:
        for category in categories:
            cps += [
                tuple(prompt_pth.split("/")[-2:]) 
                for prompt_pth in list(glob(f"{image_dir}/*/{category}/*"))
            ]

    # Remove duplicate, and sort
    cps = sorted(list(set(cps)))
    
    # Filter the prompts.
    if 'prompts' not in vars(args).keys():
        pass
    elif args.prompts is not None:
        _cps = []
        for prompt_type, prompt in cps:
            if prompt in args.prompts:
                _cps += [(prompt_type, prompt)]
        cps = _cps
        # print that filtering with 'args.prompts' has happend
        if args.verbose:
            print(f"Additionally filtered category-prompt pairs sicne '--prompts' is provided: Using {len(cps)} number of cps\n")

    # print total number of (category, prompt) pairs
    if args.verbose:
        print(f"Running for {len(cps)} number of prompts\n\n")    

    return cps


def get_unique_categories(cps):
    # list to save all categories
    all_categories = []

    # iterate for all category-prompt pairs
    for category_, prompt_ in cps: 
        all_categories.append(category_)

    # remove duplicates to retrieve unique categories
    unique_categories = sorted(list(set(all_categories)))

    return unique_categories


def get_cps_for_given_category(cps, category):
    # list to save prompts for given category 
    filtered_prompts = []

    # iterate for all category-prompt pairs
    for category_, prompt_ in cps:
        if category_ == category:
            filtered_prompts.append((category_, prompt_))
    
    # remove duplicates to retrieve cps for given category
    filtered_prompts = sorted(list(set(filtered_prompts)))
    return filtered_prompts
utils/prepare_prompts.py
jellyheadandrew/CHORUS
import os

import numpy as np
import random
import torch

max_seed_value = np.iinfo(np.uint32).max
min_seed_value = np.iinfo(np.uint32).min

def seed_everything(seed: int, workers: bool=False):
    os.environ["PL_GLOBAL_SEED"] = str(seed)
    os.environ["PL_SEED_WORKERS"] = f"{int(workers)}"
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    return seed
utils/reproducibility.py
jellyheadandrew/CHORUS
import os
from glob import glob

from constants.keypoint_detection import COCO_2DKEYPOINT_NAME2IDX

## printer used for exporting settings
class Printer(object):
    def __init__(self):
        super().__init__()

    def dict2print(self, x: dict, return_str: bool = True):
        lines = []
        for k, v in x.items():
            lines += [f"{k}:{v}"]
        towrite = "\n".join(lines)
        if return_str:
            return towrite
        else:
            print(towrite)

    def dict2txt(self, txt_pth, x: dict):
        lines = []
        for k, v in x.items():
            lines += [f"{k}:{v}"]
        towrite = "\n".join(lines)

        with open(txt_pth, "w") as f:
            f.write(towrite)

    def str2txt(self, txt_pth, x: str):
        with open(txt_pth, "w") as f:
            f.write(x)
PRINTER = Printer()


## create filtering settings
def create_filter_settings(
        ldm_model_key,
        image_dir,
        seg_dir,
        human_dir,
        save_dir,
        keypoint_dir,
        use_template,
        keypoints_which,
        filter_multi_person_object_methods,
        bbox_overlap_thres,
        keypoint_confidence_threshold,
        skip_done,
        verbose,
    ):
    # settings (to save)
    settings = dict()

    ## directory settings
    if not use_template:
        settings["image_dir"] = image_dir
        settings["seg_dir"]=seg_dir
        settings["human_dir"]=human_dir
        settings["keypoint_dir"]=keypoint_dir
    else:
        seg_thres = input("Using Template for Using Directories... Type in the Instance-Segmentation Threshold You Have Used In Previous Steps: ")
        settings["image_dir"] = image_dir
        settings["seg_dir"] = f"{image_dir}_seg_{seg_thres}"
        settings["human_dir"] = f"{image_dir}_frankmocap_seg{seg_thres}"
        settings["keypoint_dir"]=f"{image_dir}_2Dkp_COCO_HRNetDark_frank_seg{seg_thres}"

    ## multi-person, multi-object, bbox-filtering method settings
    # number of person allowed in image
    if "only-one-person-in-image" in filter_multi_person_object_methods:
        assert "multi-person-allowed-in-image" not in filter_multi_person_object_methods, "You must choose between 'only-one-person-in-image' and 'multi-person-allowed-in-image'"
        skip_multi_person_image = True
    elif "multi-person-allowed-in-image" in filter_multi_person_object_methods:
        assert "only-one-person-in-image" not in filter_multi_person_object_methods, "You must choose between 'only-one-person-in-image' and 'multi-person-allowed-in-image'"
        skip_multi_person_image = False
    else:
        assert False, "You must choose between 'only-one-person-in-image' and 'multi-person-allowed-in-image'"
    # number of objectgs allowed in image
    if "only-one-object-in-image" in filter_multi_person_object_methods:
        assert "multi-object-allowed-in-image" not in filter_multi_person_object_methods, "You must choose between 'only-one-object-in-image' and 'multi-object-allowed-in-image'"
        skip_multi_object_image = True
    elif "multi-object-allowed-in-image" in filter_multi_person_object_methods:
        assert "only-one-object-in-image" not in filter_multi_person_object_methods, "You must choose between 'only-one-object-in-image' and 'multi-object-allowed-in-image'"
        skip_multi_object_image = False
    else:       
        assert False, "You must choose between 'only-one-object-in-image' and 'multi-object-allowed-in-image'"
    # If only-one-person allowed in image: how will I count objects?
    if "1-bbox-per-person" in filter_multi_person_object_methods:
        assert "0,1-bbox-per-person" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."
        assert "n-bbox-per-person" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."
        assert "no-bbox-filtering" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."        
        allowed_bbox_num_per_human = [1]
    elif "0,1-bbox-per-person" in filter_multi_person_object_methods:
        assert "1-bbox-per-person" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."
        assert "n-bbox-per-person" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."
        assert "no-bbox-filtering" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."        
        allowed_bbox_num_per_human = [0,1]
    elif "n-bbox-per-person" in filter_multi_person_object_methods:
        assert "1-bbox-per-person" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."
        assert "0,1-bbox-per-person" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."
        assert "no-bbox-filtering" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."        
        allowed_bbox_num_per_human = [0,1,2,3,4,5]
    elif "no-bbox-filtering" in filter_multi_person_object_methods:
        assert "1-bbox-per-person" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."
        assert "0,1-bbox-per-person" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."
        assert "n-bbox-per-person" not in filter_multi_person_object_methods, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."        
        allowed_bbox_num_per_human = []
    else:
        assert False, "You must choose between '1-bbox-per-person' or '0,1-bbox-per-person' or 'n-bbox-per-person' or 'no-bbox-filtering'."    
    # save to settings
    settings["filter_multi_person_object_methods"] = sorted(filter_multi_person_object_methods)

    ## keypoint filter (coco)
    keypoints_that_must_exist, keypoints_that_must_not_exist = keypoints_which.replace(" ", "").replace("\n", "").split("|")
    # keypoints that must exist (as list[list])
    if keypoints_that_must_exist == "":
        keypoints_that_must_exist = []
    else:
        keypoints_that_must_exist = sorted(
            [
                sorted([COCO_2DKEYPOINT_NAME2IDX[name.strip()] for name in name_tuples.split(",")]) 
                for name_tuples in keypoints_that_must_exist.split("/")
            ]
        )
    # keypoints that must not exist (as list[list])
    if keypoints_that_must_not_exist == "":
        keypoints_that_must_not_exist = []
    else:
        keypoints_that_must_not_exist = sorted(
            [
                sorted([COCO_2DKEYPOINT_NAME2IDX[name.strip()] for name in name_tuples.split(",")]) 
                for name_tuples in keypoints_that_must_not_exist.split("/")
            ]
        )
    # save to settings
    settings['keypoints_that_must_exist'] = keypoints_that_must_exist
    settings['keypoints_that_must_not_exist'] = keypoints_that_must_not_exist
    settings["keypoint_confidence_threshold"] = keypoint_confidence_threshold

    ## export settings as string (to save)
    settings = dict(sorted(settings.items()))
    if verbose:
        print("\n\n===settings===")
        PRINTER.dict2print(settings, return_str=False)
    # print as txt file        
    settings_str = PRINTER.dict2print(settings, return_str=True)

    ## auxiliary-or-derived settings (not to save)
    settings["ldm_model_key"] = ldm_model_key

    settings["save_dir"] = save_dir
    settings["bbox_overlap_thres"] = bbox_overlap_thres
    
    settings["skip_done"] = skip_done
    settings["verbose"] = verbose

    settings['skip_multi_person_image'] = skip_multi_person_image
    settings['skip_multi_object_image'] = skip_multi_object_image
    settings['allowed_bbox_num_per_human'] = allowed_bbox_num_per_human

    return settings_str, settings


## finding setting number when given the filter settings
def find_filter_setting_num(settings_str, settings): # if current settings does not exist, saves as new setting
    # create saving directory for the first time
    os.makedirs(settings['save_dir'], exist_ok=True)
    
    # find all existing setting numbers
    setting_nums = sorted(
        [
            int(os.path.splitext(os.path.basename(setting_pth))[0].replace(f"settings:", "")) 
            for setting_pth in list(glob(f"{settings['save_dir']}/*.txt"))
            if os.path.splitext(os.path.basename(setting_pth))[0].replace(f"settings:", "").isdigit()
        ]
    )

    # compare setting with existing setting & find corresponding setting number
    setting_num = None
    for tmp in setting_nums:
        setting_pth = f"{settings['save_dir']}/settings:{tmp:03}.txt"
        with open(setting_pth, "r") as f:
            settings_str_ = f.read()
            if settings_str_ == settings_str:
                setting_num = tmp

    # if no same setting is found, save as new setting number (case1)
    if setting_num is None and len(setting_nums) != 0:
        # new setting number
        setting_num = max(setting_nums)+1
        # save as new .txt file
        new_save_name = f"settings:{setting_num:03}"
        txt_pth = f"{settings['save_dir']}/{new_save_name}.txt"
        PRINTER.str2txt(txt_pth, settings_str)
        # print new settings
        if settings['verbose']:
            print(f"New Settings: {setting_num:03}")
    
    # if no same setting is found, save as new setting number (case2)
    elif setting_num is None and len(setting_nums) == 0:
        # new setting number
        setting_num = 1
        # save as new .txt file
        new_save_name = f"settings:{setting_num:03}"
        txt_pth = f"{settings['save_dir']}/{new_save_name}.txt"
        PRINTER.str2txt(txt_pth, settings_str)
        # print new settings
        if settings['verbose']:
            print(f"New Settings: {setting_num:03}")
    
    # if same setting is found, use that setting number
    else:
        # print settings
        if settings['verbose']:
            print(f"Using already saved Settings: {setting_num:03}")
    
    return setting_num


## load saved settings used during filtering when setting number is given and update to args
def load_filter_settings(args, filter_setting_num):
    # load txt file with filtering settings recorded 
    with open(f"{args.filter_dir}/settings:{filter_setting_num:03}.txt", "r") as rf:
        filter_settings = [line.strip() for line in rf.read().strip().split("\n")]
    
    # 'settings to update' as dictionary
    filter_settings = {
        setting_line.split(":")[0].strip(): setting_line.split(":")[1].strip()
        for setting_line in filter_settings
    }
    filter_settings.update({'filter_setting_num': filter_setting_num})


    # update arguments to filter settings
    for k, v in sorted(vars(args).items()):
        filter_settings[k] = v

    return filter_settings


## load saved settings used during filtering when multiple setting number is given and update to args
def load_filter_settings_combination(args, filter_setting_comb, args_to_ignore=[]):
    ## final filter_settings to return
    filter_settings = dict()

    ## iterate for all filter_setting_num
    for filter_setting_num in filter_setting_comb:
        # load txt file with filtering settings recorded 
        with open(f"{args.filter_dir}/settings:{filter_setting_num:03}.txt", "r") as rf:
            _filter_settings = [line.strip() for line in rf.read().strip().split("\n")]
        
        # 'settings to update' as dictionary
        _filter_settings = {
            setting_line.split(":")[0].strip(): setting_line.split(":")[1].strip()
            for setting_line in _filter_settings
        }

        # for all new 'key-value' in _filter_settings, check for duplicate and consistency
        for k, v in _filter_settings.items():
            if k in filter_settings:
                assert filter_settings[k] == v, f"Consistency between filter_settings: '{filter_setting_comb}' not assured!!"
            else:
                filter_settings[k] = v

    # update arguments to filter settings
    for k, v in sorted(vars(args).items()):
        if k not in args_to_ignore: filter_settings[k] = v

    # update 'filter_setting_numbs' used
    filter_settings.update({'filter_setting_comb': filter_setting_comb})

    return filter_settings


## create name for 'aggregation' to use for save 
def aggr_save_name(filter_setting_comb, aggr_setting_name):
    return f"filter({','.join(str(x) for x in filter_setting_comb)})-aggr({aggr_setting_name})"
utils/settings.py
jellyheadandrew/CHORUS
import math
import torch
import torch.nn.functional as F

""" reference: https://github.com/vchoutas/smplx/blob/main/smplx/lbs.py """
def transform_mat(R, t):
    ''' Creates a batch of transformation matrices
        Args:
            - R: Bx3x3 array of a batch of rotation matrices
            - t: Bx3x1 array of a batch of translation vectors
        Returns:
            - T: Bx4x4 Transformation matrix
    '''
    # No padding left or right, only add an extra row
    return torch.cat([F.pad(R, [0, 0, 0, 1]),
                      F.pad(t, [0, 0, 0, 1], value=1)], dim=2)


""" reference: https://github.com/vchoutas/smplx/blob/main/smplx/lbs.py """
def batch_rigid_transform(
        rot_mats,
        joints,
        parents,
        dtype
    ):
    """
    Applies a batch of rigid transformations to the joints

    Parameters
    ----------
    rot_mats : torch.tensor BxNx3x3
        Tensor of rotation matrices
    joints : torch.tensor BxNx3
        Locations of joints
    parents : torch.tensor BxN
        The kinematic tree of each object
    dtype : torch.dtype, optional:
        The data type of the created tensors, the default is torch.float32

    Returns
    -------
    posed_joints : torch.tensor BxNx3
        The locations of the joints after applying the pose rotations
    rel_transforms : torch.tensor BxNx4x4
        The relative (with respect to the root joint) rigid transformations
        for all the joints
    """

    joints = torch.unsqueeze(joints, dim=-1)

    rel_joints = joints.clone()
    rel_joints[:, 1:] -= joints[:, parents[1:]]

    transforms_mat = transform_mat(
        rot_mats.reshape(-1, 3, 3),
        rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)

    transform_chain = [transforms_mat[:, 0]]
    for i in range(1, parents.shape[0]):
        # Subtract the joint location at the rest pose
        # No need for rotation, since it's identity when at rest
        curr_res = torch.matmul(transform_chain[parents[i]],
                                transforms_mat[:, i])
        transform_chain.append(curr_res)

    transforms = torch.stack(transform_chain, dim=1)

    # The last column of the transformations contains the posed joints
    posed_joints = transforms[:, :, :3, 3]

    joints_homogen = F.pad(joints, [0, 0, 0, 1])

    rel_transforms = transforms - F.pad(
        torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0])

    return posed_joints, rel_transforms


""" reference: https://github.com/nkolot/SPIN/blob/master/utils/geometry.py """
def batch_rodrigues(theta):
    """Convert axis-angle representation to rotation matrix.
    Args:
        theta: size = [B, 3]
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    l1norm = torch.norm(theta + 1e-8, p = 2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)
    return quat_to_rotmat(quat)


""" reference: https://github.com/nkolot/SPIN/blob/master/utils/geometry.py """
def quat_to_rotmat(quat):
    """Convert quaternion coefficients to rotation matrix.
    Args:
        quat: size = [B, 4] 4 <===>(w, x, y, z)
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    norm_quat = quat
    norm_quat = norm_quat/norm_quat.norm(p=2, dim=1, keepdim=True)
    w, x, y, z = norm_quat[:,0], norm_quat[:,1], norm_quat[:,2], norm_quat[:,3]

    B = quat.size(0)

    w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
    wx, wy, wz = w*x, w*y, w*z
    xy, xz, yz = x*y, x*z, y*z

    rotMat = torch.stack([w2 + x2 - y2 - z2, 2*xy - 2*wz, 2*wy + 2*xz,
                          2*wz + 2*xy, w2 - x2 + y2 - z2, 2*yz - 2*wx,
                          2*xz - 2*wy, 2*wx + 2*yz, w2 - x2 - y2 + z2], dim=1).view(B, 3, 3)
    return rotMat


""" reference: https://github.com/nkolot/SPIN/blob/master/utils/geometry.py """
def rot6d_to_rotmat(x):
    """Convert 6D rotation representation to 3x3 rotation matrix.
    Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
    Input:
        (B,6) Batch of 6-D rotation representations
    Output:
        (B,3,3) Batch of corresponding rotation matrices
    """
    x = x.view(-1,3,2)
    a1 = x[:, :, 0]
    a2 = x[:, :, 1]
    b1 = F.normalize(a1)
    b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1)
    b3 = torch.cross(b1, b2)
    return torch.stack((b1, b2, b3), dim=-1)


""" reference: https://github.com/Wallacoloo/printipi/blob/master/util/rotation_matrix.py """
def rotmat_to_axis_angle(matrix):
    """Convert the rotation matrix into the axis-angle notation.

    Conversion equations
    ====================

    From Wikipedia (http://en.wikipedia.org/wiki/Rotation_matrix), the conversion is given by::

        x = Qzy-Qyz
        y = Qxz-Qzx
        z = Qyx-Qxy
        r = hypot(x,hypot(y,z))
        t = Qxx+Qyy+Qzz
        theta = atan2(r,t-1)

    @param matrix:  The 3x3 rotation matrix to update.
    @type matrix:   3x3 numpy array
    @return:    The 3D rotation axis and angle.
    @rtype:     numpy 3D rank-1 array, float
    """

    # Axes.
    axis = torch.zeros(3, dtype=torch.float32)
    axis[0] = matrix[2,1] - matrix[1,2]
    axis[1] = matrix[0,2] - matrix[2,0]
    axis[2] = matrix[1,0] - matrix[0,1]

    # Angle.
    r = torch.hypot(axis[0], torch.hypot(axis[1], axis[2]))
    t = matrix[0,0] + matrix[1,1] + matrix[2,2]
    theta = torch.atan2(r, t-1)

    # Normalise the axis.
    axis = axis / r

    # Return the data.
    return axis, theta


def rotmat_to_axis_angle_batch(matrix):
    """Convert the rotation matrix into the axis-angle notation.

    Conversion equations
    ====================

    From Wikipedia (http://en.wikipedia.org/wiki/Rotation_matrix), the conversion is given by::

        x = Qzy-Qyz
        y = Qxz-Qzx
        z = Qyx-Qxy
        r = hypot(x,hypot(y,z))
        t = Qxx+Qyy+Qzz
        theta = atan2(r,t-1)

    @param matrix:  The 3x3 rotation matrix to update.
    @type matrix:   3x3 numpy array
    @return:    The 3D rotation axis and angle.
    @rtype:     numpy 3D rank-1 array, float
    """
    # Batch size.
    batch_size = matrix.shape[0]

    # Axes.
    axis = torch.zeros(batch_size, 3, dtype=torch.float32)
    axis[:,0] = matrix[:,2,1] - matrix[:,1,2]
    axis[:,1] = matrix[:,0,2] - matrix[:,2,0]
    axis[:,2] = matrix[:,1,0] - matrix[:,0,1]

    # Angle.
    r = torch.hypot(axis[:,0], torch.hypot(axis[:,1], axis[:,2]))
    t = matrix[:,0,0] + matrix[:,1,1] + matrix[:,2,2]
    theta = torch.atan2(r, t-1)

    # Normalise the axis.
    axis = axis / r[:,None]

    # Return the data.
    return axis, theta


def convert_Kparams_Kmat(tensor):
    K_elem = torch.split(tensor, 1)
    elem1 = torch.Tensor([1, 0, 0]).to(tensor.device) * K_elem[0]
    elem2 = torch.Tensor([0, 1, 0]).to(tensor.device) * K_elem[1]
    elem3 = torch.Tensor([1, 0, 0]).to(tensor.device) * K_elem[2] + torch.Tensor([0, 1, 0]).to(tensor.device) * K_elem[3] + torch.Tensor([0, 0, 1]).to(tensor.device)

    elem1 = elem1.unsqueeze(0).view(3, -1)
    elem2 = elem2.unsqueeze(0).view(3, -1)
    elem3 = elem3.unsqueeze(0).view(3, -1)

    result = torch.cat((elem1, elem2, elem3), dim=1)
    return result


def batch_convert_Kparams_Kmat(K_elem_batch):
    # K_elem_batch: Bx4
    assert K_elem_batch.ndim == 2
    assert K_elem_batch.shape[-1] == 4
    
    # elem1, elem2, elem3 --> Bx3
    elem1 = torch.Tensor([[1, 0, 0]]).to(K_elem_batch.device) * K_elem_batch[:,0:1]
    elem2 = torch.Tensor([[0, 1, 0]]).to(K_elem_batch.device) * K_elem_batch[:,1:2]
    elem3 = torch.Tensor([[1, 0, 0]]).to(K_elem_batch.device) * K_elem_batch[:,2:3] + \
            torch.Tensor([[0, 1, 0]]).to(K_elem_batch.device) * K_elem_batch[:,3:4] + \
            torch.Tensor([[0, 0, 1]]).to(K_elem_batch.device)

    elem1 = elem1.unsqueeze(1).permute(0,2,1) # Bx3x1
    elem2 = elem2.unsqueeze(1).permute(0,2,1) # Bx3x1
    elem3 = elem3.unsqueeze(1).permute(0,2,1) # Bx3x1

    K = torch.cat((elem1, elem2, elem3), dim=-1)
    return K


def get_focal_length(img_length, fovy, use_radian=False):
    if not use_radian:
        fovy = fovy * math.pi / 180.
    return img_length / (2.0*math.tan(fovy/2))


def get_azimuth(x, z, eps=1e-4, return_rad=False):
    if x >= eps and z >= eps:
        theta = math.atan(z/x)
    elif x < eps and x >= -eps and z >= eps:
        theta = (math.pi / 2)
    elif x < -eps and z >= eps:
        theta = math.pi + math.atan(z/x)
    elif x < -eps and z < eps and z>= - eps:
        theta = math.pi
    elif x < -eps and z < -eps:
        theta = math.pi + math.atan(z/x)
    elif x >= -eps and x < eps and z<-eps:
        theta = 3 * (math.pi / 2)
    elif x >= eps and z<-eps:
        theta = 2 * math.pi + math.atan(z/x)
    else:
        theta = 0.
    if return_rad: return theta
    else: return theta * 180 / math.pi


def get_zenith(x, y, z, eps=1e-4, return_rad=False):
    r = math.sqrt(x ** 2 + y ** 2 + z ** 2)
    cos_zenith = z / r

    if cos_zenith >= 1.-eps:
        zenith = 0.
    elif -1.+eps < cos_zenith < 1.-eps:
        zenith = math.acos(cos_zenith)
    else:
        zenith = math.pi
        
    if return_rad: return zenith
    else: return zenith * 180 / math.pi


def scale_wrt_tan(theta, scale, use_deg):
    # convert theta from 'degree' scale to 'radian' scale
    if use_deg: theta = theta * (2 * math.pi) / 360

    # scale 'tangent' value
    tan = math.tan(theta)
    tan = tan * scale

    # apply 'inverse tangent'
    theta = math.atan(tan)

    # return as 'degree' or 'radian'
    if use_deg: return theta * 360 / (2 * math.pi)
    else: return theta


def perspective_projection_grid(grid, K, R, t, grid_size=None, use_normalized_K=False, img_size=(None, None)):
    # homogeneous coordinates
    homo_grid = torch.matmul(K, torch.matmul(R, grid.T) + t.reshape(3,1)).reshape(3,*grid_size)
    
    # calculate projections
    projected_grid = homo_grid[:2].clone() # Shape: [2,*grid_size]
    projected_grid[0] /= homo_grid[2]
    projected_grid[1] /= homo_grid[2]

    # if using 'normalized-K', upscale to pixel-scale 
    if use_normalized_K:
        projected_grid *= img_size[1] # multiply W

    return projected_grid # [2,*grid_size]


def perspective_projection_mesh(mesh, K, R, t, use_normalized_K=False, img_size=(None, None), return_homo=False):
    # homogeneous coordinates
    homo_mesh = torch.matmul(K, torch.matmul(R, mesh.T) + t.reshape(3,1)).reshape(3,mesh.shape[0])
    
    # calculate projections
    projected_mesh = homo_mesh[:2].clone() # Shape: [2,N]
    projected_mesh[0] /= homo_mesh[2]
    projected_mesh[1] /= homo_mesh[2]

    # if using 'normalized-K', upscale to pixel-scale 
    if use_normalized_K:
        projected_mesh *= img_size[1] # multiply W

    if return_homo: return projected_mesh.T, homo_mesh.T # [N,2], [N,3]
    else: return projected_mesh.T # [N,2]
utils/transformations.py
jellyheadandrew/CHORUS
import os
import shutil

import torch
import numpy as np
import math
import cv2

from mayavi import mlab

from constants.mayavi_cams import QUAL_CAMS

from utils.misc import to_np_torch_recursive
from utils.transformations import get_focal_length, convert_Kparams_Kmat



def mayavi_render_mesh_heatmap_canonical(
        img_size,
        alpha,
        prob_field,
        verts,
        faces,
        color,
        thres,
        voxel_size,
        fovy,
        interactive,
        ignore_percentage,
        mesh_representation='surface',
        mesh_color=None,
        tmp_cache=None,
    ):
    # set 'mayavi.mlab' as offscreen mode
    mlab.options.offscreen = not interactive
    
    # get image shape
    H, W = img_size
    N_x, N_y, N_z = prob_field.shape

    # declare mayavi figure
    fig = mlab.figure(size=(W,H), bgcolor=(1,1,1))
    scene = fig.scene


    ## Render a 'probability field'
    # prepare 'probability field'
    prob_field = to_np_torch_recursive(prob_field, use_torch=False, device='cpu')
    prob_field = prob_field * (prob_field > ignore_percentage * prob_field.max())
    
    # add to 'mayavi.mlab' pipeline
    mlab.pipeline.volume(mlab.pipeline.scalar_field(prob_field), vmin=thres, vmax=1.0)

    ## Render a 'mesh'
    # prepare 'mesh'
    x = to_np_torch_recursive(verts[:,0], use_torch=False, device='cpu')
    y = to_np_torch_recursive(verts[:,1], use_torch=False, device='cpu')
    z = to_np_torch_recursive(verts[:,2], use_torch=False, device='cpu')
    faces = to_np_torch_recursive(faces, use_torch=False, device='cpu')
    # color
    color = to_np_torch_recursive(color, use_torch=False, device='cpu')
    if color.ndim == 1: color = np.stack([color] * x.shape[0], axis=0)
    if mesh_color is not None: color = np.ones_like(color) * mesh_color

    # add mesh to 'mayavi.mlab' pipeline
    mesh = mlab.triangular_mesh(
        x / voxel_size + N_x / 2,
        y / voxel_size + N_y / 2,
        z / voxel_size + N_z / 2,
        faces,
        representation = mesh_representation,
        scalars = np.arange(x.size)
    )

    # add texture to 'posed human'
    mesh.module_manager.scalar_lut_manager.lut.table = np.c_[
        (color * 255), 
        alpha * 255 * np.ones(verts.shape[0])
    ].astype(np.uint8)

    # update to 'mayavi.mlab' engine
    mesh.mlab_source.update()
    mesh.parent.update()

    ## Set camera for 'canonical' rendering
    renderings = []
    camera = scene.camera
    # fov
    f = get_focal_length(1.0, fovy)
    view_angle = np.arctan((H/2) / (W*f))
    
    ## iterate for all 'QUAL_CAMS'
    for i in range(2):
        # update camera of visualizer
        camera = scene.camera
        camera.position = ((np.array(QUAL_CAMS[i]['position']) / voxel_size) + np.array([N_x,N_y,N_z]) / 2).tolist()
        camera.focal_point = ((np.array(QUAL_CAMS[i]['focal_point']) / voxel_size) + np.array([N_x,N_y,N_z]) / 2).tolist()
        camera.view_angle = np.rad2deg(view_angle) * 2
        camera.view_up = QUAL_CAMS[i]['view_up']
        camera.clipping_range = [0.1, 1000000.]
        camera.compute_view_plane_normal()
        if interactive: break
        
        ## save the scene into 'np.ndarray'
        # [slight hack] save as image, and then re-open it
        os.makedirs(tmp_cache, exist_ok=True)
        mlab.savefig(f"{tmp_cache}/tmp.png")
        rendering = cv2.imread(f"{tmp_cache}/tmp.png")
        rendering = cv2.resize(rendering[:,:,::-1].copy(), dsize=(W,H))
        renderings.append(rendering)

        # remove after opening
        shutil.rmtree(tmp_cache)

    # if 'interactive', activate mlab viewer
    if interactive:
        mlab.show()
    else:
        ## close all figures
        mlab.close(all=True)
        return renderings
    
    
    

def mayavi_render_mesh_heatmap(
        img_size,
        prob_field,
        voxel_size,
        thres,
        cam_R,
        cam_t,
        cam_K,
        verts,
        faces,
        color,
        alpha,
        ignore_percentage,
        interactive,
        mesh_representation='surface',
        mesh_color=None,
        tmp_cache=None,
        **kwargs
    ):
    # set 'mayavi.mlab' as offscreen mode
    mlab.options.offscreen = not interactive
    
    # get image shape
    H, W = img_size
    N_x, N_y, N_z = prob_field.shape

    # declare mayavi figure
    fig = mlab.figure(size=(W,H), bgcolor=(1,1,1))
    scene = fig.scene

    ## Render a 'probability field'
    # prepare 'probability field'
    prob_field = to_np_torch_recursive(prob_field, use_torch=False, device='cpu')
    prob_field = prob_field * (prob_field > ignore_percentage * prob_field.max())
    
    cam_R = to_np_torch_recursive(cam_R, use_torch=False, device='cpu')
    cam_t = to_np_torch_recursive(cam_t, use_torch=False, device='cpu')
    cam_K = to_np_torch_recursive(cam_K, use_torch=False, device='cpu')

    # add to 'mayavi.mlab' pipeline
    mlab.pipeline.volume(mlab.pipeline.scalar_field(prob_field), vmin=thres, vmax=1.0)

    ## Render a 'mesh'
    # prepare 'mesh'
    x = to_np_torch_recursive(verts[:,0], use_torch=False, device='cpu')
    y = to_np_torch_recursive(verts[:,1], use_torch=False, device='cpu')
    z = to_np_torch_recursive(verts[:,2], use_torch=False, device='cpu')
    faces = to_np_torch_recursive(faces, use_torch=False, device='cpu')
    # color
    color = to_np_torch_recursive(color, use_torch=False, device='cpu')
    if color.ndim == 1: color = np.stack([color] * x.shape[0], axis=0)
    if mesh_color is not None: color = np.ones_like(color) * mesh_color

    # add mesh to 'mayavi.mlab' pipeline
    mesh = mlab.triangular_mesh(
        x / voxel_size + N_x / 2,
        y / voxel_size + N_y / 2,
        z / voxel_size + N_z / 2,
        faces,
        representation = mesh_representation,
        scalars = np.arange(x.size)
    )

    # add texture to 'posed human'
    mesh.module_manager.scalar_lut_manager.lut.table = np.c_[
        (color * 255), 
        alpha * 255 * np.ones(verts.shape[0])
    ].astype(np.uint8)

    # update to 'mayavi.mlab' engine
    mesh.mlab_source.update()
    mesh.parent.update()

    ## Set camera
    # translation
    camera_location = -cam_R.T @ cam_t
    # fov
    view_angle = np.arctan( (H / 2) / cam_K[0][0])
    # focal
    focal_point = camera_location - camera_location[2] * cam_R[2] / cam_R[2,2]
    # scale to 'pixel' scale
    camera_location = (camera_location / voxel_size) + np.array([N_x,N_y,N_z]) / 2
    focal_point = (focal_point / voxel_size) + np.array([N_x,N_y,N_z]) / 2
    # update camera of visualizer
    camera = scene.camera
    camera.position = camera_location
    camera.focal_point = focal_point
    camera.view_angle = np.rad2deg(view_angle) * 2 # 30.0 # view_angle
    camera.view_up = -cam_R[1] # negated due to openGL convention
    camera.clipping_range = [0.01, 100000.]
    camera.compute_view_plane_normal()

    # remove toolbar in the figure window
    mlab.process_ui_events()

    ## save the scene into 'np.ndarray'
    # [slight hack] save as image, and then re-open it
    os.makedirs(tmp_cache, exist_ok=True)
    mlab.savefig(f"{tmp_cache}/tmp.png", size=(W,H))
    rendering = cv2.imread(f"{tmp_cache}/tmp.png")
    rendering = cv2.resize(rendering[:,:,::-1].copy(), dsize=(W,H))

    # remove after opening
    shutil.rmtree(tmp_cache)

    # if 'interactive', activate mlab viewer
    if interactive: mlab.show()

    ## close all figures
    mlab.close(all=True)

    return rendering




def mayavi_render_mesh_heatmap_top(
        img_size,
        prob_field,
        voxel_size,
        thres,
        verts,
        faces,
        color,
        alpha,
        fovy,
        ignore_percentage,
        interactive,
        mesh_representation='surface',
        mesh_color=None,
        tmp_cache=None,
        **kwargs
    ):
    # set 'mayavi.mlab' as offscreen mode
    mlab.options.offscreen = not interactive
    
    # get image shape
    H, W = img_size
    N_x, N_y, N_z = prob_field.shape

    # declare mayavi figure
    fig = mlab.figure(size=(W,H), bgcolor=(1,1,1))
    scene = fig.scene

    ## Render a 'probability field'
    # prepare 'probability field'
    prob_field = to_np_torch_recursive(prob_field, use_torch=False, device='cpu')
    prob_field = prob_field * (prob_field > ignore_percentage * prob_field.max())    
    
    cam_R = np.array(
        [
            [1,0,0],
            [0,-1/np.sqrt(2),1/np.sqrt(2)],
            [0,-1/np.sqrt(2),-1/np.sqrt(2)]
        ], dtype=np.float32
    )
    cam_t = np.array([0,0,5])
    cam_K = to_np_torch_recursive(
        convert_Kparams_Kmat(
            W * (
                get_focal_length(1.0, fovy) * torch.Tensor([1,1,0,0])
                + torch.Tensor([0,0,0.5,0.5 * H / W])
            )
        ), 
        use_torch=False, 
        device='cpu'
    )
                             
    # add to 'mayavi.mlab' pipeline
    mlab.pipeline.volume(mlab.pipeline.scalar_field(prob_field), vmin=thres, vmax=1.0)


    ## Render a 'mesh'
    # prepare 'mesh'
    x = to_np_torch_recursive(verts[:,0], use_torch=False, device='cpu')
    y = to_np_torch_recursive(verts[:,1], use_torch=False, device='cpu')
    z = to_np_torch_recursive(verts[:,2], use_torch=False, device='cpu')
    faces = to_np_torch_recursive(faces, use_torch=False, device='cpu')
    # color
    color = to_np_torch_recursive(color, use_torch=False, device='cpu')
    if color.ndim == 1: color = np.stack([color] * x.shape[0], axis=0)
    if mesh_color is not None: color = np.ones_like(color) * mesh_color

    # add mesh to 'mayavi.mlab' pipeline
    mesh = mlab.triangular_mesh(
        x / voxel_size + N_x / 2,
        y / voxel_size + N_y / 2,
        z / voxel_size + N_z / 2,
        faces,
        representation = mesh_representation,
        scalars = np.arange(x.size)
    )

    # add texture to 'posed human'
    mesh.module_manager.scalar_lut_manager.lut.table = np.c_[
        (color * 255), 
        alpha * 255 * np.ones(verts.shape[0])
    ].astype(np.uint8)

    # update to 'mayavi.mlab' engine
    mesh.mlab_source.update()
    mesh.parent.update()


    ## Set camera
    # translation
    camera_location = -cam_R.T @ cam_t
    # fov
    view_angle = np.arctan( (H / 2) / cam_K[0][0])
    # focal
    focal_point = camera_location - camera_location[2] * cam_R[2] / cam_R[2,2]
    # scale to 'pixel' scale
    camera_location = (camera_location / voxel_size) + np.array([N_x,N_y,N_z]) / 2
    focal_point = (focal_point / voxel_size) + np.array([N_x,N_y,N_z]) / 2
    # update camera of visualizer
    camera = scene.camera
    camera.position = camera_location
    camera.focal_point = focal_point
    camera.view_angle = np.rad2deg(view_angle) * 2 # 30.0 # view_angle
    camera.view_up = -cam_R[1] # negated due to openGL convention
    camera.clipping_range = [0.01, 100000.]
    camera.compute_view_plane_normal()

    # remove toolbar in the figure window
    mlab.process_ui_events()

    ## save the scene into 'np.ndarray'
    # [slight hack] save as image, and then re-open it
    os.makedirs(tmp_cache, exist_ok=True)
    mlab.savefig(f"{tmp_cache}/tmp.png", size=(W,H))
    rendering = cv2.imread(f"{tmp_cache}/tmp.png")
    rendering = cv2.resize(rendering[:,:,::-1].copy(), dsize=(W,H))

    # remove after opening
    shutil.rmtree(tmp_cache)

    # if 'interactive', activate mlab viewer
    if interactive: mlab.show()

    ## close all figures
    mlab.close(all=True)

    return rendering
visualization/visualize_demo.py
jellyheadandrew/CHORUS

import numpy as np
import torch
import open3d as o3d


# create triangle mesh from torch.Tensor / np.ndarray input
def return_voxels(
        positions, 
        additional_translation=np.array([0,0,0], dtype=np.float32), 
        radius = 0.01,
        color = [1,0,0]
    ):
    if type(positions) == torch.Tensor:
        positions = positions.squeeze().clone().detach().cpu().numpy()
    if type(additional_translation) == torch.Tensor:
        additional_translation = additional_translation.squeeze().clone().detach().cpu().numpy()

    voxels = []
    for position in positions:
        voxel = o3d.geometry.TriangleMesh.create_sphere(radius=radius, resolution=2)
        voxel.paint_uniform_color(color)
        voxel.translate(position+additional_translation)
        voxels.append(voxel)

    return voxels


# create triangle mesh from torch.Tensor / np.ndarray input
def return_triangle_mesh(verts, faces, color = None, translation = None):
    if type(verts) == torch.Tensor:
        verts = verts.clone().detach().cpu().numpy()
    if type(faces) == torch.Tensor:
        faces = faces.clone().detach().cpu().numpy()
    if type(color) == torch.Tensor:
        color = color.clone().detach().cpu().numpy()
    if type(color) == np.ndarray or type(color) == torch.Tensor:
        color = np.squeeze(color)
    if type(translation) == torch.Tensor:
        translation = translation.clone().detach().cpu().numpy()
    if type(translation) == np.ndarray or type(translation) == torch.Tensor:
        translation = np.squeeze(translation)
    verts = np.squeeze(verts)
    faces = np.squeeze(faces)

    mesh = o3d.geometry.TriangleMesh()
    mesh.vertices = o3d.utility.Vector3dVector(verts)
    mesh.triangles = o3d.utility.Vector3iVector(faces)
    mesh.compute_vertex_normals()

    if color is not None:
        if color.shape == (3,):
            mesh.paint_uniform_color(color)
        else:
            assert color.shape == verts.shape
            mesh.vertex_colors = o3d.utility.Vector3dVector(color)

    if translation is not None:
        mesh.translate(translation)

    return mesh


# create open3d cube
def return_cube(length_x, length_y, length_z, center, color=[0.1,0.1,0.1]):
    if type(length_x) == torch.Tensor:
        length_x = length_x.clone().detach().cpu().numpy()
        length_x = np.squeeze(length_x)
    if type(length_y) == torch.Tensor:
        length_y = length_y.clone().detach().cpu().numpy()
        length_y = np.squeeze(length_y)
    if type(length_z) == torch.Tensor:
        length_z = length_z.clone().detach().cpu().numpy()
        length_z = np.squeeze(length_z)
    if type(center) == torch.Tensor:
        center = center.clone().detach().cpu().numpy()
        center = np.squeeze(center)
    if type(color) == torch.Tensor:
        color = color.clone().detach().cpu().numpy()
        color = np.squeeze(color)

    points = np.array([[length_x/2+center[0],length_y/2+center[1],length_z/2+center[2]], # 0
                                [length_x/2+center[0],length_y/2+center[1],-length_z/2+center[2]], # 1
                                [length_x/2+center[0],-length_y/2+center[1],length_z/2+center[2]], # 2
                                [length_x/2+center[0],-length_y/2+center[1],-length_z/2+center[2]], # 3
                                [-length_x/2+center[0],length_y/2+center[1],length_z/2+center[2]], # 4
                                [-length_x/2+center[0],length_y/2+center[1],-length_z/2+center[2]], # 5
                                [-length_x/2+center[0],-length_y/2+center[1],length_z/2+center[2]], # 6
                                [-length_x/2+center[0],-length_y/2+center[1],-length_z/2+center[2]]] # 7
                                , dtype=np.float32) # Shape: (8,3)
    lines = [
        [0,1],
        [0,2],
        [0,4],
        [1,3],
        [1,5],
        [2,3],
        [2,6],
        [3,7],
        [4,5],
        [4,6],
        [5,7],
        [6,7]
    ]     
    cube = o3d.geometry.LineSet(
        points=o3d.utility.Vector3dVector(points),
        lines=o3d.utility.Vector2iVector(lines)
    )
    cube.colors = o3d.utility.Vector3dVector([color]*12)
    return cube


# create open3d camera frustum
def return_camera_frustum(frustum_size, focal_length, cam_R, cam_t, color=[0.1,0.1,0.1]):
    points = [
        np.array([0,0,0], dtype=np.float32),
        np.array([0.5 * frustum_size, 0.5 * frustum_size, focal_length * frustum_size], dtype=np.float32),
        np.array([0.5 * frustum_size, -0.5 * frustum_size, focal_length * frustum_size], dtype=np.float32),
        np.array([-0.5 * frustum_size, 0.5 * frustum_size, focal_length * frustum_size], dtype=np.float32),
        np.array([-0.5 * frustum_size, -0.5 * frustum_size, focal_length * frustum_size], dtype=np.float32),
    ]
    lines = [
        [0, 1],
        [0, 2],
        [0, 3],
        [0, 4],
        [1, 2],
        [1, 3],
        [2, 4],
        [3, 4],
    ]

    # make R, t into numpy
    if type(cam_R) != np.ndarray:
        cam_R = np.array(cam_R, dtype=np.float32)
    elif type(cam_R) == torch.Tensor:
        cam_R = cam_R.detach().cpu().numpy()
    if type(cam_t) != np.ndarray:
        cam_t = np.array(cam_t, dtype=np.float32)
    elif type(cam_t) == torch.Tensor:
        cam_t = cam_t.detach().cpu().numpy()
    
    # rotate and translate via camera transformation
    final_points = []
    for point in points:
        final_point = np.matmul(cam_R.T, point - cam_t)
        final_point = final_point.tolist()
        final_points.append(final_point)

    # camera frustum
    camera_frustum = o3d.geometry.LineSet(
        points=o3d.utility.Vector3dVector(final_points), # as float 3d-vector.
        lines=o3d.utility.Vector2iVector(lines) # as int 2d-vector.
    )
    camera_frustum.colors = o3d.utility.Vector3dVector([color]*8)

    return camera_frustum


# create open3d coordinate frame (centered at world origin)
def return_coordinate_frame(size=1.0):
    # --> centered to origin (which is camera center by default)
    return o3d.geometry.TriangleMesh.create_coordinate_frame(size=size) 
visualization/visualize_open3d.py
jellyheadandrew/CHORUS
import os
from glob import glob
import shutil
import cv2


def generate_video_from_imgs(images_save_directory, fps=15.0, delete_dir=True):
    # delete videos if exists
    if os.path.exists(f"{images_save_directory}.mp4"): os.remove(f"{images_save_directory}.mp4")
    if os.path.exists(f"{images_save_directory}_before_process.mp4"): os.remove(f"{images_save_directory}_before_process.mp4")
    
    # assume there are "enumerated" images under "images_save_directory"
    assert os.path.isdir(images_save_directory)
    ImgPaths = sorted(list(glob(f"{images_save_directory}/*")))
    
    if len(ImgPaths) == 0:
        print("\tSkipping, since there must be at least one image to create mp4\n")
    else:
        # mp4 configuration
        video_path = images_save_directory + "_before_process.mp4"

        # Get height and width config
        images = sorted([ImgPath.split("/")[-1] for ImgPath in ImgPaths if ImgPath.endswith(".png")])
        frame = cv2.imread(os.path.join(images_save_directory, images[0]))
        height, width, channels = frame.shape

        # create mp4 video writer
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        video = cv2.VideoWriter(video_path, fourcc, fps, (width,height))
        for image in images:
            video.write(cv2.imread(os.path.join(images_save_directory, image)))
        cv2.destroyAllWindows()
        video.release()

        # generated video is not compatible with HTML5. Post-process and change codec of video, so that it is applicable to HTML.
        os.system(f"ffmpeg -i \"{images_save_directory}_before_process.mp4\" -vcodec libx264 -f mp4 \"{images_save_directory}.mp4\" -loglevel \"quiet\"")

    # remove group of images, and remove video before post-process.
    if delete_dir and os.path.exists(images_save_directory): shutil.rmtree(images_save_directory)
    # remove 'before-process' video
    if os.path.exists(f"{images_save_directory}_before_process.mp4"): os.remove(f"{images_save_directory}_before_process.mp4")
visualization/visualize_video.py
jellyheadandrew/CHORUS