Add preview renders and upload artifacts

Dockerfile

@@ -24,7 +24,7 @@ ARG DEBIAN_FRONTEND=noninteractive
 ENV TASKS_ROOT=/app/tasks
 
 # Keep the original Python dependency footprint
-RUN python3 -m pip install --no-cache-dir ply2splat
+RUN python3 -m pip install --no-cache-dir ply2splat plyfile
 
 WORKDIR /app
 

README.md

@@ -1,7 +1,7 @@
 # Splatter Node
 This repository contains the splatter node, part of the Auki Network. This node enables photorealistic scene rendering by training 3D Gaussian Splats.
 
-The splatter node operates in conjunction with the [reconstruction node](https://github.com/aukilabs/reconstruction-server) and the scans from the Domain Management Tool (DMT) app ([App Store](https://apps.apple.com/app/domain-management-tool/id6499270503) 🔗). The refined camera poses from the reconstruction node are used as a starting point for training the gaussian splat, making it more robust to challenging indoor environments and noisy captures.
+The splatter node operates in conjunction with the [reconstruction node](https://github.com/aukilabs/reconstruction-server) and the scans from the Domain Management Tool (DMT) app ([App Store](https://apps.apple.com/app/domain-management-tool/id6499270503) 馃敆). The refined camera poses from the reconstruction node are used as a starting point for training the gaussian splat, making it more robust to challenging indoor environments and noisy captures.
 
 For more information about the reconstruction and rendering pipeline, please refer to our [whitepaper](https://auki.gitbook.io/whitepaper/technical-overview/the-reconstruction-service).
 
@@ -26,4 +26,17 @@ This project builds upon the work of many excellent open-source projects, includ
 
 We thank their authors and contributors for making this work possible.  
 Please note that all third-party code and libraries are subject to their respective licenses, copyrights, and trademarks.
-We are not affiliated with, endorsed by, or sponsored by any of the projects or organizations mentioned above.
+We are not affiliated with, endorsed by, or sponsored by any of the projects or organizations mentioned above.
+
+## Output Artifacts
+The pipeline outputs the following artifacts under `refined/splatter`:
+- `splat.ply`
+- `splat_rot.ply`
+- `splat_rot.splat` (uploaded to DMT)
+- `preview_top.jpg`
+- `preview_angle.jpg`
+- `preview.mp4`
+- `camera_paths/preview_top.json`
+- `camera_paths/preview_angle.json`
+- `camera_paths/preview_video.json`
+- `splatfacto/` (splat model)

run.py

@@ -6,6 +6,8 @@
 import os
 import time
 import json
+import math
+import shutil
 from pathlib import Path
 
 logger = logging.getLogger("splatter-node")
@@ -129,6 +131,239 @@ def run_cmd(cmd: list):
         logger.exception(f"Unexpected error while running script: {e}")
         return 1
 
+def _normalize(vec):
+    norm = math.sqrt(sum(v * v for v in vec))
+    if norm <= 1e-8:
+        return [0.0, 0.0, 0.0]
+    return [v / norm for v in vec]
+
+def _dot(a, b):
+    return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
+
+def _cross(a, b):
+    return [
+        a[1] * b[2] - a[2] * b[1],
+        a[2] * b[0] - a[0] * b[2],
+        a[0] * b[1] - a[1] * b[0],
+    ]
+
+def _view_matrix(eye, target, up_hint):
+    lookat = [target[i] - eye[i] for i in range(3)]
+    vec2 = _normalize(lookat)
+    up = up_hint[:]
+    if abs(_dot(vec2, up)) > 0.99:
+        up = [0.0, 1.0, 0.0]
+        if abs(_dot(vec2, up)) > 0.99:
+            up = [1.0, 0.0, 0.0]
+    vec0 = _normalize(_cross(up, vec2))
+    vec1 = _normalize(_cross(vec2, vec0))
+    mat = [
+        [vec0[0], vec1[0], vec2[0], eye[0]],
+        [vec0[1], vec1[1], vec2[1], eye[1]],
+        [vec0[2], vec1[2], vec2[2], eye[2]],
+        [0.0, 0.0, 0.0, 1.0],
+    ]
+    return mat
+
+def _flatten_matrix(mat):
+    return [mat[r][c] for r in range(4) for c in range(4)]
+
+def _load_ply_bounds(ply_path: Path):
+    try:
+        from plyfile import PlyData
+        import numpy as np
+    except Exception as exc:
+        logger.warning(f"preview render skipped: missing plyfile/numpy ({exc})")
+        return None
+
+    if not ply_path.exists():
+        logger.warning(f"preview render skipped: missing ply file {ply_path}")
+        return None
+
+    try:
+        ply = PlyData.read(str(ply_path))
+        verts = ply["vertex"]
+        xs = np.asarray(verts["x"], dtype=float)
+        ys = np.asarray(verts["y"], dtype=float)
+        zs = np.asarray(verts["z"], dtype=float)
+        min_x, max_x = float(xs.min()), float(xs.max())
+        min_y, max_y = float(ys.min()), float(ys.max())
+        min_z, max_z = float(zs.min()), float(zs.max())
+    except Exception as exc:
+        logger.warning(f"preview render skipped: failed to read ply ({exc})")
+        return None
+
+    center = [
+        (min_x + max_x) / 2.0,
+        (min_y + max_y) / 2.0,
+        (min_z + max_z) / 2.0,
+    ]
+    span = [
+        max_x - min_x,
+        max_y - min_y,
+        max_z - min_z,
+    ]
+    size = max(span + [1.0])
+    return center, size
+
+def _write_camera_path(path: Path, camera_path_entries, render_width, render_height, seconds, fps):
+    payload = {
+        "camera_type": "perspective",
+        "render_height": float(render_height),
+        "render_width": float(render_width),
+        "seconds": float(seconds),
+        "fps": float(fps),
+        "is_cycle": False,
+        "smoothness_value": 0.0,
+        "camera_path": camera_path_entries,
+    }
+    path.write_text(json.dumps(payload, indent=2))
+
+def _render_previews(job_root_path: Path):
+    preview_dir = job_root_path / "refined" / "splatter"
+    config_path = preview_dir / "splatfacto" / "config.yml"
+    if not config_path.exists():
+        logger.warning(f"preview render skipped: missing config {config_path}")
+        return
+
+    if shutil.which("ns-render") is None:
+        logger.warning("preview render skipped: ns-render not found")
+        return
+    ffmpeg_available = shutil.which("ffmpeg") is not None
+    if not ffmpeg_available:
+        logger.warning("preview video skipped: ffmpeg not found")
+
+    ply_path = preview_dir / "splat.ply"
+    bounds = _load_ply_bounds(ply_path)
+    if bounds is None:
+        return
+
+    center, size = bounds
+    render_width = 1280
+    render_height = 720
+    aspect = render_width / render_height
+    fov = 60.0
+    up = [0.0, 0.0, 1.0]
+
+    camera_paths_dir = preview_dir / "camera_paths"
+    camera_paths_dir.mkdir(parents=True, exist_ok=True)
+
+    def build_entry(eye, target, render_time):
+        mat = _view_matrix(eye, target, up)
+        return {
+            "camera_to_world": _flatten_matrix(mat),
+            "fov": fov,
+            "aspect": aspect,
+            "render_time": float(render_time),
+        }
+
+    # Top-down preview image
+    top_eye = [center[0], center[1], center[2] + size * 1.5]
+    top_target = center
+    top_entry = build_entry(top_eye, top_target, 0.0)
+    top_json = camera_paths_dir / "preview_top.json"
+    _write_camera_path(top_json, [top_entry], render_width, render_height, seconds=1.0, fps=1.0)
+    top_out = preview_dir / "preview_top"
+    exit_code = run_script(
+        "ns-render",
+        "camera-path",
+        "--load-config",
+        config_path,
+        "--camera-path-filename",
+        top_json,
+        "--output-path",
+        top_out,
+        "--output-format",
+        "images",
+        "--image-format",
+        "jpeg",
+        "--jpeg-quality",
+        "90",
+    )
+    if exit_code == 0:
+        top_img_dir = top_out
+        top_img = top_img_dir / "00000.jpg"
+        if top_img.exists():
+            shutil.copy2(top_img, preview_dir / "preview_top.jpg")
+            shutil.rmtree(top_img_dir, ignore_errors=True)
+    else:
+        logger.warning("preview render failed: top-down image")
+
+    # Angled preview image
+    angle_eye = [
+        center[0] + size * 1.2,
+        center[1] + size * 1.2,
+        center[2] + size * 0.8,
+    ]
+    angle_target = center
+    angle_entry = build_entry(angle_eye, angle_target, 0.0)
+    angle_json = camera_paths_dir / "preview_angle.json"
+    _write_camera_path(angle_json, [angle_entry], render_width, render_height, seconds=1.0, fps=1.0)
+    angle_out = preview_dir / "preview_angle"
+    exit_code = run_script(
+        "ns-render",
+        "camera-path",
+        "--load-config",
+        config_path,
+        "--camera-path-filename",
+        angle_json,
+        "--output-path",
+        angle_out,
+        "--output-format",
+        "images",
+        "--image-format",
+        "jpeg",
+        "--jpeg-quality",
+        "90",
+    )
+    if exit_code == 0:
+        angle_img_dir = angle_out
+        angle_img = angle_img_dir / "00000.jpg"
+        if angle_img.exists():
+            shutil.copy2(angle_img, preview_dir / "preview_angle.jpg")
+            shutil.rmtree(angle_img_dir, ignore_errors=True)
+    else:
+        logger.warning("preview render failed: angled image")
+
+    # Preview video
+    if ffmpeg_available:
+        frames = 150
+        seconds = 5.0
+        orbit_entries = []
+        radius = size * 1.5
+        height = size * 0.6
+        for i in range(frames):
+            theta = 2.0 * math.pi * (i / frames)
+            eye = [
+                center[0] + radius * math.cos(theta),
+                center[1] + radius * math.sin(theta),
+                center[2] + height,
+            ]
+            render_time = (i / max(1, frames - 1)) * seconds
+            orbit_entries.append(build_entry(eye, center, render_time))
+
+        video_json = camera_paths_dir / "preview_video.json"
+        _write_camera_path(video_json, orbit_entries, render_width, render_height, seconds=seconds, fps=30.0)
+        preview_video = preview_dir / "preview.mp4"
+        exit_code = run_script(
+            "ns-render",
+            "camera-path",
+            "--load-config",
+            config_path,
+            "--camera-path-filename",
+            video_json,
+            "--output-path",
+            preview_video,
+            "--output-format",
+            "video",
+            "--image-format",
+            "jpeg",
+            "--jpeg-quality",
+            "90",
+        )
+        if exit_code != 0:
+            logger.warning("preview render failed: video")
+
 # Example usage
 if __name__ == "__main__":
 
@@ -188,6 +423,12 @@ def run_cmd(cmd: list):
         logger.error("failed to export gaussian splat")
         sys.exit(exit_code)
 
+    logger.info("Rendering Previews")
+    try:
+        _render_previews(args.job_root_path)
+    except Exception as exc:
+        logger.warning(f"preview render failed: {exc}")
+
     logger.info("Rotating Splat")
     exit_code = run_python_script("rotate_ply.py", 
                                   "--input", args.job_root_path / "refined/splatter/splat.ply",
@@ -204,4 +445,4 @@ def run_cmd(cmd: list):
         logger.error("failed to convert splat .ply to .splat")
         sys.exit(exit_code)
 
-    sys.exit(exit_code)
+    sys.exit(exit_code)