Spaces:
Paused
Paused
| import taichi_three as t3 | |
| import numpy as np | |
| from taichi_three.transform import * | |
| import math | |
| from pathlib import Path | |
| from tqdm import tqdm | |
| import os | |
| import cv2 | |
| import pickle | |
| import json | |
| os.environ["KMP_DUPLICATE_LIB_OK"] = "True" | |
| def extr2tf_mat(extr): | |
| """ | |
| Return C2W matrix in the shape of (4, 4). | |
| """ | |
| # flip y and z axis in the camera coordinate space, according to | |
| # [ns-data conventions](https://docs.nerf.studio/quickstart/data_conventions.html#camera-view-space) | |
| trans_mat = np.eye(4) | |
| trans_mat[1, 1] = -1 | |
| trans_mat[2, 2] = -1 | |
| tf_mat = np.stack([extr[0], extr[1], extr[2], [0, 0, 0, 1]], axis=0) | |
| tf_mat = trans_mat @ tf_mat | |
| tf_mat = np.linalg.inv(tf_mat) | |
| return tf_mat | |
| def intr2dict(intr: np.ndarray): | |
| intr_dict = { | |
| "fl_x": intr[0, 0], | |
| "fl_y": intr[1, 1], | |
| "cx": intr[0, 2], | |
| "cy": intr[1, 2], | |
| } | |
| return intr_dict | |
| def save(data_id, save_path, extrs, intrs, depths, imgs, masks): | |
| instance_path = os.path.join(save_path, data_id) | |
| img_save_path = os.path.join(save_path, data_id, 'images') | |
| depth_save_path = os.path.join(save_path, data_id, 'depth') | |
| mask_save_path = os.path.join(save_path, data_id, 'mask') | |
| Path(img_save_path).mkdir(exist_ok=True, parents=True) | |
| Path(mask_save_path).mkdir(exist_ok=True, parents=True) | |
| Path(depth_save_path).mkdir(exist_ok=True, parents=True) | |
| h, w = imgs[0].shape[:2] | |
| transforms_dict = { | |
| "w": w, | |
| "h": h, | |
| "k1": 0, "k2": 0, "p1": 0, "p2": 0, | |
| "camera_model": "OPENCV", | |
| } | |
| frames = [] | |
| num_views = len(extrs) | |
| for pid in range(num_views): | |
| # nyw note: depth is normalized to [0, 1] in taichi_three，so we need to scale it back to [0, 2^15] | |
| depth = depths[pid] * 2.0 ** 15 | |
| cv2.imwrite(os.path.join(depth_save_path, f'{pid:03d}.png'), depth.astype(np.uint16)) | |
| img = (np.clip(imgs[pid], 0, 1) * 255.0 + 0.5).astype(np.uint8)[:, :, ::-1] | |
| cv2.imwrite(os.path.join(img_save_path, f'{pid:03d}.jpg'), img) | |
| mask = (np.clip(masks[pid], 0, 1) * 255.0 + 0.5).astype(np.uint8) | |
| cv2.imwrite(os.path.join(mask_save_path, f'{pid:03d}.png'), mask[:, :, 0]) | |
| frame_dict = { | |
| "file_path": f"images/{pid:03d}.jpg", | |
| "mask_path": f"mask/{pid:03d}.png", | |
| "depth_file_path": f"depth/{pid:03d}.png", | |
| "transform_matrix": extr2tf_mat(extrs[pid]).tolist(), | |
| } | |
| intr_dict = intr2dict(intrs[pid]) | |
| frame_dict.update(intr_dict) | |
| frames.append(frame_dict) | |
| transforms_dict["frames"] = frames | |
| with open(os.path.join(instance_path, "transforms.json"), "w") as f: | |
| json.dump(transforms_dict, f, indent=4) | |
| # with open(os.path.join(instance_path, "_transforms.json"), "w") as f: | |
| # json.dump(transforms_dict, f, indent=4) | |
| def save_normal(data_id, save_path, depths, normals): | |
| normal_save_path = os.path.join(save_path, data_id, 'normals') | |
| depth_save_path = os.path.join(save_path, data_id, 'depth_smpls') | |
| Path(normal_save_path).mkdir(exist_ok=True, parents=True) | |
| Path(depth_save_path).mkdir(exist_ok=True, parents=True) | |
| num_views = len(depths) | |
| for pid in range(num_views): | |
| # nyw note: depth is normalized to [0, 1] in taichi_three，so we need to scale it back to [0, 2^15] | |
| depth = depths[pid] * 2.0 ** 15 | |
| cv2.imwrite(os.path.join(depth_save_path, f'{pid:03d}.png'), depth.astype(np.uint16)) | |
| normal = (np.clip(normals[pid], 0, 1) * 255.0 + 0.5).astype(np.uint8)[:, :, ::-1] | |
| cv2.imwrite(os.path.join(normal_save_path, f'{pid:03d}.jpg'), normal) | |
| class StaticRenderer: | |
| def __init__(self): | |
| ti.init(arch=ti.cuda, device_memory_fraction=0.8) | |
| self.scene = t3.Scene() | |
| self.N = 10 | |
| def change_all(self): | |
| save_obj = [] | |
| save_tex = [] | |
| for model in self.scene.models: | |
| save_obj.append(model.init_obj) | |
| save_tex.append(model.init_tex) | |
| ti.init(arch=ti.cuda, device_memory_fraction=0.8) | |
| print('init') | |
| self.scene = t3.Scene() | |
| for i in range(len(save_obj)): | |
| model = t3.StaticModel(self.N, obj=save_obj[i], tex=save_tex[i]) | |
| self.scene.add_model(model) | |
| def check_update(self, obj): | |
| temp_n = self.N | |
| self.N = max(obj['vi'].shape[0], self.N) | |
| self.N = max(obj['f'].shape[0], self.N) | |
| if not (obj['vt'] is None): | |
| self.N = max(obj['vt'].shape[0], self.N) | |
| if self.N > temp_n: | |
| self.N *= 2 | |
| self.change_all() | |
| self.camera_light() | |
| def add_model(self, obj, tex=None): | |
| self.check_update(obj) | |
| model = t3.StaticModel(self.N, obj=obj, tex=tex) | |
| self.scene.add_model(model) | |
| def modify_model(self, index, obj, tex=None): | |
| self.check_update(obj) | |
| self.scene.models[index].init_obj = obj | |
| self.scene.models[index].init_tex = tex | |
| self.scene.models[index]._init() | |
| def camera_light(self): | |
| camera = t3.Camera(res=(1024, 1024)) | |
| self.scene.add_camera(camera) | |
| camera_hr = t3.Camera(res=(2048, 2048)) | |
| self.scene.add_camera(camera_hr) | |
| light_dir = np.array([0, 0, 1]) | |
| light_list = [] | |
| for l in range(6): | |
| rotate = np.matmul(rotationX(math.radians(np.random.uniform(-30, 30))), | |
| rotationY(math.radians(360 // 6 * l))) | |
| dir = [*np.matmul(rotate, light_dir)] | |
| light = t3.Light(dir, color=[1.0, 1.0, 1.0]) | |
| light_list.append(light) | |
| lights = t3.Lights(light_list) | |
| self.scene.add_lights(lights) | |
| def render_data(renderer, data_path, phase, data_id, save_path, cam_nums, res, dis=1.0, is_thuman=False, is_smpl_model=False, seed_value=0): | |
| np.random.seed(seed_value) | |
| if not is_smpl_model: | |
| obj_path = os.path.join(data_path, phase, data_id, '%s.obj' % data_id) | |
| texture_path = data_path | |
| img_path = os.path.join(texture_path, phase, data_id, 'material0.jpeg') | |
| texture = cv2.imread(img_path)[:, :, ::-1] | |
| # ################ nyw add equalizeHist for texture ################ | |
| # # comment out the following lines to disable equalizeHist for texture | |
| # texture = cv2.cvtColor(texture, cv2.COLOR_RGB2HSV) | |
| # texture[:, :, 2] = cv2.equalizeHist(texture[:, :, 2]) * 0.85 # scale down the brightness by 0.85 | |
| # texture = cv2.cvtColor(texture, cv2.COLOR_HSV2RGB) | |
| # ################ nyw add equalizeHist for texture ################ | |
| texture = np.ascontiguousarray(texture) | |
| texture = texture.swapaxes(0, 1)[:, ::-1, :] | |
| else: | |
| # obj_path = '/data1/hezijian/Thuman2.1_GPS/0000.obj' | |
| obj_path = '/data1/hezijian/Thuman2.1/THuman2.0_Smpl_X_Paras/%s/mesh_smplx.obj' % data_id | |
| obj = t3.readobj(obj_path, scale=1) | |
| # height normalization | |
| vy_max = np.max(obj['vi'][:, 1]) | |
| vy_min = np.min(obj['vi'][:, 1]) | |
| human_height = 1.80 + np.random.uniform(-0.05, 0.05, 1) | |
| obj['vi'][:, :3] = obj['vi'][:, :3] / (vy_max - vy_min) * human_height | |
| obj['vi'][:, 1] -= np.min(obj['vi'][:, 1]) | |
| look_at_center = np.array([0, 0.85, 0]) | |
| base_cam_pitch = -8 | |
| ################ nyw: add multi-pitchs for better reconstruction ################ | |
| # base_cam_pitch = -8 | |
| cam_pitchs = [-8, 45, -45, 90, -90] | |
| cam_nums_for_each_pitch = [cam_nums, cam_nums//2, cam_nums//2, 1, 1] | |
| ################ nyw: add multi-pitchs for better reconstruction ################ | |
| # randomly move the scan | |
| move_range = 0.1 if human_height < 1.80 else 0.05 | |
| delta_x = np.max(obj['vi'][:, 0]) - np.min(obj['vi'][:, 0]) | |
| delta_z = np.max(obj['vi'][:, 2]) - np.min(obj['vi'][:, 2]) | |
| if delta_x > 1.0 or delta_z > 1.0: | |
| move_range = 0.01 | |
| obj['vi'][:, 0] += np.random.uniform(-move_range, move_range, 1) | |
| obj['vi'][:, 2] += np.random.uniform(-move_range, move_range, 1) | |
| if len(renderer.scene.models) >= 1: | |
| if not is_smpl_model: | |
| renderer.modify_model(0, obj, texture) | |
| else: | |
| renderer.modify_model(0, obj) | |
| else: | |
| if not is_smpl_model: | |
| renderer.add_model(obj, texture) | |
| else: | |
| renderer.add_model(obj) | |
| if is_thuman: | |
| # thuman needs a normalization of orientation | |
| smpl_path = os.path.join(data_path, 'THuman2.0_Smpl_X_Paras', data_id, 'smplx_param.pkl') | |
| with open(smpl_path, 'rb') as f: | |
| smpl_para = pickle.load(f) | |
| y_orient = smpl_para['global_orient'][0][1] | |
| angle_base = (y_orient*180.0/np.pi) | |
| # nyw note: generate one instance of thuman in this loop | |
| extrs, intrs, depths, imgs, masks, normals = [], [], [], [], [], [] | |
| for ci, cam_pitch in enumerate(cam_pitchs): | |
| for pid in range(cam_nums := cam_nums_for_each_pitch[ci]): | |
| degree_interval = 360 / cam_nums | |
| angle = angle_base + pid * degree_interval | |
| def render(dis, angle, look_at_center, p, renderer, render_2k=False, render_normal=False): | |
| ori_vec = np.array([0, 0, dis]) | |
| rotate = np.matmul(rotationY(math.radians(angle)), rotationX(math.radians(p))) | |
| fwd = np.matmul(rotate, ori_vec) | |
| cam_pos = look_at_center + fwd | |
| x_min = 0 | |
| y_min = -25 | |
| cx = res[0] * 0.5 | |
| cy = res[1] * 0.5 | |
| fx = res[0] * 0.8 | |
| fy = res[1] * 0.8 | |
| _cx = cx - x_min | |
| _cy = cy - y_min | |
| renderer.scene.cameras[0].set_intrinsic(fx, fy, _cx, _cy) | |
| renderer.scene.cameras[0].set(pos=cam_pos, target=look_at_center) | |
| renderer.scene.cameras[0]._init() | |
| if render_2k: | |
| fx = res[0] * 0.8 * 2 | |
| fy = res[1] * 0.8 * 2 | |
| _cx = (res[0] * 0.5 - x_min) * 2 | |
| _cy = (res[1] * 0.5 - y_min) * 2 | |
| renderer.scene.cameras[1].set_intrinsic(fx, fy, _cx, _cy) | |
| renderer.scene.cameras[1].set(pos=cam_pos, target=look_at_center) | |
| renderer.scene.cameras[1]._init() | |
| renderer.scene.render() | |
| camera = renderer.scene.cameras[0] | |
| camera_hr = renderer.scene.cameras[1] | |
| extrinsic = camera.export_extrinsic() | |
| intrinsic = camera.export_intrinsic() | |
| depth_map = camera.zbuf.to_numpy().swapaxes(0, 1) | |
| img = camera.img.to_numpy().swapaxes(0, 1) | |
| img_hr = camera_hr.img.to_numpy().swapaxes(0, 1) | |
| mask = camera.mask.to_numpy().swapaxes(0, 1) | |
| return extrinsic, intrinsic, depth_map, img, mask, img_hr | |
| renderer.scene.render() | |
| camera = renderer.scene.cameras[0] | |
| extrinsic = camera.export_extrinsic() | |
| intrinsic = camera.export_intrinsic() | |
| depth_map = camera.zbuf.to_numpy().swapaxes(0, 1) | |
| if not render_normal: | |
| img = camera.img.to_numpy().swapaxes(0, 1) | |
| else: | |
| img = camera.normal_map.to_numpy().swapaxes(0, 1) | |
| mask = camera.mask.to_numpy().swapaxes(0, 1) | |
| return extrinsic, intrinsic, depth_map, img, mask | |
| if not is_smpl_model: | |
| extr, intr, depth, img, mask = render(dis, angle, look_at_center, cam_pitch, renderer) | |
| extrs.append(extr) | |
| intrs.append(intr) | |
| depths.append(depth) | |
| imgs.append(img) | |
| masks.append(mask) | |
| else: | |
| extr, intr, depth, img, mask = render(dis, angle, look_at_center, cam_pitch, renderer, render_normal=True) | |
| depths.append(depth) | |
| normals.append(img) | |
| if not is_smpl_model: | |
| save(data_id, save_path, extrs, intrs, depths, imgs, masks) | |
| else: | |
| save_normal(data_id, save_path, depths, normals) | |
| if __name__ == '__main__': | |
| cam_nums = 80 | |
| scene_radius = 2.0 | |
| res = (1024, 1024) | |
| thuman_root = '/PATH/TO/Thuman2.1/' | |
| save_root = '/PATH/TO/OUTPUT' | |
| renderer = StaticRenderer() | |
| # for phase in ['train', 'val']: | |
| phase = 'all' | |
| thuman_list = sorted(os.listdir(os.path.join(thuman_root, phase))) | |
| thuman_list = thuman_list[512:513] | |
| save_path = os.path.join(save_root, phase) | |
| seed_value = np.random.randint(1,1000) | |
| for data_id in tqdm(thuman_list): | |
| render_data(renderer, thuman_root, phase, data_id, save_path, cam_nums, res, dis=scene_radius, is_thuman=True, seed_value=seed_value) | |
| render_data(renderer, thuman_root, phase, data_id, save_path, cam_nums, res, dis=scene_radius, is_thuman=True, is_smpl_model=True, seed_value=seed_value) |