pengc02's picture
all
ec9a6bc
import torch
import numpy as np
import os
import cv2
import glob
import trimesh
import pyrender
from PIL import Image
from lib.Renderer import Renderer_pyr
def get_cam_ssns(dataset_fpath):
cam_ssns = [] # sorted according to their spatial position
with open(os.path.join(dataset_fpath, 'cam_ssns.txt'), 'r') as fp:
lns = fp.readlines()
for ln in lns:
ln = ln.strip().split(' ')
if len(ln) > 0:
cam_ssns.append(ln[0])
return cam_ssns
class CalibRecorder():
def __init__(self, data_folder, param_folder):
self.data_folder = data_folder
self.camera_ids = ['01', '02', '04', '05']
self.param_folder = param_folder
os.makedirs(self.param_folder, exist_ok=True)
# self.camera_ids = ['220700191', '221501007', '222200036', '222200037', '222200038', '222200039', '222200040', '222200041',
# '222200042', '222200043', '222200044', '222200045', '222200046', '222200047', '222200048', '222200049']
def log(self, log_data):
eids = log_data['eids']
landmarks_gt = log_data['landmarks_gt']
landmarks_2d = log_data['landmarks_2d']
bfm = log_data['bfm']
intrinsics = log_data['intrinsics']
extrinsics = log_data['extrinsics']
with torch.no_grad():
vertices0, landmarks0 = log_data['bfm']()
# landmarks = torch.cat([landmarks0, vertices0[:, ::100, :]], 1)
for n, eid in enumerate(eids):
os.makedirs(os.path.join(self.param_folder, eid), exist_ok=True)
bfm.save('%s/params.npz' % (os.path.join(self.param_folder, eid)), batch_id=n)
np.save('%s/lmk_3d.npy' % (os.path.join(self.param_folder, eid)), landmarks0[n].cpu().numpy())
np.save('%s/vertices.npy' % (os.path.join(self.param_folder, eid)), vertices0[n].cpu().numpy())
# np.save('%s/keypoints_face.npy' % (os.path.join(self.param_folder, eid)), landmarks[n].cpu().numpy())
# for v in range(0, 1):
# vertices = vertices0[n].cpu().numpy()
# faces = log_data['bfm'].faces.cpu().numpy()
# mesh_trimesh = trimesh.Trimesh(vertices=vertices, faces=faces)
# mesh = pyrender.Mesh.from_trimesh(mesh_trimesh)
# renderer = Renderer_pyr(intrinsic=intrinsics[n, v], extrinsic=extrinsics[n, v])
# render_image = renderer.render(mesh)
# cv2.imwrite('%s/vis_%d.jpg' % (os.path.join(self.param_folder, eid), v+1), render_image[:,:,::-1])