This is a repository for download, preprocessing, visualizing, running evaluations on the MOYO dataset.
# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2021 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
"""
This script is used to close part meshes.
"""
from tempfile import gettempdir

import torch
import trimesh
import torch.nn as nn
import numpy as np
import pickle
# from .utils.mesh import winding_numbers

class SMPLXMesh(nn.Module):
    def __init__(self, vertices, faces):
        super(SMPLXMesh, self).__init__()

        self.vertices = vertices
        self.faces = faces

class PartVolume(nn.Module):
    def __init__(self,
                 part_name,
                 vertices,
                 faces):
        super(PartVolume, self).__init__()

        self.part_name = part_name
        self.smplx_mesh = SMPLXMesh(vertices, faces)

        self.part_triangles = None
        self.device = vertices.device

        self.new_vert_ids = []
        self.new_face_ids = []

    def close_mesh(self, boundary_vids):
        # find the center of the boundary
        mean_vert = self.smplx_mesh.vertices[:, boundary_vids, :].mean(dim=1, keepdim=True)
        self.smplx_mesh.vertices = torch.cat([self.smplx_mesh.vertices, mean_vert], dim=1)
        new_vert_idx = self.smplx_mesh.vertices.shape[1]-1
        self.new_vert_ids.append(new_vert_idx)
        # add faces
        new_faces = [[boundary_vids[i + 1], boundary_vids[i], new_vert_idx] for i in range(len(boundary_vids) - 1)]
        self.smplx_mesh.faces = torch.cat([self.smplx_mesh.faces, torch.tensor(new_faces, dtype=torch.long, device=self.device)], dim=0)
        self.new_face_ids += list(range(self.smplx_mesh.faces.shape[0]-len(new_faces), self.smplx_mesh.faces.shape[0]))

    def extract_part_triangles(self, part_vids, part_fid):
        # make sure that inputs are from a watertight part mesh
        batch_size = self.smplx_mesh.vertices.shape[0]

        part_vertices = self.smplx_mesh.vertices[:, part_vids, :]
        part_faces = self.smplx_mesh.faces[part_fid, :]

        part_mean = part_vertices.mean(dim=1, keepdim=True)

        # subtract vert mean because volume computation only applies if origin is inside the triangles
        self.smplx_mesh.vertices = self.smplx_mesh.vertices - part_mean

        # compute triangle
        if self.part_triangles is None:
            # self.part_triangles = torch.index_select(self.smplx_mesh.vertices, 1, self.smplx_mesh.faces.view(-1)).reshape(batch_size, -1, 3, 3)
            self.part_triangles = torch.index_select(self.smplx_mesh.vertices, 1, part_faces.view(-1)).reshape(batch_size, -1, 3, 3)
        else:
            self.part_triangles = torch.cat([self.part_triangles,
                                             torch.index_select(self.smplx_mesh.vertices, 1,
                                                     part_faces.view(-1)).reshape(batch_size, -1, 3, 3)], dim=1)
        # add back vert mean
        self.smplx_mesh.vertices = self.smplx_mesh.vertices + part_mean

    def part_volume(self):
        # Note: the mesh should be enclosing the origin (mean-subtracted)
        # compute volume of triangles by drawing tetrahedrons
        # https://stackoverflow.com/questions/1406029/how-to-calculate-the-volume-of-a-3d-mesh-object-the-surface-of-which-is-made-up
        x = self.part_triangles[:, :, :, 0]
        y = self.part_triangles[:, :, :, 1]
        z = self.part_triangles[:, :, :, 2]
        volume = (
                         -x[:, :, 2] * y[:, :, 1] * z[:, :, 0] +
                         x[:, :, 1] * y[:, :, 2] * z[:, :, 0] +
                         x[:, :, 2] * y[:, :, 0] * z[:, :, 1] -
                         x[:, :, 0] * y[:, :, 2] * z[:, :, 1] -
                         x[:, :, 1] * y[:, :, 0] * z[:, :, 2] +
                         x[:, :, 0] * y[:, :, 1] * z[:, :, 2]
                 ).sum(dim=1).abs() / 6.0
        return volume

class BodySegment(nn.Module):
    def __init__(self,
                 name,
                 faces,
                 segments_folder,
                 model_type='smplx',
                 append_idx=None):
        super(BodySegment, self).__init__()

        self.name = name
        self.append_idx = faces.max().item() if append_idx is None \
            else append_idx

        self.model_type = model_type
        sb_path = f'{segments_folder}/{model_type}_segments_bounds.pkl'
        sxseg = pickle.load(open(sb_path, 'rb'))

        # read mesh and find faces of segment
        segment_path = f'{segments_folder}/{model_type}_segment_{name}.ply'
        bandmesh = trimesh.load(segment_path, process=False)
        segment_vidx = torch.from_numpy(np.where(
            np.array(bandmesh.visual.vertex_colors[:,0]) == 255)[0])
        self.register_buffer('segment_vidx', segment_vidx)

        # read boundary information
        self.bands = [x for x in sxseg[name].keys()]
        self.bands_verts = [x for x in sxseg[name].values()]
        self.num_bounds = len(self.bands_verts)
        for idx, bv in enumerate(self.bands_verts):
            self.register_buffer(f'bands_verts_{idx}', torch.tensor(bv))
        self.bands_faces = self.create_band_faces()

        # read mesh and find
        segment_faces_ids = np.where(np.isin(faces.cpu().numpy(),
            segment_vidx).sum(1) == 3)[0]
        segment_faces = faces[segment_faces_ids,:]
        segment_faces = torch.cat((faces[segment_faces_ids,:],
            self.bands_faces), 0)
        self.register_buffer('segment_faces', segment_faces)

        # create vector to select vertices form faces
        tri_vidx = []
        for ii in range(faces.max().item()+1):
            tri_vidx += [torch.nonzero(faces==ii)[0].tolist()]
        self.register_buffer('tri_vidx', torch.tensor(tri_vidx))

    def create_band_faces(self):
        """
            create the faces that close the segment.
        """
        bands_faces = []
        for idx, k in enumerate(self.bands):
            new_vert_idx = self.append_idx + 1 + idx
            new_faces = [[self.bands_verts[idx][i+1], \
                self.bands_verts[idx][i], new_vert_idx] \
                for i in range(len(self.bands_verts[idx])-1)]
            bands_faces += new_faces

        bands_faces_tensor = torch.tensor(
            np.array(bands_faces).astype(np.int64), dtype=torch.long)

        return bands_faces_tensor

    def get_closed_segment(self, vertices):
        """
            create the closed segment mesh from SMPL-X vertices.
        """
        vertices = vertices.detach().clone()
        # append vertices to SMPLX, that close the segment and compute faces
        for idx in range(self.num_bounds):
            bv = eval(f'self.bands_verts_{idx}')
            close_segment_vertices = torch.mean(vertices[:, bv,:], 1,
                                    keepdim=True)
            vertices = torch.cat((vertices, close_segment_vertices), 1)
        segm_triangles = vertices[:, self.segment_faces, :]

        return segm_triangles

    # def has_self_isect_verts(self, vertices, thres=0.99):
    #     """
    #         check if segment (its vertices) are self intersecting.
    #     """
    #     segm_triangles = self.get_closed_segment(vertices)
    #     segm_verts = vertices[:,self.segment_vidx,:]
    #
    #     # do inside outside segmentation
    #     exterior = winding_numbers(segm_verts, segm_triangles) \
    #                 .le(thres)
    #
    #     return exterior
    #
    # def has_self_isect_points(self, points, triangles, thres=0.99):
    #     """
    #         check if points on segment are self intersecting.
    #     """
    #     smplx_verts = triangles[:,self.tri_vidx[:,0], self.tri_vidx[:,1],:]
    #     segm_triangles = self.get_closed_segment(smplx_verts)
    #
    #     # do inside outside segmentation
    #     exterior = winding_numbers(points, segm_triangles) \
    #                 .le(thres)
    #
    #     return exterior

class BatchBodySegment(nn.Module):
    def __init__(self,
                 names,
                 faces,
                 segments_folder,
                 model_type='smplx',
                 device='cuda'):
        super(BatchBodySegment, self).__init__()
        self.names = names
        self.num_segments = len(names)
        self.nv = faces.max().item()

        self.model_type = model_type
        sb_path = f'{segments_folder}/{model_type}_segments_bounds.pkl'
        sxseg = pickle.load(open(sb_path, 'rb'))

        self.append_idx = [len(b) for a,b in sxseg.items() \
            for c,d in b.items() if a in self.names]
        self.append_idx = np.cumsum(np.array([self.nv] + self.append_idx))

        self.segmentation = {}
        for idx, name in enumerate(names):
            self.segmentation[name] = BodySegment(name, faces, segments_folder,
                model_type).to('cuda')

    def batch_has_self_isec_verts(self, vertices):
        """
            check is mesh is intersecting with itself
        """
        exteriors = []
        for k, segm in self.segmentation.items():
            exteriors += [segm.has_self_isect_verts(vertices)]
        return exteriors

moyo/core/part_volumes.py
sha2nkt/moyo_toolkit
# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2021 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
import sys

import torch
import torch.nn as nn
import pickle as pkl
from pytorch3d.structures import Meshes
import numpy as np
from moyo.utils.mesh_utils import GMoF_unscaled, batch_face_normals, HDfier
from moyo.utils.constants import CONTACT_THRESH
from moyo.core.part_volumes import PartVolume
from pathlib import Path

SMPLX_PART_BOUNDS = '../data/essentials/yogi_segments/smplx/part_meshes_ply/smplx_segments_bounds.pkl'
FID_TO_PART = '../data/essentials/yogi_segments/smplx/part_meshes_ply/fid_to_part.pkl'
PART_VID_FID = '../data/essentials/yogi_segments/smplx/part_meshes_ply/smplx_part_vid_fid.pkl'
HD_SMPLX_MAP  = '../data/essentials/hd_model/smplx/smplx_neutral_hd_sample_from_mesh_out.pkl'

class StabilityLossCoS(nn.Module):
    def __init__(self,
                 faces,
                 cos_w = 10,
                 cos_k = 100,
                 contact_thresh=CONTACT_THRESH,
                 model_type='smplx',
    ):
        super().__init__()
        """
        Loss that ensures that the COM of the SMPL mesh is close to the center of support 
        """
        if model_type == 'smplx':
            num_faces = 20908
        if model_type == 'smpl':
            num_faces = 13776
        num_verts_hd = 20000


        assert faces is not None, 'Faces tensor is none'
        if type(faces) is not torch.Tensor:
            faces = torch.tensor(faces.astype(np.int64), dtype=torch.long)
        self.register_buffer('faces', faces)

        # self.gmof_rho = gmof_rho
        self.cos_w = cos_w
        self.cos_k = cos_k
        self.contact_thresh = contact_thresh

        self.hdfy_op = HDfier()

        with open(SMPLX_PART_BOUNDS, 'rb') as f:
            d = pkl.load(f)
            self.part_bounds = {k: d[k] for k in sorted(d)}
        self.part_order = sorted(self.part_bounds)

        with open(PART_VID_FID, 'rb') as f:
            self.part_vid_fid = pkl.load(f)

        # mapping between vid_hd and fid
        with open(HD_SMPLX_MAP, 'rb') as f:
            faces_vert_is_sampled_from = pkl.load(f)['faces_vert_is_sampled_from']
        index_row_col = torch.stack(
            [torch.LongTensor(np.arange(0, num_verts_hd)), torch.LongTensor(faces_vert_is_sampled_from)], dim=0)
        values = torch.ones(num_verts_hd, dtype=torch.float)
        size = torch.Size([num_verts_hd, num_faces])
        hd_vert_on_fid = torch.sparse.FloatTensor(index_row_col, values, size)

        # mapping between fid and part label
        with open(FID_TO_PART, 'rb') as f:
            fid_to_part_dict = pkl.load(f)
        fid_to_part = torch.zeros([len(fid_to_part_dict.keys()), len(self.part_order)], dtype=torch.float32)
        for fid, partname in fid_to_part_dict.items():
            part_idx = self.part_order.index(partname)
            fid_to_part[fid, part_idx] = 1.

        # mapping between vid_hd and part label
        self.hd_vid_in_part = self.vertex_id_to_part_mapping(hd_vert_on_fid, fid_to_part)

    def compute_triangle_area(self, triangles):
        ### Compute the area of each triangle in the mesh
        # Compute the cross product of the two vectors of each triangle
        # Then compute the length of the cross product
        # Finally, divide by 2 to get the area of each triangle

        vectors = torch.diff(triangles, dim=2)
        crosses = torch.cross(vectors[:, :, 0], vectors[:, :, 1])
        area = torch.norm(crosses, dim=2) / 2
        return area

    def compute_per_part_volume(self, vertices):
        """
        Compute the volume of each part in the reposed mesh
        """
        part_volume = []
        for part_name, part_bounds in self.part_bounds.items():
            # get part vid and fid
            part_vid = torch.LongTensor(self.part_vid_fid[part_name]['vert_id']).to(vertices.device)
            part_fid = torch.LongTensor(self.part_vid_fid[part_name]['face_id']).to(vertices.device)
            pv = PartVolume(part_name, vertices, self.faces)
            for bound_name, bound_vids in part_bounds.items():
                pv.close_mesh(bound_vids)
            # add extra vids and fids to original part ids
            new_vert_ids = torch.LongTensor(pv.new_vert_ids).to(vertices.device)
            new_face_ids = torch.LongTensor(pv.new_face_ids).to(vertices.device)
            part_vid = torch.cat((part_vid, new_vert_ids), dim=0)
            part_fid = torch.cat((part_fid, new_face_ids), dim=0)
            pv.extract_part_triangles(part_vid, part_fid)
            part_volume.append(pv.part_volume())
        return torch.vstack(part_volume).permute(1,0).to(vertices.device)

    def vertex_id_to_part_volume_mapping(self, per_part_volume, device):
        batch_size = per_part_volume.shape[0]
        self.hd_vid_in_part = self.hd_vid_in_part.to(device)
        hd_vid_in_part = self.hd_vid_in_part[None, :, :].repeat(batch_size, 1, 1)
        vid_to_vol = torch.bmm(hd_vid_in_part, per_part_volume[:, :, None])
        return vid_to_vol

    def vertex_id_to_part_mapping(self, hd_vert_on_fid, fid_to_part):
        vid_to_part = torch.mm(hd_vert_on_fid, fid_to_part)
        return vid_to_part

    def forward(self, vertices):
        # Note: the vertices should be aligned along y-axis and in world coordinates
        batch_size = vertices.shape[0]
        # calculate per part volume
        per_part_volume = self.compute_per_part_volume(vertices)
        # sample 20k vertices uniformly on the smpl mesh
        vertices_hd = self.hdfy_op.hdfy_mesh(vertices)
        # get volume per vertex id in the hd mesh
        volume_per_vert_hd = self.vertex_id_to_part_volume_mapping(per_part_volume, vertices.device)
        # calculate com using volume weighted mean
        com = torch.sum(vertices_hd * volume_per_vert_hd, dim=1) / torch.sum(volume_per_vert_hd, dim=1)

        # # get COM of the SMPLX mesh
        # triangles = torch.index_select(vertices, 1, self.faces.view(-1)).reshape(batch_size, -1, 3, 3)
        # triangle_centroids = torch.mean(triangles, dim=2)
        # triangle_area = self.compute_triangle_area(triangles)
        # com_naive = torch.einsum('bij,bi->bj', triangle_centroids, triangle_area) / torch.sum(triangle_area, dim=1)

        # pressure based center of support
        ground_plane_height = 0.0
        eps = 1e-6
        vertex_height = (vertices_hd[:, :, 1] - ground_plane_height)
        inside_mask = (vertex_height < 0.0).float()
        outside_mask = (vertex_height >= 0.0).float()
        pressure_weights = inside_mask * (1-self.cos_k*vertex_height) + outside_mask *  torch.exp(-self.cos_w * vertex_height)
        cos = torch.sum(vertices_hd * pressure_weights.unsqueeze(-1), dim=1) / (torch.sum(pressure_weights, dim=1, keepdim=True) +eps)

        # naive center of support
        # vertex_height_robustified = GMoF_unscaled(rho=self.gmof_rho)(vertex_height)
        contact_confidence = torch.sum(pressure_weights, dim=1)
        # contact_mask = (vertex_height < self.contact_thresh).float()
        # num_contact_verts = torch.sum(contact_mask, dim=1)
        # contact_centroid_naive = torch.sum(vertices_hd * contact_mask[:, :, None], dim=1) / (torch.sum(contact_mask, dim=1) + eps)

        # project com, cos to ground plane (x-z plane)
        # weight loss by number of contact vertices to zero out if zero vertices in contact
        com_xz = torch.stack([com[:, 0], torch.zeros_like(com)[:, 0], com[:, 2]], dim=1)
        contact_centroid_xz = torch.stack([cos[:, 0], torch.zeros_like(cos)[:, 0], cos[:, 2]], dim=1)
        # stability_loss = (contact_confidence * torch.norm(com_xz - contact_centroid_xz, dim=1)).sum(dim=-1)
        stability_loss = (torch.norm(com_xz - contact_centroid_xz, dim=1))
        return stability_loss


class StabilityLossPoS(nn.Module):
    def __init__(self,
                 faces,
                 gmof_rho=0.02,
                 contact_thresh=CONTACT_THRESH,
                 use_hd=False,
    ):
        super().__init__()
        """
        Loss that ensures that the COM of the SMPL mesh is close to the plane of support 
        """
        assert faces is not None, 'Faces tensor is none'
        if type(faces) is not torch.Tensor:
            faces = torch.tensor(faces.astype(np.int64), dtype=torch.long)
        self.register_buffer('faces', faces)

        # self.gmof_rho = gmof_rho
        self.contact_thresh = contact_thresh

    def compute_triangle_area(self, triangles):
        ### Compute the area of each triangle in the mesh
        # Compute the cross product of the two vectors of each triangle
        # Then compute the length of the cross product
        # Finally, divide by 2 to get the area of each triangle

        vectors = torch.diff(triangles, dim=2)
        crosses = torch.cross(vectors[:, :, 0], vectors[:, :, 1])
        area = torch.norm(crosses, dim=2) / 2
        return area

    def forward(self, vertices):

        batch_size = vertices.shape[0]

        # get vertices in contact
        ground_plane_height = 0.0  # obtained by visualization on the presented pose
        vertex_height = (vertices[:, :, 1] - ground_plane_height)
        # vertex_height_robustified = GMoF_unscaled(rho=self.gmof_rho)(vertex_height)
        contact_mask = (vertex_height_robustified < self.contact_thresh).float()

        # get plane of support
        # ToDo: fit to the contact vertices to minimize distance instead
        contact_centroid = torch.sum(vertices * contact_mask[:, :, None], dim=1) / torch.sum(contact_mask, dim=1)
        support_plane_offset = contact_centroid[:, 0]

        # get COM of the SMPLX mesh
        triangles = torch.index_select(vertices, 1, self.faces.view(-1)).reshape(batch_size, -1, 3, 3)
        triangle_centroids = torch.mean(triangles, dim=2)
        triangle_area = self.compute_triangle_area(triangles)
        com = torch.einsum('bij,bi->bj', triangle_centroids, triangle_area) / torch.sum(triangle_area, dim=1)

        # get distance of COM to plane of support
        stability_loss = torch.abs(com[:, 0] - support_plane_offset).sum(dim=-1)
        return stability_loss

class GroundLoss(nn.Module):
    def __init__(self,
                 faces,
                 device='cuda',
                 model_type='smplx',
                 out_alpha1=1,
                 out_alpha2=0.5,
                 in_alpha1=1,
                 in_alpha2=0.15,
    ):
        super().__init__()
        assert faces is not None, 'Faces tensor is none'
        if type(faces) is not torch.Tensor:
            faces = torch.tensor(faces.astype(np.int64), dtype=torch.long)
        self.register_buffer('faces', faces)

        self.device = device
        self.model_type = model_type

        # loss specifications

        # push, pull, and contact weights
        # self.inside_w = inside_loss_weight
        # self.outside_w = outside_loss_weight

        # hyper params
        self.a1 = out_alpha1
        self.a2 = out_alpha2
        self.b1 = in_alpha1
        self.b2 = in_alpha2

        self.hdfy_op = HDfier()



    def forward(self, vertices):
        """
                Loss that ensures that the ground plane pulls body mesh vertices towards it till contact and resolves
                ground-plane intersections
                Using assymetric loss such that the
                - loss is 0 when the body is in contact with the ground plane
                - loss is >> 0 when the plane interpenetrates the body (completely implausible)
                - loss >0 if the plane is close to nearby vertices that are not touching
                vertices:the vertices should be aligned along y-axis and in world coordinates
        """
        bs = vertices.shape[0]
        vertices_hd = self.hdfy_op.hdfy_mesh(vertices)
        # get vertices under the ground plane
        ground_plane_height = 0.0  # obtained by visualization on the presented pose
        vertex_height = (vertices_hd[:, :, 1] - ground_plane_height)
        inside_mask = vertex_height < 0.00
        outside_mask = vertex_height >= 0.00
        # pull closeby outside vertices
        v2v_pull = (self.a1 * torch.tanh((vertex_height * outside_mask) / self.a2) ** 2)
        # # apply loss to inside vertices to remove intersection
        v2v_push = (self.b1 * torch.tanh((vertex_height * inside_mask) / self.b2)**2)
        return v2v_pull, v2v_push

class SupportRegLoss(nn.Module):
    def __init__(self,
                 faces,
                 inside_loss_weight=1.0,
                 outside_loss_weight=1.0,
                 contact_loss_weight=2.5,
                 align_faces=False,
                 gmof_rho=0.02,
                 use_hd=True,
                 device='cuda',
                 model_type='smplx',
                 alpha1=0.005,
                 alpha2=0.005,
                 beta1=600,
    ):
        super().__init__()
        assert faces is not None, 'Faces tensor is none'
        if type(faces) is not torch.Tensor:
            faces = torch.tensor(faces.astype(np.int64), dtype=torch.long)
        self.register_buffer('faces', faces)

        self.device = device
        self.model_type = model_type

        # loss specifications
        self.use_hd = use_hd
        self.align_faces = align_faces
        # self.gmof_rho = gmof_rho

        # push, pull, and contact weights
        self.inside_w = inside_loss_weight
        self.contact_w = contact_loss_weight

        # hyper params
        self.a1 = alpha1
        self.a2 = alpha2
        self.b1 = beta1



    def forward(self, vertices):
        """
                Loss that ensures that the ground plane pulls body mesh vertices towards it till contact and resolves
                ground-plane intersections
                Using assymetric loss such that the
                - loss is 0 when the body is in contact with the ground plane
                - loss is >> 0 when the plane interpenetrates the body (completely implausible)
                - loss >0 if the plane is close to nearby vertices that are not touching
        """
        bs = vertices.shape[0]

        import ipdb; ipdb.set_trace()
        # get vertices under the ground plane
        ground_plane_height = 0.0  # obtained by visualization on the presented pose
        vertex_height = (vertices[:, :, 1] - ground_plane_height)
        inside_mask = (vertex_height < 0.00).float()

        # pull closeby outside vertices
        v2v_pull = self.contact_w * (self.a1 * torch.tanh(vertex_height * [~inside_mask] / self.a2)**2).sum()
        # # apply loss to inside vertices to remove intersection
        v2v_push = self.inside_w * (self.b1 * (vertex_height * inside_mask)**2).sum()
        contactloss = v2v_push + v2v_pull

        # now align faces that are close
        # dot product should be -1 for faces in contact
        if self.align_faces:
            triangles = torch.index_select(vertices, 1, self.faces.view(-1)).reshape(bs, -1, 3, 3)
            face_normals = batch_face_normals(triangles)

        face_angle_loss = torch.zeros(bs, device=vertices.device)
        # Todo: Implement normal angle contact loss
        if self.align_faces:
            if self.use_hd:
                if hd_faces_in_contact[idx] is not None:
                    hd_fn_in_contact = face_normals[idx][hd_faces_in_contact[idx]]
                    dotprod_normals = 1 + (hd_fn_in_contact[0] * hd_fn_in_contact[1]).sum(1)
                    face_angle_loss[idx] = dotprod_normals.sum()
            else:
                sys.exit('You can only align vertices when use_hd=True')

        return contactloss, face_angle_loss
moyo/core/support.py
sha2nkt/moyo_toolkit
"""
Usage python eval/com_evaluation.py --img_folder ../data/moyo/20221004_with_com/images/val/ --pp_folder ../data/moyo//20221004_with_com/mosh/val/ --nexus_com_c3d_folder ../data/moyo//20221004_with_com/com/val

"""

import argparse
import copy
import glob
import json
import os
import os.path as osp
from pathlib import Path

import cv2
import numpy as np
import torch
import trimesh
import c3d

from tqdm import tqdm
import pickle as pkl

from moyo.utils.misc import colors, copy2cpu as c2c

import smplx
from moyo.utils.geometry import matrix_to_axis_angle, ea2rm, axis_angle_to_matrix
from moyo.utils.constants import frame_select_dict_nexus as frame_select_dict
from moyo.utils.mesh_utils import smplx_breakdown
from moyo.utils.biomech_eval_utils import BiomechanicalEvaluator

VICON_FPS = 60
PRESSURE_MAP_FPS = 60
IOI_FPS = 30


def read_keypoints(keypoint_fn):
    with open(keypoint_fn) as keypoint_file:
        data = json.load(keypoint_file)

    keypoints = []
    for idx, person_data in enumerate(data['people']):
        body_keypoints = np.array(
            person_data['pose_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])

        left_hand_keyp = np.array(
            person_data['hand_left_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])

        right_hand_keyp = np.array(
            person_data['hand_right_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])

        face_keypoints = np.array(
            person_data['face_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])[17: 17 + 51, :]

        contour_keyps = np.array(
            person_data['face_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])[:17, :]

        body_keypoints = np.concatenate([body_keypoints, left_hand_keyp,
                                         right_hand_keyp, face_keypoints, contour_keyps], axis=0)

        keypoints.append(body_keypoints)
    return keypoints


def select_keypoints(keypoints, img):
    img_center = np.array(img.shape[:2]) / 2  # height, width
    # select keypoints closest to image center weighted by inverse confidence
    if len(keypoints) > 1:
        kpselect = np.inf * np.ones(len(keypoints))
        # major joints openpose
        op_to_12 = [9, 10, 11, 12, 13, 14, 2, 3, 4, 5, 6, 7]
        for idx, personkpts in enumerate(keypoints):
            kpdist = personkpts[op_to_12, :2] - img_center
            kpdist = np.linalg.norm(kpdist, axis=1)
            kpconf = np.dot(kpdist, (- personkpts[op_to_12, 2] + 1))
            kpselect[idx] = kpconf
        kpselidx = np.argmin(kpselect)
    elif len(keypoints) == 1:
        kpselidx = 0
    else:
        keypoints = None
    keypoints = np.stack(keypoints[kpselidx:kpselidx + 1])
    return keypoints


def visualize_mesh(bm_output, faces, frame_id=0, display=False):
    imw, imh = 1600, 1600

    body_mesh = trimesh.Trimesh(vertices=c2c(bm_output.vertices[frame_id]),
                                faces=faces,
                                vertex_colors=np.tile(colors['grey'], (6890, 1)),
                                process=False,
                                maintain_order=True)
    return body_mesh


def get_trans_offset(pelvis, smplx_params, trans, body_model):
    '''
    SOMA fits are in AMASS format (xy plane is ground) but SMPLify-XMC uses Pyrender format (xz plane is ground). global_orient is corrected by 270 rotation along x-axis to match the two formats.
    While the ground plane is z=0 in soma, since global_orient is wrt pelvis, this would mean the ground plane is not at y=0 after rotation.
    Fix: distance of pelvis to ground in preserved. Use it to push the mesh up after rotation such that the ground plane is at y=0.
    Args:
        pelvis: pelvis joint before rotation
    '''
    bs = trans.shape[0]
    init_root_orient = smplx_params['global_orient']
    pelvis_height = pelvis[:, 2]  # this is always conserved

    new_smplx_params = copy.deepcopy(smplx_params)
    # Rotate AMASS root orient to smplify-xmc format
    R_init = axis_angle_to_matrix(init_root_orient)
    R1 = ea2rm(torch.tensor([[np.radians(270)]]), torch.tensor([[np.radians(0)]]),
               torch.tensor([[np.radians(0)]])).float().to(R_init.device)
    R = torch.bmm(R1.expand(bs, -1, -1), R_init)
    new_smplx_params['global_orient'] = matrix_to_axis_angle(R)

    # posed body with hand, with global orient
    body_model_output = body_model(
        global_orient=new_smplx_params['global_orient'],
        body_pose=new_smplx_params['body_pose'])

    new_pelvis = body_model_output.joints[:, 0]
    new_ground_plane_height = new_pelvis[:, 1] - pelvis_height
    trans_offset = -new_ground_plane_height
    return trans_offset


def smplx_to_mesh(body_params, model_folder, model_type, gender='neutral'):
    with torch.no_grad():
        device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
        num_betas = 10

        smplx_params = smplx_breakdown(body_params, device)

        trans = torch.from_numpy(body_params['trans']).float().to(device)

        betas = torch.from_numpy(body_params['betas']).float().to(device).unsqueeze(0)
        betas = betas[:, :num_betas]  # only using 10 betas

        body_model_params = dict(model_path=model_folder,
                                 model_type=model_type,
                                 gender=gender,
                                 v_template=smplx_params['v_template'],
                                 # joint_mapper=joint_mapper,
                                 batch_size=trans.shape[0],
                                 create_global_orient=True,
                                 create_body_pose=True,
                                 create_betas=True,
                                 num_betas=num_betas,
                                 create_left_hand_pose=True,
                                 create_right_hand_pose=True,
                                 create_expression=True,
                                 create_jaw_pose=True,
                                 create_leye_pose=True,
                                 create_reye_pose=True,
                                 create_transl=True,
                                 use_pca=False,
                                 flat_hand_mean=True,
                                 dtype=torch.float32)

        body_model = smplx.create(**body_model_params).to(device)
        body_model_output = body_model(transl=trans,
                                       global_orient=smplx_params['global_orient'],
                                       body_pose=smplx_params['body_pose'],
                                       left_hand_pose=smplx_params['left_hand_pose'],
                                       right_hand_pose=smplx_params['right_hand_pose'])

        pelvis = body_model_output.joints[:, 0]
        trans_offset = get_trans_offset(pelvis, smplx_params, trans, body_model)
        zero_vec = torch.zeros_like(trans_offset)
        transl = torch.stack([zero_vec, trans_offset, zero_vec], dim=-1)

    return body_model_output, smplx_params, transl, body_model.faces

def extract_com_from_c3d(c3d_file):
    with open(c3d_file, 'rb') as f:
        data = c3d.Reader(f)  # data contains the header of the c3d files, info such as the label name
        com_idx = np.where(data.point_labels == 'CentreOfMass                  ')
        com_idx_floor = np.where(data.point_labels == 'CentreOfMassFloor             ')
        coms, coms_floor = [], []
        for frame in data.read_frames():  # call read_frames() to iterate across the marker positions
            frame = frame[1][:, :3]  # this contains the marker positions
            coms.append(frame[com_idx, :])
            coms_floor.append(frame[com_idx_floor, :])

    coms = np.concatenate(coms, axis=0)
    coms_floor = np.concatenate(coms_floor, axis=0)
    return coms, coms_floor



def yoga82_extra(img_folder, pp_folder, nexus_com_c3d_folder, model_folder,
                 model_type='smplx', subsampling_factor=1, save_outputs=False):
    imgpaths_ = []
    imgnames_ = []
    subject_ids_ = []
    relpaths_ = []
    heights_ = []
    weights_ = []
    genders_ = []
    gt_coms_ = []
    pred_coms_ = []
    com_errors_ = []
    pp_body_poses_ = []
    pp_left_hand_poses_ = []
    pp_right_hand_poses_ = []
    pp_global_orients_ = []
    pp_transls_ = []

    # presented poses
    pp_names = os.listdir(pp_folder)
    pp_names = [os.path.splitext(x)[0] for x in pp_names if x[0] != '.' and '.pkl' in x]

    all_pose_names = [pp_name.replace('_stageii', '') for pp_name in pp_names]
    all_pose_names = sorted(all_pose_names)

    pp_pose_cache = {}

    no_images_found = set()
    no_pp_found = set()

    smplx_faces = smplx.create(model_path=model_folder, model_type=model_type).faces
    smplx_faces = torch.tensor(smplx_faces.astype(np.int64), dtype=torch.long)
    com_evaluator = BiomechanicalEvaluator(faces=smplx_faces)

    for i, pname in enumerate(tqdm(sorted(all_pose_names))):
        if 'female_stagei' in pname: # do not analyze female stage1 file
            continue
        # get all possible variations
        pp_pkl = osp.join(pp_folder, pname + '_stageii.pkl')
        pp_params = pkl.load(open(pp_pkl, 'rb'))
        # read height, weight, gender meta
        height = None
        weight = None
        gender = 'female'
        subj_id = 'unknown'
        pp_body_model_output, smplx_params, transls, faces = smplx_to_mesh(pp_params, model_folder,
                                                                           model_type, gender=gender)
        # get pressure xml
        com_c3d = osp.join(nexus_com_c3d_folder, pname + '.c3d')
        gt_coms, gt_coms_floor = extract_com_from_c3d(com_c3d)

        # pick selected frame
        selected_frame = frame_select_dict['_'.join(pname.split('_')[6:])]

        # get all images for this pose
        img_dir = Path(osp.join(img_folder, pname))
        for cam in range(1, 7):
            cam_name = Path('YOGI_Cam_' + str(cam).zfill(2))
            img_dir_cam = Path.joinpath(img_dir, cam_name)
            types = ('*.jpg', '*.png')  # the tuple of file types
            images = []
            for files in types:
                images.extend(glob.glob(osp.join(img_dir_cam, files)))
            if len(images) == 0:
                print(f'No images found for {pname}')
                no_images_found.add(pname)

            # subsample frames
            images.sort()

            # folders to log outputs
            if cam == 1 and save_outputs:
                out_img_dir = osp.join(pp_folder, 'images_small', pname, cam_name)
                os.makedirs(out_img_dir, exist_ok=True)
                out_obj_dir = osp.join(pp_folder, 'meshes', pname, cam_name)
                os.makedirs(out_obj_dir, exist_ok=True)

            for img_path in images[
                            :-1:subsampling_factor]:  # sometimes, the last frame is missing in pressure or vicon data
                # get frame correspondences
                img_fnum = int(img_path.split('_')[-1].split('.')[0])
                vicon_fnum = img_fnum * (VICON_FPS // IOI_FPS)

                # only run on selected frames
                if img_fnum != selected_frame:
                    continue

                # load image
                img_fn, _ = osp.splitext(osp.split(img_path)[1])
                relpath = osp.dirname(img_path).replace(str(img_dir), '').strip('/')

                if cam == 1 and save_outputs:
                    img = cv2.imread(img_path).astype(np.float32)
                    # resize image to half size
                    img = cv2.resize(img, (img.shape[1] // 2, img.shape[0] // 2))
                    cv2.imwrite(osp.join(out_img_dir, f'{pname}_{vicon_fnum}.jpg'), img)
                    print('[[Image Saved]]\n', osp.join(out_img_dir, f'{pname}_{vicon_fnum}.jpg'))

                transl = transls[vicon_fnum]
                pp_mesh = visualize_mesh(pp_body_model_output, faces, frame_id=vicon_fnum)
                pp_verts = pp_mesh.vertices
                # pp_verts = pp_body_model_output.vertices[vicon_fnum]
                vertices = torch.tensor(pp_verts).unsqueeze(0)
                if cam == 1 and save_outputs:
                    out_obj = osp.join(out_obj_dir, f'{pname}_{vicon_fnum}.obj')
                    pp_mesh.export(out_obj)
                    print('[[Mesh]]\n', out_obj)

                # get com for current frame
                gt_com = gt_coms[vicon_fnum, :]
                gt_com = gt_com/1000.0 # convert to meters
                print('[[GT COM]]\n', gt_com)

                # get com_error
                pred_com, com_error = com_evaluator.evaluate_com(gt_com, vertices)
                pred_com = pred_com.numpy()
                com_error = com_error.squeeze().numpy()

                body_pose = smplx_params['body_pose'].squeeze().cpu()
                left_hand_pose = smplx_params['left_hand_pose'].squeeze().cpu()
                right_hand_pose = smplx_params['right_hand_pose'].squeeze().cpu()
                global_orient = smplx_params['global_orient'].squeeze().cpu()
                transl = transls.squeeze().cpu()

                # add to dict if keypoints were detected
                imgpaths_.append(img_path)
                imgnames_.append(img_fn)
                subject_ids_.append(subj_id)
                relpaths_.append(relpath)
                heights_.append(height)
                weights_.append(weight)
                genders_.append(gender)
                gt_coms_.append(gt_com)
                pred_coms_.append(pred_com)
                com_errors_.append(com_error)
                pp_body_poses_.append(np.array(body_pose))
                pp_left_hand_poses_.append(np.array(left_hand_pose))
                pp_right_hand_poses_.append(np.array(right_hand_pose))
                pp_global_orients_.append(np.array(global_orient))
                pp_transls_.append(np.array(transl))

    # get mean com errro
    print('Mean COM error: ', np.mean(com_errors_))
    print(f'Folders with no images: \n {no_images_found}')
    print(f'Folders with no presented poses: \n {no_pp_found}')

    # save as db file
    outfile = osp.join('../data', 'dbs_biomech', 'newmoyo_gt_50_nexus_221004.npz')
    os.makedirs(osp.dirname(outfile), exist_ok=True)
    np.savez(outfile,
             imgpath=imgpaths_,
             subject_id=subject_ids_,
             imgname=imgnames_,
             relpath=relpaths_,
             height=heights_,
             weight=weights_,
             gender=genders_,
             gt_com=gt_coms_,
             pred_com=pred_coms_,
             com_error=com_errors_,
             body_pose=pp_body_poses_,
             left_hand_pose=pp_left_hand_poses_,
             right_hand_pose=pp_right_hand_poses_,
             global_orient=pp_global_orients_,
             transl=pp_transls_
             )

    # Average com_errors_
    com_errors_ = np.array(com_errors_)
    print('Average COM error: ', np.mean(com_errors_))


if __name__ == "__main__":
    """
    yoga-82 dir /ps/scratch/ps_shared/4Paul/Yoga-82
    """

    parser = argparse.ArgumentParser()
    parser.add_argument('--img_folder', required=True,
                        help='folder containing yoga 82 images')
    parser.add_argument('--nexus_com_c3d_folder', required=True,
                        help='folder containing GT pressure CSVs')
    parser.add_argument('--pp_folder', required=True,
                        help='presented poses folder')
    parser.add_argument('--model_folder', required=False, default='/ps/project/common/smplifyx/models/',
                        help='path to SMPL/SMPL-X model folder')
    parser.add_argument('--save_outputs', action='store_true', default=False,
                         help='if enabled, saves outputs to disk, but slows down eval')

    args = parser.parse_args()
    yoga82_extra(args.img_folder, args.pp_folder, args.nexus_com_c3d_folder, args.model_folder)
moyo/eval/com_evaluation.py
sha2nkt/moyo_toolkit
"""
Usage
python eval/pressure_map_evaluation.py --img_folder ../data/moyo/20220923_20220926_with_hands/images/val/ --pp_folder ../data/moyo/20220923_20220926_with_hands/mosh/val/ --pressure_xml_folder ../data/moyo/20220923_20220926_with_hands/pressure/val/xml --pressure_csv_folder ../data/moyo/20220923_20220926_with_hands/pressure/val/single_csv
"""

import argparse
import copy
import glob
import json
import os
import os.path as osp

import cv2
import numpy as np
import torch
import trimesh

from tqdm import tqdm
import pickle as pkl
from scipy import stats


import sys

from moyo.utils.misc import colors, copy2cpu as c2c


import smplx
from moyo.utils.geometry import matrix_to_axis_angle, ea2rm, axis_angle_to_matrix
from moyo.utils.constants import frame_select_dict_combined as frame_select_dict
from moyo.utils.biomech_eval_utils import BiomechanicalEvaluator
from moyo.utils.pressure_mat_utils import PressureMat
from psbody.mesh import Mesh

VICON_FPS = 60
PRESSURE_MAP_FPS = 60
IOI_FPS = 30


def read_keypoints(keypoint_fn):
    with open(keypoint_fn) as keypoint_file:
        data = json.load(keypoint_file)

    keypoints = []
    for idx, person_data in enumerate(data['people']):
        body_keypoints = np.array(
            person_data['pose_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])

        left_hand_keyp = np.array(
            person_data['hand_left_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])

        right_hand_keyp = np.array(
            person_data['hand_right_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])

        face_keypoints = np.array(
            person_data['face_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])[17: 17 + 51, :]

        contour_keyps = np.array(
            person_data['face_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])[:17, :]

        body_keypoints = np.concatenate([body_keypoints, left_hand_keyp,
                                         right_hand_keyp, face_keypoints, contour_keyps], axis=0)

        keypoints.append(body_keypoints)
    return keypoints


def select_keypoints(keypoints, img):
    img_center = np.array(img.shape[:2]) / 2  # height, width
    # select keypoints closest to image center weighted by inverse confidence
    if len(keypoints) > 1:
        kpselect = np.inf * np.ones(len(keypoints))
        # major joints openpose
        op_to_12 = [9, 10, 11, 12, 13, 14, 2, 3, 4, 5, 6, 7]
        for idx, personkpts in enumerate(keypoints):
            kpdist = personkpts[op_to_12, :2] - img_center
            kpdist = np.linalg.norm(kpdist, axis=1)
            kpconf = np.dot(kpdist, (- personkpts[op_to_12, 2] + 1))
            kpselect[idx] = kpconf
        kpselidx = np.argmin(kpselect)
    elif len(keypoints) == 1:
        kpselidx = 0
    else:
        keypoints = None
    keypoints = np.stack(keypoints[kpselidx:kpselidx + 1])
    return keypoints


def visualize_mesh(bm_output, faces, frame_id=0, display=False):
    imw, imh = 1600, 1600

    body_mesh = trimesh.Trimesh(vertices=c2c(bm_output.vertices[frame_id]),
                                faces=faces,
                                vertex_colors=np.tile(colors['grey'], (6890, 1)),
                                process=False,
                                maintain_order=True)

    if display:
        mv = MeshViewer(width=imw, height=imh, use_offscreen=True)
        mv.set_static_meshes([body_mesh])
        body_image = mv.render(render_wireframe=False)
        show_image(body_image)

    return body_mesh


def get_trans_offset(pelvis, smplx_params, trans, body_model):
    '''
    SOMA fits are in AMASS format (xy plane is ground) but SMPLify-XMC uses Pyrender format (xz plane is ground). global_orient is corrected by 270 rotation along x-axis to match the two formats.
    While the ground plane is z=0 in soma, since global_orient is wrt pelvis, this would mean the ground plane is not at y=0 after rotation.
    Fix: distance of pelvis to ground in preserved. Use it to push the mesh up after rotation such that the ground plane is at y=0.
    Args:
        pelvis: pelvis joint before rotation
    '''
    bs = trans.shape[0]
    init_root_orient = smplx_params['global_orient']
    pelvis_height = pelvis[:, 2]  # this is always conserved

    new_smplx_params = copy.deepcopy(smplx_params)
    # Rotate AMASS root orient to smplify-xmc format
    R_init = axis_angle_to_matrix(init_root_orient)
    R1 = ea2rm(torch.tensor([[np.radians(270)]]), torch.tensor([[np.radians(0)]]),
               torch.tensor([[np.radians(0)]])).float().to(R_init.device)
    R = torch.bmm(R1.expand(bs, -1, -1), R_init)
    new_smplx_params['global_orient'] = matrix_to_axis_angle(R)

    # posed body with hand, with global orient
    body_model_output = body_model(
        global_orient=new_smplx_params['global_orient'],
        body_pose=new_smplx_params['body_pose'])

    new_pelvis = body_model_output.joints[:, 0]
    new_ground_plane_height = new_pelvis[:, 1] - pelvis_height
    trans_offset = -new_ground_plane_height
    return trans_offset

def smplx_breakdown(bdata, device):
    num_frames = len(bdata['trans'])

    bdata['poses'] = bdata['fullpose']

    global_orient = torch.from_numpy(bdata['poses'][:, :3]).float().to(device)
    body_pose = torch.from_numpy(bdata['poses'][:, 3:66]).float().to(device)
    jaw_pose = torch.from_numpy(bdata['poses'][:, 66:69]).float().to(device)
    leye_pose = torch.from_numpy(bdata['poses'][:, 69:72]).float().to(device)
    reye_pose = torch.from_numpy(bdata['poses'][:, 72:75]).float().to(device)
    left_hand_pose = torch.from_numpy(bdata['poses'][:, 75:120]).float().to(device)
    right_hand_pose = torch.from_numpy(bdata['poses'][:, 120:]).float().to(device)

    # if bdata['v_template_fname'] == '/ps/project/amass/MOCAP/PS_MoCaps/YOGI/subject_data/subject_v_template/210727_03596/smplx_female.ply':
    v_template_fname = '/is/cluster/scratch/stripathi/pycharm_remote/yogi/data/v_templates/220923_yogi_03596_minimal_simple_female/mesh.ply'
    # else:
    #     print("Something weird happening with V_template loading!!!!")
    #     import ipdb; ipdb.set_trace()
    v_template = Mesh(filename=v_template_fname)

    body_params = {'global_orient': global_orient, 'body_pose': body_pose,
                  'jaw_pose': jaw_pose, 'leye_pose': leye_pose, 'reye_pose': reye_pose,
                  'left_hand_pose': left_hand_pose, 'right_hand_pose': right_hand_pose,
                   'v_template': torch.Tensor(v_template.v).to(device),}
    return body_params


def smplx_to_mesh(body_params, model_folder, model_type, gender='neutral'):
    with torch.no_grad():
        device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
        num_betas = 10

        smplx_params = smplx_breakdown(body_params, device)

        trans = torch.from_numpy(body_params['trans']).float().to(device)

        betas = torch.from_numpy(body_params['betas']).float().to(device).unsqueeze(0)
        betas = betas[:, :num_betas]  # only using 10 betas

        body_model_params = dict(model_path=model_folder,
                                 model_type=model_type,
                                 gender=gender,
                                 v_template=smplx_params['v_template'],
                                 # joint_mapper=joint_mapper,
                                 batch_size=trans.shape[0],
                                 create_global_orient=True,
                                 create_body_pose=True,
                                 create_betas=True,
                                 num_betas=num_betas,
                                 create_left_hand_pose=True,
                                 create_right_hand_pose=True,
                                 create_expression=True,
                                 create_jaw_pose=True,
                                 create_leye_pose=True,
                                 create_reye_pose=True,
                                 create_transl=True,
                                 use_pca=False,
                                 flat_hand_mean=True,
                                 dtype=torch.float32)

        body_model = smplx.create(**body_model_params).to(device)
        body_model_output = body_model(transl=trans,
                                       global_orient=smplx_params['global_orient'],
                                       body_pose=smplx_params['body_pose'],
                                       left_hand_pose=smplx_params['left_hand_pose'],
                                       right_hand_pose=smplx_params['right_hand_pose'])

        pelvis = body_model_output.joints[:, 0]
        trans_offset = get_trans_offset(pelvis, smplx_params, trans, body_model)
        zero_vec = torch.zeros_like(trans_offset)
        transl = torch.stack([zero_vec, trans_offset, zero_vec], dim=-1)

    return body_model_output, smplx_params, transl, body_model.faces

def yoga82_extra(img_folder, pp_folder, pressure_xml_folder, pressure_csv_folder, model_folder, save_outputs=False,
                 model_type='smplx', subsampling_factor=1):
    imgpaths_ = []
    imgnames_ = []
    subject_ids_ = []
    relpaths_ = []
    heights_ = []
    weights_ = []
    genders_ = []
    pressures_ = []
    pressure_heatmaps_ = []
    pp_body_poses_ = []
    pp_left_hand_poses_ = []
    pp_right_hand_poses_ = []
    pp_global_orients_ = []
    pp_transls_ = []

    # presented poses
    pp_names = os.listdir(pp_folder)
    pp_names = [os.path.splitext(x)[0] for x in pp_names if x[0] != '.' and '.pkl' in x]

    all_pose_names = [pp_name.replace('_stageii', '') for pp_name in pp_names]
    all_pose_names = sorted(all_pose_names)

    pp_pose_cache = {}

    no_images_found = set()
    no_pp_found = set()

    smplx_faces = smplx.create(model_path=model_folder, model_type=model_type).faces
    smplx_faces = torch.tensor(smplx_faces.astype(np.int64), dtype=torch.long)
    pressure_evaluator = BiomechanicalEvaluator(faces=smplx_faces)

    for i, pname in enumerate(tqdm(all_pose_names)):
        print(pname)
        if 'Cow_Face_Pose_or_Gomukhasana_-c' in pname: # pressure mat was broken for this pose
            continue
        if 'female_stagei' in pname: # do not analyze female stage1 file
            continue
        # get all possible variations
        pp_pkl = osp.join(pp_folder, pname + '_stageii.pkl')
        pp_params = pkl.load(open(pp_pkl, 'rb'))
        # read height, weight, gender meta
        height = None
        weight = None
        gender = 'female'
        subj_id = 'unknown'
        pp_body_model_output, smplx_params, transls, faces = smplx_to_mesh(pp_params, model_folder,
                                                       model_type, gender=gender)
        # get pressure xml
        pressure_xml = osp.join(pressure_xml_folder, pname + '.xml')
        pressure_csv = osp.join(pressure_csv_folder, pname + '.csv')
        pressure_mat_c3d = pressure_xml.replace('/xml/', '/pressure_mat_c3d/').replace('.xml', '.c3d')
        # initialize pressure mat
        pressure_mat = PressureMat(pressure_csv, pressure_xml, pressure_mat_c3d)

        # pick selected frame
        selected_frame = frame_select_dict['_'.join(pname.split('_')[5:])]

        # get all images for this pose
        img_dir = osp.join(img_folder, pname)
        for cam in [1]:
            cam_name = 'YOGI_Cam_' + str(cam).zfill(2)
            img_dir_cam = osp.join(img_dir, cam_name)
            types = ('*.jpg', '*.png')  # the tuple of file types
            images = []
            for files in types:
                images.extend(glob.glob(osp.join(img_dir_cam, files)))
            if len(images) == 0:
                print(f'No images found for {pname}')
                no_images_found.add(pname)

            # subsample frames
            images.sort()

            # folders to log outputs
            if cam == 1 and save_outputs:
                out_img_dir = osp.join(pp_folder, 'images_small', pname, cam_name)
                os.makedirs(out_img_dir, exist_ok=True)
                out_pressure_heatmap_dir = osp.join(pp_folder, 'pressure_heatmaps', pname, cam_name)
                os.makedirs(out_pressure_heatmap_dir, exist_ok=True)
                out_obj_dir = osp.join(pp_folder, 'meshes', pname, cam_name)
                os.makedirs(out_obj_dir, exist_ok=True)

            for img_path in images[
                            :-1:subsampling_factor]:  # sometimes, the last frame is missing in pressure or vicon data
                # get frame correspondences
                img_fnum = int(img_path.split('_')[-1].split('.')[0])

                #only run on selected frames
                if img_fnum != selected_frame:
                    continue
                vicon_fnum = img_fnum * (VICON_FPS // IOI_FPS)
                pressure_fnum = img_fnum * (PRESSURE_MAP_FPS // IOI_FPS)
                assert vicon_fnum == pressure_fnum

                # load image
                img_fn, _ = osp.splitext(osp.split(img_path)[1])
                relpath = osp.dirname(img_path).replace(img_dir, '').strip('/')

                if cam == 1 and save_outputs:
                    img = cv2.imread(img_path).astype(np.float32)
                    # resize image to half size
                    img = cv2.resize(img, (img.shape[1] // 2, img.shape[0] // 2))
                    cv2.imwrite(osp.join(out_img_dir, f'{pname}_{pressure_fnum}.jpg'), img)
                    print('[[Image Saved]]\n', osp.join(out_img_dir, f'{pname}_{pressure_fnum}.jpg'))
                if cam != 1:
                    continue

                # get pressure for current frame
                # increase_heatmap size before saving
                gt_pressure = pressure_mat.gt_pressures[pressure_fnum]
                gt_heatmap = pressure_mat.gt_heatmaps[pressure_fnum]
                gt_cop_relative = pressure_mat.gt_cops_relative[pressure_fnum]
                if np.isnan(gt_cop_relative).any():
                    continue
                # gt_cop_metric = pressure_mat.gt_cops_metric[pressure_fnum]
                if cam == 1 and save_outputs:
                    pressure_name = str(img_fnum).zfill(4) + '.png'
                    heatmap_large = cv2.resize(gt_heatmap, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_NEAREST)
                    # add gt_cop point to heatmap
                    cv2.drawMarker(heatmap_large, (
                        int(gt_cop_relative[0] / pressure_mat.mat_size_metric[0] * heatmap_large.shape[1]),
                        int(gt_cop_relative[1] / pressure_mat.mat_size_metric[1] * heatmap_large.shape[0])),
                                   [0, 255, 0],
                                   markerType=cv2.MARKER_STAR, markerSize=40, thickness=5)
                    # heatmap_large = cv2.flip(heatmap_large, 0) # to have better visualization wrt to blender mat configuration
                    cv2.imwrite(os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_name}'),
                                heatmap_large)
                    print('[[GT Pressure Heatmap Saved]]\n',
                          os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_name}'))



                transl = transls[vicon_fnum]
                pp_mesh = visualize_mesh(pp_body_model_output, faces, frame_id=img_fnum)
                pp_verts = pp_mesh.vertices
                # pp_verts = pp_body_model_output.vertices[img_fnum]
                vertices = torch.tensor(pp_verts).unsqueeze(0)
                if cam == 1 and save_outputs:
                    mesh_name = str(img_fnum).zfill(4) + '.obj'
                    out_obj = osp.join(out_obj_dir, f'{pname}_{mesh_name}')
                    pp_mesh.export(out_obj)
                    print('[[Mesh]]\n', out_obj)


                # get per-vertex pressure
                pressure_evaluator.evaluate_pressure(gt_pressure,
                                                   gt_cop_relative,
                                                   vertices,
                                                   pressure_mat.mat_size,
                                                   pressure_mat.mat_size_metric,
                                                   pressure_mat.marker_positions_metric)
                if cam == 1 and save_outputs:
                    # iou first visualization
                    pred_heatmap = pressure_evaluator.iou_first.pred_heatmaps[-1]
                    pred_cop_relative = pressure_evaluator.iou_first.pred_cops_relative[-1]
                    cop_error = pressure_evaluator.iou_first.cop_errors[-1]
                    iou = pressure_evaluator.iou_first.ious[-1]
                    frame_diff = pressure_evaluator.iou_first.frame_diffs[-1]
                    best_cop_w = pressure_evaluator.iou_first.best_cop_w[-1]
                    best_cop_k = pressure_evaluator.iou_first.best_cop_k[-1]

                    pred_heatmap_large = cv2.resize(pred_heatmap, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_NEAREST)
                    cv2.drawMarker(pred_heatmap_large, (
                        int(pred_cop_relative[0] / pressure_mat.mat_size_metric[0] * pred_heatmap_large.shape[1]),
                        int(pred_cop_relative[1] / pressure_mat.mat_size_metric[1] * pred_heatmap_large.shape[0])),
                                   [0, 255, 255],
                                   markerType=cv2.MARKER_STAR, markerSize=40, thickness=5)
                    cv2.putText(pred_heatmap_large, f'{cop_error*1000:.2f} mm', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (0, 255, 255), 2)
                    cv2.putText(pred_heatmap_large, f'{iou:.2f}', (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255,255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'{frame_diff:.2f}', (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_w: {best_cop_w:.2f}', (10, 200), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_k: {best_cop_k:.2f}', (10, 250), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.imwrite(os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_iou_first.png'),
                                pred_heatmap_large)
                    print('[[PRED Pressure Heatmap Saved]]\n',
                          os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_iou_first.png'))

                    # cop error first visualizations
                    pred_heatmap = pressure_evaluator.cop_error_first.pred_heatmaps[-1]
                    pred_cop_relative = pressure_evaluator.cop_error_first.pred_cops_relative[-1]
                    cop_error = pressure_evaluator.cop_error_first.cop_errors[-1]
                    iou = pressure_evaluator.cop_error_first.ious[-1]
                    frame_diff = pressure_evaluator.cop_error_first.frame_diffs[-1]
                    best_cop_w = pressure_evaluator.cop_error_first.best_cop_w[-1]
                    best_cop_k = pressure_evaluator.cop_error_first.best_cop_k[-1]

                    pred_heatmap_large = cv2.resize(pred_heatmap, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_NEAREST)
                    cv2.drawMarker(pred_heatmap_large, (
                        int(pred_cop_relative[0] / pressure_mat.mat_size_metric[0] * pred_heatmap_large.shape[1]),
                        int(pred_cop_relative[1] / pressure_mat.mat_size_metric[1] * pred_heatmap_large.shape[0])),
                                   [0, 255, 255],
                                   markerType=cv2.MARKER_STAR, markerSize=40, thickness=5)
                    cv2.putText(pred_heatmap_large, f'{cop_error * 1000:.2f} mm', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (0, 255, 255), 2)
                    cv2.putText(pred_heatmap_large, f'{iou:.2f}', (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'{frame_diff:.2f}', (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_w: {best_cop_w:.2f}', (10, 200), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_k: {best_cop_k:.2f}', (10, 250), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.imwrite(os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_cop_error_first.png'),
                                pred_heatmap_large)
                    print('[[PRED Pressure Heatmap Saved]]\n',
                          os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_cop_error_first.png'))

                    # frame_idd first visualizations
                    pred_heatmap = pressure_evaluator.frame_diff_first.pred_heatmaps[-1]
                    pred_cop_relative = pressure_evaluator.frame_diff_first.pred_cops_relative[-1]
                    cop_error = pressure_evaluator.frame_diff_first.cop_errors[-1]
                    iou = pressure_evaluator.frame_diff_first.ious[-1]
                    frame_diff = pressure_evaluator.frame_diff_first.frame_diffs[-1]
                    best_cop_w = pressure_evaluator.frame_diff_first.best_cop_w[-1]
                    best_cop_k = pressure_evaluator.frame_diff_first.best_cop_k[-1]

                    pred_heatmap_large = cv2.resize(pred_heatmap, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_NEAREST)
                    cv2.drawMarker(pred_heatmap_large, (
                        int(pred_cop_relative[0] / pressure_mat.mat_size_metric[0] * pred_heatmap_large.shape[1]),
                        int(pred_cop_relative[1] / pressure_mat.mat_size_metric[1] * pred_heatmap_large.shape[0])),
                                   [0, 255, 255],
                                   markerType=cv2.MARKER_STAR, markerSize=40, thickness=5)
                    cv2.putText(pred_heatmap_large, f'{cop_error * 1000:.2f} mm', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (0, 255, 255), 2)
                    cv2.putText(pred_heatmap_large, f'{iou:.2f}', (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'{frame_diff:.2f}', (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_w: {best_cop_w:.2f}', (10, 200), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_k: {best_cop_k:.2f}', (10, 250), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.imwrite(
                        os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_frame_diff_first.png'),
                        pred_heatmap_large)
                    print('[[PRED Pressure Heatmap Saved]]\n',
                          os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_frame_diff_first.png'))



                body_pose = smplx_params['body_pose'][vicon_fnum].squeeze().cpu()
                left_hand_pose = smplx_params['left_hand_pose'][vicon_fnum].squeeze().cpu()
                right_hand_pose = smplx_params['right_hand_pose'][vicon_fnum].squeeze().cpu()
                global_orient = smplx_params['global_orient'][vicon_fnum].squeeze().cpu()
                transl = transl.squeeze().cpu()

                # add to dict if keypoints were detected
                imgpaths_.append(img_path)
                imgnames_.append(img_fn)
                subject_ids_.append(subj_id)
                relpaths_.append(relpath)
                heights_.append(height)
                weights_.append(weight)
                genders_.append(gender)
                pressures_.append(gt_pressure)
                pressure_heatmaps_.append(gt_heatmap)
                pp_body_poses_.append(np.array(body_pose))
                pp_left_hand_poses_.append(np.array(left_hand_pose))
                pp_right_hand_poses_.append(np.array(right_hand_pose))
                pp_global_orients_.append(np.array(global_orient))
                pp_transls_.append(np.array(transl))

    print('Mean COP error (iou first): ', np.mean(pressure_evaluator.iou_first.cop_errors))
    print('Mean Pressure IOU (iou first): ', np.mean(pressure_evaluator.iou_first.ious))
    print('Mean Pressure Frame Diff (iou first): ', np.mean(pressure_evaluator.iou_first.frame_diffs))
    # print mode of cop_w and cop_k
    print('Mode of cop_w (iou first): ', stats.mode(pressure_evaluator.iou_first.best_cop_w))
    print('Mode of cop_k (iou first): ', stats.mode(pressure_evaluator.iou_first.best_cop_k))
    print('---------------------------------------------')
    print('Mean COP error (cop error first): ', np.mean(pressure_evaluator.cop_error_first.cop_errors))
    print('Mean Pressure IOU (cop error first): ', np.mean(pressure_evaluator.cop_error_first.ious))
    print('Mean Pressure Frame Diff (cop error first): ', np.mean(pressure_evaluator.cop_error_first.frame_diffs))
    # print mode of cop_w and cop_k
    print('Mode of cop_w (cop error first): ', stats.mode(pressure_evaluator.cop_error_first.best_cop_w))
    print('Mode of cop_k (cop error first): ', stats.mode(pressure_evaluator.cop_error_first.best_cop_k))
    print('---------------------------------------------')
    print('Mean COP error (frame diff first): ', np.mean(pressure_evaluator.frame_diff_first.cop_errors))
    print('Mean Pressure IOU (frame diff first): ', np.mean(pressure_evaluator.frame_diff_first.ious))
    print('Mean Pressure Frame Diff (frame diff first): ', np.mean(pressure_evaluator.frame_diff_first.frame_diffs))
    # print mode of cop_w and cop_k
    print('Mode of cop_w (frame diff first): ', stats.mode(pressure_evaluator.frame_diff_first.best_cop_w))
    print('Mode of cop_k (frame diff first): ', stats.mode(pressure_evaluator.frame_diff_first.best_cop_k))
    print('---------------------------------------------')
    print(f'Folders with no images: \n {no_images_found}')
    print(f'Folders with no presented poses: \n {no_pp_found}')

    # save as db file
    outfile = osp.join('../data', 'dbs_biomech', 'newmoyo_gt_first_day_vicon_220923.npz')
    os.makedirs(osp.dirname(outfile), exist_ok=True)
    np.savez(outfile,
             imgpath=imgpaths_,
             subject_id=subject_ids_,
             imgname=imgnames_,
             relpath=relpaths_,
             height=heights_,
             weight=weights_,
             gender=genders_,
             pressure=pressures_,
             pressure_heatmap=pressure_heatmaps_,
             body_pose=pp_body_poses_,
             left_hand_pose=pp_left_hand_poses_,
             right_hand_pose=pp_right_hand_poses_,
             global_orient=pp_global_orients_,
             transl=pp_transls_
             )


if __name__ == "__main__":
    """
    yoga-82 dir /ps/scratch/ps_shared/4Paul/Yoga-82
    """

    parser = argparse.ArgumentParser()
    parser.add_argument('--img_folder', required=True,
                        help='folder containing yoga 82 images')
    # parser.add_argument('--openpose_folder', required=True,
    #                     help='folder containing yoga 82 keypoints')
    parser.add_argument('--pressure_xml_folder', required=True,
                        help='folder containing GT pressure XMLs')
    parser.add_argument('--pressure_csv_folder', required=True,
                        help='folder containing GT pressure CSVs')
    parser.add_argument('--pp_folder', required=True,
                        help='presented poses folder')
    # parser.add_argument('--essentials_folder', required=True,
    #                     help='geodesics, ... essentials for smplify-xmc')
    parser.add_argument('--model_folder', required=False, default='/ps/project/common/smplifyx/models/',
                        help='path to SMPL/SMPL-X model folder')
    parser.add_argument('--save_outputs', action='store_true', default=False,
                        help='if enabled, saves outputs to disk, but slows down eval')

    args = parser.parse_args()
    yoga82_extra(args.img_folder, args.pp_folder, args.pressure_xml_folder, args.pressure_csv_folder, args.model_folder, args.save_outputs)
moyo/eval/pressure_map_evaluation.py
sha2nkt/moyo_toolkit
import argparse
import json
import os
import os.path as op
import sys
import traceback
import warnings
from glob import glob
from hashlib import sha256

import requests
from clint.textui import progress
from loguru import logger
from tqdm import tqdm

warnings.filterwarnings("ignore", message="Unverified HTTPS request")


def download(url_files_dict, post_data):
    # Download urls
    logger.warning(
        f"Note: Download Images? = {args.download_images}. To get images, rerun with -i flag in the bash script")
    session = requests.Session()

    for url_tag, url_file in url_files_dict.items():
        if not op.exists(url_file):
            raise FileNotFoundError(f"File {url_file} not found")

        if args.download_images == "False":
            if "image" in url_tag:
                continue
        # Read the URLs from the file
        with open(url_file, "r") as f:
            urls = f.readlines()
        # Strip newline characters from the URLs
        urls = [url.strip() for url in urls]

        # Loop through the URLs and download the files
        logger.info(f"Start downloading from {url_file}\n")

        pbar = tqdm(urls)
        for url in pbar:
            # Get the filename from the URL
            filepath = url.split("=")[-1]
            if op.exists(op.join(args.out_dir, filepath)):
                logger.info(f"File {filepath} already exists. Skipping...")
                continue

            pbar.set_description(f"Downloading {url}")
            # Make a POST request with the username and password
            response = session.post(
                url,
                data=post_data,
                stream=True,
                verify=False,
                allow_redirects=True,
            )

            if response.status_code == 401:
                logger.warning(
                    f"Authentication failed for URLs in {url_file}. Username/password correct?"
                )
                sys.exit(1)


            # Write the contents of the response to a file
            out_p = op.join(args.out_dir, filepath)
            os.makedirs(op.dirname(out_p), exist_ok=True)
            total_length = int(response.headers.get("content-length"))
            with open(out_p, "wb") as f:
                for chunk in progress.bar(
                        response.iter_content(chunk_size=1024*1024*100),
                        expected_size=(total_length / (1024*1024*100)) + 1,
                ):
                    if chunk:
                        f.write(chunk)
                        f.flush()

def verify_checksum(out_dir):
    logger.info("Verifying checksums for downloaded files")

    fnames = glob(op.join(out_dir, "**/*"), recursive=True)
    pbar = tqdm(fnames)

    with open("./moyo/bash/assets/checksum.json", "r") as f:
        gt_checksum = json.load(f)

    hash_dict = {}
    for fname in pbar:
        if op.isdir(fname):
            continue
        if ".zip" not in fname:
            continue

        try:
            with open(fname, "rb") as f:
                pbar.set_description(f"Reading {fname}")
                data = f.read()
                hashcode = sha256(data).hexdigest()
                key = fname.replace("./", "")

                hash_dict[key] = hashcode
                if hashcode != gt_checksum[key]:
                    print(f"Error: {fname} has different checksum!")
                else:
                    pbar.set_description(f"Hashcode of {fname} is correct!")
                    # print(f'Hashcode of {fname} is correct!')
        except:
            print(f"Error processing {fname}")
            traceback.print_exc()
            continue

    out_p = op.join(out_dir, "checksum.json")
    with open(out_p, "w") as f:
        json.dump(hash_dict, f, indent=4, sort_keys=True)
    print(f"Checksum file saved to {out_p}!")

def unzip(url_files_dict):
    for url_tag, url_file in url_files_dict.items():
        if not op.exists(url_file):
            raise FileNotFoundError(f"File {url_file} not found")

        if args.download_images == "False":
            logger.warning("Note: Skipping image download. To get images, rerun with -i flag in the bash script")
            if "image" in url_tag:
                continue
        # Read the URLs from the file
        with open(url_file, "r") as f:
            urls = f.readlines()
        # Strip newline characters from the URLs
        urls = [url.strip() for url in urls]

        # Loop through the URLs and download the files
        logger.info(f"Start downloading from {url_file}")
        pbar = tqdm(urls)
        for url in pbar:
            pbar.set_description(f"Unzipping {url[-40:]}")

            # Get the filename from the URL
            filepath = url.split("=")[-1]

            # Write the contents of the response to a file
            out_p = op.join(args.out_dir, filepath)
            if out_p.endswith(".zip"):
                logger.info(f"Unzipping {out_p}")
                os.system(f"unzip -o {out_p} -d {args.out_dir}")


def delete_zip(url_files_dict):
    for url_tag, url_file in url_files_dict.items():
        if not op.exists(url_file):
            raise FileNotFoundError(f"File {url_file} not found")

        if args.download_images == "False":
            logger.warning("Note: Skipping image download. To get images, rerun with -i flag in the bash script")
            if "image" in url_tag:
                continue
        # Read the URLs from the file
        with open(url_file, "r") as f:
            urls = f.readlines()
        # Strip newline characters from the URLs
        urls = [url.strip() for url in urls]

        # Loop through the URLs and download the files
        logger.info(f"Start downloading from {url_file}")
        pbar = tqdm(urls)
        for url in pbar:
            pbar.set_description(f"Unzipping {url[-40:]}")

            # Get the filename from the URL
            filepath = url.split("=")[-1]

            # Write the contents of the response to a file
            out_p = op.join(args.out_dir, filepath)
            os.makedirs(op.dirname(out_p), exist_ok=True)
            logger.info(f"Unzipping {out_p}")
            os.system(f"unzip -o {out_p} -d {args.out_dir}")


def download_data(args):
    # Define the username and password
    flag = "MOYO"

    username = args.username
    password = args.password
    password_fake = "*" * len(password)

    logger.info(f"Username: {username}")
    logger.info(f"Password: {password_fake}")

    post_data = {"username": username, "password": password}

    # Get URL paths
    url_files_dict = {
        "cameras_urls": op.join(args.url_dir, "cameras.txt"),
        "image_train_urls": op.join(args.url_dir, "images_train.txt"),
        "image_val_urls": op.join(args.url_dir, "images_val.txt"),
        "com_urls": op.join(args.url_dir, "coms.txt"),
        "pressure_urls": op.join(args.url_dir, "pressures.txt"),
        "vicon_urls": op.join(args.url_dir, "vicon.txt"),
        "smplx_fits_urls": op.join(args.url_dir, "smplx_fits.txt"),
    }

    # Download urls
    download(url_files_dict, post_data)
    # Verify checksums
    verify_checksum(args.out_dir)

    # Unzip files
    if args.unzip == "True":
        unzip(url_files_dict)
    # Delete zip files
    if args.delete == "True":
        delete_zip(url_files_dict)
    logger.info("Done")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Download files from a list of URLs")
    parser.add_argument(
        "--url_dir",
        type=str,
        help="Path to file directory containing URLs",
        required=True,
    )
    parser.add_argument(
        "--out_dir",
        type=str,
        help="Path to folder to store downloaded files",
        required=True,
    )
    parser.add_argument(
        "--username", type=str, help="Username for the MOYO data download, register at moyo.is.mpg.de", required=True
    )
    parser.add_argument(
        "--password", type=str, help="Password for the MOYO data download, register at moyo.is.mpg.de", required=True
    )
    parser.add_argument(
        "--download_images", type=str, help="Whether to download the images", default="True"
    )
    parser.add_argument(
        "--unzip", type=str, help="Whether to unzip the downloaded zip files", default="True"
    )
    parser.add_argument(
        "--delete", type=str, help="Whether to delete the downloaded zip files", default="True"
    )

    args = parser.parse_args()

    download_data(args)

moyo/scripts/download_data.py
sha2nkt/moyo_toolkit
import argparse
import glob
import hashlib
import json
import os

from tqdm import tqdm


def calculate_checksum(file_path):
    hasher = hashlib.sha256()
    with open(file_path, 'rb') as f:
        for byte_block in iter(lambda: f.read(4096), b''):
            hasher.update(byte_block)
    return hasher.hexdigest()


def generate_checksums(data_dir):
    checksums = {}
    zip_files = glob.glob(os.path.join(data_dir, '**/*.zip'), recursive=True)

    for file_path in tqdm(zip_files):
        checksum = calculate_checksum(file_path)
        checksums[file_path] = checksum

    with open('./moyo/bash/assets/checksum.json', 'w') as json_file:
        json.dump(checksums, json_file, indent=4)


def main():
    parser = argparse.ArgumentParser(description='Generate checksums for zip files in a folder recursively.')
    parser.add_argument('--data_dir', help='Path to the folder containing zip files')
    args = parser.parse_args()

    generate_checksums(args.data_dir)


if __name__ == '__main__':
    main()

moyo/scripts/generate_gt_checksums.py
sha2nkt/moyo_toolkit
"""
Script to get 2D projects from mosh_fits.

python scripts/ioi_vicon_frame_sync.py --img_folder ../data/moyo/20220923_20220926_with_hands/images/ --c3d_folder ../data/moyo/20220923_20220926_with_hands/vicon --cam_folder_first ../data/moyo/20220923_20220926_with_hands/cameras/20220923/220923_Afternoon_PROCESSED_CAMERA_PARAMS/cameras_param.json --cam_folder_second ../data/moyo/20220923_20220926_with_hands/cameras/20220926/220926_Morning_PROCESSED_CAMERA_PARAMS/cameras_param.json --output_dir ../data/moyo_images_mocap_projected --frame_offset 1 --split val"""

import argparse
import glob
import json
import os
import os.path as osp

import cv2
import ipdb
import numpy as np
import torch
import trimesh
from ezc3d import c3d as ezc3d
from tqdm import tqdm

from moyo.utils.constants import frame_select_dict_combined as frame_select_dict
from moyo.utils.misc import colors, copy2cpu as c2c


class Opt:
    def __init__(self):
        self.pose_track = False
        self.tracking = False
        self.showbox = False


def visualize_mesh(bm_output, faces, frame_id=0, display=False):
    imw, imh = 1600, 1600

    body_mesh = trimesh.Trimesh(vertices=c2c(bm_output.vertices[frame_id]),
                                faces=faces,
                                vertex_colors=np.tile(colors['grey'], (6890, 1)),
                                process=False,
                                maintain_order=True)

    if display:
        mv = MeshViewer(width=imw, height=imh, use_offscreen=True)
        mv.set_static_meshes([body_mesh])
        body_image = mv.render(render_wireframe=False)
        show_image(body_image)

    return body_mesh


def project2d(j3d, cam_params, downsample_factor=1.0):
    """
    Project 3D points to 2D
    Args:
        j3d: (N, 3) 3D joints
        cam_params: dict
        downsample_factor: resize factor

    Returns:
        j2d : 2D joint locations
    """
    mm2m = 1000
    j3d = torch.tensor(j3d, dtype=torch.float32)

    # cam intrinsics
    f = cam_params['focal'] * downsample_factor
    cx = cam_params['princpt'][0] * downsample_factor
    cy = cam_params['princpt'][1] * downsample_factor

    # cam extrinsics
    R = torch.tensor(cam_params['rotation'])
    t = -torch.mm(R, torch.tensor(cam_params['position'])[:, None]).squeeze()  # t= -RC

    # cam matrix
    K = torch.tensor([[f, 0, cx],
                      [0, f, cy],
                      [0, 0, 1]]).to(j3d.device)

    Rt = torch.cat([R, t[:, None]], dim=1).to(j3d.device)

    # apply extrinsics
    bs = j3d.shape[0]
    j3d_cam = torch.bmm(Rt[None, :, :].expand(bs, -1, -1), j3d[:, :, None])
    j2d = torch.bmm(K[None, :].expand(bs, -1, -1), j3d_cam)
    j2d = j2d / j2d[:, [-1]]
    return j2d[:, :-1, :].squeeze()


def visualize_on_img(j2d, img_name, out_dir):
    # Visualize the joints
    pose_name = img_name.split('/')[-2]
    img_num = img_name.split('/')[-1]
    img = cv2.imread(img_name)
    fname = img_name.split('/')[-1]
    os.makedirs(out_dir, exist_ok=True)
    ext = osp.splitext(fname)[1]
    for n in range(j2d.shape[0]):
        # check if nan
        if np.any(np.isnan(j2d[n, :])):
            continue
        cor_x, cor_y = int(j2d[n, 0]), int(j2d[n, 1])
        cv2.circle(img, (cor_x, cor_y), 1, (0, 255, 0), 5)

    out_img_path = osp.join(out_dir, fname).replace(f'{ext}', '_markers.png')
    cv2.imwrite(out_img_path, img)
    print(f'{out_img_path} is saved')


def main(img_folder, c3d_folder, model_folder, output_dir, cam_folders, frame_offset, split, downsample_factor):
    # presented poses
    c3d_folder = os.path.join(c3d_folder, split, 'c3d')
    c3d_names = os.listdir(c3d_folder)
    c3d_names = [os.path.splitext(x)[0] for x in c3d_names if x[0] != '.' and '.c3d' in x]
    c3d_names = sorted(c3d_names)

    img_folder = os.path.join(img_folder, split)
    # list all folders in img_folder
    img_pose_folders = os.listdir(img_folder)
    img_pose_folders = [item for item in img_pose_folders if os.path.isdir(os.path.join(img_folder, item))]

    for i, c3d_name in enumerate(tqdm(c3d_names)):
        print(f'Processing {i}: {c3d_name}')
        # load images, loop through and read smplx and keypoints
        pose_name = '-'.join('_'.join(c3d_name.split('_')[5:]).split('-')[:-1])
        if pose_name == "Happy_Baby_Pose":
            pose_name = "Happy_Baby_Pose_or_Ananda_Balasana_"  # this name was changed during mocap

        # get soma fit
        try:
            c3d_path = glob.glob(osp.join(c3d_folder, f'{c3d_name}.c3d'))[0]
        except:
            print(f'{c3d_folder}/{c3d_name}_stageii.pkl does not exist. SKIPPING!!!')
            continue

        c3d = ezc3d(c3d_path)

        markers3d = c3d['data']['points'].transpose(2, 1,
                                                    0)  # Frames x NumPoints x 4 (x,y,z,1) in homogenous coordinates

        try:
            c3d_var = '_'.join(c3d_name.split('_')[5:])
            selected_frame = frame_select_dict[c3d_var]
        except:
            print(f'{c3d_var} does not exist in frame_selection_dict. SKIPPING!!!')
            continue

        j3d = markers3d[selected_frame]

        # load images belonging to presented pose
        img_dir = [dir for dir in img_pose_folders if c3d_var in dir]

        if len(img_dir) == 0:
            print(f'{pose_name} does not exist in {img_folder}. SKIPPING!!!')
            continue
        if len(img_dir) > 1:
            print(f'{pose_name} has more than one folder. Choose one:')
            for i, dir in enumerate(img_dir):
                print(f'{i}. {dir}')
            select_idx = int(input())
            if select_idx > len(img_dir) - 1:
                print(f'{select_idx} is out of range. SKIPPING!!!')
                continue
            img_dir = img_dir[select_idx]
        if len(img_dir) == 1:
            img_dir = img_dir[0]

        img_dir = osp.join(img_folder, img_dir)
        image_paths = glob.glob(osp.join(img_dir, '*/*.jpg'))

        # select a particular image frame
        selected_img_frame = (selected_frame // 2) + frame_offset
        image_paths = [path for path in image_paths if
                       int(selected_img_frame) == int(osp.splitext(path)[0].split('_')[-1])]

        # only take Cam_01
        image_paths = [path for path in image_paths if 'Cam_01' in path]

        if len(image_paths) == 0:
            print(f'No images found for {img_dir}')

        out_dir = osp.join(output_dir, pose_name)

        for img_path in image_paths:
            img_name = osp.basename(img_path)
            sessionid = img_name.split('_')[0]
            if '220923' in sessionid:
                cam_path = cam_folders['220923']
            elif '220926' in sessionid:
                cam_path = cam_folders['220926']
            else:
                print(f'Invalid {sessionid}. Check!!!')
                ipdb.set_trace()
            # Load camera matrix from json file
            with open(cam_path, 'rb') as fp:
                cameras = json.load(fp)

            name_splits = img_name.split('_')
            cam_num = int(name_splits[name_splits.index('Cam') + 1])
            cam_id = f'cam_{cam_num}'

            cam_params = cameras[cam_id]
            j2d = project2d(j3d, cam_params, downsample_factor=downsample_factor)

            # convert to openpose format
            visualize_on_img(j2d.cpu().numpy(), img_path, out_dir)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_folder', required=True,
                        help='folder_containing_ioi_data')
    parser.add_argument('--c3d_folder', required=True,
                        help='folder containing raw c3d mocap files')
    parser.add_argument('--cam_folder_first', required=True,
                        help='folder containing camera matrix for first session 210727')
    parser.add_argument('--cam_folder_second', required=True,
                        help='folder containing camera matrix for second session 211117')
    parser.add_argument('--split', required=True, choices=['train', 'val', 'test'],
                        help='split to process')
    parser.add_argument('--model_folder', required=False, default='/ps/project/common/smplifyx/models/',
                        help='path to SMPL/SMPL-X model folder')
    parser.add_argument('--output_dir', type=str, help='Path to the output directory')
    parser.add_argument('--downsample_factor', type=float, default=0.5,
                        help='Downsample factor for the images. Release images are downsampled by 0.5')
    parser.add_argument('--frame_offset', type=int,
                        help='To get the frame offset, divide the vicon fnum by 2 and add this offset')

    args = parser.parse_args()

    cam_folders = {'220923': args.cam_folder_first, '220926': args.cam_folder_second}

    main(args.img_folder, args.c3d_folder, args.model_folder, args.output_dir, cam_folders, args.frame_offset,
         args.split, args.downsample_factor)

moyo/scripts/ioi_vicon_frame_sync.py
sha2nkt/moyo_toolkit
import sys
import torch
import numpy as np
import os
import os.path as osp

sys.path.append(os.path.join(os.path.dirname(__file__), '../..'))
sys.path.append(os.path.join(os.path.dirname(__file__), '../../..'))

from moyo.core.support import StabilityLossCoS
from moyo.utils.constants import CONTACT_THRESH
from moyo.utils.mesh_utils import HDfier
from moyo.utils.pressure_mat_utils import vis_heatmap

class MetricsCollector:
    def __init__(self):
        self.ious = []
        self.cop_errors = []
        self.frame_diffs = []
        self.binary_ths = []
        self.best_cop_w = []
        self.best_cop_k = []
        self.pred_heatmaps = []
        self.pred_cops_relative = []

    def assign(self, *, iou, cop_error, frame_diff, binary_th, best_cop_w, best_cop_k, pred_heatmap, pred_cop_relative):
        self.ious.append(iou)
        self.cop_errors.append(cop_error)
        self.frame_diffs.append(frame_diff)
        self.binary_ths.append(binary_th)
        self.best_cop_w.append(best_cop_w)
        self.best_cop_k.append(best_cop_k)
        self.pred_heatmaps.append(pred_heatmap)
        self.pred_cops_relative.append(pred_cop_relative)


class BiomechanicalEvaluator(StabilityLossCoS):
    def __init__(self,
                 faces,
                 cop_w=10,
                 cop_k=100,
                 contact_thresh=CONTACT_THRESH,
                 model_type='smplx',
                 ):
        super().__init__(faces, cop_w, cop_k, contact_thresh, model_type)

        self.iou_first = MetricsCollector()
        self.cop_error_first = MetricsCollector()
        self.frame_diff_first = MetricsCollector()

        self.cop_k_range = np.linspace(10, 200, 10).astype(np.float64)
        self.cop_w_range = np.linspace(10, 5000, 15).astype(np.float64)

    def generate_aligned_heatmap(self, vertices, mat_size, marker_positions,
                                 ground_plane_height=0.0, heatmap_res=512, cop_w=10, cop_k=100, vp=0.1,
                                 com_type='volume', cop_type='pressure', debug=False):
        """
        get vertices height along *z-axis* from ground-plane and pass it through function e(-wx) to get values for the heatmap.
        align and crop the heatmap so that it is it matches the Gt pressure map
        Args:
            mat_size: size of the pressure mat
            marker_positions: bottom left, top left, top right, bottom right
            cop_w: paramaeter for contact function
            vp: view padding empty space around contact region (in m)
        """
        # map 4 x 4 m area to a 512x512 image
        pressure_map = np.zeros((mat_size[1], mat_size[0]), dtype=np.float32)
        heatmap_point = np.zeros((heatmap_res, heatmap_res), dtype=np.float32)

        # uniformly sample vertices from the SMPL mesh so that the contact vertices are not biased to hands and face
        vertices = torch.tensor(vertices)[None, :, :]
        vertices_hd = HDfier().hdfy_mesh(vertices).cpu().numpy()[0]
        vertices = vertices.cpu().numpy()[0]

        # calculate values for heatmap
        vertex_height = vertices_hd[:, 2] - ground_plane_height
        vertex_height = vertex_height.astype(np.float64)

        # Get metric range. x-z plane is ground
        mat_bl_corner = marker_positions[0, :2] / 1000
        mat_tl_corner = marker_positions[1, :2] / 1000
        mat_tr_corner = marker_positions[2, :2] / 1000
        mat_br_corner = marker_positions[3, :2] / 1000
        m_x = mat_br_corner[0], mat_bl_corner[0]
        m_y = mat_tr_corner[1], mat_br_corner[1]
        m_range_x = m_x[1] - m_x[0]
        m_range_y = m_y[1] - m_y[0]

        # filter out vertices outside of the mat
        mask = (vertices_hd[:, 0] >= m_x[0]) & (vertices_hd[:, 0] <= m_x[1]) & (vertices_hd[:, 1] >= m_y[0]) & (
                vertices_hd[:, 1] <= m_y[1])
        vertex_height = vertex_height[mask]
        vertices_hd = vertices_hd[mask]

        # Normalize metric values to image range.
        max_mat_idx_x = mat_size[0] - 1
        max_mat_idx_y = mat_size[1] - 1
        v_x = np.rint((m_x[1] - vertices_hd[:, 0]) / m_range_x * (max_mat_idx_x)).astype(int) # flipped because global x axis is opposite of pressure mat x-axis
        v_y = np.rint((m_y[1] - vertices_hd[:, 1]) / m_range_y * (max_mat_idx_y)).astype(
            int)  # flipped because mat goes from bottom to up while image goes from top to bottom

        # assymetric function for inside and outside vertices
        inside_mask = (vertex_height < 0)
        outside_mask = (vertex_height >= 0)
        # v_z = inside_mask * np.exp(-4 * cop_w * vertex_height) + outside_mask * np.exp(-cop_w * vertex_height)
        v_z = inside_mask * (1 - cop_k * vertex_height) + outside_mask * np.exp(-cop_w * vertex_height)
        # v_z = np.exp(-cop_w*vertex_height)

        # normalize v_z to [0, 1]
        v_z = (v_z - np.min(v_z)) / (np.max(v_z) - np.min(v_z))

        # foe each overlapping vertex, add the value to the heatmap
        for i in range(len(v_x)):
            pressure_map[v_y[i], v_x[i]] += v_z[i]
        # normalize heatmap to [0, 1]
        pressure_map = (pressure_map - np.min(pressure_map)) / (np.max(pressure_map) - np.min(pressure_map))

        # flip the pressure_map because the GT is inverted
        # pressure_map = np.flipud(pressure_map)
        heatmap = vis_heatmap(pressure_map)
        return pressure_map, heatmap

    def iou(self, gt_pressure, pred_pressure):
        """
        Intersection over union
        :return:
        """
        binary_threholds = np.linspace(0, 1, 10)
        ious = []
        for th in binary_threholds:
            pred_binary = pred_pressure > th
            gt_binary = gt_pressure > th
            intersection = np.logical_and(pred_binary, gt_binary)
            union = np.logical_or(pred_binary, gt_binary)
            eps = 1e-13
            iou_score = np.sum(intersection) / (np.sum(union) + eps)
            ious.append(iou_score)
        max_iou_score = np.max(ious)
        max_th_idx = np.argmax(ious)
        max_th = binary_threholds[max_th_idx]
        return max_iou_score, max_th

    def draw_ious_graph(self):
        from numpy import exp, arange
        from pylab import meshgrid, cm, imshow, contour, clabel, colorbar, axis, title, show

        # the function that I'm going to plot
        x = arange(-3.0, 3.0, 0.1)
        y = arange(-3.0, 3.0, 0.1)
        X, Y = meshgrid(x, y)  # grid of point
        Z = z_func(X, Y)  # evaluation of the function on the grid

        im = imshow(Z, cmap=cm.RdBu)  # drawing the function
        # adding the Contour lines with labels
        cset = contour(Z, arange(-1, 1.5, 0.2), linewidths=2, cmap=cm.Set2)
        clabel(cset, inline=True, fmt='%1.1f', fontsize=10)
        colorbar(im)  # adding the colobar on the right
        # latex fashion title
        title('$z=(1-x^2+y^3) e^{-(x^2+y^2)/2}$')
        show()

    def evaluate_com(self, gt_com, vertices):
        # Note: the vertices should be aligned along y-axis and in world coordinates
        batch_size = vertices.shape[0]
        vertices = vertices.float()
        # calculate per part volume
        per_part_volume = self.compute_per_part_volume(vertices)
        # sample 20k vertices uniformly on the smpl mesh
        vertices_hd = self.hdfy_op.hdfy_mesh(vertices)
        # get volume per vertex id in the hd mesh
        volume_per_vert_hd = self.vertex_id_to_part_volume_mapping(per_part_volume, vertices.device)
        # calculate com using volume weighted mean
        com = torch.sum(vertices_hd * volume_per_vert_hd, dim=1) / torch.sum(volume_per_vert_hd, dim=1)
        com_error = torch.norm(com - gt_com, dim=1)
        return com, com_error

    def evaluate_pressure(self, gt_pressure, gt_cop_relative, vertices, mat_size, mat_size_global, marker_positions):
        """
        Evaluate the predicted pressure map against the ground truth pressure map.
        Args:
            gt_pressure:
            gt_cop_relative: cop relative to the pressure mat in mms
            vertices:
            mat_size: the resolution of the mat image
            mat_bbox: [tl_x, tl_y, br_x, br_y] of the mat in the image

        Returns:

        """
        # Note: the vertices should be aligned along y-axis and in world coordinates
        batch_size = vertices.shape[0]
        vertices = vertices.float()
        # calculate per part volume
        per_part_volume = self.compute_per_part_volume(vertices)
        # sample 20k vertices uniformly on the smpl mesh
        vertices_hd = self.hdfy_op.hdfy_mesh(vertices)
        # # get volume per vertex id in the hd mesh
        # volume_per_vert_hd = self.vertex_id_to_part_volume_mapping(per_part_volume, vertices.device)
        # # calculate com using volume weighted mean
        # com = torch.sum(vertices_hd * volume_per_vert_hd, dim=1) / torch.sum(volume_per_vert_hd, dim=1)

        # pressure based center of support
        ground_plane_height = 0.0
        eps = 1e-6
        vertex_height = (vertices_hd[:, :, 2] - ground_plane_height).double()
        inside_mask = (vertex_height < 0.0).float()
        outside_mask = (vertex_height >= 0.0).float()

        iou_first_best_iou = 0.0 # best iou score
        cop_error_first_best_cop_error = np.inf # best cop error
        frame_diff_first_best_frame_diff = np.inf
        for cop_k in self.cop_k_range:
            for cop_w in self.cop_w_range:
                pressure_weights = inside_mask * (1 - cop_k * vertex_height) + outside_mask * torch.exp(
                    - cop_w * vertex_height)
                pred_cos_global = torch.sum(vertices_hd * pressure_weights.unsqueeze(-1), dim=1) / (
                        torch.sum(pressure_weights, dim=1, keepdim=True) + eps)

                ## convert pred_cos relative to the pressure mat
                # Get metric range. x-z plane is ground
                mat_bl_corner = marker_positions[0, :2] / 1000
                mat_tl_corner = marker_positions[1, :2] / 1000
                mat_tr_corner = marker_positions[2, :2] / 1000
                mat_br_corner = marker_positions[3, :2] / 1000
                m_x = mat_br_corner[0], mat_bl_corner[0]
                m_y = mat_tr_corner[1], mat_br_corner[1]
                m_range_x = m_x[1] - m_x[0]
                m_range_y = m_y[1] - m_y[0]

                # convert all vertices relative to the pressure mat
                mat_size_global_x = mat_size_global[0]
                mat_size_global_y = mat_size_global[1]
                v_x_relative = (m_x[1] - vertices_hd[:, :, 0]) / m_range_x * (mat_size_global_x) # flipped because global x axis is opposite of pressure mat x-axis
                v_y_relative = (m_y[1] - vertices_hd[:, :, 1]) / m_range_y * (mat_size_global_y) # flipped because mat goes from bottom to up while image goes from top to bottom
                vertices_hd_relative = torch.stack([v_x_relative, v_y_relative, vertices_hd[:, :, 2]], dim=2)
                pred_cos_relative = torch.sum(vertices_hd_relative * pressure_weights.unsqueeze(-1), dim=1) / (
                        torch.sum(pressure_weights, dim=1, keepdim=True) + eps)
                pred_cos_relative = pred_cos_relative[0, :2].cpu().numpy()

                # compute l2 distance between gt and pred cop
                cop_error = np.linalg.norm(pred_cos_relative - gt_cop_relative, axis=0) # in metres

                # get pressure heatmap
                pred_pressure, pred_heatmap = self.generate_aligned_heatmap(vertices, mat_size, marker_positions,
                                                                            cop_w=cop_w, cop_k=cop_k)

                # calculate intersection-over-union between gt_pressure and predicted pressure map
                iou, binary_th = self.iou(gt_pressure, pred_pressure)
                # calculate mean frame diff
                frame_diff = np.mean(np.abs(gt_pressure - pred_pressure))
                if frame_diff <= frame_diff_first_best_frame_diff:
                    frame_diff_first_best_iou = iou
                    frame_diff_first_best_th = binary_th
                    frame_diff_first_best_cop_error = cop_error
                    frame_diff_first_best_frame_diff = frame_diff
                    frame_diff_first_best_pred_heatmap = pred_heatmap
                    frame_diff_first_best_pred_cop_relative = pred_cos_relative
                    frame_diff_first_best_cop_w = cop_w
                    frame_diff_first_best_cop_k = cop_k
                if iou >= iou_first_best_iou:
                    iou_first_best_iou = iou
                    iou_first_best_cop_error = cop_error
                    iou_first_best_frame_diff = frame_diff
                    iou_first_best_th = binary_th
                    iou_first_best_pred_heatmap = pred_heatmap
                    iou_first_best_pred_cop_relative = pred_cos_relative
                    iou_first_best_cop_w = cop_w
                    iou_first_best_cop_k = cop_k
                if cop_error <= cop_error_first_best_cop_error:
                    cop_error_first_best_iou = iou
                    cop_error_first_best_cop_error = cop_error
                    cop_error_first_best_frame_diff = frame_diff
                    cop_error_first_best_th = binary_th
                    cop_error_first_best_pred_heatmap = pred_heatmap
                    cop_error_first_best_pred_cop_relative = pred_cos_relative
                    cop_error_first_best_cop_w = cop_w
                    cop_error_first_best_cop_k = cop_k

        self.frame_diff_first.assign(iou=frame_diff_first_best_iou,
                                     cop_error=frame_diff_first_best_cop_error,
                                     frame_diff=frame_diff_first_best_frame_diff,
                                     binary_th=frame_diff_first_best_th,
                                     best_cop_w=frame_diff_first_best_cop_w,
                                     best_cop_k=frame_diff_first_best_cop_k,
                                     pred_heatmap=frame_diff_first_best_pred_heatmap,
                                     pred_cop_relative=frame_diff_first_best_pred_cop_relative)

        self.iou_first.assign(iou=iou_first_best_iou,
                                      cop_error=iou_first_best_cop_error,
                                      frame_diff = iou_first_best_frame_diff,
                                      binary_th=iou_first_best_th,
                                        best_cop_w=iou_first_best_cop_w,
                                        best_cop_k=iou_first_best_cop_k,
                                      pred_heatmap=iou_first_best_pred_heatmap,
                                      pred_cop_relative=iou_first_best_pred_cop_relative)
        self.cop_error_first.assign(iou=cop_error_first_best_iou,
                                            cop_error=cop_error_first_best_cop_error,
                                            frame_diff=cop_error_first_best_frame_diff,
                                            binary_th=cop_error_first_best_th,
                                            best_cop_w=cop_error_first_best_cop_w,
                                            best_cop_k=cop_error_first_best_cop_k,
                                            pred_heatmap=cop_error_first_best_pred_heatmap,
                                            pred_cop_relative=cop_error_first_best_pred_cop_relative)

moyo/utils/biomech_eval_utils.py
sha2nkt/moyo_toolkit
MOYO_V_TEMPLATE = '../data/v_templates/220923_yogi_03596_minimal_simple_female/mesh.ply'
IOI_PATH = '/is/cluster/scratch/stripathi/pycharm_remote/yogi/data/Yogi-IOI/yoga_dataset_images'
ESSENTIALS_DIR = '/is/cluster/scratch/stripathi/pycharm_remote/yogi/smplify-xmc-support/data/essentials'
SMPLX_VERT_SEG = '/is/cluster/scratch/stripathi/pycharm_remote/yogi/utils/meshscripts/smplx_vert_segmentation.json'
SMPL_VERT_SEG = '/is/cluster/scratch/stripathi/pycharm_remote/yogi/utils/meshscripts/smpl_vert_segmentation.json'
SEGMENTS_DIR = '/is/cluster/scratch/stripathi/pycharm_remote/yogi/smplify-xmc-support/data/essentials/yogi_segments'

CONTACT_THRESH = 0.15  # gmof_rho in overlay vis = contact thresh

# for vicon captures, frame_num for ioi at 30 fps
frame_select_dict_combined = {
    'Boat_Pose_or_Paripurna_Navasana_-a': 147,
    'Boat_Pose_or_Paripurna_Navasana_-b': 168,
    'Boat_Pose_or_Paripurna_Navasana_-c': 185,
    'Boat_Pose_or_Paripurna_Navasana_-d': 147,
    'Boat_Pose_or_Paripurna_Navasana_-e': 205,
    'Boat_Pose_or_Paripurna_Navasana_-f': 165,
    'Bound_Angle_Pose_or_Baddha_Konasana_-b': 465,
    'Bound_Angle_Pose_or_Baddha_Konasana_-c': 237,
    'Bound_Angle_Pose_or_Baddha_Konasana_-d': 481,
    'Bound_Angle_Pose_or_Baddha_Konasana_-e': 438,
    'Bound_Angle_Pose_or_Baddha_Konasana_-f': 160,
    'Bound_Angle_Pose_or_Baddha_Konasana_-g': 240,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-a': 470,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-b': 644,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-c': 422,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-d': 506,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-e': 675,
    'Camel_Pose_or_Ustrasana_-a': 582,
    'Camel_Pose_or_Ustrasana_-b': 566,
    'Camel_Pose_or_Ustrasana_-c': 218,
    'Camel_Pose_or_Ustrasana_-d': 288,
    'Cat_Cow_Pose_or_Marjaryasana_-a': 239,
    'Cat_Cow_Pose_or_Marjaryasana_-b': 172,
    'Cat_Cow_Pose_or_Marjaryasana_-c': 256,
    'Chair_Pose_or_Utkatasana_-a': 151,
    'Chair_Pose_or_Utkatasana_-b': 203,
    'Chair_Pose_or_Utkatasana_-c': 231,
    'Child_Pose_or_Balasana_-a': 426,
    'Child_Pose_or_Balasana_-b': 301,
    'Cobra_Pose_or_Bhujangasana_-a': 278,
    'Cobra_Pose_or_Bhujangasana_-b': 294,
    'Cobra_Pose_or_Bhujangasana_-c': 551,
    'Cobra_Pose_or_Bhujangasana_-d': 281,
    'Cockerel_Pose-a': 1377,
    'Cockerel_Pose-b': 1202,
    'Corpse_Pose_or_Savasana_-a': 249,
    'Cow_Face_Pose_or_Gomukhasana_-b': 481,
    'Cow_Face_Pose_or_Gomukhasana_-c': 482,
    'Cow_Face_Pose_or_Gomukhasana_-e': 429,
    'Crane_(Crow)_Pose_or_Bakasana_-a': 258,
    'Crane_(Crow)_Pose_or_Bakasana_-b': 314,
    'Dolphin_Plank_Pose_or_Makara_Adho_Mukha_Svanasana_-a': 297,
    'Dolphin_Pose_or_Ardha_Pincha_Mayurasana_-a': 246,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-a': 158,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-b': 257,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-c': 165,
    'Extended_Puppy_Pose_or_Uttana_Shishosana_-a': 581,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-a': 246,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-b': 276,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-d': 353,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-e': 278,
    'Extended_Revolved_Triangle_Pose_or_Utthita_Trikonasana_-a': 312,
    'Extended_Revolved_Triangle_Pose_or_Utthita_Trikonasana_-b': 277,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-a': 788,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-b': 649,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-c': 832,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-d': 664,
    'Firefly_Pose_or_Tittibhasana_-a': 344,
    'Firefly_Pose_or_Tittibhasana_-b': 275,
    'Fish_Pose_or_Matsyasana_-c': 830,
    'Fish_Pose_or_Matsyasana_-d': 376,
    'Four-Limbed_Staff_Pose_or_Chaturanga_Dandasana_-a': 197,
    'Four-Limbed_Staff_Pose_or_Chaturanga_Dandasana_-b': 140,
    'Garland_Pose_or_Malasana_-a': 136,
    'Gate_Pose_or_Parighasana_-a': 393,
    'Half_Lord_of_the_Fishes_Pose_or_Ardha_Matsyendrasana_-a': 376,
    'Half_Moon_Pose_or_Ardha_Chandrasana_-a': 348,
    'Half_Moon_Pose_or_Ardha_Chandrasana_-b': 464,
    'Handstand_pose_or_Adho_Mukha_Vrksasana_-a': 1270,
    'Head-to-Knee_Forward_Bend_pose_or_Janu_Sirsasana_-a': 376,
    'Intense_Side_Stretch_Pose_or_Parsvottanasana_-a': 328,
    'Legs-Up-the-Wall_Pose_or_Viparita_Karani_-a': 264,
    'Locust_Pose_or_Salabhasana_-a': 276,
    'Locust_Pose_or_Salabhasana_-b': 268,
    'Low_Lunge_pose_or_Anjaneyasana_-a': 545,
    'Low_Lunge_pose_or_Anjaneyasana_-b': 432,
    'Low_Lunge_pose_or_Anjaneyasana_-c': 260,
    'Low_Lunge_pose_or_Anjaneyasana_-d': 380,
    'Noose_Pose_or_Pasasana_-b': 317,
    'Peacock_Pose_or_Mayurasana_-a': 603,
    'Peacock_Pose_or_Mayurasana_-b': 911,
    'Plank_Pose_or_Kumbhakasana_-a': 231,
    'Plow_Pose_or_Halasana_-a': 390,
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-a': 608,
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-b': 528,
    'Revolved_Head-to-Knee_Pose_or_Parivrtta_Janu_Sirsasana_-b': 404,
    'Scale_Pose_or_Tolasana_-a': 477,
    'Scorpion_pose_or_vrischikasana-a': 1203,
    'Scorpion_pose_or_vrischikasana-b': 589,
    'Scorpion_pose_or_vrischikasana-c': 653,
    'Shoulder-Pressing_Pose_or_Bhujapidasana_-a': 237,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-a': 397,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-b': 725,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-c': 730,
    'Side_Plank_Pose_or_Vasisthasana_-a': 409,
    'Side_Plank_Pose_or_Vasisthasana_-b': 361,
    'Side_Plank_Pose_or_Vasisthasana_-c': 276,
    'Side_Plank_Pose_or_Vasisthasana_-d': 361,
    'Side_Plank_Pose_or_Vasisthasana_-e': 386,
    'Sitting_pose_1_(normal)-a': 309,
    'Sitting_pose_1_(normal)-b': 223,
    'Sitting_pose_1_(normal)-c': 216,
    'Split_pose-a': 499,
    'Split_pose-b': 415,
    'Split_pose-c': 687,
    'Staff_Pose_or_Dandasana_-a': 282,
    'Staff_Pose_or_Dandasana_-b': 235,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-a': 422,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-b': 317,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-c': 333,
    'Standing_Forward_Bend_pose_or_Uttanasana_-a': 168,
    'Standing_Forward_Bend_pose_or_Uttanasana_-b': 135,
    'Standing_Split_pose_or_Urdhva_Prasarita_Eka_Padasana_-a': 224,
    'Standing_Split_pose_or_Urdhva_Prasarita_Eka_Padasana_-b': 227,
    'Supported_Headstand_pose_or_Salamba_Sirsasana_-b': 486,
    'Supported_Shoulderstand_pose_or_Salamba_Sarvangasana_-a': 415,
    'Supta_Baddha_Konasana_-b': 552,
    'Supta_Virasana_Vajrasana-c': 591,
    'Tree_Pose_or_Vrksasana_-a': 379,
    'Tree_Pose_or_Vrksasana_-b': 219,
    'Akarna_Dhanurasana-a': 569,
    'Akarna_Dhanurasana-b': 426,
    'Akarna_Dhanurasana-c': 323,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-a': 596,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-b': 490,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-c': 526,
    'Bound_Angle_Pose_or_Baddha_Konasana_-a': 480,
    'Bow_Pose_or_Dhanurasana_-a': 327,
    'Bow_Pose_or_Dhanurasana_-b': 590,
    'Cow_Face_Pose_or_Gomukhasana_-a': 565,
    'Cow_Face_Pose_or_Gomukhasana_-d': 537,
    'Eagle_Pose_or_Garudasana_-a': 372,
    'Eight-Angle_Pose_or_Astavakrasana_-a': 419,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-c': 235,
    'Fish_Pose_or_Matsyasana_-a': 830,
    'Fish_Pose_or_Matsyasana_-b': 757,
    'Frog_Pose_or_Bhekasana-a': 332,
    'Frog_Pose_or_Bhekasana-b': 559,
    'Frog_Pose_or_Bhekasana-c': 372,
    'Garland_Pose_or_Malasana_-b': 295,
    'Garland_Pose_or_Malasana_-c': 265,
    'Gate_Pose_or_Parighasana_-b': 364,
    'Half_Lord_of_the_Fishes_Pose_or_Ardha_Matsyendrasana_-b': 1032,
    'Half_Lord_of_the_Fishes_Pose_or_Ardha_Matsyendrasana_-c': 589,
    'Happy_Baby_Pose_or_Ananda_Balasana_-a': 443,
    'Head-to-Knee_Forward_Bend_pose_or_Janu_Sirsasana_-b': 432,
    'Heron_Pose_or_Krounchasana_-a': 320,
    'Heron_Pose_or_Krounchasana_-b': 505,
    'Heron_Pose_or_Krounchasana_-c': 370,
    'Intense_Side_Stretch_Pose_or_Parsvottanasana_-b': 460,
    'Locust_Pose_or_Salabhasana_-c': 477,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-a': 422,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-b': 1573,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-c': 679,
    'Noose_Pose_or_Pasasana_-a': 602,
    'Pigeon_Pose_or_Kapotasana_-a': 1360,
    'Plow_Pose_or_Halasana_-b': 794,
    'Rajakapotasana-a': 940,
    'Rajakapotasana-b': 1299,
    'Rajakapotasana-c': 642,
    'Rajakapotasana-d': 631,
    'Reclining_Hand-to-Big-Toe_Pose_or_Supta_Padangusthasana_-a': 321,
    'Reclining_Hand-to-Big-Toe_Pose_or_Supta_Padangusthasana_-b': 410,
    'Revolved_Head-to-Knee_Pose_or_Parivrtta_Janu_Sirsasana_-a': 628,
    'Seated_Forward_Bend_pose_or_Paschimottanasana_-a': 357,
    'Side-Reclining_Leg_Lift_pose_or_Anantasana_-a': 690,
    'Sun_Salutation_Flow ': 1000,
    'Supported_Headstand_pose_or_Salamba_Sirsasana_-a': 490,
    'Supta_Baddha_Konasana_-a': 563,
    'Supta_Virasana_Vajrasana-a': 495,
    'Supta_Virasana_Vajrasana-b': 562,
    'Supta_Virasana_Vajrasana-d': 526,
    'Tortoise_Pose-a': 1370,
    'Tortoise_Pose-b': 492,
    'Upward_Bow_(Wheel)_Pose_or_Urdhva_Dhanurasana_-a': 709,
    'Upward_Bow_(Wheel)_Pose_or_Urdhva_Dhanurasana_-b': 502,
    'Upward_Facing_Two-Foot_Staff_Pose_or_Dwi_Pada_Viparita_Dandasana_-a': 550,
    'Upward_Plank_Pose_or_Purvottanasana_-a': 695,
    'Upward_Plank_Pose_or_Purvottanasana_-b': 414,
    'viparita_virabhadrasana_or_reverse_warrior_pose-a': 372,
    'viparita_virabhadrasana_or_reverse_warrior_pose-b': 337,
    'Virasana_or_Vajrasana-a': 311,
    'Virasana_or_Vajrasana-b': 451,
    'Virasana_or_Vajrasana-c': 257,
    'Warrior_I_Pose_or_Virabhadrasana_I_-a': 289,
    'Warrior_II_Pose_or_Virabhadrasana_II_-a': 223,
    'Warrior_III_Pose_or_Virabhadrasana_III_-a': 423,
    'Warrior_III_Pose_or_Virabhadrasana_III_-b': 393,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-a': 599,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-b': 497,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-c': 1028,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-d': 1112,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-a': 566,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-b': 456,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-c': 463,
    'Wild_Thing_pose_or_Camatkarasana_-a': 395,
    'Wind_Relieving_pose_or_Pawanmuktasana-a': 394,
    'Wind_Relieving_pose_or_Pawanmuktasana-b': 443,
    'Yogic_sleep_pose-a': 1293,
    'Yogic_sleep_pose-b': 1217
}

# ioi frame nums at 30 fps
frame_select_dict_first_session = {
    'Boat_Pose_or_Paripurna_Navasana_-a': 147,
    'Boat_Pose_or_Paripurna_Navasana_-b': 168,
    'Boat_Pose_or_Paripurna_Navasana_-c': 185,
    'Boat_Pose_or_Paripurna_Navasana_-d': 147,
    'Boat_Pose_or_Paripurna_Navasana_-e': 205,
    'Boat_Pose_or_Paripurna_Navasana_-f': 165,
    'Bound_Angle_Pose_or_Baddha_Konasana_-b': 465,
    'Bound_Angle_Pose_or_Baddha_Konasana_-c': 237,
    'Bound_Angle_Pose_or_Baddha_Konasana_-d': 481,
    'Bound_Angle_Pose_or_Baddha_Konasana_-e': 438,
    'Bound_Angle_Pose_or_Baddha_Konasana_-f': 160,
    'Bound_Angle_Pose_or_Baddha_Konasana_-g': 240,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-a': 470,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-b': 644,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-c': 422,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-d': 506,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-e': 675,
    'Camel_Pose_or_Ustrasana_-a': 582,
    'Camel_Pose_or_Ustrasana_-b': 566,
    'Camel_Pose_or_Ustrasana_-c': 218,
    'Camel_Pose_or_Ustrasana_-d': 288,
    'Cat_Cow_Pose_or_Marjaryasana_-a': 239,
    'Cat_Cow_Pose_or_Marjaryasana_-b': 172,
    'Cat_Cow_Pose_or_Marjaryasana_-c': 256,
    'Chair_Pose_or_Utkatasana_-a': 151,
    'Chair_Pose_or_Utkatasana_-b': 203,
    'Chair_Pose_or_Utkatasana_-c': 231,
    'Child_Pose_or_Balasana_-a': 426,
    'Child_Pose_or_Balasana_-b': 301,
    'Cobra_Pose_or_Bhujangasana_-a': 278,
    'Cobra_Pose_or_Bhujangasana_-b': 294,
    'Cobra_Pose_or_Bhujangasana_-c': 551,
    'Cobra_Pose_or_Bhujangasana_-d': 281,
    'Cockerel_Pose-a': 1377,
    'Cockerel_Pose-b': 1202,
    'Corpse_Pose_or_Savasana_-a': 249,
    'Cow_Face_Pose_or_Gomukhasana_-b': 481,
    'Cow_Face_Pose_or_Gomukhasana_-c': 782,
    'Cow_Face_Pose_or_Gomukhasana_-e': 429,
    'Crane_(Crow)_Pose_or_Bakasana_-a': 258,
    'Crane_(Crow)_Pose_or_Bakasana_-b': 314,
    'Dolphin_Plank_Pose_or_Makara_Adho_Mukha_Svanasana_-a': 297,
    'Dolphin_Pose_or_Ardha_Pincha_Mayurasana_-a': 246,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-a': 158,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-b': 257,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-c': 165,
    'Extended_Puppy_Pose_or_Uttana_Shishosana_-a': 581,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-a': 246,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-b': 276,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-d': 353,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-e': 278,
    'Extended_Revolved_Triangle_Pose_or_Utthita_Trikonasana_-a': 312,
    'Extended_Revolved_Triangle_Pose_or_Utthita_Trikonasana_-b': 277,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-a': 788,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-b': 649,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-c': 832,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-d': 664,
    'Firefly_Pose_or_Tittibhasana_-a': 344,
    'Firefly_Pose_or_Tittibhasana_-b': 275,
    'Fish_Pose_or_Matsyasana_-c': 830,
    'Fish_Pose_or_Matsyasana_-d': 376,
    'Four-Limbed_Staff_Pose_or_Chaturanga_Dandasana_-a': 197,
    'Four-Limbed_Staff_Pose_or_Chaturanga_Dandasana_-b': 140,
    'Garland_Pose_or_Malasana_-a': 136,
    'Gate_Pose_or_Parighasana_-a': 393,
    'Half_Lord_of_the_Fishes_Pose_or_Ardha_Matsyendrasana_-a': 376,
    'Half_Moon_Pose_or_Ardha_Chandrasana_-a': 348,
    'Half_Moon_Pose_or_Ardha_Chandrasana_-b': 464,
    'Handstand_pose_or_Adho_Mukha_Vrksasana_-a': 1270,
    'Head-to-Knee_Forward_Bend_pose_or_Janu_Sirsasana_-a': 376,
    'Intense_Side_Stretch_Pose_or_Parsvottanasana_-a': 328,
    'Legs-Up-the-Wall_Pose_or_Viparita_Karani_-a': 264,
    'Locust_Pose_or_Salabhasana_-a': 276,
    'Locust_Pose_or_Salabhasana_-b': 268,
    'Low_Lunge_pose_or_Anjaneyasana_-a': 545,
    'Low_Lunge_pose_or_Anjaneyasana_-b': 432,
    'Low_Lunge_pose_or_Anjaneyasana_-c': 260,
    'Low_Lunge_pose_or_Anjaneyasana_-d': 380,
    'Noose_Pose_or_Pasasana_-b': 317,
    'Peacock_Pose_or_Mayurasana_-a': 603,
    'Peacock_Pose_or_Mayurasana_-b': 911,
    'Plank_Pose_or_Kumbhakasana_-a': 231,
    'Plow_Pose_or_Halasana_-a': 390,
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-a': 608,
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-b': 528,
    'Revolved_Head-to-Knee_Pose_or_Parivrtta_Janu_Sirsasana_-b': 404,
    'Scale_Pose_or_Tolasana_-a': 477,
    'Scorpion_pose_or_vrischikasana-a': 1203,
    'Scorpion_pose_or_vrischikasana-b': 589,
    'Scorpion_pose_or_vrischikasana-c': 653,
    'Shoulder-Pressing_Pose_or_Bhujapidasana_-a': 237,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-a': 397,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-b': 725,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-c': 730,
    'Side_Plank_Pose_or_Vasisthasana_-a': 409,
    'Side_Plank_Pose_or_Vasisthasana_-b': 361,
    'Side_Plank_Pose_or_Vasisthasana_-c': 276,
    'Side_Plank_Pose_or_Vasisthasana_-d': 361,
    'Side_Plank_Pose_or_Vasisthasana_-e': 386,
    'Sitting_pose_1_(normal)-a': 309,
    'Sitting_pose_1_(normal)-b': 223,
    'Sitting_pose_1_(normal)-c': 216,
    'Split_pose-a': 499,
    'Split_pose-b': 415,
    'Split_pose-c': 687,
    'Staff_Pose_or_Dandasana_-a': 282,
    'Staff_Pose_or_Dandasana_-b': 235,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-a': 422,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-b': 317,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-c': 333,
    'Standing_Forward_Bend_pose_or_Uttanasana_-a': 168,
    'Standing_Forward_Bend_pose_or_Uttanasana_-b': 135,
    'Standing_Split_pose_or_Urdhva_Prasarita_Eka_Padasana_-a': 224,
    'Standing_Split_pose_or_Urdhva_Prasarita_Eka_Padasana_-b': 227,
    'Supported_Headstand_pose_or_Salamba_Sirsasana_-b': 486,
    'Supported_Shoulderstand_pose_or_Salamba_Sarvangasana_-a': 415,
    'Supta_Baddha_Konasana_-b': 552,
    'Supta_Virasana_Vajrasana-c': 591,
    'Tree_Pose_or_Vrksasana_-a': 379,
    'Tree_Pose_or_Vrksasana_-b': 219
}

# ioi frame nums at 30 fps
frame_select_dict_second_session = {
    'Akarna_Dhanurasana-a': 569,
    'Akarna_Dhanurasana-b': 426,
    'Akarna_Dhanurasana-c': 323,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-a': 596,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-b': 490,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-c': 526,
    'Bound_Angle_Pose_or_Baddha_Konasana_-a': 480,
    'Bow_Pose_or_Dhanurasana_-a': 327,
    'Bow_Pose_or_Dhanurasana_-b': 590,
    'Cow_Face_Pose_or_Gomukhasana_-a': 565,
    'Cow_Face_Pose_or_Gomukhasana_-c': 482,
    'Cow_Face_Pose_or_Gomukhasana_-d': 537,
    'Eagle_Pose_or_Garudasana_-a': 372,
    'Eight-Angle_Pose_or_Astavakrasana_-a': 419,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-c': 235,
    'Fish_Pose_or_Matsyasana_-a': 830,
    'Fish_Pose_or_Matsyasana_-b': 757,
    'Frog_Pose_or_Bhekasana-a': 332,
    'Frog_Pose_or_Bhekasana-b': 559,
    'Frog_Pose_or_Bhekasana-c': 372,
    'Garland_Pose_or_Malasana_-b': 295,
    'Garland_Pose_or_Malasana_-c': 265,
    'Gate_Pose_or_Parighasana_-b': 364,
    'Half_Lord_of_the_Fishes_Pose_or_Ardha_Matsyendrasana_-b': 1032,
    'Half_Lord_of_the_Fishes_Pose_or_Ardha_Matsyendrasana_-c': 589,
    'Happy_Baby_Pose_or_Ananda_Balasana_-a': 443,
    'Head-to-Knee_Forward_Bend_pose_or_Janu_Sirsasana_-b': 432,
    'Heron_Pose_or_Krounchasana_-a': 320,
    'Heron_Pose_or_Krounchasana_-b': 505,
    'Heron_Pose_or_Krounchasana_-c': 370,
    'Intense_Side_Stretch_Pose_or_Parsvottanasana_-b': 460,
    'Locust_Pose_or_Salabhasana_-c': 477,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-a': 422,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-b': 1573,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-c': 679,
    'Noose_Pose_or_Pasasana_-a': 602,
    'Pigeon_Pose_or_Kapotasana_-a': 1360,
    'Plow_Pose_or_Halasana_-b': 794,
    'Rajakapotasana-a': 940,
    'Rajakapotasana-b': 1299,
    'Rajakapotasana-c': 642,
    'Rajakapotasana-d': 631,
    'Reclining_Hand-to-Big-Toe_Pose_or_Supta_Padangusthasana_-a': 321,
    'Reclining_Hand-to-Big-Toe_Pose_or_Supta_Padangusthasana_-b': 410,
    'Revolved_Head-to-Knee_Pose_or_Parivrtta_Janu_Sirsasana_-a': 628,
    'Seated_Forward_Bend_pose_or_Paschimottanasana_-a': 357,
    'Side-Reclining_Leg_Lift_pose_or_Anantasana_-a': 690,
    'Sun_Salutation_Flow ': 1000,
    'Supported_Headstand_pose_or_Salamba_Sirsasana_-a': 490,
    'Supta_Baddha_Konasana_-a': 563,
    'Supta_Virasana_Vajrasana-a': 495,
    'Supta_Virasana_Vajrasana-b': 562,
    'Supta_Virasana_Vajrasana-d': 526,
    'Tortoise_Pose-a': 1370,
    'Tortoise_Pose-b': 492,
    'Upward_Bow_(Wheel)_Pose_or_Urdhva_Dhanurasana_-a': 709,
    'Upward_Bow_(Wheel)_Pose_or_Urdhva_Dhanurasana_-b': 502,
    'Upward_Facing_Two-Foot_Staff_Pose_or_Dwi_Pada_Viparita_Dandasana_-a': 550,
    'Upward_Plank_Pose_or_Purvottanasana_-a': 695,
    'Upward_Plank_Pose_or_Purvottanasana_-b': 414,
    'viparita_virabhadrasana_or_reverse_warrior_pose-a': 372,
    'viparita_virabhadrasana_or_reverse_warrior_pose-b': 337,
    'Virasana_or_Vajrasana-a': 311,
    'Virasana_or_Vajrasana-b': 451,
    'Virasana_or_Vajrasana-c': 257,
    'Warrior_I_Pose_or_Virabhadrasana_I_-a': 289,
    'Warrior_II_Pose_or_Virabhadrasana_II_-a': 223,
    'Warrior_III_Pose_or_Virabhadrasana_III_-a': 423,
    'Warrior_III_Pose_or_Virabhadrasana_III_-b': 393,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-a': 599,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-b': 497,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-c': 1028,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-d': 1112,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-a': 566,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-b': 456,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-c': 463,
    'Wild_Thing_pose_or_Camatkarasana_-a': 395,
    'Wind_Relieving_pose_or_Pawanmuktasana-a': 394,
    'Wind_Relieving_pose_or_Pawanmuktasana-b': 443,
    'Yogic_sleep_pose-a': 1293,
    'Yogic_sleep_pose-b': 1217
}

# ioi frame nums at 30 fps. Note: for dynamic movements, we choose an arbitrary frame = 1000 (e.g. Sun_Salutation_Flow, 'Front-Back Walking')
frame_select_dict_nexus = {
    'Akarna_Dhanurasana-a': 321,
    'Akarna_Dhanurasana-b': 312,
    'Akarna_Dhanurasana-c': 320,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-a': 738,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-b': 464,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-c': 458,
    'Boat_Pose_or_Paripurna_Navasana_-a': 156,
    'Boat_Pose_or_Paripurna_Navasana_-b': 140,
    'Boat_Pose_or_Paripurna_Navasana_-c': 192,
    'Boat_Pose_or_Paripurna_Navasana_-d': 171,
    'Boat_Pose_or_Paripurna_Navasana_-e': 244,
    'Boat_Pose_or_Paripurna_Navasana_-f': 195,
    'Bound_Angle_Pose_or_Baddha_Konasana_-a': 423,
    'Bound_Angle_Pose_or_Baddha_Konasana_-b': 464,
    'Bound_Angle_Pose_or_Baddha_Konasana_-c': 379,
    'Bound_Angle_Pose_or_Baddha_Konasana_-d': 457,
    'Bound_Angle_Pose_or_Baddha_Konasana_-e': 388,
    'Bound_Angle_Pose_or_Baddha_Konasana_-f': 209,
    'Bound_Angle_Pose_or_Baddha_Konasana_-g': 238,
    'Bow_Pose_or_Dhanurasana_-a': 320,
    'Bow_Pose_or_Dhanurasana_-b': 316,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-a': 454,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-b': 434,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-c': 408,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-d': 545,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-e': 660,
    'Camel_Pose_or_Ustrasana_-a': 425,
    'Camel_Pose_or_Ustrasana_-b': 406,
    'Camel_Pose_or_Ustrasana_-c': 398,
    'Camel_Pose_or_Ustrasana_-d': 467,
    'Cat_Cow_Pose_or_Marjaryasana_-a': 275,
    'Cat_Cow_Pose_or_Marjaryasana_-b': 260,
    'Cat_Cow_Pose_or_Marjaryasana_-c': 1000,
    'Chair_Pose_or_Utkatasana_-a': 235,
    'Chair_Pose_or_Utkatasana_-b': 337,
    'Chair_Pose_or_Utkatasana_-c': 295,
    'Child_Pose_or_Balasana_-a': 351,
    'Child_Pose_or_Balasana_-b': 238,
    'Circle_Around': 1000,
    'Cobra_Pose_or_Bhujangasana_-a': 338,
    'Cobra_Pose_or_Bhujangasana_-b': 348,
    'Cobra_Pose_or_Bhujangasana_-c': 334,
    'Cobra_Pose_or_Bhujangasana_-d': 283,
    'Cockerel_Pose-b': 1317,
    'Corpse_Pose_or_Savasana_-a': 327,
    'Cow_Face_Pose_or_Gomukhasana_-a': 641,
    'Cow_Face_Pose_or_Gomukhasana_-b': 348,
    'Cow_Face_Pose_or_Gomukhasana_-c': 524,
    'Cow_Face_Pose_or_Gomukhasana_-d': 636,
    'Cow_Face_Pose_or_Gomukhasana_-e': 379,
    'Crane_(Crow)_Pose_or_Bakasana_-a': 542,
    'Crane_(Crow)_Pose_or_Bakasana_-b': 290,
    'Dolphin_Plank_Pose_or_Makara_Adho_Mukha_Svanasana_-a': 307,
    'Dolphin_Pose_or_Ardha_Pincha_Mayurasana_-a': 266,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-a': 188,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-b': 299,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-c': 280,
    'Eagle_Pose_or_Garudasana_-a': 499,
    'Eight-Angle_Pose_or_Astavakrasana_-a': 258,
    'Extended_Puppy_Pose_or_Uttana_Shishosana_-a': 491,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-a': 321,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-b': 361,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-c': 288,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-d': 628,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-e': 448,
    'Extended_Revolved_Triangle_Pose_or_Utthita_Trikonasana_-a': 497,
    'Extended_Revolved_Triangle_Pose_or_Utthita_Trikonasana_-b': 417,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-a': 769,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-b': 703,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-c': 981,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-d': 482,
    'Firefly_Pose_or_Tittibhasana_-a': 324,
    'Firefly_Pose_or_Tittibhasana_-b': 245,
    'Fish_Pose_or_Matsyasana_-a': 643,
    'Fish_Pose_or_Matsyasana_-b': 824,
    'Fish_Pose_or_Matsyasana_-c': 695,
    'Fish_Pose_or_Matsyasana_-d': 355,
    'Four-Limbed_Staff_Pose_or_Chaturanga_Dandasana_-a': 254,
    'Four-Limbed_Staff_Pose_or_Chaturanga_Dandasana_-b': 200,
    'Frog_Pose_or_Bhekasana-a': 400,
    'Frog_Pose_or_Bhekasana-b': 437,
    'Frog_Pose_or_Bhekasana-c': 458,
    'Front_Back_Walking': 1000,
    'Garland_Pose_or_Malasana_-a': 180,
    'Garland_Pose_or_Malasana_-b': 256,
    'Garland_Pose_or_Malasana_-c': 240,
    'Gate_Pose_or_Parighasana_-a': 352,
    'Gate_Pose_or_Parighasana_-b': 468,
    'Half_Lord_of_the_Fishes_Pose_or_Ardha_Matsyendrasana_-a': 374,
    'Half_Lord_of_the_Fishes_Pose_or_Ardha_Matsyendrasana_-b': 481,
    'Half_Lord_of_the_Fishes_Pose_or_Ardha_Matsyendrasana_-c': 491,
    'Half_Moon_Pose_or_Ardha_Chandrasana_-a': 393,
    'Half_Moon_Pose_or_Ardha_Chandrasana_-b': 351,
    'Happy_Baby_Pose_or_Ananda_Balasana_-a': 314,
    'Head-to-Knee_Forward_Bend_pose_or_Janu_Sirsasana_-a': 426,
    'Heron_Pose_or_Krounchasana_-a': 480,
    'Heron_Pose_or_Krounchasana_-b': 413,
    'Heron_Pose_or_Krounchasana_-c': 350,
    'Intense_Side_Stretch_Pose_or_Parsvottanasana_-a': 343,
    'Intense_Side_Stretch_Pose_or_Parsvottanasana_-b': 337,
    'Legs-Up-the-Wall_Pose_or_Viparita_Karani_-a': 310,
    'Locust_Pose_or_Salabhasana_-a': 256,
    'Locust_Pose_or_Salabhasana_-b': 248,
    'Locust_Pose_or_Salabhasana_-c': 507,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-a': 488,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-b': 1121,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-c': 617,
    'Low_Lunge_pose_or_Anjaneyasana_-a': 379,
    'Low_Lunge_pose_or_Anjaneyasana_-b': 282,
    'Low_Lunge_pose_or_Anjaneyasana_-c': 345,
    'Low_Lunge_pose_or_Anjaneyasana_-d': 411,
    'Noose_Pose_or_Pasasana_-a': 678,
    'Noose_Pose_or_Pasasana_-b': 403,
    'Peacock_Pose_or_Mayurasana_-a': 699,
    'Peacock_Pose_or_Mayurasana_-b': 856,
    'Pigeon_Pose_or_Kapotasana_-a': 1468,
    'Plank_Pose_or_Kumbhakasana_-a': 229,
    'Plow_Pose_or_Halasana_-a': 489,
    'Plow_Pose_or_Halasana_-b': 593,
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-a': 711,
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-b': 314,
    'Rajakapotasana-a': 1012,
    'Rajakapotasana-b': 1200,
    'Rajakapotasana-c': 675,
    'Rajakapotasana-d': 652,
    'Random_Movements_Jumps': 1000,
    'Random_Movements_Squats': 1000,
    'Reclining_Hand-to-Big-Toe_Pose_or_Supta_Padangusthasana_-a': 368,
    'Reclining_Hand-to-Big-Toe_Pose_or_Supta_Padangusthasana_-b': 465,
    'Revolved_Head-to-Knee_Pose_or_Parivrtta_Janu_Sirsasana_-a': 470,
    'Revolved_Head-to-Knee_Pose_or_Parivrtta_Janu_Sirsasana_-b': 477,
    'Scale_Pose_or_Tolasana_-a': 487,
    'Scorpion_pose_or_vrischikasana-a': 839,
    'Scorpion_pose_or_vrischikasana-c': 664,
    'Seated_Forward_Bend_pose_or_Paschimottanasana_-a': 292,
    'Shoulder-Pressing_Pose_or_Bhujapidasana_-a': 299,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-a': 470,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-b': 414,
    'Side_Plank_Pose_or_Vasisthasana_-a': 365,
    'Side_Plank_Pose_or_Vasisthasana_-b': 427,
    'Side_Plank_Pose_or_Vasisthasana_-c': 276,
    'Side_Plank_Pose_or_Vasisthasana_-d': 388,
    'Side_Plank_Pose_or_Vasisthasana_-e': 367,
    'Side_Step': 1000,
    'Side-Reclining_Leg_Lift_pose_or_Anantasana_-a': 506,
    'Sitting_pose_1_(normal)-a': 201,
    'Sitting_pose_1_(normal)-b': 179,
    'Sitting_pose_1_(normal)-c': 223,
    'Split_pose-a': 699,
    'Split_pose-b': 404,
    'Split_pose-c': 800,
    'Staff_Pose_or_Dandasana_-a': 484,
    'Staff_Pose_or_Dandasana_-b': 262,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-a': 589,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-b': 408,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-c': 386,
    'Standing_Forward_Bend_pose_or_Uttanasana_-a': 253,
    'Standing_Forward_Bend_pose_or_Uttanasana_-b': 231,
    'Standing_Split_pose_or_Urdhva_Prasarita_Eka_Padasana_-a': 198,
    'Standing_Split_pose_or_Urdhva_Prasarita_Eka_Padasana_-b': 246,
    'Sun_Salutation_Flow': 1000,
    'Supported_Headstand_pose_or_Salamba_Sirsasana_-a': 408,
    'Supported_Headstand_pose_or_Salamba_Sirsasana_-b': 463,
    'Supported_Shoulderstand_pose_or_Salamba_Sarvangasana_-a': 432,
    'Supta_Baddha_Konasana_-a': 387,
    'Supta_Baddha_Konasana_-b': 661,
    'Supta_Virasana_Vajrasana-a': 467,
    'Supta_Virasana_Vajrasana-c': 601,
    'Tree_Pose_or_Vrksasana_-a': 320,
    'Tree_Pose_or_Vrksasana_-b': 266,
    'Upward_Bow_(Wheel)_Pose_or_Urdhva_Dhanurasana_-a': 475,
    'Upward_Bow_(Wheel)_Pose_or_Urdhva_Dhanurasana_-b': 446,
    'Upward_Facing_Two-Foot_Staff_Pose_or_Dwi_Pada_Viparita_Dandasana_-a': 702,
    'Upward_Plank_Pose_or_Purvottanasana_-a': 433,
    'Upward_Plank_Pose_or_Purvottanasana_-b': 270,
    'viparita_virabhadrasana_or_reverse_warrior_pose-a': 376,
    'viparita_virabhadrasana_or_reverse_warrior_pose-b': 368,
    'Virasana_or_Vajrasana-a': 276,
    'Virasana_or_Vajrasana-c': 207,
    'Warm_Up': 1000,
    'Warrior_I_Pose_or_Virabhadrasana_I_-a': 344,
    'Warrior_II_Pose_or_Virabhadrasana_II_-a': 322,
    'Warrior_III_Pose_or_Virabhadrasana_III_-a': 332,
    'Warrior_III_Pose_or_Virabhadrasana_III_-b': 328,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-a': 591,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-b': 547,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-c': 835,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-d': 706,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-a': 457,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-c': 552,
    'Wild_Thing_pose_or_Camatkarasana_-a': 498,
    'Wind_Relieving_pose_or_Pawanmuktasana-a': 302
}

# list of mismatch names: vicon_name:ioi_name
name_changes_first_session = {
    'Akarna_Dhanurasana-a': 'Akarna_Dhanurasana-b_1',
    'Akarna_Dhanurasana-b': 'Akarna_Dhanurasana-b_2',
    'Side-Reclining_Leg_Lift_pose_or_Anantasana_-a': 'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-a',
    'Rajakapotasana-c': 'check_name',
    'Side-Reclining_Leg_Lift_pose_or_Anantasana_-a': 'CoreView2',
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-a': 'Pose_Dedicated_to_the_Sage_Ko_or_Eka_Pada_Kound_I_and_II-a'

}

# list of rejected names with frame_nums useful in matching IOI with Vicon
reject_names_first_session = []

moyo/utils/constants.py
sha2nkt/moyo_toolkit
# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2021 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de

import torch
import cv2
import numpy as np
import math
import torch.nn.functional as F

def ea2rm(x, y, z):
    cos_x, sin_x = torch.cos(x), torch.sin(x)
    cos_y, sin_y = torch.cos(y), torch.sin(y)
    cos_z, sin_z = torch.cos(z), torch.sin(z)

    R = torch.stack(
            [torch.cat([cos_y*cos_z, sin_x*sin_y*cos_z - cos_x*sin_z, cos_x*sin_y*cos_z + sin_x*sin_z], dim=1),
            torch.cat([cos_y*sin_z, sin_x*sin_y*sin_z + cos_x*cos_z, cos_x*sin_y*sin_z - sin_x*cos_z], dim=1),
            torch.cat([-sin_y, sin_x*cos_y, cos_x*cos_y], dim=1)], dim=1)
    return R

def quaternion_to_axis_angle(quaternions):
    """
    Convert rotations given as quaternions to axis/angle.
    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).
    Returns:
        Rotations given as a vector in axis angle form, as a tensor
            of shape (..., 3), where the magnitude is the angle
            turned anticlockwise in radians around the vector's
            direction.
    """
    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
    half_angles = torch.atan2(norms, quaternions[..., :1])
    angles = 2 * half_angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    return quaternions[..., 1:] / sin_half_angles_over_angles

def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
    """
    Returns torch.sqrt(torch.max(0, x))
    but with a zero subgradient where x is 0.
    """
    ret = torch.zeros_like(x)
    positive_mask = x > 0
    ret[positive_mask] = torch.sqrt(x[positive_mask])
    return ret

def matrix_to_quaternion(matrix: torch.Tensor):
    batch_dim = matrix.shape[:-2]
    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(matrix.reshape(*batch_dim, 9), dim=-1)
    q_abs = _sqrt_positive_part(torch.stack([ 1.0 + m00 + m11 + m22, 1.0 + m00 - m11 - m22,  1.0 - m00 + m11 - m22, 1.0 - m00 - m11 + m22,],dim=-1,))
    # we produce the desired quaternion multiplied by each of r, i, j, k
    quat_by_rijk = torch.stack([torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),],dim=-2,)
    # clipping is not important here; if q_abs is small, the candidate won't be picked
    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None])
    # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
    # forall i; we pick the best-conditioned one (with the largest denominator)
    bla= quat_candidates[F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :].reshape(*batch_dim, 4)
    return bla

def matrix_to_axis_angle(matrix):
    """
    Convert rotations given as rotation matrices to axis/angle.
    Args:
        matrix: Rotation matrices as tensor of shape (..., 3, 3).
    Returns:
        Rotations given as a vector in axis angle form, as a tensor
            of shape (..., 3), where the magnitude is the angle
            turned anticlockwise in radians around the vector's
            direction.
    """
    return quaternion_to_axis_angle(matrix_to_quaternion(matrix))

def _angle_from_tan(
    axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool
):
    """
    Extract the first or third Euler angle from the two members of
    the matrix which are positive constant times its sine and cosine.

    Args:
        axis: Axis label "X" or "Y or "Z" for the angle we are finding.
        other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
            convention.
        data: Rotation matrices as tensor of shape (..., 3, 3).
        horizontal: Whether we are looking for the angle for the third axis,
            which means the relevant entries are in the same row of the
            rotation matrix. If not, they are in the same column.
        tait_bryan: Whether the first and third axes in the convention differ.

    Returns:
        Euler Angles in radians for each matrix in data as a tensor
        of shape (...).
    """

    i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
    if horizontal:
        i2, i1 = i1, i2
    even = (axis + other_axis) in ["XY", "YZ", "ZX"]
    if horizontal == even:
        return torch.atan2(data[..., i1], data[..., i2])
    if tait_bryan:
        return torch.atan2(-data[..., i2], data[..., i1])
    return torch.atan2(data[..., i2], -data[..., i1])


def _index_from_letter(letter: str):
    if letter == "X":
        return 0
    if letter == "Y":
        return 1
    if letter == "Z":
        return 2



def matrix_to_euler_angles(matrix, convention: str):
    """
    Convert rotations given as rotation matrices to Euler angles in radians.

    Args:
        matrix: Rotation matrices as tensor of shape (..., 3, 3).
        convention: Convention string of three uppercase letters.

    Returns:
        Euler angles in radians as tensor of shape (..., 3).
    """
    if len(convention) != 3:
        raise ValueError("Convention must have 3 letters.")
    if convention[1] in (convention[0], convention[2]):
        raise ValueError(f"Invalid convention {convention}.")
    for letter in convention:
        if letter not in ("X", "Y", "Z"):
            raise ValueError(f"Invalid letter {letter} in convention string.")
    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
        raise ValueError(f"Invalid rotation matrix  shape f{matrix.shape}.")
    i0 = _index_from_letter(convention[0])
    i2 = _index_from_letter(convention[2])
    tait_bryan = i0 != i2
    if tait_bryan:
        central_angle = torch.asin(
            matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0)
        )
    else:
        central_angle = torch.acos(matrix[..., i0, i0])

    o = (
        _angle_from_tan(
            convention[0], convention[1], matrix[..., i2], False, tait_bryan
        ),
        central_angle,
        _angle_from_tan(
            convention[2], convention[1], matrix[..., i0, :], True, tait_bryan
        ),
    )
    return torch.stack(o, -1)

def axis_angle_to_quaternion(axis_angle):
    """
    Convert rotations given as axis/angle to quaternions.

    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.

    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
    half_angles = 0.5 * angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions

def quaternion_to_matrix(quaternions):
    """
    Convert rotations given as quaternions to rotation matrices.

    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).

    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    r, i, j, k = torch.unbind(quaternions, -1)
    two_s = 2.0 / (quaternions * quaternions).sum(-1)

    o = torch.stack(
        (
            1 - two_s * (j * j + k * k),
            two_s * (i * j - k * r),
            two_s * (i * k + j * r),
            two_s * (i * j + k * r),
            1 - two_s * (i * i + k * k),
            two_s * (j * k - i * r),
            two_s * (i * k - j * r),
            two_s * (j * k + i * r),
            1 - two_s * (i * i + j * j),
        ),
        -1,
    )
    return o.reshape(quaternions.shape[:-1] + (3, 3))


def axis_angle_to_matrix(axis_angle):
    """
    Convert rotations given as axis/angle to rotation matrices.

    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.

    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))

def rotmat_from_euler_angles(pitch, roll, yaw):
    x = pitch
    y = yaw
    z = roll
    cos_x, sin_x = torch.cos(x), torch.sin(x)
    cos_y, sin_y = torch.cos(y), torch.sin(y)
    cos_z, sin_z = torch.cos(z), torch.sin(z)

    R = torch.stack(
            [torch.cat([cos_y*cos_z, sin_x*sin_y*cos_z - cos_x*sin_z, cos_x*sin_y*cos_z + sin_x*sin_z], dim=1),
            torch.cat([cos_y*sin_z, sin_x*sin_y*sin_z + cos_x*cos_z, cos_x*sin_y*sin_z - sin_x*cos_z], dim=1),
            torch.cat([-sin_y, sin_x*cos_y, cos_x*cos_y], dim=1)], dim=1)

    return R

def euler_angles_from_rotmat(R):
    """
    computer euler angles for rotation around x, y, z axis
    from rotation amtrix
    R: 4x4 rotation matrix
    https://www.gregslabaugh.net/publications/euler.pdf
    """
    r21 = np.round(R[:,2,0].item(), 4)
    if abs(r21) != 1:
        y_angle1 = -1 * torch.asin(R[:,2,0])
        y_angle2 = math.pi + torch.asin(R[:,2,0])
        cy1, cy2 = torch.cos(y_angle1), torch.cos(y_angle2)

        x_angle1 = torch.atan2(R[:,2,1] / cy1, R[:,2,2] / cy1)
        x_angle2 = torch.atan2(R[:,2,1] / cy2, R[:,2,2] / cy2)
        z_angle1 = torch.atan2(R[:,1,0] / cy1, R[:,0,0] / cy1)
        z_angle2 = torch.atan2(R[:,1,0] / cy2, R[:,0,0] / cy2)

        s1 = (x_angle1, y_angle1, z_angle1)
        s2 = (x_angle2, y_angle2, z_angle2)
        s = (s1, s2)

    else:
        z_angle = torch.tensor([0], device=R.device).float()
        if r21 == -1:
            y_angle = torch.tensor([math.pi / 2], device=R.device).float()
            x_angle = z_angle + torch.atan2(R[:,0,1], R[:,0,2])
        else:
            y_angle = -torch.tensor([math.pi / 2], device=R.device).float()
            x_angle = -z_angle + torch.atan2(-R[:,0,1], R[:,0,2])
        s = ((x_angle, y_angle, z_angle), )
    return s



def estimate_translation_np_old(S, joints_2d, joints_conf, focal_length_x=5000, focal_length_y=5000, W=224, H=224):
    """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
    Taken from: https://github.com/nkolot/SPIN/blob/master/utils/geometry.py
    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_x, focal_length_y])
    # optical center
    center = np.array([W/2., H/2.])
    # 2d joints mean
    p2d = joints_2d.mean(0)
    # 3d joints mean
    P3d = S.mean(0)[:2]

    trans_xy = (p2d-center) / f - P3d

    # 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)
    #trans[:2] = trans_xy

    return trans


def euler_angles_from_rotmat(R):
    """
    computer euler angles for rotation around x, y, z axis
    from rotation amtrix
    R: 4x4 rotation matrix
    https://www.gregslabaugh.net/publications/euler.pdf
    """
    r21 = np.round(R[:,2,0].item(), 4)
    if abs(r21) != 1:
        y_angle1 = -1 * torch.asin(R[:,2,0])
        y_angle2 = math.pi + torch.asin(R[:,2,0])
        cy1, cy2 = torch.cos(y_angle1), torch.cos(y_angle2)

        x_angle1 = torch.atan2(R[:,2,1] / cy1, R[:,2,2] / cy1)
        x_angle2 = torch.atan2(R[:,2,1] / cy2, R[:,2,2] / cy2)
        z_angle1 = torch.atan2(R[:,1,0] / cy1, R[:,0,0] / cy1)
        z_angle2 = torch.atan2(R[:,1,0] / cy2, R[:,0,0] / cy2)

        s1 = (x_angle1, y_angle1, z_angle1)
        s2 = (x_angle2, y_angle2, z_angle2)
        s = (s1, s2)

    else:
        z_angle = torch.tensor([0], device=R.device).float()
        if r21 == -1:
            y_angle = torch.tensor([math.pi / 2], device=R.device).float()
            x_angle = z_angle + torch.atan2(R[:,0,1], R[:,0,2])
        else:
            y_angle = -torch.tensor([math.pi / 2], device=R.device).float()
            x_angle = -z_angle + torch.atan2(-R[:,0,1], R[:,0,2])
        s = ((x_angle, y_angle, z_angle), )
    return s



def estimate_translation_np(S, joints_2d, joints_conf, 
    fx=5000, fy=5000, cx=112, cy=112, R=None, eqtype='optimize'):
    """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
    Taken from: https://github.com/nkolot/SPIN/blob/master/utils/geometry.py
    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([fx, fy])
    # optical center
    center = np.array([cx, cy])
    # rotation
    if R is None:
        R = np.eye(3)

    if eqtype == 'iter1':
        # camera matrix 
        K = np.array(
            [[fx, 0, cx], 
            [0, fy, cy],
            [0, 0, 1]]
        )
        K_inv = np.linalg.inv(K)

        RS = np.einsum('ij,bi->bj', R, S)
        RS = RS[:,:2] / RS[:,2, np.newaxis]

        Kx = np.einsum('ij,bi->bj', K_inv[:2, :2], joints_2d) + K_inv[:2,2]

        Q = Kx - RS

        t = Q.sum(0) / num_joints

        # get scale factor 
        sfs = []
        for idx in range(joints_2d.shape[0]):
            sf = joints_2d[idx] / (K[:2, :2] @ ((R @ S[idx])[:2] + t) + K[:2, 2])
            sfs += [sf]
        z = np.array(sfs).mean()
        trans = np.array([t[0] * z, t[1] * z, 1 / z])

    elif eqtype == 'iter2':
        RS = np.einsum('ij,bi->bj', R, S)

        # transformations
        Z = np.reshape(np.tile(RS[:,2],(2,1)).T,-1)
        XY = np.reshape(RS[:,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)

    else:
        # 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 optimize_translation(S, joints_2d, joints_conf, 
    fx=5000, fy=5000, cx=112, cy=112, R=None, eqtype='optimize'):
    """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
    Taken from: https://github.com/nkolot/SPIN/blob/master/utils/geometry.py
    Input:
        S: (25, 3) 3D joint locations
        joints: (25, 3) 2D joint locations and confidence
    Returns:
        (3,) camera translation vector
    """

    trans = None

    return trans

def estimate_translation_np_(S, joints_2d, joints_conf, focal_length_x=5000, focal_length_y=5000, W=224, H=224):
    """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
    Taken from: https://github.com/nkolot/SPIN/blob/master/utils/geometry.py
    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_x, focal_length_y])
    # optical center
    center = np.array([W/2., H/2.])
    # 2d joints mean
    p2d = joints_2d.mean(0)
    # 3d joints mean
    P3d = S.mean(0)[:2]

    trans_xy = (p2d-center) / f - P3d

    # 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)
    #trans[:2] = trans_xy

    return trans
moyo/utils/geometry.py
sha2nkt/moyo_toolkit
# borrowed from https://github.com/CZ-Wu/GPNet/blob/d8df3b1489800626d4f503f62d2ae8daffe8b686/tools/visualization.py

import os.path as osp

import numpy as np
import torch
import torch.nn as nn
import trimesh

from moyo.utils.constants import ESSENTIALS_DIR, MOYO_V_TEMPLATE


def smplx_breakdown(bdata, device):
    num_frames = len(bdata['trans'])

    bdata['poses'] = bdata['fullpose']

    global_orient = torch.from_numpy(bdata['poses'][:, :3]).float().to(device)
    body_pose = torch.from_numpy(bdata['poses'][:, 3:66]).float().to(device)
    jaw_pose = torch.from_numpy(bdata['poses'][:, 66:69]).float().to(device)
    leye_pose = torch.from_numpy(bdata['poses'][:, 69:72]).float().to(device)
    reye_pose = torch.from_numpy(bdata['poses'][:, 72:75]).float().to(device)
    left_hand_pose = torch.from_numpy(bdata['poses'][:, 75:120]).float().to(device)
    right_hand_pose = torch.from_numpy(bdata['poses'][:, 120:]).float().to(device)

    v_template = trimesh.load(MOYO_V_TEMPLATE, process=False)

    body_params = {'global_orient': global_orient, 'body_pose': body_pose,
                   'jaw_pose': jaw_pose, 'leye_pose': leye_pose, 'reye_pose': reye_pose,
                   'left_hand_pose': left_hand_pose, 'right_hand_pose': right_hand_pose,
                   'v_template': torch.Tensor(v_template.vertices).to(device), }
    return body_params


def batch_face_normals(triangles):
    # Calculate the edges of the triangles
    # Size: BxFx3
    edge0 = triangles[:, :, 1] - triangles[:, :, 0]
    edge1 = triangles[:, :, 2] - triangles[:, :, 0]
    # Compute the cross product of the edges to find the normal vector of
    # the triangle
    aCrossb = torch.cross(edge0, edge1, dim=2)
    # Normalize the result to get a unit vector
    normals = aCrossb / torch.norm(aCrossb, 2, dim=2, keepdim=True)

    return normals


def sparse_batch_mm(m1, m2):
    """
    https://github.com/pytorch/pytorch/issues/14489

    m1: sparse matrix of size N x M
    m2: dense matrix of size B x M x K
    returns m1@m2 matrix of size B x N x K
    """

    batch_size = m2.shape[0]
    # stack m2 into columns: (B x N x K) -> (N, B, K) -> (N, B * K)
    m2_stack = m2.transpose(0, 1).reshape(m1.shape[1], -1)
    result = m1.mm(m2_stack).reshape(m1.shape[0], batch_size, -1) \
        .transpose(1, 0)
    return result


class HDfier():
    def __init__(self, model_type='smplx'):
        hd_operator_path = osp.join(ESSENTIALS_DIR, 'hd_model', model_type,
                                    f'{model_type}_neutral_hd_vert_regressor_sparse.npz')
        hd_operator = np.load(hd_operator_path)
        self.hd_operator = torch.sparse.FloatTensor(
            torch.tensor(hd_operator['index_row_col']),
            torch.tensor(hd_operator['values']),
            torch.Size(hd_operator['size']))

    def hdfy_mesh(self, vertices, model_type='smplx'):
        """
        Applies a regressor that maps SMPL vertices to uniformly distributed vertices
        """
        # device = body.vertices.device
        # check if vertices ndim are 3, if not , add a new axis
        if vertices.ndim != 3:
            # batchify the vertices
            vertices = vertices[None, :, :]

        # check if vertices are an ndarry, if yes, make pytorch tensor
        if isinstance(vertices, np.ndarray):
            vertices = torch.from_numpy(vertices).to(self.device)

        vertices = vertices.to(torch.double)

        if self.hd_operator.device != vertices.device:
            self.hd_operator = self.hd_operator.to(vertices.device)
        hd_verts = sparse_batch_mm(self.hd_operator, vertices).to(torch.float)
        return hd_verts


def get_submesh(verts, faces, verts_retained=None, faces_retained=None, min_vert_in_face=2):
    '''
        Given a mesh, create a (smaller) submesh
        indicate faces or verts to retain as indices or boolean

        @return new_verts: the new array of 3D vertices
                new_faces: the new array of faces
                bool_faces: the faces indices wrt the input mesh
                vetex_ids: the vertex_ids wrt the input mesh
        '''

    if verts_retained is not None:
        # Transform indices into bool array
        if verts_retained.dtype != 'bool':
            vert_mask = np.zeros(len(verts), dtype=bool)
            vert_mask[verts_retained] = True
        else:
            vert_mask = verts_retained

        # Faces with at least min_vert_in_face vertices
        bool_faces = np.sum(vert_mask[faces.ravel()].reshape(-1, 3), axis=1) > min_vert_in_face

    elif faces_retained is not None:
        # Transform indices into bool array
        if faces_retained.dtype != 'bool':
            bool_faces = np.zeros(len(faces_retained), dtype=bool)
        else:
            bool_faces = faces_retained

    new_faces = faces[bool_faces]
    # just in case additional vertices are added
    vertex_ids = list(set(new_faces.ravel()))

    oldtonew = -1 * np.ones([len(verts)])
    oldtonew[vertex_ids] = range(0, len(vertex_ids))

    new_verts = verts[vertex_ids]
    new_faces = oldtonew[new_faces].astype('int32')

    return (new_verts, new_faces, bool_faces, vertex_ids)


def get_world_mesh_list(planeWidth=4, axisHeight=0.7, axisRadius=0.02, add_plane=True):
    groundColor = [220, 220, 220, 255]  # face_colors: [R, G, B, transparency]
    xColor = [255, 0, 0, 128]
    yColor = [0, 255, 0, 128]
    zColor = [0, 0, 255, 128]

    if add_plane:
        ground = trimesh.primitives.Box(
            center=[0, 0, -0.0001],
            extents=[planeWidth, planeWidth, 0.0002]
        )
        ground.visual.face_colors = groundColor

    xAxis = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    xAxis.apply_transform(matrix=np.mat(
        ((0, 0, 1, axisHeight / 2),
         (0, 1, 0, 0),
         (-1, 0, 0, 0),
         (0, 0, 0, 1))
    ))
    xAxis.visual.face_colors = xColor
    yAxis = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    yAxis.apply_transform(matrix=np.mat(
        ((1, 0, 0, 0),
         (0, 0, -1, axisHeight / 2),
         (0, 1, 0, 0),
         (0, 0, 0, 1))
    ))
    yAxis.visual.face_colors = yColor
    zAxis = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    zAxis.apply_transform(matrix=np.mat(
        ((1, 0, 0, 0),
         (0, 1, 0, 0),
         (0, 0, 1, axisHeight / 2),
         (0, 0, 0, 1))
    ))
    zAxis.visual.face_colors = zColor
    xBox = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    xBox.apply_translation((axisHeight, 0, 0))
    xBox.visual.face_colors = xColor
    yBox = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    yBox.apply_translation((0, axisHeight, 0))
    yBox.visual.face_colors = yColor
    zBox = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    zBox.apply_translation((0, 0, axisHeight))
    zBox.visual.face_colors = zColor
    if add_plane:
        worldMeshList = [ground, xAxis, yAxis, zAxis, xBox, yBox, zBox]
    else:
        worldMeshList = [xAxis, yAxis, zAxis, xBox, yBox, zBox]
    return worldMeshList


def get_checkerboard_plane(plane_width=4, num_boxes=9, center=True):
    pw = plane_width / num_boxes
    white = [220, 220, 220, 255]
    black = [35, 35, 35, 255]

    meshes = []
    for i in range(num_boxes):
        for j in range(num_boxes):
            c = i * pw, j * pw
            ground = trimesh.primitives.Box(
                center=[0, 0, -0.0001],
                extents=[pw, pw, 0.0002]
            )
            if center:
                c = c[0] + (pw / 2) - (plane_width / 2), c[1] + (pw / 2) - (plane_width / 2)
            # trans = trimesh.transformations.scale_and_translate(scale=1, translate=[c[0], c[1], 0])
            ground.apply_translation([c[0], c[1], 0])
            # ground.apply_transform(trimesh.transformations.rotation_matrix(np.rad2deg(-120), direction=[1,0,0]))
            ground.visual.face_colors = black if ((i + j) % 2) == 0 else white
            meshes.append(ground)

    # orient mesh be ground in pyrender axis system
    pose = trimesh.transformations.rotation_matrix(np.radians(90), [1, 0, 0])
    meshes = [mesh.apply_transform(pose) for mesh in meshes]
    return meshes


def get_meshed_plane(num_boxes=9):
    def makeVertexGrid(n, m):
        # vertices are incremented in X and Y but remain at zero in Z
        x = np.repeat(np.arange(-n / 2, n / 2), m)
        y = np.tile(np.arange(-m / 2, m / 2), n)
        z = np.zeros(n * m)
        grid = np.column_stack((x, y, z))
        return grid

    def makeFaceGrid(n, m):
        # Each face references three vertices by their index numbers, in counterclockwise order. for every vertex i not
        # in the right column or the top row, triangles [i, i+m,i+m+1] and [i, i+m+1, i+1] are created, where m is the
        # number of rows in the array.
        grid = np.zeros((2 * (m - 1) * (n - 1), 3))
        counter = 0
        for i in range((n - 1) * m):
            if i % m == m - 1:
                pass
            else:
                grid[counter] = [i, i + m, i + m + 1]
                grid[counter + 1] = [i, i + m + 1, i + 1]
                counter += 2
        return grid

    vertices = makeVertexGrid(num_boxes, num_boxes)
    faces = makeFaceGrid(num_boxes, num_boxes)

    # Add faces at the top of the mesh to avoid backface culling during visualization
    faces = np.vstack((faces, np.fliplr(faces)))

    white = [220, 220, 220, 255]
    black = [35, 35, 35, 255]
    face_colors = [black if i % 2 == 0 else white for i, _ in enumerate(faces)]
    ground = trimesh.Trimesh(vertices=vertices, faces=faces, face_colors=face_colors)
    # ground.visual.face_colors = white
    return ground


class GMoF(nn.Module):
    def __init__(self, rho=1):
        super(GMoF, self).__init__()
        self.rho = rho

    def extra_repr(self):
        return 'rho = {}'.format(self.rho)

    def forward(self, residual):
        squared_res = residual ** 2
        dist = torch.div(squared_res, squared_res + self.rho ** 2)
        return self.rho ** 2 * dist


class GMoF_unscaled(nn.Module):
    def __init__(self, rho=1):
        super(GMoF_unscaled, self).__init__()
        self.rho = rho

    def extra_repr(self):
        return 'rho = {}'.format(self.rho)

    def forward(self, residual):
        squared_res = residual ** 2
        if type(residual) is torch.Tensor:
            dist = torch.div(squared_res, squared_res + self.rho ** 2)
        else:
            dist = squared_res / (squared_res + self.rho ** 2)
        return dist


def getNewCoordinate(axisHeight=0.05, axisRadius=0.001):
    xColor = [200, 50, 0, 128]
    yColor = [0, 200, 50, 128]
    zColor = [50, 0, 200, 128]

    xAxis2 = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    xAxis2.apply_transform(matrix=np.mat(
        ((0, 0, 1, axisHeight / 2),
         (0, 1, 0, 0),
         (-1, 0, 0, 0),
         (0, 0, 0, 1))
    ))
    xAxis2.visual.face_colors = xColor
    yAxis2 = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    yAxis2.apply_transform(matrix=np.mat(
        ((1, 0, 0, 0),
         (0, 0, -1, axisHeight / 2),
         (0, 1, 0, 0),
         (0, 0, 0, 1))
    ))
    yAxis2.visual.face_colors = yColor
    zAxis2 = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    zAxis2.apply_transform(matrix=np.mat(
        ((1, 0, 0, 0),
         (0, 1, 0, 0),
         (0, 0, 1, axisHeight / 2),
         (0, 0, 0, 1))
    ))
    zAxis2.visual.face_colors = zColor
    xBox2 = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    xBox2.apply_translation((axisHeight, 0, 0))
    xBox2.visual.face_colors = xColor
    yBox2 = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    yBox2.apply_translation((0, axisHeight, 0))
    yBox2.visual.face_colors = yColor
    zBox2 = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    zBox2.apply_translation((0, 0, axisHeight))
    zBox2.visual.face_colors = zColor

    return 1


def meshVisualization(mesh):
    worldMeshList = get_world_mesh_list(planeWidth=0.2, axisHeight=0.05, axisRadius=0.001)
    mesh.visual.face_colors = [255, 128, 255, 200]
    worldMeshList.append(mesh)
    scene = trimesh.Scene(worldMeshList)
    scene.show()


def meshPairVisualization(mesh1, mesh2):
    worldMeshList = get_world_mesh_list(planeWidth=0.2, axisHeight=0.05, axisRadius=0.001)
    mesh1.visual.face_colors = [255, 128, 255, 200]
    mesh2.visual.face_colors = [255, 255, 128, 200]

    worldMeshList.append((mesh1, mesh2))
    scene = trimesh.Scene(worldMeshList)
    scene.show()


if __name__ == '__main__':
    # meshVisualization(trimesh.convex.convex_hull(np.array([[0.1, 0, 0], [0, 0.1, 0], [0, 0, 0.1], [0, 0, 0]])))
    scene = trimesh.Scene(get_checkerboard_plane())
    scene.add_geometry(get_world_mesh_list(add_plane=False))
    scene.show()

moyo/utils/mesh_utils.py
sha2nkt/moyo_toolkit
import numpy as np

colors = {
    'pink': [.6, .0, .4],
    'purple': [.9, .7, .7],
    'cyan': [.7, .75, .5],
    'red': [1.0, 0.0, 0.0],

    'green': [.0, 1., .0],
    'yellow': [1., 1., 0],
    'brown': [.5, .2, .1],
    'brown-light': [0.654, 0.396, 0.164],
    'blue': [.0, .0, 1.],

    'offwhite': [.8, .9, .9],
    'white': [1., 1., 1.],
    'orange': [1., .2, 0],

    'grey': [.7, .7, .7],
    'grey-blue': [0.345, 0.580, 0.713],
    'black': np.zeros(3),
    'white': np.ones(3),

    'yellowg': [0.83, 1, 0],
}

def copy2cpu(tensor):
    if isinstance(tensor, np.ndarray): return tensor
    return tensor.detach().cpu().numpy()
moyo/utils/misc.py
sha2nkt/moyo_toolkit
"""
Usage
python tests/biomechanical_evaluations/pressure_map_evaluation.py --img_folder /ps/data/YOGI/YOGI_220923_03596/IOI/20220923_Yogi_capture_PNG_MH/ --pp_folder /is/cluster/scratch/gbecherini/mosh/YOGI_2/YOGI_2_def-vt_type-smplx_grab_yogi_head-corr-no_wt-fingers-15000.0_wt-feet-300.0_wt-head-10000.0_seed-100_OFFSET/220923/03596 --pressure_xml_folder /ps/project/vicondata/ViconDataCaptures/OfficialCaptures/YOGI_2/pressure/220923/xml/ --pressure_csv_folder /ps/project/vicondata/ViconDataCaptures/OfficialCaptures/YOGI_2/pressure/220923/single_csv/YOGI/03596/

"""

import argparse
import copy
import glob
import json
import os
import os.path as osp
import shutil

import cv2
import numpy as np
import torch
import trimesh
import pandas as pd

from tqdm import tqdm
import pickle as pkl
from scipy import stats


import sys

sys.path.append(os.path.join(os.path.dirname(__file__), '../..'))
sys.path.append(os.path.join(os.path.dirname(__file__), '../../..'))

from body_visualizer.tools.vis_tools import colors

from human_body_prior.tools.omni_tools import copy2cpu as c2c

import smplx
from lib.utils.geometry import matrix_to_axis_angle, ea2rm, axis_angle_to_matrix
from utils.constants import frame_select_dict_combined as frame_select_dict
from utils.constants import YOGA82_PATH
from biomech_eval_utils import BiomechanicalEvaluator
from pressure_mat_utils import PressureMat
from psbody.mesh import Mesh

VICON_FPS = 60
PRESSURE_MAP_FPS = 60
IOI_FPS = 30


def read_keypoints(keypoint_fn):
    with open(keypoint_fn) as keypoint_file:
        data = json.load(keypoint_file)

    keypoints = []
    for idx, person_data in enumerate(data['people']):
        body_keypoints = np.array(
            person_data['pose_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])

        left_hand_keyp = np.array(
            person_data['hand_left_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])

        right_hand_keyp = np.array(
            person_data['hand_right_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])

        face_keypoints = np.array(
            person_data['face_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])[17: 17 + 51, :]

        contour_keyps = np.array(
            person_data['face_keypoints_2d'],
            dtype=np.float32).reshape([-1, 3])[:17, :]

        body_keypoints = np.concatenate([body_keypoints, left_hand_keyp,
                                         right_hand_keyp, face_keypoints, contour_keyps], axis=0)

        keypoints.append(body_keypoints)
    return keypoints


def select_keypoints(keypoints, img):
    img_center = np.array(img.shape[:2]) / 2  # height, width
    # select keypoints closest to image center weighted by inverse confidence
    if len(keypoints) > 1:
        kpselect = np.inf * np.ones(len(keypoints))
        # major joints openpose
        op_to_12 = [9, 10, 11, 12, 13, 14, 2, 3, 4, 5, 6, 7]
        for idx, personkpts in enumerate(keypoints):
            kpdist = personkpts[op_to_12, :2] - img_center
            kpdist = np.linalg.norm(kpdist, axis=1)
            kpconf = np.dot(kpdist, (- personkpts[op_to_12, 2] + 1))
            kpselect[idx] = kpconf
        kpselidx = np.argmin(kpselect)
    elif len(keypoints) == 1:
        kpselidx = 0
    else:
        keypoints = None
    keypoints = np.stack(keypoints[kpselidx:kpselidx + 1])
    return keypoints


def visualize_mesh(bm_output, faces, frame_id=0, display=False):
    imw, imh = 1600, 1600

    body_mesh = trimesh.Trimesh(vertices=c2c(bm_output.vertices[frame_id]),
                                faces=faces,
                                vertex_colors=np.tile(colors['grey'], (6890, 1)),
                                process=False,
                                maintain_order=True)

    if display:
        mv = MeshViewer(width=imw, height=imh, use_offscreen=True)
        mv.set_static_meshes([body_mesh])
        body_image = mv.render(render_wireframe=False)
        show_image(body_image)

    return body_mesh


def get_trans_offset(pelvis, smplx_params, trans, body_model):
    '''
    SOMA fits are in AMASS format (xy plane is ground) but SMPLify-XMC uses Pyrender format (xz plane is ground). global_orient is corrected by 270 rotation along x-axis to match the two formats.
    While the ground plane is z=0 in soma, since global_orient is wrt pelvis, this would mean the ground plane is not at y=0 after rotation.
    Fix: distance of pelvis to ground in preserved. Use it to push the mesh up after rotation such that the ground plane is at y=0.
    Args:
        pelvis: pelvis joint before rotation
    '''
    bs = trans.shape[0]
    init_root_orient = smplx_params['global_orient']
    pelvis_height = pelvis[:, 2]  # this is always conserved

    new_smplx_params = copy.deepcopy(smplx_params)
    # Rotate AMASS root orient to smplify-xmc format
    R_init = axis_angle_to_matrix(init_root_orient)
    R1 = ea2rm(torch.tensor([[np.radians(270)]]), torch.tensor([[np.radians(0)]]),
               torch.tensor([[np.radians(0)]])).float().to(R_init.device)
    R = torch.bmm(R1.expand(bs, -1, -1), R_init)
    new_smplx_params['global_orient'] = matrix_to_axis_angle(R)

    # posed body with hand, with global orient
    body_model_output = body_model(
        global_orient=new_smplx_params['global_orient'],
        body_pose=new_smplx_params['body_pose'])

    new_pelvis = body_model_output.joints[:, 0]
    new_ground_plane_height = new_pelvis[:, 1] - pelvis_height
    trans_offset = -new_ground_plane_height
    return trans_offset

def smplx_breakdown(bdata, device):
    num_frames = len(bdata['trans'])

    bdata['poses'] = bdata['fullpose']

    global_orient = torch.from_numpy(bdata['poses'][:, :3]).float().to(device)
    body_pose = torch.from_numpy(bdata['poses'][:, 3:66]).float().to(device)
    jaw_pose = torch.from_numpy(bdata['poses'][:, 66:69]).float().to(device)
    leye_pose = torch.from_numpy(bdata['poses'][:, 69:72]).float().to(device)
    reye_pose = torch.from_numpy(bdata['poses'][:, 72:75]).float().to(device)
    left_hand_pose = torch.from_numpy(bdata['poses'][:, 75:120]).float().to(device)
    right_hand_pose = torch.from_numpy(bdata['poses'][:, 120:]).float().to(device)

    # if bdata['v_template_fname'] == '/ps/project/amass/MOCAP/PS_MoCaps/YOGI/subject_data/subject_v_template/210727_03596/smplx_female.ply':
    v_template_fname = '/is/cluster/scratch/stripathi/pycharm_remote/yogi/data/v_templates/220923_yogi_03596_minimal_simple_female/mesh.ply'
    # else:
    #     print("Something weird happening with V_template loading!!!!")
    #     import ipdb; ipdb.set_trace()
    v_template = Mesh(filename=v_template_fname)

    body_params = {'global_orient': global_orient, 'body_pose': body_pose,
                  'jaw_pose': jaw_pose, 'leye_pose': leye_pose, 'reye_pose': reye_pose,
                  'left_hand_pose': left_hand_pose, 'right_hand_pose': right_hand_pose,
                   'v_template': torch.Tensor(v_template.v).to(device),}
    return body_params


def smplx_to_mesh(body_params, model_folder, model_type, gender='neutral'):
    with torch.no_grad():
        device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
        num_betas = 10

        smplx_params = smplx_breakdown(body_params, device)

        trans = torch.from_numpy(body_params['trans']).float().to(device)

        betas = torch.from_numpy(body_params['betas']).float().to(device).unsqueeze(0)
        betas = betas[:, :num_betas]  # only using 10 betas

        body_model_params = dict(model_path=model_folder,
                                 model_type=model_type,
                                 gender=gender,
                                 v_template=smplx_params['v_template'],
                                 # joint_mapper=joint_mapper,
                                 batch_size=trans.shape[0],
                                 create_global_orient=True,
                                 create_body_pose=True,
                                 create_betas=True,
                                 num_betas=num_betas,
                                 create_left_hand_pose=True,
                                 create_right_hand_pose=True,
                                 create_expression=True,
                                 create_jaw_pose=True,
                                 create_leye_pose=True,
                                 create_reye_pose=True,
                                 create_transl=True,
                                 use_pca=False,
                                 flat_hand_mean=True,
                                 dtype=torch.float32)

        body_model = smplx.create(**body_model_params).to(device)
        body_model_output = body_model(transl=trans,
                                       global_orient=smplx_params['global_orient'],
                                       body_pose=smplx_params['body_pose'],
                                       left_hand_pose=smplx_params['left_hand_pose'],
                                       right_hand_pose=smplx_params['right_hand_pose'])

        pelvis = body_model_output.joints[:, 0]
        trans_offset = get_trans_offset(pelvis, smplx_params, trans, body_model)
        zero_vec = torch.zeros_like(trans_offset)
        transl = torch.stack([zero_vec, trans_offset, zero_vec], dim=-1)

    return body_model_output, smplx_params, transl, body_model.faces

def yoga82_extra(img_folder, pp_folder, pressure_xml_folder, pressure_csv_folder, model_folder,
                 model_type='smplx', subsampling_factor=1):
    imgpaths_ = []
    imgnames_ = []
    subject_ids_ = []
    relpaths_ = []
    heights_ = []
    weights_ = []
    genders_ = []
    pressures_ = []
    pressure_heatmaps_ = []
    pp_body_poses_ = []
    pp_left_hand_poses_ = []
    pp_right_hand_poses_ = []
    pp_global_orients_ = []
    pp_transls_ = []

    # presented poses
    pp_names = os.listdir(pp_folder)
    pp_names = [os.path.splitext(x)[0] for x in pp_names if x[0] != '.' and '.pkl' in x]

    all_pose_names = [pp_name.replace('_stageii', '') for pp_name in pp_names]
    all_pose_names = sorted(all_pose_names)

    pp_pose_cache = {}

    no_images_found = set()
    no_pp_found = set()

    smplx_faces = smplx.create(model_path=model_folder, model_type=model_type).faces
    smplx_faces = torch.tensor(smplx_faces.astype(np.int64), dtype=torch.long)
    pressure_evaluator = BiomechanicalEvaluator(faces=smplx_faces)

    for i, pname in enumerate(tqdm(all_pose_names)):
        if 'Cow_Face_Pose_or_Gomukhasana_-c' in pname: # pressure mat was broken for this pose
            continue
        if 'female_stagei' in pname: # do not analyze female stage1 file
            continue
        # get all possible variations
        pp_pkl = osp.join(pp_folder, pname + '_stageii.pkl')
        pp_params = pkl.load(open(pp_pkl, 'rb'))
        # read height, weight, gender meta
        height = None
        weight = None
        gender = 'female'
        subj_id = 'unknown'
        pp_body_model_output, smplx_params, transls, faces = smplx_to_mesh(pp_params, model_folder,
                                                       model_type, gender=gender)
        # get pressure xml
        pressure_xml = osp.join(pressure_xml_folder, pname + '.xml')
        pressure_csv = osp.join(pressure_csv_folder, pname + '.csv')
        pressure_mat_c3d = pressure_xml.replace('pressure', 'vicon').replace('/xml/', '/3_Gaps_filled/').replace('.xml', '.c3d')
        # initialize pressure mat
        pressure_mat = PressureMat(pressure_csv, pressure_xml, pressure_mat_c3d)

        # pick selected frame
        selected_frame = frame_select_dict['_'.join(pname.split('_')[5:])]

        # get all images for this pose
        img_dir = osp.join(img_folder, pname)
        for cam in range(1, 4):
            cam_name = 'YOGI_Cam_' + str(cam).zfill(2)
            img_dir_cam = osp.join(img_dir, cam_name)
            types = ('*.jpg', '*.png')  # the tuple of file types
            images = []
            for files in types:
                images.extend(glob.glob(osp.join(img_dir_cam, files)))
            if len(images) == 0:
                print(f'No images found for {pname}')
                no_images_found.add(pname)

            # subsample frames
            images.sort()

            # folders to log outputs
            if cam == 1 or cam == 3:
                out_img_dir = osp.join(pp_folder, 'images_small', pname, cam_name)
                os.makedirs(out_img_dir, exist_ok=True)
                out_pressure_heatmap_dir = osp.join(pp_folder, 'pressure_heatmaps', pname, cam_name)
                os.makedirs(out_pressure_heatmap_dir, exist_ok=True)
                out_obj_dir = osp.join(pp_folder, 'meshes', pname, cam_name)
                os.makedirs(out_obj_dir, exist_ok=True)

            for img_path in images[
                            :-1:subsampling_factor]:  # sometimes, the last frame is missing in pressure or vicon data
                # get frame correspondences
                img_fnum = int(img_path.split('_')[-1].split('.')[0])

                #only run on selected frames
                if img_fnum != selected_frame:
                    continue
                vicon_fnum = img_fnum * (VICON_FPS // IOI_FPS)
                pressure_fnum = img_fnum * (PRESSURE_MAP_FPS // IOI_FPS)
                assert vicon_fnum == pressure_fnum

                # load image
                img_fn, _ = osp.splitext(osp.split(img_path)[1])
                relpath = osp.dirname(img_path).replace(img_dir, '').strip('/')

                if cam == 1 or cam == 3:
                    img = cv2.imread(img_path).astype(np.float32)
                    # resize image to half size
                    img = cv2.resize(img, (img.shape[1] // 2, img.shape[0] // 2))
                    cv2.imwrite(osp.join(out_img_dir, f'{pname}_{pressure_fnum}.jpg'), img)
                    print('[[Image Saved]]\n', osp.join(out_img_dir, f'{pname}_{pressure_fnum}.jpg'))
                if cam != 1:
                    continue

                # get pressure for current frame
                # increase_heatmap size before saving
                gt_pressure = pressure_mat.gt_pressures[pressure_fnum]
                gt_heatmap = pressure_mat.gt_heatmaps[pressure_fnum]
                gt_cop_relative = pressure_mat.gt_cops_relative[pressure_fnum]
                if np.isnan(gt_cop_relative).any():
                    continue
                # gt_cop_metric = pressure_mat.gt_cops_metric[pressure_fnum]
                if cam == 1:
                    heatmap_large = cv2.resize(gt_heatmap, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_NEAREST)
                    # add gt_cop point to heatmap
                    cv2.drawMarker(heatmap_large, (
                        int(gt_cop_relative[0] / pressure_mat.mat_size_metric[0] * heatmap_large.shape[1]),
                        int(gt_cop_relative[1] / pressure_mat.mat_size_metric[1] * heatmap_large.shape[0])),
                                   [0, 255, 0],
                                   markerType=cv2.MARKER_STAR, markerSize=40, thickness=5)
                    # heatmap_large = cv2.flip(heatmap_large, 0) # to have better visualization wrt to blender mat configuration
                    cv2.imwrite(os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_gt.png'),
                                heatmap_large)
                    print('[[GT Pressure Heatmap Saved]]\n',
                          os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_gt.png'))



                transl = transls[vicon_fnum]
                pp_mesh = visualize_mesh(pp_body_model_output, faces, frame_id=vicon_fnum)
                pp_verts = pp_mesh.vertices
                # pp_verts = pp_body_model_output.vertices[vicon_fnum]
                vertices = torch.tensor(pp_verts).unsqueeze(0)
                if cam == 1:
                    out_obj = osp.join(out_obj_dir, f'{pname}_{vicon_fnum}.obj')
                    pp_mesh.export(out_obj)
                    print('[[Mesh]]\n', out_obj)


                # get per-vertex pressure
                pressure_evaluator.evaluate_pressure(gt_pressure,
                                                   gt_cop_relative,
                                                   vertices,
                                                   pressure_mat.mat_size,
                                                   pressure_mat.mat_size_metric,
                                                   pressure_mat.marker_positions_metric)
                if cam == 1:
                    # iou first visualization
                    pred_heatmap = pressure_evaluator.iou_first.pred_heatmaps[-1]
                    pred_cop_relative = pressure_evaluator.iou_first.pred_cops_relative[-1]
                    cop_error = pressure_evaluator.iou_first.cop_errors[-1]
                    iou = pressure_evaluator.iou_first.ious[-1]
                    frame_diff = pressure_evaluator.iou_first.frame_diffs[-1]
                    best_cop_w = pressure_evaluator.iou_first.best_cop_w[-1]
                    best_cop_k = pressure_evaluator.iou_first.best_cop_k[-1]

                    pred_heatmap_large = cv2.resize(pred_heatmap, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_NEAREST)
                    cv2.drawMarker(pred_heatmap_large, (
                        int(pred_cop_relative[0] / pressure_mat.mat_size_metric[0] * pred_heatmap_large.shape[1]),
                        int(pred_cop_relative[1] / pressure_mat.mat_size_metric[1] * pred_heatmap_large.shape[0])),
                                   [0, 255, 255],
                                   markerType=cv2.MARKER_STAR, markerSize=40, thickness=5)
                    cv2.putText(pred_heatmap_large, f'{cop_error*1000:.2f} mm', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (0, 255, 255), 2)
                    cv2.putText(pred_heatmap_large, f'{iou:.2f}', (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255,255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'{frame_diff:.2f}', (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_w: {best_cop_w:.2f}', (10, 200), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_k: {best_cop_k:.2f}', (10, 250), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.imwrite(os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_iou_first.png'),
                                pred_heatmap_large)
                    print('[[PRED Pressure Heatmap Saved]]\n',
                          os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_iou_first.png'))

                    # cop error first visualizations
                    pred_heatmap = pressure_evaluator.cop_error_first.pred_heatmaps[-1]
                    pred_cop_relative = pressure_evaluator.cop_error_first.pred_cops_relative[-1]
                    cop_error = pressure_evaluator.cop_error_first.cop_errors[-1]
                    iou = pressure_evaluator.cop_error_first.ious[-1]
                    frame_diff = pressure_evaluator.cop_error_first.frame_diffs[-1]
                    best_cop_w = pressure_evaluator.cop_error_first.best_cop_w[-1]
                    best_cop_k = pressure_evaluator.cop_error_first.best_cop_k[-1]

                    pred_heatmap_large = cv2.resize(pred_heatmap, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_NEAREST)
                    cv2.drawMarker(pred_heatmap_large, (
                        int(pred_cop_relative[0] / pressure_mat.mat_size_metric[0] * pred_heatmap_large.shape[1]),
                        int(pred_cop_relative[1] / pressure_mat.mat_size_metric[1] * pred_heatmap_large.shape[0])),
                                   [0, 255, 255],
                                   markerType=cv2.MARKER_STAR, markerSize=40, thickness=5)
                    cv2.putText(pred_heatmap_large, f'{cop_error * 1000:.2f} mm', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (0, 255, 255), 2)
                    cv2.putText(pred_heatmap_large, f'{iou:.2f}', (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'{frame_diff:.2f}', (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_w: {best_cop_w:.2f}', (10, 200), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_k: {best_cop_k:.2f}', (10, 250), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.imwrite(os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_cop_error_first.png'),
                                pred_heatmap_large)
                    print('[[PRED Pressure Heatmap Saved]]\n',
                          os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_cop_error_first.png'))

                    # frame_idd first visualizations
                    pred_heatmap = pressure_evaluator.frame_diff_first.pred_heatmaps[-1]
                    pred_cop_relative = pressure_evaluator.frame_diff_first.pred_cops_relative[-1]
                    cop_error = pressure_evaluator.frame_diff_first.cop_errors[-1]
                    iou = pressure_evaluator.frame_diff_first.ious[-1]
                    frame_diff = pressure_evaluator.frame_diff_first.frame_diffs[-1]
                    best_cop_w = pressure_evaluator.frame_diff_first.best_cop_w[-1]
                    best_cop_k = pressure_evaluator.frame_diff_first.best_cop_k[-1]

                    pred_heatmap_large = cv2.resize(pred_heatmap, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_NEAREST)
                    cv2.drawMarker(pred_heatmap_large, (
                        int(pred_cop_relative[0] / pressure_mat.mat_size_metric[0] * pred_heatmap_large.shape[1]),
                        int(pred_cop_relative[1] / pressure_mat.mat_size_metric[1] * pred_heatmap_large.shape[0])),
                                   [0, 255, 255],
                                   markerType=cv2.MARKER_STAR, markerSize=40, thickness=5)
                    cv2.putText(pred_heatmap_large, f'{cop_error * 1000:.2f} mm', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (0, 255, 255), 2)
                    cv2.putText(pred_heatmap_large, f'{iou:.2f}', (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'{frame_diff:.2f}', (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_w: {best_cop_w:.2f}', (10, 200), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.putText(pred_heatmap_large, f'cop_k: {best_cop_k:.2f}', (10, 250), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                (255, 255, 0), 2)
                    cv2.imwrite(
                        os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_frame_diff_first.png'),
                        pred_heatmap_large)
                    print('[[PRED Pressure Heatmap Saved]]\n',
                          os.path.join(out_pressure_heatmap_dir, f'{pname}_{pressure_fnum}_pred_frame_diff_first.png'))



                body_pose = smplx_params['body_pose'][vicon_fnum].squeeze().cpu()
                left_hand_pose = smplx_params['left_hand_pose'][vicon_fnum].squeeze().cpu()
                right_hand_pose = smplx_params['right_hand_pose'][vicon_fnum].squeeze().cpu()
                global_orient = smplx_params['global_orient'][vicon_fnum].squeeze().cpu()
                transl = transl.squeeze().cpu()

                # add to dict if keypoints were detected
                imgpaths_.append(img_path)
                imgnames_.append(img_fn)
                subject_ids_.append(subj_id)
                relpaths_.append(relpath)
                heights_.append(height)
                weights_.append(weight)
                genders_.append(gender)
                pressures_.append(gt_pressure)
                pressure_heatmaps_.append(gt_heatmap)
                pp_body_poses_.append(np.array(body_pose))
                pp_left_hand_poses_.append(np.array(left_hand_pose))
                pp_right_hand_poses_.append(np.array(right_hand_pose))
                pp_global_orients_.append(np.array(global_orient))
                pp_transls_.append(np.array(transl))

    print('Mean COP error (iou first): ', np.mean(pressure_evaluator.iou_first.cop_errors))
    print('Mean Pressure IOU (iou first): ', np.mean(pressure_evaluator.iou_first.ious))
    print('Mean Pressure Frame Diff (iou first): ', np.mean(pressure_evaluator.iou_first.frame_diffs))
    # print mode of cop_w and cop_k
    print('Mode of cop_w (iou first): ', stats.mode(pressure_evaluator.iou_first.best_cop_w))
    print('Mode of cop_k (iou first): ', stats.mode(pressure_evaluator.iou_first.best_cop_k))
    print('---------------------------------------------')
    print('Mean COP error (cop error first): ', np.mean(pressure_evaluator.cop_error_first.cop_errors))
    print('Mean Pressure IOU (cop error first): ', np.mean(pressure_evaluator.cop_error_first.ious))
    print('Mean Pressure Frame Diff (cop error first): ', np.mean(pressure_evaluator.cop_error_first.frame_diffs))
    # print mode of cop_w and cop_k
    print('Mode of cop_w (cop error first): ', stats.mode(pressure_evaluator.cop_error_first.best_cop_w))
    print('Mode of cop_k (cop error first): ', stats.mode(pressure_evaluator.cop_error_first.best_cop_k))
    print('---------------------------------------------')
    print('Mean COP error (frame diff first): ', np.mean(pressure_evaluator.frame_diff_first.cop_errors))
    print('Mean Pressure IOU (frame diff first): ', np.mean(pressure_evaluator.frame_diff_first.ious))
    print('Mean Pressure Frame Diff (frame diff first): ', np.mean(pressure_evaluator.frame_diff_first.frame_diffs))
    # print mode of cop_w and cop_k
    print('Mode of cop_w (frame diff first): ', stats.mode(pressure_evaluator.frame_diff_first.best_cop_w))
    print('Mode of cop_k (frame diff first): ', stats.mode(pressure_evaluator.frame_diff_first.best_cop_k))
    print('---------------------------------------------')
    print(f'Folders with no images: \n {no_images_found}')
    print(f'Folders with no presented poses: \n {no_pp_found}')

    # save as db file
    outfile = osp.join('data', 'dbs_biomech', 'newmoyo_gt_first_day_vicon_220923.npz')
    np.savez(outfile,
             imgpath=imgpaths_,
             subject_id=subject_ids_,
             imgname=imgnames_,
             relpath=relpaths_,
             height=heights_,
             weight=weights_,
             gender=genders_,
             pressure=pressures_,
             pressure_heatmap=pressure_heatmaps_,
             body_pose=pp_body_poses_,
             left_hand_pose=pp_left_hand_poses_,
             right_hand_pose=pp_right_hand_poses_,
             global_orient=pp_global_orients_,
             transl=pp_transls_
             )


if __name__ == "__main__":
    """
    yoga-82 dir /ps/scratch/ps_shared/4Paul/Yoga-82
    """

    parser = argparse.ArgumentParser()
    parser.add_argument('--img_folder', required=True,
                        help='folder containing yoga 82 images')
    # parser.add_argument('--openpose_folder', required=True,
    #                     help='folder containing yoga 82 keypoints')
    parser.add_argument('--pressure_xml_folder', required=True,
                        help='folder containing GT pressure XMLs')
    parser.add_argument('--pressure_csv_folder', required=True,
                        help='folder containing GT pressure CSVs')
    parser.add_argument('--pp_folder', required=True,
                        help='presented poses folder')
    # parser.add_argument('--essentials_folder', required=True,
    #                     help='geodesics, ... essentials for smplify-xmc')
    parser.add_argument('--model_folder', required=False, default='/ps/project/common/smplifyx/models/',
                        help='path to SMPL/SMPL-X model folder')

    args = parser.parse_args()
    yoga82_extra(args.img_folder, args.pp_folder, args.pressure_xml_folder, args.pressure_csv_folder, args.model_folder)
moyo/utils/pressure_map_evaluation.py
sha2nkt/moyo_toolkit
import re
import c3d
import cv2
import numpy as np
import pandas as pd
from bs4 import BeautifulSoup as bs

VICON_FPS = 60
PRESSURE_MAP_FPS = 60
IOI_FPS = 30

class PressureMat:
    def __init__(self, csv_path, xml_path, c3d_file):
        self.csv_path = csv_path
        self.xml_path = xml_path
        self.c3d_file = c3d_file

        self.marker_positions_metric = self.extract_pressure_mat_markers()

        self.gt_cops_relative, gt_pressure_frameids = self.process_csv()
        # self.gt_cops_metric = self.compute_gt_cop_metric()
        self.gt_pressures, self.gt_heatmaps, self.mat_size, self.mat_size_metric = self.process_xml()
        # filter out pressure frames that are not in the csv file
        self.gt_pressures = self.gt_pressures[gt_pressure_frameids]
        self.gt_heatmaps = [self.gt_heatmaps[i] for i in gt_pressure_frameids]


    def parse_pressure_crops(self, pressure_crop):
        '''
            Takes raw pressure data from xml which is a list of strings for every frame containing pressure
            pressure_crops: N x 1 list of strings
            '''
        pressure_crop = [pressure_crop[i].get_text() for i in range(len(pressure_crop))]
        pressure_crop = [re.split(r'\n\t+', pressure_crop[i])[1:-1] for i in range(len(pressure_crop))]
        for i in range(len(pressure_crop)):
            for j in range(len(pressure_crop[i])):
                row = re.split(' ', pressure_crop[i][j])
                row = [float(el) for el in row]
                pressure_crop[i][j] = row
        return pressure_crop

    def unnormalize_pressure(self, pressure_crop, cell_begin, cell_end, mat_size, frame_count):
        '''
        Unnormalize pressure data to the full size of the pressure mat
        pressure_crop: N x 1 list of lists of lists (N frames, 1 pressure crop, pressure_crop[i] is a list of lists of floats (pressure values))
        cell_begin: N x 2 list of lists (N frames, 2 coordinates of the beginning of the pressure crop)
        cell_end: N x 2 list of lists (N frames, 2 coordinates of the end of the pressure crop)
        mat_size: 2 x 1 list (2 coordinates of the full size of the pressure mat)
        '''
        pressure_unnorm = np.zeros([frame_count, mat_size[1], mat_size[0]])
        for i in range(frame_count):
            pressure_crop_i = np.flipud(pressure_crop[i])
            pressure_unnorm[i, cell_begin[i][1]:cell_end[i][1],
            cell_begin[i][0]:cell_end[i][0]] = pressure_crop_i
        return pressure_unnorm

    def process_xml(self):
        # read xml file and parse
        with open(self.xml_path, 'r') as f:
            xml = f.read()
        xml_soup = bs(xml, 'lxml')
        mat_size = [int(xml_soup.movements.clips.clip.cell_count.x.get_text()),
                    int(xml_soup.movements.clips.clip.cell_count.y.get_text())]  # number of cells in x and y direction
        sensel_size = [float(xml_soup.movements.clips.clip.cell_size.x.get_text()),
                       float(xml_soup.movements.clips.clip.cell_size.y.get_text())]  # size of every cell in mm
        mat_size_metric = [mat_size[0] * sensel_size[0] / 1000,
                           mat_size[1] * sensel_size[1] / 1000]  # size of the pressure mat in m
        fps = float(xml_soup.movements.clips.clip.frequency.get_text())
        frame_count = int(xml_soup.movements.clips.clip.count.get_text())

        # the pressure mat crops the pressure data, uncrop it to the full size
        cell_begin_raw = xml_soup.movements.clips.clip.data.find_all('cell_begin')
        assert len(
            cell_begin_raw) == frame_count, 'Number of frames in xml file does not match the number of frames in the pressure data'
        cell_begin = [[int(cell_begin_raw[i].x.get_text()), int(cell_begin_raw[i].y.get_text())] for i in
                      range(len(cell_begin_raw))]
        cell_extent_raw = xml_soup.movements.clips.clip.data.find_all('cell_count')
        assert len(
            cell_extent_raw) == frame_count, 'Number of frames in xml file does not match the number of frames in the pressure data'
        cell_end = [[cell_begin[i][0] + int(cell_extent_raw[i].x.get_text()),
                     cell_begin[i][1] + int(cell_extent_raw[i].y.get_text())] for i in range(len(cell_extent_raw))]
        pressure_crop_raw = xml_soup.movements.clips.clip.data.find_all('cells')
        assert len(
            pressure_crop_raw) == frame_count, 'Number of frames in xml file does not match the number of frames in the pressure data'
        # parse pressure crops
        pressure_crop = self.parse_pressure_crops(pressure_crop_raw)
        # unnormalize pressure crops
        pressure_unnorm = self.unnormalize_pressure(pressure_crop, cell_begin, cell_end, mat_size, frame_count)
        # visualize pressure as heatmap
        heatmaps = vis_heatmap_seq(pressure_unnorm, normalize=True)
        return pressure_unnorm, heatmaps, mat_size, mat_size_metric


    def process_csv(self):
        # read csv file and parse
        pressure_md = pd.read_csv(self.csv_path, header=[0], nrows=1)
        pressure_df = pd.read_csv(self.csv_path, header=[2])

        # extract important pressure md from csv
        pressure_fps = pressure_md['frequency'].values[0]
        assert pressure_fps == PRESSURE_MAP_FPS, 'Pressure fps is not equal to default'
        pressure_frame_count = pressure_md['count'].values[0]

        # get the cop coordinates
        cop_x = pressure_df['Pressure, Raw Pressure-distribution-Pressure center-x (mm)'].values / 1000  # mm to m
        cop_y = pressure_df['Pressure, Raw Pressure-distribution-Pressure center-y (mm)'].values / 1000
        cop = np.array([cop_x, cop_y]).T

        pressure_timestamps = pressure_df['time']
        # convert pressure df to numpy array
        pressure_timestamps = pressure_timestamps.values
        assert pressure_timestamps.shape[
                   0] == pressure_frame_count, 'Number of frames in md does not match the number of frames in the timestamp'

        # convert timestamps to frame numbers and round to lower
        pressure_fids = np.rint(pressure_timestamps * pressure_fps).astype(int)
        assert pressure_fids.shape[0] == pressure_fids[-1] - pressure_fids[
            0] + 1, 'Some frames are missing in the timestamps'
        return cop, pressure_fids


    def extract_pressure_mat_markers(self):
        with open(self.c3d_file, 'rb') as f:
            data = c3d.Reader(f)  # data contains the header of the c3d files, info such as the label name
            for frame in data.read_frames():  # call read_frames() to iterate across the marker positions
                frame = frame[1][:, :3]  # this contains the marker positions
                break

        marker_id_one = np.where(data.point_labels == 'SensorMat:M_1                                                   ')[0][0]
        marker_id_two = np.where(data.point_labels == 'SensorMat:M_2                                                   ')[0][0]
        marker_id_three = np.where(data.point_labels == 'SensorMat:M_3                                                   ')[0][0]
        marker_id_four = np.where(data.point_labels == 'SensorMat:M_4                                                   ')[0][0]
        marker_ids_unordered = [marker_id_one, marker_id_two, marker_id_three, marker_id_four]

        # Put the markers in the correct order
        for i, marker_id in enumerate(marker_ids_unordered):
            assert marker_id is not None, f'Marker {i} not found'
            if frame[marker_id, :][0] > 0 and frame[marker_id, :][1] > 0:
                marker_id_bl = marker_id
            if frame[marker_id, :][0] > 0 and frame[marker_id, :][1] < 0:
                marker_id_tl = marker_id
            if frame[marker_id, :][0] < 0 and frame[marker_id, :][1] < 0:
                marker_id_tr = marker_id
            if frame[marker_id, :][0] < 0 and frame[marker_id, :][1] > 0:
                marker_id_br = marker_id

        marker_ids = [marker_id_bl, marker_id_tl, marker_id_tr, marker_id_br]      # (bottom left, top left, top right, bottom right)

        marker_positions = frame[marker_ids, :]
        return marker_positions

    # def compute_gt_cop_metric(self):
    #     """
    #     Convert relative gt_cop to metric gt_cop in the world coordinates
    #     Args:
    #         gt_cop: relative to the pressure mat
    #         mat_size_metric:
    #         marker_positions:
    #
    #     Returns:
    #
    #     """
    #     import ipdb; ipdb.set_trace()
    #     mat_marker_bl = self.marker_positions_metric[0]
    #     mat_marker_br = self.marker_positions_metric[1]
    #
    #     # convert gt_cop to metric
    #     gt_cop_metric = self.gt_cops_relative.copy()
    #     gt_cop_metric[:, 0] = mat_marker_br[0] / 1000 - self.gt_cops_relative[:,
    #                                                     0]  # because the x-axis is to your right in mat coordinate but to your left in image coordinates
    #     gt_cop_metric[:, 1] = self.gt_cops_relative[:, 1] + mat_marker_bl[1] / 1000
    #     return gt_cop_metric


def vis_heatmap_seq(pressure, normalize=True):
    # normalize pressure value over all frames
    # if normalize:
    #     pressure = (pressure - np.min(pressure)) / (np.max(pressure) - np.min(pressure))
    # visualize pressure as heatmap
    heatmaps = []
    for i in range(len(pressure)):
        if normalize:
            pressure_i = (pressure[i] - np.min(pressure[i])) / (np.max(pressure[i]) - np.min(pressure[i]))
        else:
            pressure_i = pressure[i]
        heatmap = cv2.applyColorMap(np.uint8(255 * pressure_i), cv2.COLORMAP_JET)
        heatmaps.append(heatmap)
    return heatmaps


def vis_heatmap(pressure, normalize=True):
    if normalize:
        pressure = (pressure - np.min(pressure)) / (np.max(pressure) - np.min(pressure))
    heatmap = cv2.applyColorMap(np.uint8(255 * pressure), cv2.COLORMAP_JET)
    return heatmap
moyo/utils/pressure_mat_utils.py
sha2nkt/moyo_toolkit