Update app.py

app.py

@@ -30,7 +30,57 @@ from src.utils.image import process_image
 
 os.environ["ANYSPLAT_PROCESSED"] = f"{os.getcwd()}/proprocess_results"
 
+from plyfile import PlyData
+import numpy as np
+import argparse
+from io import BytesIO
+
+
+def process_ply_to_splat(ply_file_path):
+    plydata = PlyData.read(ply_file_path)
+    vert = plydata["vertex"]
+    sorted_indices = np.argsort(
+        -np.exp(vert["scale_0"] + vert["scale_1"] + vert["scale_2"])
+        / (1 + np.exp(-vert["opacity"]))
+    )
+    buffer = BytesIO()
+    for idx in sorted_indices:
+        v = plydata["vertex"][idx]
+        position = np.array([v["x"], v["y"], v["z"]], dtype=np.float32)
+        scales = np.exp(
+            np.array(
+                [v["scale_0"], v["scale_1"], v["scale_2"]],
+                dtype=np.float32,
+            )
+        )
+        rot = np.array(
+            [v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]],
+            dtype=np.float32,
+        )
+        SH_C0 = 0.28209479177387814
+        color = np.array(
+            [
+                0.5 + SH_C0 * v["f_dc_0"],
+                0.5 + SH_C0 * v["f_dc_1"],
+                0.5 + SH_C0 * v["f_dc_2"],
+                1 / (1 + np.exp(-v["opacity"])),
+            ]
+        )
+        buffer.write(position.tobytes())
+        buffer.write(scales.tobytes())
+        buffer.write((color * 255).clip(0, 255).astype(np.uint8).tobytes())
+        buffer.write(
+            ((rot / np.linalg.norm(rot)) * 128 + 128)
+            .clip(0, 255)
+            .astype(np.uint8)
+            .tobytes()
+        )
 
+    return buffer.getvalue()
+
+def save_splat_file(splat_data, output_path):
+    with open(output_path, "wb") as f:
+        f.write(splat_data)
 
 def get_reconstructed_scene(outdir, model, device):
 
@@ -63,6 +113,7 @@ def get_reconstructed_scene(outdir, model, device):
     )
 
     plyfile = os.path.join(outdir, "gaussians.ply")
+    splatfile = os.path.join(outdir, "gaussians.splat")
     
     export_ply(
         gaussians.means[0],
@@ -73,10 +124,12 @@ def get_reconstructed_scene(outdir, model, device):
         Path(plyfile),
         save_sh_dc_only=True,
     )
+    splat_data = process_ply_to_splat(plyfile)
+    save_splat_file(splat_data, splatfile)
 
     # Clean up
     torch.cuda.empty_cache()
-    return plyfile, video, depth_colored
+    return splatfile, video, depth_colored
 
 
 # 2) Handle uploaded video/images --> produce target_dir + images