| import torch | |
| import math | |
| from pathlib import Path | |
| import time | |
| def euler_rotation(rx=0, ry=0, rz=0, backwards=False): | |
| Ms = [] | |
| if rz != 0: | |
| cosz = math.cos(rz) | |
| sinz = math.sin(rz) | |
| Ms.append(torch.tensor( | |
| [[cosz, -sinz, 0], | |
| [sinz, cosz, 0], | |
| [0, 0, 1]])) | |
| if ry != 0: | |
| cosy = math.cos(ry) | |
| siny = math.sin(ry) | |
| Ms.append(torch.tensor( | |
| [[cosy, 0, siny], | |
| [0, 1, 0], | |
| [-siny, 0, cosy]])) | |
| if rx != 0: | |
| cosx = math.cos(rx) | |
| sinx = math.sin(rx) | |
| Ms.append(torch.tensor( | |
| [[1, 0, 0], | |
| [0, cosx, -sinx], | |
| [0, sinx, cosx]])) | |
| rotation = torch.eye(3) | |
| if backwards and len(Ms) > 0: | |
| Ms = Ms[::-1] | |
| for mat in Ms[::-1]: | |
| rotation = torch.matmul(rotation, mat) | |
| return rotation | |
| def export(file, vs, faces, vn=None, color=None): | |
| with open(file, 'w+') as f: | |
| for vi, v in enumerate(vs): | |
| if color is None: | |
| f.write("v %f %f %f\n" % (v[0], v[1], v[2])) | |
| else: | |
| f.write("v %f %f %f %f %f %f\n" % (v[0], v[1], v[2], color[vi][0], color[vi][1], color[vi][2])) | |
| if vn is not None: | |
| f.write("vn %f %f %f\n" % (vn[vi, 0], vn[vi, 1], vn[vi, 2])) | |
| for face in faces: | |
| f.write("f %d %d %d\n" % (face[0] + 1, face[1] + 1, face[2] + 1)) | |
| def xyz2tensor(txt, append_normals=False): | |
| pts = [] | |
| for line in txt.split('\n'): | |
| line = line.strip() | |
| line = line.lstrip('v ') | |
| spt = line.split(' ') | |
| if 'nan' in line: | |
| continue | |
| if len(spt) == 6: | |
| pts.append(torch.tensor([float(x) for x in spt])) | |
| if len(spt) == 3: | |
| t = [float(x) for x in spt] | |
| if append_normals: | |
| t += [0.0 for _ in range(3)] | |
| pts.append(torch.tensor(t)) | |
| rtn = torch.stack(pts, dim=0) | |
| return rtn | |
| def read_xyz_file(path: Path): | |
| with open(path, 'r') as file: | |
| return xyz2tensor(file.read(), append_normals=True) | |
| def embed_color(img: torch.Tensor, color, box_size=70): | |
| shp = img.shape | |
| D2 = [shp[2] - box_size, shp[2]] | |
| D3 = [shp[3] - box_size, shp[3]] | |
| img = img.clone() | |
| img[:, :3, D2[0]:D2[1], D3[0]:D3[1]] = color[:, :, None, None] | |
| if img.shape[1] == 4: | |
| img[:, -1, D2[0]:D2[1], D3[0]:D3[1]] = 1 | |
| return img | |
| def get_n_params(model): | |
| pp=0 | |
| for p in list(model.parameters()): | |
| nn=1 | |
| for s in list(p.size()): | |
| nn = nn*s | |
| pp += nn | |
| return pp | |
| def xyz2tensor(txt, append_normals=False): | |
| pts = [] | |
| for line in txt.split('\n'): | |
| line = line.strip() | |
| line = line.lstrip('v ') | |
| spt = line.split(' ') | |
| if 'nan' in line: | |
| continue | |
| if len(spt) == 6: | |
| pts.append(torch.tensor([float(x) for x in spt])) | |
| if len(spt) == 3: | |
| t = [float(x) for x in spt] | |
| if append_normals: | |
| t += [0.0 for _ in range(3)] | |
| pts.append(torch.tensor(t)) | |
| rtn = torch.stack(pts, dim=0) | |
| return rtn | |
| def read_xyz_file(path: Path): | |
| with open(path, 'r') as file: | |
| return xyz2tensor(file.read(), append_normals=True) | |