For when the situation calms down

app.py

@@ -1,6 +1,7 @@
 from pulsar_clip import PulsarCLIP, CONFIG_SPEC
 from datetime import datetime
 import gradio as gr
+import utils
 
 
 def generate(*args):
@@ -14,7 +15,8 @@ def generate(*args):
     from tqdm.auto import tqdm
     from subprocess import Popen, PIPE
     fps = 30
-    video_path = f"{datetime.strftime(datetime.now(), '%Y-%m-%d-%H-%M-%S')}.mp4"
+    filename = datetime.strftime(datetime.now(), "%Y-%m-%d-%H-%M-%S")
+    video_path = f"{filename}.mp4"
     if frames:
         p = Popen((f"ffmpeg -y -f image2pipe -vcodec png -r {fps} -i - -vcodec libx264 -r {fps} "
                    f"-pix_fmt yuv420p -crf 17 -preset fast ").split() + [str(video_path)], stdin=PIPE)
@@ -22,7 +24,9 @@ def generate(*args):
             im.save(p.stdin, "PNG")
         p.stdin.close()
         p.wait()
-    model_path = None  # TODO
+    model_path = f"{filename}.ply"
+    model.save_ply(model_path)
+    # model_path = None  # TODO
     return [video_path, model_path]
 
 

pulsar_clip.py

@@ -224,6 +224,10 @@ class PulsarCLIP(object):
                     #yield img
         except KeyboardInterrupt:
             pass
+    
+    
+    def save_ply(self, fn):
+        utils.save_ply(self.get_points(), fn)
 
 
 class DotDict(dict):

utils.py

@@ -1,3 +1,4 @@
+import numpy as np
 import random
 import torch
 import math
@@ -69,3 +70,83 @@ def load_clip(model_name="ViT-B/16", device="cuda:0" if torch.cuda.is_available(
     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")