189 lines
7.0 KiB
Python
189 lines
7.0 KiB
Python
'''
|
|
@ 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 |