Spaces:
Build error
Build error
| import numpy as np | |
| import random | |
| import torch | |
| import math | |
| def rotate_axis(x, add_angle=0, axis=1): # TODO Replace with a rotation matrix # But this is more fun | |
| axes = list(range(3)) | |
| axes.remove(axis) | |
| ax1, ax2 = axes | |
| angle = torch.atan2(x[..., ax1], x[..., ax2]) | |
| if isinstance(add_angle, torch.Tensor): | |
| while add_angle.ndim < angle.ndim: | |
| add_angle = add_angle.unsqueeze(-1) | |
| angle = angle + add_angle | |
| dist = x.norm(dim=-1) | |
| t = [] | |
| _, t = zip(*sorted([ | |
| (axis, x[..., axis]), | |
| (ax1, torch.sin(angle) * dist), | |
| (ax2, torch.cos(angle) * dist), | |
| ])) | |
| return torch.stack(t, dim=-1) | |
| noise_level = 0.5 | |
| # stolen from https://gist.github.com/ac1b097753f217c5c11bc2ff396e0a57 | |
| # ported from https://github.com/pvigier/perlin-numpy/blob/master/perlin2d.py | |
| def rand_perlin_2d(shape, res, fade=lambda t: 6 * t ** 5 - 15 * t ** 4 + 10 * t ** 3): | |
| delta = (res[0] / shape[0], res[1] / shape[1]) | |
| d = (shape[0] // res[0], shape[1] // res[1]) | |
| grid = torch.stack(torch.meshgrid(torch.arange(0, res[0], delta[0]), torch.arange(0, res[1], delta[1])), dim=-1) % 1 | |
| angles = 2 * math.pi * torch.rand(res[0] + 1, res[1] + 1) | |
| gradients = torch.stack((torch.cos(angles), torch.sin(angles)), dim=-1) | |
| tile_grads = lambda slice1, slice2: gradients[slice1[0]:slice1[1], slice2[0]:slice2[1]].repeat_interleave(d[0], | |
| 0).repeat_interleave( | |
| d[1], 1) | |
| dot = lambda grad, shift: ( | |
| torch.stack((grid[:shape[0], :shape[1], 0] + shift[0], grid[:shape[0], :shape[1], 1] + shift[1]), | |
| dim=-1) * grad[:shape[0], :shape[1]]).sum(dim=-1) | |
| n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0]) | |
| n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0]) | |
| n01 = dot(tile_grads([0, -1], [1, None]), [0, -1]) | |
| n11 = dot(tile_grads([1, None], [1, None]), [-1, -1]) | |
| t = fade(grid[:shape[0], :shape[1]]) | |
| return math.sqrt(2) * torch.lerp(torch.lerp(n00, n10, t[..., 0]), torch.lerp(n01, n11, t[..., 0]), t[..., 1]) | |
| def rand_perlin_2d_octaves(shape, res, octaves=1, persistence=0.5): | |
| noise = torch.zeros(shape) | |
| frequency = 1 | |
| amplitude = 1 | |
| for _ in range(octaves): | |
| noise += amplitude * rand_perlin_2d(shape, (frequency * res[0], frequency * res[1])) | |
| frequency *= 2 | |
| amplitude *= persistence | |
| noise *= random.random() - noise_level # haha | |
| noise += random.random() - noise_level # haha x2 | |
| return noise | |
| def load_clip(model_name="ViT-B/16", device="cuda:0" if torch.cuda.is_available() else "cpu"): | |
| import clip | |
| model, preprocess = clip.load(model_name, device=device, jit=False) | |
| if len(preprocess.transforms) > 4: | |
| preprocess.transforms = preprocess.transforms[-1:] | |
| return model, preprocess | |
| # http://blog.andreaskahler.com/2009/06/creating-icosphere-mesh-in-code.html | |
| def ico(): | |
| phi = (1 + 5 ** 0.5) / 2 | |
| return ( | |
| np.array([ | |
| [-1, phi, 0], | |
| [1, phi, 0], | |
| [-1, -phi, 0], | |
| [1, -phi, 0], | |
| [0, -1, phi], | |
| [0, 1, phi], | |
| [0, -1, -phi], | |
| [0, 1, -phi], | |
| [phi, 0, -1], | |
| [phi, 0, 1], | |
| [-phi, 0, -1], | |
| [-phi, 0, 1] | |
| ]) / phi, | |
| [ | |
| [0, 11, 5], | |
| [0, 5, 1], | |
| [0, 1, 7], | |
| [0, 7, 10], | |
| [0, 10, 11], | |
| [1, 5, 9], | |
| [5, 11, 4], | |
| [11, 10, 2], | |
| [10, 7, 6], | |
| [7, 1, 8], | |
| [3, 9, 4], | |
| [3, 4, 2], | |
| [3, 2, 6], | |
| [3, 6, 8], | |
| [3, 8, 9], | |
| [4, 9, 5], | |
| [2, 4, 11], | |
| [6, 2, 10], | |
| [8, 6, 7], | |
| [9, 8, 1] | |
| ] | |
| ) | |
| def ico_at(xyz=np.array([0, 0, 0]), radius=1.0, i=0): | |
| vert, idx = ico() | |
| return vert * radius + xyz, [[y + i for y in x] for x in idx] | |
| def save_ply(points, out_path): | |
| with torch.inference_mode(): | |
| vert_pos, vert_col, vert_rad, vert_opa = (x.detach().cpu().numpy() for x in points) | |
| verts = [] | |
| faces = [] | |
| for xyz, radius, (r, g, b), a in zip(vert_pos, vert_rad, vert_col, vert_opa): | |
| v, i = ico_at(xyz, radius, len(verts)) | |
| for x, y, z in v: | |
| verts.append((x, y, z, int(r * 255), int(g * 255), int(b * 255), int(a * 255))) | |
| faces += i | |
| with open(out_path, "w") as out_file: | |
| out_file.write("ply\n") | |
| out_file.write("format ascii 1.0\n") | |
| out_file.write(f"element vertex {len(verts)}\n") | |
| out_file.write("property float x\n") | |
| out_file.write("property float y\n") | |
| out_file.write("property float z\n") | |
| out_file.write("property uchar red\n") | |
| out_file.write("property uchar green\n") | |
| out_file.write("property uchar blue\n") | |
| out_file.write("property uchar alpha\n") | |
| out_file.write(f"element face {len(faces)}\n") | |
| out_file.write("property list uchar int vertex_index\n") | |
| out_file.write("end_header\n") | |
| for v in verts: | |
| out_file.write(" ".join(map(str, v)) + "\n") | |
| for f in faces: | |
| out_file.write(" ".join(map(str, [len(f)] + f)) + "\n") | |