gsplat 1.0.0

data/examples/fit_image.rb

@@ -0,0 +1,35 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require "optparse"
+require "gsplat"
+
+options = {
+  width: 64,
+  height: 64,
+  n_gaussians: 5_000,
+  steps: 300,
+  learning_rate: 20.0
+}
+parser = OptionParser.new do |options_parser|
+  options_parser.banner = "Usage: bundle exec ruby examples/fit_image.rb [options]"
+  options_parser.on("--width N", Integer) { |value| options[:width] = value }
+  options_parser.on("--height N", Integer) { |value| options[:height] = value }
+  options_parser.on("--gaussians N", Integer) { |value| options[:n_gaussians] = value }
+  options_parser.on("--steps N", Integer) { |value| options[:steps] = value }
+  options_parser.on("--learning-rate X", Float) { |value| options[:learning_rate] = value }
+  options_parser.on_tail("-h", "--help", "Show this help") do
+    puts options_parser
+    exit
+  end
+end
+parser.parse!
+
+steps = options.delete(:steps)
+result = Gsplat::Training::ImageFitter.new(**options).fit(steps: steps)
+warn format(
+  "fit complete: PSNR %<initial>.2f dB -> %<final>.2f dB (%<steps>d steps)",
+  initial: result.initial_psnr,
+  final: result.final_psnr,
+  steps: steps
+)

data/examples/generate_sample_data.rb

@@ -0,0 +1,147 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require "fileutils"
+require "optparse"
+require "gsplat"
+
+# Builds the deterministic miniature scene used by the runnable examples.
+module SampleData
+  WIDTH = 16
+  HEIGHT = 16
+  CAMERA_TRANSLATIONS = [-0.12, 0.0, 0.12].freeze
+
+  module_function
+
+  def generate(root)
+    sparse_directory = File.join(root, "colmap", "sparse", "0")
+    image_directory = File.join(root, "colmap", "images")
+    FileUtils.mkdir_p([sparse_directory, image_directory])
+
+    parameters, colors = gaussian_parameters
+    views, intrinsics = cameras
+    images, = Gsplat.rasterization(
+      means: parameters.fetch(:means),
+      quats: parameters.fetch(:quats),
+      scales: Numo::NMath.exp(parameters.fetch(:scales)),
+      opacities: sigmoid(parameters.fetch(:opacities)),
+      colors: parameters.fetch(:sh0),
+      viewmats: views,
+      ks: intrinsics,
+      width: WIDTH,
+      height: HEIGHT,
+      sh_degree: 0
+    )
+
+    write_colmap_model(sparse_directory, colors, parameters.fetch(:means))
+    write_images(image_directory, images)
+    Gsplat::IO::Ply.write(File.join(root, "splats.ply"), parameters)
+  end
+
+  def gaussian_parameters
+    means = Numo::SFloat.zeros(16, 3)
+    colors = Numo::SFloat.zeros(16, 3)
+    16.times do |index|
+      x_coord = ((index % 4) * 2) + 0.5
+      y_coord = ((index / 4) * 2) + 0.5
+      fraction = index.fdiv(15)
+      means[index, true] = [(x_coord - 4) / 4.0, (y_coord - 4) / 4.0, 2]
+      colors[index, true] = [fraction, 0.2 + (0.5 * fraction), 1 - fraction]
+    end
+
+    parameters = {
+      means: means,
+      quats: identity_quaternions(16),
+      scales: Numo::SFloat.ones(16, 3) * Math.log(0.2),
+      opacities: Numo::SFloat.ones(16) * Math.log(0.85 / 0.15),
+      sh0: Gsplat::Utils.rgb_to_sh(colors).reshape(16, 1, 3),
+      shN: Numo::SFloat.zeros(16, 0, 3)
+    }
+    [parameters, colors]
+  end
+  private_class_method :gaussian_parameters
+
+  def cameras
+    count = CAMERA_TRANSLATIONS.length
+    views = Numo::SFloat.zeros(count, 4, 4)
+    intrinsics = Numo::SFloat.zeros(count, 3, 3)
+    count.times do |index|
+      views[index, true, true] = Numo::SFloat.eye(4)
+      views[index, 0, 3] = CAMERA_TRANSLATIONS.fetch(index)
+      intrinsics[index, true, true] = Numo::SFloat[[16, 0, 8], [0, 16, 8], [0, 0, 1]]
+    end
+    [views, intrinsics]
+  end
+  private_class_method :cameras
+
+  def identity_quaternions(count)
+    quaternions = Numo::SFloat.zeros(count, 4)
+    quaternions[true, 0] = 1
+    quaternions
+  end
+  private_class_method :identity_quaternions
+
+  def sigmoid(values)
+    1.0 / (1.0 + Numo::NMath.exp(-values))
+  end
+  private_class_method :sigmoid
+
+  def write_colmap_model(directory, colors, means)
+    File.write(
+      File.join(directory, "cameras.txt"),
+      "# One PINHOLE camera shared by all sample views.\n1 PINHOLE #{WIDTH} #{HEIGHT} 16 16 8 8\n"
+    )
+    image_records = CAMERA_TRANSLATIONS.each_index.map do |index|
+      format(
+        "%<id>d 1 0 0 0 %<tx>.6f 0 0 1 view_%<frame>03d.png\n",
+        id: index + 1,
+        tx: CAMERA_TRANSLATIONS[index],
+        frame: index
+      )
+    end
+    File.write(
+      File.join(directory, "images.txt"),
+      "# Each registered image is followed by an empty POINTS2D line.\n#{image_records.join("\n")}"
+    )
+    point_records = means.shape[0].times.map do |index|
+      rgb = colors[index, true].to_a.map { |value| (value * 255).round }
+      format(
+        "%<id>d %<x>.9g %<y>.9g %<z>.9g %<r>d %<g>d %<b>d 0",
+        id: index + 1,
+        x: means[index, 0],
+        y: means[index, 1],
+        z: means[index, 2],
+        r: rgb[0],
+        g: rgb[1],
+        b: rgb[2]
+      )
+    end
+    File.write(File.join(directory, "points3D.txt"), "# Synthetic colored seed points.\n#{point_records.join("\n")}\n")
+  end
+  private_class_method :write_colmap_model
+
+  def write_images(directory, images)
+    CAMERA_TRANSLATIONS.each_index do |index|
+      path = File.join(directory, format("view_%03d.png", index))
+      Gsplat::IO::Image.write(path, images[index, true, true, true], backend: :chunky_png)
+    end
+  end
+  private_class_method :write_images
+end
+
+options = { output: File.expand_path("data", __dir__) }
+parser = OptionParser.new do |options_parser|
+  options_parser.banner = "Usage: bundle exec ruby examples/generate_sample_data.rb [options]"
+  options_parser.on("--output PATH", "Destination directory (default: examples/data)") do |value|
+    options[:output] = value
+  end
+  options_parser.on_tail("-h", "--help", "Show this help") do
+    puts options_parser
+    exit
+  end
+end
+parser.parse!
+
+Gsplat.backend = :ruby
+SampleData.generate(File.expand_path(options.fetch(:output)))
+warn "sample data written to #{File.expand_path(options.fetch(:output))}"

data/examples/render_path.rb

@@ -0,0 +1,117 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require "fileutils"
+require "optparse"
+require "gsplat"
+
+sample_ply = File.expand_path("data/splats.ply", __dir__)
+options = { ply: sample_ply, output: nil, frames: nil, width: nil, height: nil, dry_run: false }
+parser = OptionParser.new do |options_parser|
+  options_parser.banner = "Usage: bundle exec ruby examples/render_path.rb [options]"
+  options_parser.on("--ply FILE", "Inria PLY file (default: bundled sample)") { |value| options[:ply] = value }
+  options_parser.on("--output PATH") { |value| options[:output] = value }
+  options_parser.on("--frames N", Integer) { |value| options[:frames] = value }
+  options_parser.on("--width N", Integer) { |value| options[:width] = value }
+  options_parser.on("--height N", Integer) { |value| options[:height] = value }
+  options_parser.on("--dry-run", "Validate the PLY and camera path without writing images") do
+    options[:dry_run] = true
+  end
+  options_parser.on_tail("-h", "--help", "Show this help") do
+    puts options_parser
+    exit
+  end
+end
+parser.parse!
+
+using_sample = File.expand_path(options.fetch(:ply)) == sample_ply
+options[:output] ||= using_sample ? "renders/sample" : "renders"
+options[:frames] ||= using_sample ? 12 : 60
+options[:width] ||= using_sample ? 64 : 512
+options[:height] ||= using_sample ? 64 : 512
+
+def normalize(vector)
+  length = Math.sqrt(vector.sum { |value| value * value })
+  vector.map { |value| value / length }
+end
+
+def cross(left, right)
+  [
+    (left[1] * right[2]) - (left[2] * right[1]),
+    (left[2] * right[0]) - (left[0] * right[2]),
+    (left[0] * right[1]) - (left[1] * right[0])
+  ]
+end
+
+def look_at(position, target)
+  forward = normalize(target.zip(position).map { |target_value, position_value| target_value - position_value })
+  right = normalize(cross([0.0, 1.0, 0.0], forward))
+  up = cross(forward, right)
+  rotation = [right, up, forward]
+  view = Numo::SFloat.eye(4)
+  view[0...3, 0...3] = Numo::SFloat.cast(rotation)
+  view[0...3, 3] = Numo::SFloat.cast(rotation.map { |row| -row.zip(position).sum { |a, b| a * b } })
+  view
+end
+
+def sh_coefficients(params)
+  direct = params.fetch(:sh0)
+  rest = params.fetch(:shN)
+  output = direct.class.zeros(direct.shape[0], direct.shape[1] + rest.shape[1], 3)
+  output[true, 0...direct.shape[1], true] = direct
+  output[true, direct.shape[1]...output.shape[1], true] = rest unless rest.shape[1].zero?
+  output
+end
+
+params = Gsplat::IO::Ply.read(options.fetch(:ply))
+center = params.fetch(:means).mean(axis: 0).to_a
+distances = Numo::NMath.sqrt(((params.fetch(:means) - Numo::SFloat.cast(center))**2).sum(axis: 1))
+radius = [distances.max.to_f * 2.0, 1.0].max
+colors = sh_coefficients(params)
+degree = [Math.sqrt(colors.shape[1]).to_i - 1, 4].min
+focal = [options.fetch(:width), options.fetch(:height)].max.to_f
+intrinsics = Numo::SFloat[
+  [[focal, 0, options.fetch(:width) / 2.0],
+   [0, focal, options.fetch(:height) / 2.0],
+   [0, 0, 1]]
+]
+if options.fetch(:dry_run)
+  warn format(
+    "render path valid: gaussians=%<gaussians>d frames=%<frames>d degree=%<degree>d",
+    gaussians: params.fetch(:means).shape[0],
+    frames: options.fetch(:frames),
+    degree: degree
+  )
+  exit
+end
+
+FileUtils.mkdir_p(options.fetch(:output))
+
+options.fetch(:frames).times do |frame|
+  angle = 2 * Math::PI * frame / options.fetch(:frames)
+  position = [
+    center[0] + (radius * Math.cos(angle)),
+    center[1] + (radius * 0.25),
+    center[2] + (radius * Math.sin(angle))
+  ]
+  view = look_at(position, center).reshape(1, 4, 4)
+  rendered, = Gsplat.rasterization(
+    means: params.fetch(:means),
+    quats: params.fetch(:quats),
+    scales: Numo::NMath.exp(params.fetch(:scales)),
+    opacities: 1.0 / (1.0 + Numo::NMath.exp(-params.fetch(:opacities))),
+    colors: colors,
+    viewmats: view,
+    ks: intrinsics,
+    width: options.fetch(:width),
+    height: options.fetch(:height),
+    sh_degree: degree
+  )
+  path = File.join(options.fetch(:output), format("frame_%04d.png", frame))
+  Gsplat::IO::Image.write(path, rendered[0, true, true, true])
+end
+warn format(
+  "render complete: frames=%<frames>d output=%<output>s",
+  frames: options.fetch(:frames),
+  output: options.fetch(:output)
+)

data/examples/simple_trainer.rb

@@ -0,0 +1,76 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require "fileutils"
+require "optparse"
+require "gsplat"
+
+sample_data = File.expand_path("data/colmap", __dir__)
+options = {
+  data: sample_data,
+  output_dir: nil,
+  max_steps: nil,
+  data_factor: 1,
+  strategy: "default",
+  dry_run: false
+}
+parser = OptionParser.new do |options_parser|
+  options_parser.banner = "Usage: bundle exec ruby examples/simple_trainer.rb [options]"
+  options_parser.on("--data PATH", "COLMAP dataset root (default: bundled sample)") { |value| options[:data] = value }
+  options_parser.on("--output PATH") { |value| options[:output_dir] = value }
+  options_parser.on("--steps N", Integer) { |value| options[:max_steps] = value }
+  options_parser.on("--data-factor N", Integer) { |value| options[:data_factor] = value }
+  options_parser.on("--strategy NAME", %w[default mcmc]) { |value| options[:strategy] = value }
+  options_parser.on("--dry-run", "Validate the sparse model without loading images") { options[:dry_run] = true }
+  options_parser.on_tail("-h", "--help", "Show this help") do
+    puts options_parser
+    exit
+  end
+end
+parser.parse!
+
+using_sample = File.expand_path(options.fetch(:data)) == sample_data
+options[:output_dir] ||= using_sample ? "results/sample" : "results"
+options[:max_steps] ||= using_sample ? 10 : 30_000
+
+if options.fetch(:dry_run)
+  dataset = Gsplat::IO::Colmap.read(options.fetch(:data), data_factor: options.fetch(:data_factor))
+  warn format(
+    "dataset valid: cameras=%<cameras>d images=%<images>d points=%<points>d",
+    cameras: dataset.cameras.length,
+    images: dataset.images.length,
+    points: dataset.points3d.length
+  )
+  exit
+end
+
+scene = Gsplat::Training::Scene.from_colmap(
+  options.fetch(:data),
+  data_factor: options.fetch(:data_factor)
+)
+config_options = {
+  max_steps: options.fetch(:max_steps),
+  output_dir: options.fetch(:output_dir),
+  eval_steps: [options.fetch(:max_steps)],
+  save_steps: [options.fetch(:max_steps)]
+}
+config_options.merge!(opacity_reg: 0.01, scale_reg: 0.01) if options.fetch(:strategy) == "mcmc"
+config = Gsplat::Training::Config.new(**config_options)
+strategy = if options.fetch(:strategy) == "mcmc"
+             Gsplat::Strategy::MCMC.new
+           else
+             Gsplat::Strategy::Default.new
+           end
+trainer = Gsplat::Training::Trainer.new(scene: scene, config: config, strategy: strategy)
+result = trainer.train
+FileUtils.mkdir_p(options.fetch(:output_dir))
+Gsplat::IO::Ply.write(
+  File.join(options.fetch(:output_dir), "splats.ply"),
+  trainer.params
+)
+warn format(
+  "training complete: step=%<step>d PSNR=%<psnr>.2f SSIM=%<ssim>.4f",
+  step: result.step,
+  psnr: result.final_metrics.fetch(:psnr),
+  ssim: result.final_metrics.fetch(:ssim)
+)

data/ext/gsplat_native/common.h

@@ -0,0 +1,54 @@
+#ifndef GSPLAT_COMMON_H
+#define GSPLAT_COMMON_H
+
+#include <ruby.h>
+#include <ruby/thread.h>
+#include <numo/narray.h>
+#include <numo/intern.h>
+
+static inline narray_t *
+gs_narray(VALUE value)
+{
+    narray_t *array;
+    GetNArray(value, array);
+    return array;
+}
+
+static inline void
+gs_require_sfloat(VALUE value, const char *name)
+{
+    if (!rb_obj_is_kind_of(value, numo_cSFloat)) {
+        rb_raise(rb_eTypeError, "%s must be Numo::SFloat", name);
+    }
+}
+
+static inline void
+gs_require_int32(VALUE value, const char *name)
+{
+    if (!rb_obj_is_kind_of(value, numo_cInt32)) {
+        rb_raise(rb_eTypeError, "%s must be Numo::Int32", name);
+    }
+}
+
+static inline void
+gs_require_same_shape(VALUE left, VALUE right)
+{
+    narray_t *left_array = gs_narray(left);
+    narray_t *right_array = gs_narray(right);
+    if (left_array->ndim != right_array->ndim ||
+        left_array->size != right_array->size) {
+        rb_raise(rb_eArgError, "NArray shapes differ");
+    }
+    for (int axis = 0; axis < left_array->ndim; ++axis) {
+        if (left_array->shape[axis] != right_array->shape[axis]) {
+            rb_raise(rb_eArgError, "NArray shapes differ");
+        }
+    }
+}
+
+#define GS_SFLOAT_READ(value) ((const float *)na_get_pointer_for_read(value))
+#define GS_SFLOAT_WRITE(value) ((float *)na_get_pointer_for_write(value))
+#define GS_INT32_READ(value) ((const int32_t *)na_get_pointer_for_read(value))
+#define GS_INT32_WRITE(value) ((int32_t *)na_get_pointer_for_write(value))
+
+#endif

data/ext/gsplat_native/extconf.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+require "mkmf"
+require "rubygems"
+
+# mkmf configures extension builds through these documented globals.
+# rubocop:disable Style/GlobalVars
+numo_spec = Gem::Specification.find_by_name("numo-narray")
+numo_include = File.join(numo_spec.full_gem_path, "lib", "numo")
+$INCFLAGS = "-I#{numo_include} #{$INCFLAGS}"
+$CFLAGS = "#{$CFLAGS} -O3"
+
+abort "numo/narray.h is required" unless have_header("numo/narray.h")
+
+openmp_enabled = try_compile("int main(void) { return 0; }", "-fopenmp")
+if openmp_enabled
+  $CFLAGS = "#{$CFLAGS} -fopenmp"
+  $LDFLAGS = "#{$LDFLAGS} -fopenmp"
+elsif RUBY_PLATFORM.include?("darwin")
+  prefixes = [
+    ENV.fetch("LIBOMP_PREFIX", nil),
+    "/opt/homebrew/opt/libomp",
+    "/usr/local/opt/libomp"
+  ].compact
+  prefix = prefixes.find { |candidate| File.file?(File.join(candidate, "include", "omp.h")) }
+  if prefix
+    original_flags = [$INCFLAGS, $CFLAGS, $LDFLAGS]
+    $INCFLAGS = "-I#{File.join(prefix, 'include')} #{$INCFLAGS}"
+    $CFLAGS = "#{$CFLAGS} -Xpreprocessor -fopenmp"
+    $LDFLAGS = "#{$LDFLAGS} -L#{File.join(prefix, 'lib')} -lomp"
+    openmp_enabled = try_link("#include <omp.h>\nint main(void) { return omp_get_max_threads() < 1; }")
+    $INCFLAGS, $CFLAGS, $LDFLAGS = original_flags unless openmp_enabled
+  end
+end
+
+$defs << "-DGSPLAT_OPENMP=1" if openmp_enabled
+# rubocop:enable Style/GlobalVars
+
+create_makefile("gsplat/gsplat_native")

data/ext/gsplat_native/gsplat_native.c

@@ -0,0 +1,60 @@
+#include "common.h"
+
+typedef struct {
+    const float *left;
+    const float *right;
+    float *output;
+    size_t count;
+} gs_add_args;
+
+void gs_init_spherical_harmonics(VALUE native);
+void gs_init_projection(VALUE native);
+void gs_init_intersections(VALUE native);
+void gs_init_rasterization(VALUE native);
+
+static void *
+gs_add_without_gvl(void *opaque)
+{
+    gs_add_args *args = (gs_add_args *)opaque;
+#ifdef GSPLAT_OPENMP
+#pragma omp parallel for
+#endif
+    for (size_t index = 0; index < args->count; ++index) {
+        args->output[index] = args->left[index] + args->right[index];
+    }
+    return NULL;
+}
+
+static VALUE
+gs_native_add(VALUE self, VALUE left, VALUE right)
+{
+    (void)self;
+    gs_require_sfloat(left, "left");
+    gs_require_sfloat(right, "right");
+    gs_require_same_shape(left, right);
+
+    VALUE output = rb_funcall(left, rb_intern("dup"), 0);
+    gs_add_args args = {
+        GS_SFLOAT_READ(left),
+        GS_SFLOAT_READ(right),
+        GS_SFLOAT_WRITE(output),
+        gs_narray(left)->size
+    };
+    rb_thread_call_without_gvl(gs_add_without_gvl, &args, RUBY_UBF_IO, NULL);
+    RB_GC_GUARD(left);
+    RB_GC_GUARD(right);
+    RB_GC_GUARD(output);
+    return output;
+}
+
+void
+Init_gsplat_native(void)
+{
+    VALUE gsplat = rb_const_get(rb_cObject, rb_intern("Gsplat"));
+    VALUE native = rb_define_module_under(gsplat, "Native");
+    rb_define_singleton_method(native, "add", gs_native_add, 2);
+    gs_init_spherical_harmonics(native);
+    gs_init_projection(native);
+    gs_init_intersections(native);
+    gs_init_rasterization(native);
+}