EasyMocap/easymocap/dataset/mv1pmf_mirror.py

'''
  @ Date: 2021-01-12 17:12:50
  @ Author: Qing Shuai
  @ LastEditors: Qing Shuai
  @ LastEditTime: 2021-04-14 11:26:36
  @ FilePath: /EasyMocapRelease/easymocap/dataset/mv1pmf_mirror.py
'''
import os
from os.path import join
import numpy as np
import cv2
from .base import ImageFolder
from .mv1pmf import MVBase
from .mirror import calc_mirror_transform, flipSMPLParams, mirrorPoint3D, flipPoint2D, mirror_Rh
from ..mytools.file_utils import get_bbox_from_pose, read_json

class MV1PMF_Mirror(MVBase):
    def __init__(self, root, cams=[], pid=0, out=None, config={}, 
        image_root='images', annot_root='annots', kpts_type='body15',
        undis=True, no_img=False,
        verbose=False) -> None:
        self.mirror = np.array([[0., 1., 0., 0.]])
        super().__init__(root=root, cams=cams, out=out, config=config, 
            image_root=image_root, annot_root=annot_root, 
            kpts_type=kpts_type, undis=undis, no_img=no_img)
        self.pid = pid
        self.verbose = False

    def __str__(self) -> str:
        return 'Dataset for MultiMirror: {} views'.format(len(self.cams))
    
    def write_keypoints3d(self, keypoints3d, nf):
        results = []
        M = self.Mirror[0]
        pid = self.pid
        val = {'id': pid, 'keypoints3d': keypoints3d}
        results.append(val)
        kpts = keypoints3d
        kpts3dm = (M[:3, :3] @ kpts[:, :3].T + M[:3, 3:]).T
        kpts3dm = np.hstack([kpts3dm, kpts[:, 3:]])
        kpts3dm = flipPoint2D(kpts3dm)
        val1 = {'id': pid + 1, 'keypoints3d': kpts3dm}
        results.append(val1)
        super().write_keypoints3d(results, nf)

    def write_smpl(self, params, nf):
        outname = join(self.out, 'smpl', '{:06d}.json'.format(nf))
        results = []
        M = self.Mirror[0]
        pid = self.pid
        val = {'id': pid}
        val.update(params)
        results.append(val)
        # 增加镜子里的人的
        val = {'id': pid + 1}
        val.update(flipSMPLParams(params, self.mirror))
        results.append(val)
        self.writer.write_smpl(results, outname)

    def vis_smpl(self, vertices, faces, images, nf, sub_vis=[], 
        mode='smpl', extra_data=[], add_back=True):
        outname = join(self.out, 'smpl', '{:06d}.jpg'.format(nf))
        render_data = {}
        if len(vertices.shape) == 3:
            vertices = vertices[0]
        pid = self.pid
        render_data[pid] = {'vertices': vertices, 'faces': faces, 
            'vid': pid, 'name': 'human_{}_{}'.format(nf, pid)}
        vertices_m = mirrorPoint3D(vertices, self.Mirror[0])
        render_data[pid+1] = {'vertices': vertices_m, 'faces': faces, 
            'vid': pid, 'name': 'human_mirror_{}_{}'.format(nf, pid)}
        
        cameras = {'K': [], 'R':[], 'T':[]}
        if len(sub_vis) == 0:
            sub_vis = self.cams
        for key in cameras.keys():
            cameras[key] = [self.cameras[cam][key] for cam in sub_vis]
        images = [images[self.cams.index(cam)] for cam in sub_vis]
        self.writer.vis_smpl(render_data, images, cameras, outname, add_back=add_back)
    
    def vis_detections(self, images, annots, nf, to_img=True, sub_vis=[]):
        outname = join(self.out, 'detec', '{:06d}.jpg'.format(nf))
        lDetections = []
        nViews = len(images)
        for nv in range(len(images)):
            det = {
                'id': self.pid,
                'bbox': annots['bbox'][nv],
                'keypoints2d': annots['keypoints'][nv]
            }
            det_m = {
                'id': self.pid + 1,
                'bbox': annots['bbox'][nv+nViews],
                'keypoints2d': annots['keypoints'][nv+nViews]
            }
            lDetections.append([det, det_m])
        if len(sub_vis) != 0:
            valid_idx = [self.cams.index(i) for i in sub_vis]
            images = [images[i] for i in valid_idx]
            lDetections = [lDetections[i] for i in valid_idx]
        return self.writer.vis_keypoints2d_mv(images, lDetections, outname=outname, vis_id=False)

    def vis_repro(self, images, kpts_repro, nf, to_img=True, sub_vis=[]):
        outname = join(self.out, 'repro', '{:06d}.jpg'.format(nf))
        lDetections = []
        for nv in range(len(images)):
            det = {
                'id': -1,
                'keypoints2d': kpts_repro[nv],
                'bbox': get_bbox_from_pose(kpts_repro[nv], images[nv])
            }
            det_mirror = {
                'id': -1,
                'keypoints2d': kpts_repro[nv+len(images)],
                'bbox': get_bbox_from_pose(kpts_repro[nv+len(images)], images[nv])
            }
            lDetections.append([det, det_mirror])
        if len(sub_vis) != 0:
            valid_idx = [self.cams.index(i) for i in sub_vis]
            images = [images[i] for i in valid_idx]
            lDetections = [lDetections[i] for i in valid_idx]
        return self.writer.vis_keypoints2d_mv(images, lDetections, outname=outname, vis_id=False)

    @property
    def Mirror(self):
        M = calc_mirror_transform(self.mirror)
        return M

    @property
    def Pall(self):
        return self.Pall_

    @Pall.setter
    def Pall(self, value):
        M = self.Mirror
        if M.shape[0] == 1 and M.shape[0] != value.shape[0]:
            M = M.repeat(value.shape[0], 0)
        Pall_mirror = np.einsum('bmn,bno->bmo', value, M)
        Pall = np.vstack((value, Pall_mirror))
        self.Pall_ = Pall
        
    def __getitem__(self, index: int):
        images, annots_all = super().__getitem__(index)
        annots0 = self.select_person(annots_all, index, self.pid)
        annots1 = self.select_person(annots_all, index, self.pid + 1)
        # flip points
        # stack it as only one person
        annots = {
            'bbox': np.vstack([annots0['bbox'], annots1['bbox']]),
            'keypoints': np.vstack([annots0['keypoints'], flipPoint2D(annots1['keypoints'])]),
        }
        return images, annots

class ImageFolderMirror(ImageFolder):
    def normal(self, nf):
        annname = join(self.annot_root, self.annotlist[nf])
        data = read_json(annname)
        if 'vanish_point' in data.keys():
            vp1 = np.array(data['vanish_point'][1])
            vp1[2] = 1
            K = self.camera(nf)['K']
            normal = np.linalg.inv(K) @ vp1.reshape(3, 1)
            normal = normal.T / np.linalg.norm(normal)
        else:
            normal = None
        # normal: (1, 3)
        return normal

    def normal_all(self, start, end):
        normals = []
        for nf in range(start, end):
            annname = join(self.annot_root, self.annotlist[nf])
            data = read_json(annname)
            if 'vanish_point' in data.keys():
                vp1 = np.array(data['vanish_point'][1])
                vp1[2] = 1
                K = self.camera(nf)['K']
                normal = np.linalg.inv(K) @ vp1.reshape(3, 1)
                normal = normal.T / np.linalg.norm(normal)
                normals.append(normal)
        # nFrames, 1, 3
        if len(normals) > 0:
            normals = np.stack(normals)
        else:
            normals = None
        return normals

if __name__ == "__main__":
    pass