Spaces:
Runtime error
Runtime error
| import math | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| import cv2 as cv | |
| import trimesh | |
| def pos_map_to_mesh(pos_map: torch.Tensor): | |
| pd = (0, 0, 1 if pos_map.shape[1] % 2 == 1 else 0, 0, 1 if pos_map.shape[0] % 2 == 1 else 0, 0) | |
| pos_map = F.pad(pos_map, pd, 'constant', 0) | |
| mask = torch.linalg.norm(pos_map, dim = -1) > 0.1 | |
| # cv.imshow('mask', mask.cpu().numpy().astype(np.uint8) * 255) | |
| mask = cv.erode(mask.cpu().numpy().astype(np.uint8), (5, 5), iterations = 20) | |
| mask = torch.from_numpy(mask > 0).to(pos_map.device) | |
| # cv.imshow('mask_eroded', mask.cpu().numpy().astype(np.uint8) * 255) | |
| # cv.waitKey(0) | |
| v0 = pos_map[:-1, :-1].reshape(-1, 3) | |
| v1 = pos_map[1:, :-1].reshape(-1, 3) | |
| v2 = pos_map[:-1, 1:].reshape(-1, 3) | |
| v3 = pos_map[1:, 1:].reshape(-1, 3) | |
| m0 = mask[:-1, :-1].reshape(-1) | |
| m1 = mask[1:, :-1].reshape(-1) | |
| m2 = mask[:-1, 1:].reshape(-1) | |
| m3 = mask[1:, 1:].reshape(-1) | |
| vertices = torch.cat([v0, v1, v2, v3], 0) | |
| masks = torch.cat([m0, m1, m2, m3], 0) | |
| pnum = v0.shape[0] | |
| a = torch.arange(0, pnum).to(torch.int64).to(pos_map.device) | |
| f1 = torch.stack([a, a + pnum, a + pnum * 2], 1) | |
| f2 = torch.stack([a + pnum, a + pnum * 3, a + pnum * 2], 1) | |
| faces = torch.cat([f1, f2], 0) | |
| # remove invalid faces | |
| face_mask = masks[faces.reshape(-1)].reshape(-1, 3).sum(1) == 3 | |
| face_mask = torch.logical_and(face_mask, torch.linalg.norm(vertices[faces[:, 0]] - vertices[faces[:, 1]], dim = 1) < 0.02) | |
| face_mask = torch.logical_and(face_mask, torch.linalg.norm(vertices[faces[:, 1]] - vertices[faces[:, 2]], dim = 1) < 0.02) | |
| face_mask = torch.logical_and(face_mask, torch.linalg.norm(vertices[faces[:, 0]] - vertices[faces[:, 2]], dim = 1) < 0.02) | |
| valid_faces = faces[face_mask] | |
| # debug | |
| mesh = trimesh.Trimesh(vertices = vertices.cpu().numpy(), faces = valid_faces.cpu().numpy()) | |
| mesh.export('./debug/mesh.obj') | |
| exit(1) | |
| def to_HSV(c: torch.Tensor): | |
| """ | |
| :param c: (N, 1) or (N,) | |
| :return: (N, 3) | |
| """ | |
| h = (1 - c) * 240. / 60. | |
| x = 1 - torch.abs(h.to(torch.int64) % 2 + h - h.to(torch.int64) - 1.) | |
| rgb = torch.zeros((c.shape[0], 3)).to(c).to(torch.int64) | |
| cond_1 = torch.logical_and(h >= 0, h < 1) | |
| rgb[cond_1, 0] = 255 | |
| rgb[cond_1, 1] = (x[cond_1] * 255).to(torch.int64) | |
| cond_2 = torch.logical_and(h >= 1, h < 2) | |
| rgb[cond_2, 0] = (x[cond_2] * 255).to(torch.int64) | |
| rgb[cond_2, 1] = 255 | |
| cond_3 = torch.logical_and(h >= 2, h < 3) | |
| rgb[cond_3, 1] = 255 | |
| rgb[cond_3, 2] = (x[cond_3] * 255).to(torch.int64) | |
| cond_4 = h >= 3 | |
| rgb[cond_4, 1] = (x[cond_4] * 255).to(torch.int64) | |
| rgb[cond_4, 2] = 255 | |
| rgb.clip_(0, 255) | |
| return rgb.to(torch.uint8) | |
| # def calc_back_mv(dist): | |
| # rot_center = np.array([0, 0, dist], np.float32) | |
| # trans_mat = np.identity(4, np.float32) | |
| # trans_mat[:3, :3] = cv.Rodrigues(np.array([0, math.pi, 0]))[0] | |
| # trans_mat[:3, 3] = (np.identity(3) - trans_mat[:3, :3]) @ rot_center | |
| # | |
| # return trans_mat | |
| def calc_front_mv(object_center, tar_pos = np.array([0, 0, 2.0])): | |
| """ | |
| calculate an extrinsic matrix for rendering the front of a 3D object | |
| under the assumption of fx,fy=550, cx,cy=256, img_h,img_w=512 | |
| :param object_center: np.ndarray, (3,): the original center of the 3D object | |
| :param tar_pos: np.ndarray, (3,): the target center of the 3D object | |
| :return: extr_mat: np.ndarray, (4, 4) | |
| """ | |
| mat_2origin = np.identity(4, np.float32) | |
| mat_2origin[:3, 3] = -object_center | |
| mat_rotX = np.identity(4, np.float32) | |
| mat_rotX[:3, :3] = cv.Rodrigues(np.array([math.pi, 0, 0]))[0] | |
| mat_2tarPos = np.identity(4, np.float32) | |
| mat_2tarPos[:3, 3] = tar_pos | |
| extr_mat = mat_2tarPos @ mat_rotX @ mat_2origin | |
| return extr_mat | |
| def calc_back_mv(object_center, tar_pos = np.array([0, 0, 2.0])): | |
| """ | |
| calculate an extrinsic matrix for rendering the back of a 3D object | |
| under the assumption of fx,fy=550, cx,cy=256, img_h,img_w=512 | |
| :param object_center: np.ndarray, (3,): the original center of the 3D object | |
| :param tar_pos: np.ndarray, (3,): the target center of the 3D object | |
| :return: extr_mat: np.ndarray, (4, 4) | |
| """ | |
| mat_2origin = np.identity(4, np.float32) | |
| mat_2origin[:3, 3] = -object_center | |
| mat_rotX = np.identity(4, np.float32) | |
| mat_rotX[:3, :3] = cv.Rodrigues(np.array([math.pi, 0, 0]))[0] | |
| mat_rotY = np.identity(4, np.float32) | |
| mat_rotY[:3, :3] = cv.Rodrigues(np.array([0, math.pi, 0]))[0] | |
| mat_2tarPos = np.identity(4, np.float32) | |
| mat_2tarPos[:3, 3] = tar_pos | |
| extr_mat = mat_2tarPos @ mat_rotY @ mat_rotX @ mat_2origin | |
| return extr_mat | |
| def calc_free_mv(object_center, tar_pos = np.array([0, 0, 2.0]), rot_Y = 0., rot_X = 0., global_orient = None): | |
| """ | |
| calculate an extrinsic matrix for rendering the back of a 3D object | |
| under the assumption of fx,fy=550, cx,cy=256, img_h,img_w=512 | |
| :param object_center: np.ndarray, (3,): the original center of the 3D object | |
| :param tar_pos: np.ndarray, (3,): the target center of the 3D object | |
| :param rot_Y: float, rotation angle along Y axis | |
| :param global_orient: np.ndarray, global orientation of the 3D object | |
| :return: extr_mat: np.ndarray, (4, 4) | |
| """ | |
| # import ipdb; ipdb.set_trace() | |
| mat_2origin = np.identity(4, np.float32) | |
| mat_2origin[:3, 3] = -object_center | |
| mat_invGlobalOrient = np.identity(4, np.float32) | |
| if global_orient is not None: | |
| mat_invGlobalOrient[:3, :3] = cv.Rodrigues(np.array([math.pi, 0., 0.]))[0] @ np.linalg.inv(global_orient) | |
| else: | |
| mat_invGlobalOrient[:3, :3] = cv.Rodrigues(np.array([math.pi, 0., 0.]))[0] | |
| mat_rotY = np.identity(4, np.float32) | |
| mat_rotY[:3, :3] = cv.Rodrigues(np.array([0, rot_Y, 0]))[0] | |
| mat_rotX = np.identity(4, np.float32) | |
| mat_rotX[:3, :3] = cv.Rodrigues(np.array([rot_X, 0, 0]))[0] | |
| mat_2tarPos = np.identity(4, np.float32) | |
| mat_2tarPos[:3, 3] = tar_pos | |
| extr_mat = mat_2tarPos @ mat_rotX @ mat_rotY @ mat_invGlobalOrient @ mat_2origin | |
| return extr_mat | |
| def calculate_cano_front_mv(mesh_center): | |
| if isinstance(mesh_center, torch.Tensor): | |
| mesh_center = mesh_center.cpu().numpy() | |
| front_mv = np.identity(4, np.float32) | |
| front_mv[:3, 3] = -mesh_center + np.array([0, 0, -10], np.float32) | |
| front_mv[1:3] *= -1 | |
| return front_mv | |
| def calculate_cano_back_mv(mesh_center): | |
| if isinstance(mesh_center, torch.Tensor): | |
| mesh_center = mesh_center.cpu().numpy() | |
| back_mv = np.identity(4, np.float32) | |
| rot_y = cv.Rodrigues(np.array([0, np.pi, 0], np.float32))[0] | |
| back_mv[:3, :3] = rot_y | |
| back_mv[:3, 3] = -rot_y @ mesh_center + np.array([0, 0, -10], np.float32) | |
| back_mv[1:3] *= -1 | |
| return back_mv | |
| def paper_visualize_pos_map(pos_map): | |
| mask = np.linalg.norm(pos_map, axis = -1) > 1e-6 | |
| valid_pos = pos_map[mask] | |
| min_xyz = valid_pos.min(0)[None] | |
| max_xyz = valid_pos.max(0)[None] | |
| normalized_pos = (valid_pos - min_xyz) / (max_xyz - min_xyz) | |
| pos_map[mask] = normalized_pos | |
| pos_map[~mask] = np.array([0.5, 0.5, 0.5]) | |
| return pos_map | |
| def paper_visualize_gaussian_map(gaussian_map): | |
| mask = np.linalg.norm(gaussian_map, axis = -1) > 1e-6 | |
| valid_gaussians = gaussian_map[mask] | |
| u, s, v = np.linalg.svd(valid_gaussians.transpose()) | |
| print(u, s, v) | |