320 lines
12 KiB
Python
320 lines
12 KiB
Python
'''
|
||
@ Date: 2021-11-22 15:16:14
|
||
@ Author: Qing Shuai
|
||
@ LastEditors: Qing Shuai
|
||
@ LastEditTime: 2022-09-28 21:40:17
|
||
@ FilePath: /EasyMocapPublic/easymocap/visualize/render_base.py
|
||
'''
|
||
from glob import glob
|
||
from os.path import join
|
||
import numpy as np
|
||
|
||
from ..mytools.file_utils import read_json
|
||
from ..mytools.debug_utils import log
|
||
from ..mytools.reader import read_keypoints3d, read_smpl
|
||
import os
|
||
from ..mytools.camera_utils import read_cameras, Undistort
|
||
import cv2
|
||
from ..mytools.vis_base import merge, plot_keypoints_auto
|
||
from ..config.baseconfig import load_object
|
||
from .geometry import load_sphere
|
||
|
||
def imwrite(imgname, img):
|
||
if not os.path.exists(os.path.dirname(imgname)):
|
||
os.makedirs(os.path.dirname(imgname))
|
||
if img.shape[0] % 2 == 1 or img.shape[1] % 2 == 1:
|
||
img = cv2.resize(img, (img.shape[1]//2*2, img.shape[0]//2*2))
|
||
cv2.imwrite(imgname, img)
|
||
|
||
def compute_normals(vertices, faces):
|
||
normal = np.zeros_like(vertices)
|
||
|
||
# compute normal per triangle
|
||
normal_faces = np.cross(vertices[faces[:,1]] - vertices[faces[:,0]], vertices[faces[:,2]] - vertices[faces[:,0]])
|
||
|
||
# sum normals at vtx
|
||
normal[faces[:, 0]] += normal_faces[:]
|
||
normal[faces[:, 1]] += normal_faces[:]
|
||
normal[faces[:, 2]] += normal_faces[:]
|
||
|
||
# compute norms
|
||
normal = normal / np.linalg.norm(normal, axis=-1, keepdims=True)
|
||
|
||
return normal
|
||
|
||
def get_dilation_of_mesh(delta):
|
||
def func(info):
|
||
vertices = info['vertices']
|
||
normals = compute_normals(info['vertices'], info['faces'])
|
||
vertices += delta * normals
|
||
return info
|
||
return func
|
||
|
||
class Results:
|
||
def __init__(self, body_model, path, rend_type,
|
||
operation='none') -> None:
|
||
self.path = path
|
||
self.body_model = body_model
|
||
self.skelnames = sorted(glob(join(path, '*.json')))
|
||
self.ismulti = False
|
||
if len(self.skelnames) == 0:
|
||
# 尝试找多视角的结果
|
||
subs = sorted(os.listdir(path))
|
||
assert len(subs) > 0, path
|
||
self.ismulti = True
|
||
self.subs = subs
|
||
self.skelnames = {}
|
||
for sub in subs:
|
||
skelnames = sorted(glob(join(path, sub, '*.json')))
|
||
self.skelnames[sub] = skelnames
|
||
self.rend_type = rend_type
|
||
self.read_func = {'skel': read_keypoints3d, 'mesh': read_smpl}[rend_type]
|
||
if operation.startswith('dilation'):
|
||
# TODO: 暂时直接解析
|
||
opfunc = get_dilation_of_mesh(float(operation.replace('dilation:', '')))
|
||
else:
|
||
opfunc = lambda x:x
|
||
self.operation = opfunc
|
||
|
||
def process(self, info):
|
||
return info
|
||
|
||
def read(self, skelname):
|
||
results = self.read_func(skelname)
|
||
render_data = {}
|
||
trans = np.array([[0., 0., 0., 0.]])
|
||
for info in results:
|
||
info['vertices'] = self.body_model(return_verts=True, return_tensor=False, **info)[0]
|
||
info['keypoints3d'] = self.body_model(return_verts=False, return_tensor=False, **info)[0]
|
||
info['vertices'] += trans[:, :3]
|
||
info['faces'] = self.body_model.faces
|
||
d = self.operation(info)
|
||
render_data[d['id']] = {
|
||
'vertices': d['vertices'], 'keypoints3d': d['keypoints3d'],
|
||
'faces': info['faces'],
|
||
'vid': d['id'], 'name': 'human_{}'.format(d['id'])}
|
||
if self.rend_type == 'skel':
|
||
render_data[d['id']]['smooth'] = False
|
||
return render_data
|
||
|
||
def get_multi(self, index):
|
||
datas = {}
|
||
for sub in self.subs:
|
||
skelname = self.skelnames[sub][index]
|
||
basename = os.path.basename(skelname).replace('.json', '')
|
||
render_data = self.read(skelname)
|
||
datas[sub] = render_data
|
||
return basename, datas
|
||
|
||
def __getitem__(self, index):
|
||
if self.ismulti:
|
||
return self.get_multi(index)
|
||
else:
|
||
skelname = self.skelnames[index]
|
||
basename = os.path.basename(skelname).replace('.json', '')
|
||
return basename, self.read(skelname)
|
||
|
||
def __len__(self):
|
||
if self.ismulti:
|
||
return len(self.skelnames[self.subs[0]])
|
||
else:
|
||
return len(self.skelnames)
|
||
|
||
class ResultsObjects(Results):
|
||
def __init__(self, object,**kwargs) -> None:
|
||
super().__init__(**kwargs)
|
||
self.object = object
|
||
|
||
def __getitem__(self, index):
|
||
basename, datas = super().__getitem__(index)
|
||
objectname = join(self.path, '..', '..', self.object, basename+'.json')
|
||
objects = read_json(objectname)
|
||
ret_objects = {}
|
||
for obj in objects:
|
||
pid = obj['id']
|
||
vertices, faces = load_sphere()
|
||
vertices *= 0.12
|
||
vertices += np.array(obj['keypoints3d'])[:, :3]
|
||
ret_objects[1000+pid] = {
|
||
'vertices': vertices,
|
||
'faces': faces,
|
||
'vid': pid,
|
||
'name': 'object_{}'.format(pid)
|
||
}
|
||
if self.ismulti:
|
||
for key, data in datas.items():
|
||
data.update(ret_objects)
|
||
else:
|
||
datas.update(ret_objects)
|
||
return basename, datas
|
||
|
||
class Images:
|
||
def __init__(self, path, subs, image_args) -> None:
|
||
if path == 'none':
|
||
self.images = path
|
||
# no need for images
|
||
assert len(subs) > 0, '{} must non-empty'.format(subs)
|
||
else:
|
||
self.images = join(path, 'images')
|
||
if len(subs) == 0:
|
||
subs = sorted(os.listdir(self.images))
|
||
if subs[0].isdigit():
|
||
subs.sort(key=lambda x:int(x))
|
||
self.cameras = read_cameras(path)
|
||
self.subs = subs
|
||
self.image_args = image_args
|
||
self.cameras_vis = {sub:cam.copy() for sub, cam in self.cameras.items()}
|
||
# rescale the camera
|
||
for cam in self.cameras_vis.values():
|
||
K = cam['K'].copy()
|
||
cam['K'][:2, :] *= image_args.scale
|
||
self.distortMap = {}
|
||
|
||
def __call__(self, basename):
|
||
if self.images == 'none':
|
||
# 返回空的图像
|
||
imgs = {sub: self.blank.copy() for sub in self.subs}
|
||
import ipdb; ipdb.set_trace()
|
||
else:
|
||
imgs = {}
|
||
for sub in self.subs:
|
||
imgname = join(self.images, sub, basename+self.image_args.ext)
|
||
if not os.path.exists(imgname):
|
||
for ext in ['.jpg', '.png']:
|
||
imgname_ = imgname.replace(self.image_args.ext, ext)
|
||
if os.path.exists(imgname_):
|
||
self.image_args.ext = ext
|
||
imgname = imgname_
|
||
assert os.path.exists(imgname), imgname
|
||
img = cv2.imread(imgname)
|
||
if self.image_args.scale != 1:
|
||
img = cv2.resize(img, None, fx=self.image_args.scale, fy=self.image_args.scale)
|
||
if self.image_args.undis:
|
||
camera = self.cameras_vis[sub]
|
||
K, D = camera['K'], camera['dist']
|
||
if sub not in self.distortMap.keys():
|
||
h, w = img.shape[:2]
|
||
mapx, mapy = cv2.initUndistortRectifyMap(camera['K'], camera['dist'], None, camera['K'], (w,h), 5)
|
||
self.distortMap[sub] = (mapx, mapy)
|
||
mapx, mapy = self.distortMap[sub]
|
||
img = cv2.remap(img, mapx, mapy, cv2.INTER_LINEAR)
|
||
imgs[sub] = img
|
||
return imgs, self.cameras_vis
|
||
|
||
class Repro:
|
||
def __init__(self, out, merge='none') -> None:
|
||
self.out = out
|
||
os.makedirs(self.out, exist_ok=True)
|
||
self.merge = merge
|
||
|
||
@staticmethod
|
||
def repro(P, keypoints3d, img):
|
||
if keypoints3d.shape[1] == 3:
|
||
keypoints3d = np.hstack((keypoints3d, np.ones((keypoints3d.shape[0], 1))))
|
||
kcam = np.hstack([keypoints3d[:, :3], np.ones((keypoints3d.shape[0], 1))]) @ P.T
|
||
kcam = kcam[:, :]/kcam[:, 2:]
|
||
k2d = np.hstack((kcam, (keypoints3d[:, 3:]>0.)&(kcam[:, 2:] >0.1)))
|
||
from ..estimator.wrapper_base import bbox_from_keypoints
|
||
bbox = bbox_from_keypoints(k2d)
|
||
return k2d, bbox
|
||
|
||
def __call__(self, images, results, cameras, basename):
|
||
imgsout = {}
|
||
cams = list(images.keys())
|
||
for nv, cam in enumerate(cams):
|
||
img = images[cam]
|
||
outname = join(self.out, basename+'_' + cam +'.jpg')
|
||
# K可能缩放过了，所以需要重新计算
|
||
P = cameras[cam]['K'] @ cameras[cam]['RT']
|
||
for pid, info in results.items():
|
||
keypoints3d = info['keypoints3d']
|
||
k2d, bbox = self.repro(P, keypoints3d, img)
|
||
lw = int(max(bbox[2] - bbox[0], bbox[3] - bbox[1])/50)
|
||
# plot_bbox(img, bbox, pid=pid, vis_id=pid)
|
||
plot_keypoints_auto(img, k2d, pid=pid, use_limb_color=False)
|
||
imgsout[outname] = img
|
||
if self.merge == 'none':
|
||
for outname, img in imgsout.items():
|
||
cv2.imwrite(outname, img)
|
||
else:
|
||
outname = join(self.out, basename+'.jpg')
|
||
out = merge(list(imgsout.values()), square=True)
|
||
cv2.imwrite(outname, out)
|
||
|
||
class Outputs:
|
||
def __init__(self, out, mode, backend, scene={}) -> None:
|
||
self.out = out
|
||
os.makedirs(self.out, exist_ok=True)
|
||
from .render_func import get_render_func, get_ext
|
||
self.render_func = get_render_func(mode, backend)
|
||
self.mode = mode
|
||
self.ext = get_ext(mode)
|
||
self.extra_mesh = {}
|
||
self.scene = []
|
||
for key, val in scene.items():
|
||
mesh = load_object(val.module, val.args)
|
||
log('[vis] Load extra mesh {}'.format(key))
|
||
self.scene.append(mesh)
|
||
|
||
def __call__(self, images, results, cameras, basename):
|
||
for i, mesh in enumerate(self.scene):
|
||
results[10000+i] = mesh
|
||
render_results = self.render_func(images, results, cameras, self.extra_mesh)
|
||
output = merge(list(render_results.values()), square=True)
|
||
if output.shape[0] > 10000:
|
||
scale = 5000./output.shape[0]
|
||
output = cv2.resize(output, None, fx=scale, fy=scale)
|
||
outname = join(self.out, basename+self.ext)
|
||
imwrite(outname, output)
|
||
return 0
|
||
|
||
class MIOutputs(Outputs):
|
||
def __init__(self, out, mode, backend, merge=True, scene={}) -> None:
|
||
super().__init__(out, mode, backend, scene=scene)
|
||
self.merge = merge
|
||
|
||
def __call__(self, images, results, cameras, basename):
|
||
# 传个subs进来
|
||
# save the results to individual folders
|
||
subs = list(images.keys())
|
||
outputs = {}
|
||
for nv, sub in enumerate(subs):
|
||
if sub in results.keys():
|
||
result = results[sub]
|
||
# consider the case that the result is not in the results
|
||
# especially in multiview results and render novel view
|
||
else:
|
||
result = results
|
||
value0 = list(result.values())[0]
|
||
if 'vertices' not in value0.keys():
|
||
continue
|
||
for i, mesh in enumerate(self.scene):
|
||
result[10000+i] = mesh
|
||
if self.mode == 'instance-mask' or self.mode == 'instance-depth':
|
||
# TODO: use depth to render instance mask and consider occlusion
|
||
for pid, value in result.items():
|
||
if 'vertices' not in value.keys():
|
||
print('[ERROR] vis view {}, pid {}'.format(sub, pid))
|
||
output = self.render_func({sub:images[sub]},
|
||
{pid:value},
|
||
{sub:cameras[sub]}, self.extra_mesh)
|
||
outname = join(self.out, sub, basename + '_{}'.format(pid) + self.ext)
|
||
imwrite(outname, output[sub])
|
||
continue
|
||
else:
|
||
output = self.render_func({sub:images[sub]},
|
||
result,
|
||
{sub:cameras[sub]}, self.extra_mesh)
|
||
outputs[sub] = output[sub]
|
||
if len(outputs.keys()) == 0:
|
||
return 0
|
||
if self.merge:
|
||
outname = join(self.out, basename + self.ext)
|
||
outputs = merge(list(outputs.values()), square=True)
|
||
imwrite(outname, outputs)
|
||
else:
|
||
for nv, sub in enumerate(subs):
|
||
if sub not in outputs.keys():continue
|
||
outname = join(self.out, sub, basename + self.ext)
|
||
imwrite(outname, outputs[sub])
|
||
return 0 |