Spaces:
Paused
Paused
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| import os | |
| import cv2 | |
| # borrowed from https://github.com/daniilidis-group/neural_renderer/blob/master/neural_renderer/vertices_to_faces.py | |
| def face_vertices(vertices, faces): | |
| """ | |
| :param vertices: [batch size, number of vertices, 3] | |
| :param faces: [batch size, number of faces, 3] | |
| :return: [batch size, number of faces, 3, 3] | |
| """ | |
| assert (vertices.ndimension() == 3) | |
| assert (faces.ndimension() == 3) | |
| assert (vertices.shape[0] == faces.shape[0]) | |
| assert (vertices.shape[2] == 3) | |
| assert (faces.shape[2] == 3) | |
| bs, nv = vertices.shape[:2] | |
| bs, nf = faces.shape[:2] | |
| device = vertices.device | |
| faces = faces + (torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None] | |
| vertices = vertices.reshape((bs * nv, 3)) | |
| # pytorch only supports long and byte tensors for indexing | |
| return vertices[faces.long()] | |
| def vertex_normals(vertices, faces): | |
| """ | |
| :param vertices: [batch size, number of vertices, 3] | |
| :param faces: [batch size, number of faces, 3] | |
| :return: [batch size, number of vertices, 3] | |
| """ | |
| assert (vertices.ndimension() == 3) | |
| assert (faces.ndimension() == 3) | |
| assert (vertices.shape[0] == faces.shape[0]) | |
| assert (vertices.shape[2] == 3) | |
| assert (faces.shape[2] == 3) | |
| bs, nv = vertices.shape[:2] | |
| bs, nf = faces.shape[:2] | |
| device = vertices.device | |
| normals = torch.zeros(bs * nv, 3).to(device) | |
| faces = faces + (torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None] # expanded faces | |
| vertices_faces = vertices.reshape((bs * nv, 3))[faces.long()] | |
| faces = faces.reshape(-1, 3) | |
| vertices_faces = vertices_faces.reshape(-1, 3, 3) | |
| normals.index_add_(0, faces[:, 1].long(), | |
| torch.cross(vertices_faces[:, 2] - vertices_faces[:, 1], vertices_faces[:, 0] - vertices_faces[:, 1])) | |
| normals.index_add_(0, faces[:, 2].long(), | |
| torch.cross(vertices_faces[:, 0] - vertices_faces[:, 2], vertices_faces[:, 1] - vertices_faces[:, 2])) | |
| normals.index_add_(0, faces[:, 0].long(), | |
| torch.cross(vertices_faces[:, 1] - vertices_faces[:, 0], vertices_faces[:, 2] - vertices_faces[:, 0])) | |
| normals = F.normalize(normals, eps=1e-6, dim=1) | |
| normals = normals.reshape((bs, nv, 3)) | |
| # pytorch only supports long and byte tensors for indexing | |
| return normals | |
| def batch_orth_proj(X, camera): | |
| ''' orthgraphic projection | |
| X: 3d vertices, [bz, n_point, 3] | |
| camera: scale and translation, [bz, 3], [scale, tx, ty] | |
| ''' | |
| camera = camera.clone().view(-1, 1, 3) | |
| X_trans = X[:, :, :2] + camera[:, :, 1:] | |
| X_trans = torch.cat([X_trans, X[:,:,2:]], 2) | |
| Xn = (camera[:, :, 0:1] * X_trans) | |
| return Xn | |
| from pytorch3d.renderer import look_at_view_transform | |
| def weak_cam2persp_cam(wcam, focal_length=12, z_dist=10): | |
| """_summary_ | |
| Args: | |
| wcam (torch.Tensor): In shape Bx3, for each [s, x, y] | |
| focal_length (int, optional): perspective camera focal length. Defaults to 12. | |
| z_dist (int, optional): perspective camera at (0, 0, z). Defaults to 10. | |
| Returns: | |
| R, T: Rotation matrix and translation vector | |
| """ | |
| bz = wcam.shape[0] | |
| R, T = look_at_view_transform(dist=z_dist, device=wcam.device) | |
| R = R.repeat(bz, 1, 1) | |
| T = T.repeat(bz, 1) | |
| T[:, 2] = focal_length / wcam[:, 0] | |
| T[:, 1] = wcam[:, 2] | |
| T[:, 0] = -wcam[:, 1] | |
| return R, T | |
| def cam2persp_cam_fov(wcam, tanfov=1/12): | |
| """ | |
| scale by changing zdist,unchange fov | |
| Returns: w2c | |
| R, T: Rotation matrix and translation vector | |
| """ | |
| # blender coord to pytorch3d coord | |
| flx=-1.0 | |
| fly= 1.0 | |
| flz=-1.0 | |
| bz = wcam.shape[0] | |
| z_dist=1/tanfov | |
| R=torch.tensor([[flx,0,0], | |
| [0,fly,0], | |
| [0,0,flz]], device=wcam.device, dtype=torch.float32) | |
| T=torch.tensor([0,0,z_dist], device=wcam.device, dtype=torch.float32) | |
| R = R.repeat(bz, 1, 1) | |
| T = T.repeat(bz, 1) | |
| T[:, 2] = T[:, 2] / wcam[:, 0] | |
| T[:, 1] = wcam[:, 2] *fly | |
| T[:, 0] = wcam[:, 1] *flx | |
| return R, T | |
| def cam2persp_cam_fov_body(wcam, tanfov=1/12): | |
| """ | |
| scale by changing zdist,unchange fov | |
| Returns: w2c | |
| R, T: Rotation matrix and translation vector | |
| """ | |
| #image coord to pytorch3d coord | |
| flx=-1.0 | |
| fly=-1.0 | |
| flz=1.0 | |
| bz = wcam.shape[0] | |
| z_dist=1/tanfov | |
| R=torch.tensor([[flx,0,0], | |
| [0,fly,0], | |
| [0,0,flz]],device=wcam.device,dtype=torch.float32) | |
| T=torch.tensor([0,0,z_dist],device=wcam.device,dtype=torch.float32) | |
| R = R.repeat(bz, 1, 1) | |
| T = T.repeat(bz, 1) | |
| T[:, 2] = T[:, 2] / wcam[:, 0] | |
| T[:, 1] = wcam[:, 2]*fly # not needed | |
| T[:, 0] = wcam[:, 1]*flx # not needed | |
| return R, T | |
| # def cam2persp_cam_fov(wcam, tanfov=1): | |
| # """ | |
| # scale by changing zdist,unchange fov | |
| # Returns: w2c | |
| # R, T: Rotation matrix and translation vector | |
| # """ | |
| # bz = wcam.shape[0] | |
| # z_dist=2/tanfov | |
| # R=torch.tensor([[1,0,0], | |
| # [0,1,0], | |
| # [0,0,1]],device=wcam.device,dtype=torch.float32) | |
| # T=torch.tensor([0,0,-z_dist],device=wcam.device,dtype=torch.float32) | |
| # R = R.repeat(bz, 1, 1) | |
| # T = T.repeat(bz, 1) | |
| # T[:, 2] = T[:, 2] / wcam[:, 0] | |
| # T[:, 1] = wcam[:, 2] | |
| # T[:, 0] = wcam[:, 1] | |
| # return R, T |