diff --git a/Dockerfile b/Dockerfile index 0f6859d..2450bca 100644 --- a/Dockerfile +++ b/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 diff --git a/README.md b/README.md index 95538f5..5915892 100644 --- a/README.md +++ b/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. \ No newline at end of file +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) diff --git a/run.py b/run.py index e49d601..247cd46 100644 --- a/run.py +++ b/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) \ No newline at end of file + sys.exit(exit_code) diff --git a/server/rust/runner/src/lib.rs b/server/rust/runner/src/lib.rs index 217ceec..bf1b186 100644 --- a/server/rust/runner/src/lib.rs +++ b/server/rust/runner/src/lib.rs @@ -748,18 +748,22 @@ impl compute_runner_api::Runner for HelloRunner { return Err(anyhow!("expected output missing: {}", splat_abs.display())); } - let upload_key = if let Some(suffix) = - refined_suffix.as_deref().filter(|s| !s.is_empty()) - { - if suffix.starts_with('_') { - format!("refined_splat{suffix}") + let suffix = refined_suffix.as_deref().filter(|s| !s.is_empty()); + if suffix.is_none() { + warn!("refined manifest suffix missing; uploading as splat_data without timestamp"); + } + let build_upload_key = |base: &str| -> String { + if let Some(suffix) = suffix { + if suffix.starts_with('_') { + format!("{base}{suffix}") + } else { + format!("{base}_{suffix}") + } } else { - format!("refined_splat_{suffix}") + base.to_string() } - } else { - warn!("refined manifest suffix missing; uploading as splat_data without timestamp"); - "refined_splat".to_string() }; + let upload_key = build_upload_key("refined_splat"); ctx.ctrl .progress(json!({ @@ -799,6 +803,87 @@ impl compute_runner_api::Runner for HelloRunner { "uploaded": upload_key.as_str(), })) .await; + + let preview_specs = [ + ( + "preview_top.jpg", + "refined_splat_preview_top", + "splat_preview_top", + ), + ( + "preview_angle.jpg", + "refined_splat_preview_angle", + "splat_preview_angle", + ), + ( + "preview.mp4", + "refined_splat_preview_video", + "splat_preview_video", + ), + ]; + ensure_task_not_cancelled(&ctx, "before preview upload").await?; + for (filename, key_base, data_type) in preview_specs { + let rel_path = PathBuf::from("refined").join("splatter").join(filename); + let abs_path = job_root.join(&rel_path); + if !abs_path.exists() { + let _ = ctx + .ctrl + .log_event(json!({ + "level": "info", + "stage": "upload", + "message": "preview missing; skipping upload", + "file": abs_path.display().to_string(), + "data_type": data_type, + })) + .await; + continue; + } + let upload_name = build_upload_key(key_base); + let upload_result = ctx + .output + .put_domain_artifact(compute_runner_api::runner::DomainArtifactRequest { + rel_path: upload_name.as_str(), + name: upload_name.as_str(), + data_type, + existing_id: None, + content: compute_runner_api::runner::DomainArtifactContent::File(&abs_path), + }) + .await; + match upload_result { + Ok(_) => { + let _ = ctx + .ctrl + .log_event(json!({ + "level": "info", + "stage": "upload", + "message": "preview uploaded", + "uploaded": upload_name.as_str(), + "file": abs_path.display().to_string(), + "data_type": data_type, + })) + .await; + } + Err(err) => { + warn!( + %err, + file = %abs_path.display(), + data_type = %data_type, + "preview upload failed" + ); + let _ = ctx + .ctrl + .log_event(json!({ + "level": "warn", + "stage": "upload", + "message": "preview upload failed", + "error": err.to_string(), + "file": abs_path.display().to_string(), + "data_type": data_type, + })) + .await; + } + } + } ctx.ctrl .progress(json!({ "progress": 100,