| 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) |
|
|
