Spaces:
Sleeping
Sleeping
| import torch | |
| import numpy as np | |
| import pyrender | |
| import cv2 | |
| from pytorch3d.renderer import look_at_view_transform | |
| def Rotate_y_180(X, pos='right'): | |
| R = torch.eye(3).to(X.device) | |
| R[0,0] = -1.0 | |
| R[2,2] = -1.0 | |
| if pos == 'right': | |
| X = torch.matmul(X, R) | |
| else: | |
| X = torch.matmul(R, X) | |
| return X | |
| def Rotate_z_180(X, pos='right'): | |
| R = torch.eye(3).to(X.device) | |
| R[0,0] = -1.0 | |
| R[1,1] = -1.0 | |
| if pos == 'right': | |
| X = torch.matmul(X, R) | |
| else: | |
| X = torch.matmul(R, X) | |
| return X | |
| class Renderer_pyr(): | |
| def __init__(self, intrinsic, extrinsic, image_size=2048): | |
| self.name = 'pyrender' | |
| self.image_size = image_size | |
| self.R = extrinsic[0:3, 0:3] | |
| self.T = extrinsic[0:3, 3:4] | |
| self.K = intrinsic | |
| R_pyrender = Rotate_y_180(torch.inverse(self.R), 'right').float() | |
| T_pyrender = -torch.matmul(torch.inverse(self.R), self.T)[:,0].float() | |
| self.camera = pyrender.IntrinsicsCamera(self.K[0,0], self.K[1,1], self.image_size - self.K[0,2], self.image_size - self.K[1,2]) | |
| self.camera_pose = np.eye(4) | |
| self.camera_pose[0:3,0:3] = R_pyrender.cpu().numpy() | |
| self.camera_pose[0:3,3] = T_pyrender.cpu().numpy() | |
| self.lights = [] | |
| light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=100.0) | |
| light_pose1 = np.array([[1.0, 0.0, 0.0, -3.0], | |
| [0.0, 1.0, 0.0, -3.0], | |
| [0.0, 0.0, 1.0, -1.0], | |
| [0.0, 0.0, 0.0, 1.0],]) | |
| light_pose2 = np.array([[1.0, 0.0, 0.0, 0.0], | |
| [0.0, 1.0, 0.0, -6.0], | |
| [0.0, 0.0, 1.0, -1.0], | |
| [0.0, 0.0, 0.0, 1.0],]) | |
| light_pose3 = np.array([[1.0, 0.0, 0.0, 3.0], | |
| [0.0, 1.0, 0.0, 3.0], | |
| [0.0, 0.0, 1.0, -1.0], | |
| [0.0, 0.0, 0.0, 1.0],]) | |
| light_pose4 = np.array([[1.0, 0.0, 0.0, 0.0], | |
| [0.0, 1.0, 0.0, 1.0], | |
| [0.0, 0.0, 1.0, -5.0], | |
| [0.0, 0.0, 0.0, 1.0],]) | |
| self.lights.append((light, light_pose1)) | |
| self.lights.append((light, light_pose2)) | |
| self.lights.append((light, light_pose3)) | |
| self.lights.append((light, light_pose4)) | |
| def render(self, mesh, return_mask=False): | |
| scene = pyrender.Scene(ambient_light=[0.1, 0.1, 0.1], bg_color=[0.0, 0.0, 0.0]) | |
| scene.add(mesh) | |
| for light in self.lights: | |
| scene.add(light[0], pose=light[1]) | |
| scene.add(self.camera, pose=self.camera_pose) | |
| osr = pyrender.OffscreenRenderer(self.image_size, self.image_size) | |
| color, depth = osr.render(scene) | |
| color = cv2.flip(color, -1) | |
| depth = cv2.flip(depth, -1) | |
| if return_mask: | |
| return color, (depth > 0).astype(np.uint8) * 255 | |
| else: | |
| return color | |
| def save(self, path): | |
| intrinsic = self.K.numpy() | |
| extrinsic = np.hstack((self.R.numpy(), self.T.numpy())) | |
| np.savez(path, intrinsic=intrinsic, extrinsic=extrinsic) |