gsplat 1.0.0

data/lib/gsplat/io/checkpoint.rb

@@ -0,0 +1,145 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Gsplat
+  module IO
+    # Portable NPZ checkpoint for parameters, Adam moments, step, and config.
+    module Checkpoint
+      # Current checkpoint schema version.
+      VERSION = 1
+      # Separator reserved for structured NPZ entry names.
+      SEPARATOR = "___"
+      # Loaded checkpoint payload.
+      #
+      # @!attribute params
+      #   @return [Hash{Symbol=>Numo::NArray}]
+      # @!attribute optimizer_states
+      #   @return [Hash] optimizer moments grouped by optimizer and parameter
+      # @!attribute step
+      #   @return [Integer]
+      # @!attribute config
+      #   @return [Hash]
+      Snapshot = Data.define(:params, :optimizer_states, :step, :config)
+
+      module_function
+
+      def save(target, params:, optimizers:, step:, config: {})
+        raise ArgumentError, "step must be a non-negative integer" unless step.is_a?(Integer) && !step.negative?
+
+        arrays = {
+          "checkpoint_version" => Numo::Int32[VERSION],
+          "checkpoint_step" => Numo::Int64[step],
+          "checkpoint_config_json" => encode_json(config)
+        }
+        params.each do |name, value|
+          arrays[key("param", name)] = Ops::TensorOps.data(value)
+        end
+        optimizers.each do |optimizer_name, optimizer|
+          optimizer.groups.each_key do |group_name|
+            state = optimizer.state(group_name)
+            prefix = key("optimizer", optimizer_name, group_name)
+            arrays["#{prefix}#{SEPARATOR}step"] = Numo::Int64[state.step]
+            arrays["#{prefix}#{SEPARATOR}exp_avg"] = state.exp_avg
+            arrays["#{prefix}#{SEPARATOR}exp_avg_sq"] = state.exp_avg_sq
+          end
+        end
+        Npy.write_npz(target, arrays)
+      end
+
+      # Loads parameters and optimizer state without mutating live objects.
+      #
+      # @param source [String, #read] NPZ checkpoint path or IO
+      # @return [Snapshot]
+      def load(source)
+        arrays = Npy.read_npz(source)
+        version = arrays.fetch("checkpoint_version")[0].to_i
+        raise NotSupportedError, "unsupported checkpoint version #{version}" unless version == VERSION
+
+        Snapshot.new(
+          params: extract_params(arrays),
+          optimizer_states: extract_optimizer_states(arrays),
+          step: arrays.fetch("checkpoint_step")[0].to_i,
+          config: decode_json(arrays.fetch("checkpoint_config_json"))
+        )
+      rescue KeyError => e
+        raise Gsplat::Error, "invalid checkpoint: #{e.message}"
+      end
+
+      # Loads a checkpoint into existing Variables and optimizers.
+      #
+      # @param source [String, #read] NPZ checkpoint path or IO
+      # @param params [Hash{Symbol=>Autograd::Variable}]
+      # @param optimizers [Hash{Symbol=>Optim::Adam}]
+      # @return [Snapshot]
+      def restore!(source, params:, optimizers:)
+        snapshot = load(source)
+        snapshot.params.each do |name, value|
+          params.fetch(name).replace_data!(value.dup)
+        end
+        snapshot.optimizer_states.each do |optimizer_name, groups|
+          optimizer = optimizers.fetch(optimizer_name)
+          groups.each do |group_name, state|
+            optimizer.load_state!(
+              group_name,
+              step: state.fetch(:step),
+              exp_avg: state.fetch(:exp_avg),
+              exp_avg_sq: state.fetch(:exp_avg_sq)
+            )
+          end
+        end
+        snapshot
+      rescue KeyError => e
+        raise Gsplat::Error, "checkpoint target mismatch: #{e.message}"
+      end
+
+      def key(*parts)
+        values = parts.map(&:to_s)
+        invalid = values.find { |value| value.empty? || value.include?(SEPARATOR) || !value.match?(/\A\w+\z/) }
+        raise ArgumentError, "invalid checkpoint name #{invalid.inspect}" if invalid
+
+        values.join(SEPARATOR)
+      end
+      private_class_method :key
+
+      def encode_json(config)
+        value = config.respond_to?(:to_h) ? config.to_h : config
+        Numo::UInt8.cast(JSON.generate(value).bytes)
+      rescue JSON::GeneratorError => e
+        raise ArgumentError, "config is not JSON serializable: #{e.message}"
+      end
+      private_class_method :encode_json
+
+      def decode_json(bytes)
+        JSON.parse(bytes.to_a.pack("C*"), symbolize_names: true)
+      rescue JSON::ParserError => e
+        raise Gsplat::Error, "invalid checkpoint config JSON: #{e.message}"
+      end
+      private_class_method :decode_json
+
+      def extract_params(arrays)
+        prefix = "param#{SEPARATOR}"
+        arrays.filter_map do |name, value|
+          next unless name.start_with?(prefix)
+
+          [name.delete_prefix(prefix).to_sym, value]
+        end.to_h
+      end
+      private_class_method :extract_params
+
+      def extract_optimizer_states(arrays)
+        output = Hash.new { |hash, name| hash[name] = Hash.new { |groups, group| groups[group] = {} } }
+        pattern = /\Aoptimizer#{SEPARATOR}(\w+)#{SEPARATOR}(\w+)#{SEPARATOR}(step|exp_avg|exp_avg_sq)\z/
+        arrays.each do |name, value|
+          match = pattern.match(name)
+          next unless match
+
+          state_value = match[3] == "step" ? value[0].to_i : value
+          output[match[1].to_sym][match[2].to_sym][match[3].to_sym] = state_value
+        end
+        output.to_h { |name, groups| [name, groups.to_h] }
+      end
+      private_class_method :extract_optimizer_states
+    end
+  end
+end

data/lib/gsplat/io/colmap.rb

@@ -0,0 +1,175 @@
+# frozen_string_literal: true
+
+require_relative "colmap_binary"
+require_relative "colmap_text"
+
+module Gsplat
+  module IO
+    # COLMAP sparse reconstruction reader and camera conversion utilities.
+    module Colmap
+      # Parsed camera calibration record.
+      Camera = Data.define(:id, :model, :width, :height, :params, :intrinsics, :distortion)
+      # Parsed registered-image pose and feature observations.
+      ImageRecord = Data.define(
+        :id, :qvec, :tvec, :camera_id, :name, :points2d, :point3d_ids, :rotation, :world_to_camera
+      )
+      # Parsed sparse 3D point and visibility track.
+      Point3D = Data.define(:id, :xyz, :rgb, :error, :track)
+      # Complete sparse model returned by {.read}.
+      Dataset = Data.define(:cameras, :images, :points3d, :path)
+
+      # Supported COLMAP numeric camera model IDs.
+      MODELS = {
+        0 => "SIMPLE_PINHOLE",
+        1 => "PINHOLE",
+        2 => "SIMPLE_RADIAL",
+        4 => "OPENCV"
+      }.freeze
+
+      module_function
+
+      # Reads cameras, registered images, and points from a sparse model directory.
+      def read(path, data_factor: 1)
+        directory = model_directory(path)
+        extension = model_extension(directory)
+        Dataset.new(
+          cameras: read_cameras(File.join(directory, "cameras.#{extension}"), data_factor: data_factor),
+          images: read_images(File.join(directory, "images.#{extension}")),
+          points3d: read_points3d(File.join(directory, "points3D.#{extension}")),
+          path: directory
+        )
+      end
+
+      # Reads and scales camera calibration records from a bin/txt file.
+      #
+      # @param path [String]
+      # @param data_factor [Numeric] image downsampling factor
+      # @return [Hash{Integer=>Camera}]
+      def read_cameras(path, data_factor: 1)
+        validate_factor!(data_factor)
+        raw_records(path, :cameras).transform_values { |record| camera_record(record, data_factor) }
+      end
+
+      # Reads registered image poses from a bin/txt file.
+      #
+      # @param path [String]
+      # @return [Hash{Integer=>ImageRecord}]
+      def read_images(path)
+        raw_records(path, :images).transform_values { |record| image_record(record) }
+      end
+
+      # Reads sparse points from a bin/txt file.
+      #
+      # @param path [String]
+      # @return [Hash{Integer=>Point3D}]
+      def read_points3d(path)
+        raw_records(path, :points3d).transform_values do |record|
+          Point3D.new(
+            id: record.fetch(:id),
+            xyz: Numo::DFloat.cast(record.fetch(:xyz)),
+            rgb: Numo::UInt8.cast(record.fetch(:rgb)),
+            error: record.fetch(:error),
+            track: record.fetch(:track).freeze
+          )
+        end
+      end
+
+      # Converts a COLMAP wxyz quaternion to a 3x3 world-to-camera rotation.
+      def qvec_to_rotmat(qvec)
+        quaternion = Numo::DFloat.cast(qvec).reshape(1, 4)
+        Math::Quaternion.to_rotmat(quaternion)[0, true, true].dup
+      end
+
+      def camera_record(record, factor)
+        model = model_name(record.fetch(:model))
+        params = record.fetch(:params).map(&:to_f)
+        focal, distortion = intrinsics_for(model, params)
+        intrinsics = Numo::DFloat[
+          [focal[0] / factor, 0, focal[2] / factor],
+          [0, focal[1] / factor, focal[3] / factor],
+          [0, 0, 1]
+        ]
+        Camera.new(
+          id: record.fetch(:id),
+          model: model,
+          width: record.fetch(:width).fdiv(factor).to_i,
+          height: record.fetch(:height).fdiv(factor).to_i,
+          params: Numo::DFloat.cast(params),
+          intrinsics: intrinsics,
+          distortion: Numo::DFloat.cast(distortion)
+        )
+      end
+      private_class_method :camera_record
+
+      def image_record(record)
+        qvec = Numo::DFloat.cast(record.fetch(:qvec))
+        tvec = Numo::DFloat.cast(record.fetch(:tvec))
+        rotation = qvec_to_rotmat(qvec)
+        view = Numo::DFloat.eye(4)
+        view[0...3, 0...3] = rotation
+        view[0...3, 3] = tvec
+        coordinates = record.fetch(:points2d)
+        points2d = coordinates.empty? ? Numo::DFloat.zeros(0, 2) : Numo::DFloat.cast(coordinates)
+        ImageRecord.new(
+          id: record.fetch(:id),
+          qvec: qvec,
+          tvec: tvec,
+          camera_id: record.fetch(:camera_id),
+          name: record.fetch(:name),
+          points2d: points2d,
+          point3d_ids: Numo::Int64.cast(record.fetch(:point3d_ids)),
+          rotation: rotation,
+          world_to_camera: view
+        )
+      end
+      private_class_method :image_record
+
+      def intrinsics_for(model, params)
+        case model
+        when "SIMPLE_PINHOLE" then [[params[0], params[0], params[1], params[2]], []]
+        when "PINHOLE" then [params.first(4), []]
+        when "SIMPLE_RADIAL" then [[params[0], params[0], params[1], params[2]], [params[3]]]
+        when "OPENCV" then [params.first(4), params.drop(4)]
+        end
+      end
+      private_class_method :intrinsics_for
+
+      def model_name(value)
+        return value if MODELS.value?(value)
+
+        MODELS.fetch(value) { raise NotSupportedError, "unsupported COLMAP camera model #{value.inspect}" }
+      end
+      private_class_method :model_name
+
+      def raw_records(path, kind)
+        data = File.binread(path)
+        reader = File.extname(path) == ".bin" ? ColmapBinary : ColmapText
+        reader.public_send(kind, data)
+      end
+      private_class_method :raw_records
+
+      def model_directory(path)
+        candidates = [path, File.join(path, "sparse", "0"), File.join(path, "sparse")]
+        directory = candidates.find do |candidate|
+          File.file?(File.join(candidate, "cameras.bin")) ||
+            File.file?(File.join(candidate, "cameras.txt"))
+        end
+        directory || raise(Gsplat::Error, "COLMAP sparse model not found under #{path}")
+      end
+      private_class_method :model_directory
+
+      def model_extension(directory)
+        return "bin" if %w[cameras images points3D].all? { |name| File.file?(File.join(directory, "#{name}.bin")) }
+        return "txt" if %w[cameras images points3D].all? { |name| File.file?(File.join(directory, "#{name}.txt")) }
+
+        raise Gsplat::Error, "COLMAP sparse model files are incomplete in #{directory}"
+      end
+      private_class_method :model_extension
+
+      def validate_factor!(factor)
+        raise ArgumentError, "data_factor must be positive" unless factor.is_a?(Numeric) && factor.positive?
+      end
+      private_class_method :validate_factor!
+    end
+  end
+end

data/lib/gsplat/io/colmap_binary.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module IO
+    # Decoder for COLMAP sparse-model binary files.
+    module ColmapBinary
+      CAMERA_PARAM_COUNTS = {
+        0 => 3, 1 => 4, 2 => 4, 3 => 5, 4 => 8, 5 => 8,
+        6 => 12, 7 => 5, 8 => 4, 9 => 5, 10 => 12
+      }.freeze
+
+      # Bounds-checked little-endian binary cursor.
+      class Cursor
+        def initialize(data)
+          @data = data
+          @offset = 0
+        end
+
+        def scalar(format, width)
+          bytes = @data.byteslice(@offset, width)
+          raise Gsplat::Error, "truncated COLMAP binary file" unless bytes&.bytesize == width
+
+          @offset += width
+          bytes.unpack1(format)
+        end
+
+        def array(format, width, count)
+          Array.new(count) { scalar(format, width) }
+        end
+
+        def cstring
+          finish = @data.index("\0", @offset)
+          raise Gsplat::Error, "unterminated COLMAP image name" unless finish
+
+          value = @data.byteslice(@offset, finish - @offset)
+          @offset = finish + 1
+          value.force_encoding(Encoding::UTF_8)
+        end
+      end
+
+      module_function
+
+      def cameras(data)
+        cursor = Cursor.new(data)
+        Array.new(cursor.scalar("Q<", 8)).to_h do
+          id = cursor.scalar("l<", 4)
+          model = cursor.scalar("l<", 4)
+          count = CAMERA_PARAM_COUNTS.fetch(model) do
+            raise NotSupportedError, "unknown COLMAP camera model id #{model}"
+          end
+          width = cursor.scalar("Q<", 8)
+          height = cursor.scalar("Q<", 8)
+          [id, { id: id, model: model, width: width, height: height, params: cursor.array("E", 8, count) }]
+        end
+      end
+
+      def images(data)
+        cursor = Cursor.new(data)
+        Array.new(cursor.scalar("Q<", 8)).to_h do
+          id = cursor.scalar("l<", 4)
+          qvec = cursor.array("E", 8, 4)
+          tvec = cursor.array("E", 8, 3)
+          camera_id = cursor.scalar("l<", 4)
+          name = cursor.cstring
+          count = cursor.scalar("Q<", 8)
+          points2d, point3d_ids = image_points(cursor, count)
+          [id, { id: id, qvec: qvec, tvec: tvec, camera_id: camera_id, name: name,
+                 points2d: points2d, point3d_ids: point3d_ids }]
+        end
+      end
+
+      def points3d(data)
+        cursor = Cursor.new(data)
+        Array.new(cursor.scalar("Q<", 8)).to_h do
+          id = cursor.scalar("Q<", 8)
+          xyz = cursor.array("E", 8, 3)
+          rgb = cursor.array("C", 1, 3)
+          error = cursor.scalar("E", 8)
+          track_count = cursor.scalar("Q<", 8)
+          track = Array.new(track_count) { [cursor.scalar("l<", 4), cursor.scalar("l<", 4)] }
+          [id, { id: id, xyz: xyz, rgb: rgb, error: error, track: track }]
+        end
+      end
+
+      def image_points(cursor, count)
+        coordinates = []
+        ids = []
+        count.times do
+          coordinates << cursor.array("E", 8, 2)
+          raw_id = cursor.scalar("Q<", 8)
+          ids << (raw_id == ((1 << 64) - 1) ? -1 : raw_id)
+        end
+        [coordinates, ids]
+      end
+      private_class_method :image_points
+    end
+  end
+end

data/lib/gsplat/io/colmap_text.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module IO
+    # Decoder for COLMAP sparse-model text files.
+    module ColmapText
+      module_function
+
+      def cameras(data)
+        content_lines(data).to_h do |line|
+          fields = line.split
+          id = Integer(fields.shift, 10)
+          model = fields.shift
+          width = Integer(fields.shift, 10)
+          height = Integer(fields.shift, 10)
+          [id, { id: id, model: model, width: width, height: height, params: fields.map { |item| Float(item) } }]
+        end
+      end
+
+      def images(data)
+        lines = data.lines(chomp: true)
+        output = {}
+        until lines.empty?
+          metadata = next_metadata_line(lines)
+          break unless metadata
+
+          fields = metadata.split
+          id = Integer(fields.shift, 10)
+          qvec = fields.shift(4).map { |item| Float(item) }
+          tvec = fields.shift(3).map { |item| Float(item) }
+          camera_id = Integer(fields.shift, 10)
+          name = fields.join(" ")
+          points_line = lines.shift.to_s.strip
+          points2d, point3d_ids = parse_image_points(points_line)
+          output[id] = {
+            id: id, qvec: qvec, tvec: tvec, camera_id: camera_id,
+            name: name, points2d: points2d, point3d_ids: point3d_ids
+          }
+        end
+        output
+      end
+
+      def points3d(data)
+        content_lines(data).to_h do |line|
+          fields = line.split
+          id = Integer(fields.shift, 10)
+          xyz = fields.shift(3).map { |item| Float(item) }
+          rgb = fields.shift(3).map { |item| Integer(item, 10) }
+          error = Float(fields.shift)
+          track = fields.each_slice(2).map { |image_id, point_index| [Integer(image_id, 10), Integer(point_index, 10)] }
+          [id, { id: id, xyz: xyz, rgb: rgb, error: error, track: track }]
+        end
+      end
+
+      def content_lines(data)
+        data.lines(chomp: true).map(&:strip).reject { |line| line.empty? || line.start_with?("#") }
+      end
+      private_class_method :content_lines
+
+      def next_metadata_line(lines)
+        until lines.empty?
+          line = lines.shift.strip
+          return line unless line.empty? || line.start_with?("#")
+        end
+        nil
+      end
+      private_class_method :next_metadata_line
+
+      def parse_image_points(line)
+        fields = line.split
+        coordinates = []
+        ids = []
+        fields.each_slice(3) do |x_coord, y_coord, point_id|
+          raise Gsplat::Error, "invalid COLMAP image point row" unless point_id
+
+          coordinates << [Float(x_coord), Float(y_coord)]
+          ids << Integer(point_id, 10)
+        end
+        [coordinates, ids]
+      end
+      private_class_method :parse_image_points
+    end
+  end
+end

data/lib/gsplat/io/image.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require_relative "image_backends"
+
+module Gsplat
+  module IO
+    # RGB image input/output with vips-first optional backends.
+    module Image
+      # Optional image backend implementations in preference order.
+      BACKENDS = {
+        vips: VipsImageBackend,
+        chunky_png: ChunkyPngImageBackend
+      }.freeze
+
+      module_function
+
+      # Reads an image as float32 [H,W,3] in the range 0..1.
+      def read(path, backend: :auto)
+        resolve_backend(backend).read(path)
+      end
+
+      # Writes a float [H,W,3] image, clamping values to 0..1.
+      def write(path, array, backend: :auto)
+        validate_array!(array)
+        resolve_backend(backend).write(path, array)
+      end
+
+      # Returns image codecs currently loadable in this process.
+      #
+      # @return [Array<Symbol>] subset of `:vips` and `:chunky_png`
+      def available_backends
+        BACKENDS.filter_map { |name, implementation| name if implementation.available? }
+      end
+
+      def resolve_backend(name)
+        if name == :auto
+          selected = BACKENDS.values.find(&:available?)
+          return selected if selected
+
+          raise NotSupportedError, "image IO requires ruby-vips or chunky_png"
+        end
+        implementation = BACKENDS.fetch(name.to_sym) do
+          raise ArgumentError, "unknown image backend #{name.inspect}"
+        end
+        return implementation if implementation.available?
+
+        raise NotSupportedError, "image backend #{name} is unavailable"
+      end
+      private_class_method :resolve_backend
+
+      def validate_array!(array)
+        valid = array.is_a?(Numo::NArray) &&
+                [Numo::SFloat, Numo::DFloat].include?(array.class) &&
+                array.ndim == 3 && array.shape[2] == 3
+        return if valid
+
+        actual = array.respond_to?(:shape) ? array.shape.inspect : array.class
+        raise ShapeError, "expected image [H,W,3] floating array, got #{actual}"
+      end
+      private_class_method :validate_array!
+    end
+  end
+end

data/lib/gsplat/io/image_backends.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module IO
+    # Optional ruby-vips implementation.
+    module VipsImageBackend
+      module_function
+
+      def available?
+        require "vips"
+        true
+      rescue LoadError, StandardError
+        false
+      end
+
+      def read(path)
+        image = Vips::Image.new_from_file(path.to_s, access: :sequential)
+        scale = { uchar: 255.0, ushort: 65_535.0 }.fetch(image.format.to_sym, 1.0)
+        rgb = rgb_image(image).cast(:float)
+        Numo::SFloat.from_binary(rgb.write_to_memory).reshape(rgb.height, rgb.width, 3) / scale
+      end
+
+      def write(path, array)
+        bytes = quantized_bytes(array)
+        image = Vips::Image.new_from_memory(bytes, array.shape[1], array.shape[0], 3, :uchar)
+        image.write_to_file(path.to_s)
+        path
+      end
+
+      def rgb_image(image)
+        return image.extract_band(0, n: 3) if image.bands >= 3
+
+        band = image.extract_band(0)
+        band.bandjoin([band, band])
+      end
+      private_class_method :rgb_image
+
+      def quantized_bytes(array)
+        array.to_a.flatten.map { |value| (value.to_f.clamp(0.0, 1.0) * 255).round }.pack("C*")
+      end
+      private_class_method :quantized_bytes
+    end
+
+    # Portable chunky_png implementation.
+    module ChunkyPngImageBackend
+      module_function
+
+      def available?
+        require "chunky_png"
+        true
+      rescue LoadError
+        false
+      end
+
+      def read(path)
+        image = ChunkyPNG::Image.from_file(path.to_s)
+        values = image.height.times.flat_map do |y_coord|
+          image.width.times.flat_map do |x_coord|
+            pixel = image[x_coord, y_coord]
+            [ChunkyPNG::Color.r(pixel), ChunkyPNG::Color.g(pixel), ChunkyPNG::Color.b(pixel)]
+          end
+        end
+        Numo::SFloat.cast(values).reshape(image.height, image.width, 3) / 255.0
+      end
+
+      def write(path, array)
+        image = ChunkyPNG::Image.new(array.shape[1], array.shape[0])
+        array.shape[0].times do |y_coord|
+          array.shape[1].times do |x_coord|
+            rgb = array[y_coord, x_coord, true].to_a.map do |value|
+              (value.to_f.clamp(0.0, 1.0) * 255).round
+            end
+            image[x_coord, y_coord] = ChunkyPNG::Color.rgb(*rgb)
+          end
+        end
+        image.save(path.to_s)
+        path
+      end
+    end
+  end
+end