gsplat 1.0.0

data/lib/gsplat/compression/grid_sort.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Compression
+    # Deterministic Morton ordering for square Gaussian parameter grids.
+    module GridSort
+      QUANTIZATION_MAX = 1023
+
+      module_function
+
+      def prepare(parameters, use_sort: true)
+        values = parameters.to_h { |name, value| [name.to_s, data(value).dup] }
+        validate_lengths!(values)
+        side = ::Math.sqrt(values.fetch("means").shape[0]).floor
+        raise ArgumentError, "at least one Gaussian is required" if side.zero?
+
+        values = crop_to_square(values, side * side)
+        values = reorder(values, morton_order(values.fetch("means"))) if use_sort
+        [values, side]
+      end
+
+      def crop_to_square(values, count)
+        return values if values.fetch("means").shape[0] == count
+
+        opacities = values.fetch("opacities").reshape(values.fetch("opacities").shape[0])
+        keep = (0...opacities.size).sort_by { |index| [-opacities[index].to_f, index] }.first(count)
+        reorder(values, keep)
+      end
+      private_class_method :crop_to_square
+
+      def morton_order(means)
+        coordinates = means.to_a
+        mins = 3.times.map { |axis| coordinates.map { |point| point[axis] }.min }
+        maxs = 3.times.map { |axis| coordinates.map { |point| point[axis] }.max }
+        (0...coordinates.length).sort_by do |index|
+          quantized = 3.times.map do |axis|
+            span = maxs[axis] - mins[axis]
+            span.zero? ? 0 : (((coordinates[index][axis] - mins[axis]) / span) * QUANTIZATION_MAX).round
+          end
+          [morton3(*quantized), index]
+        end
+      end
+
+      def morton3(x_coord, y_coord, z_coord)
+        code = 0
+        10.times do |bit|
+          code |= ((x_coord >> bit) & 1) << (3 * bit)
+          code |= ((y_coord >> bit) & 1) << ((3 * bit) + 1)
+          code |= ((z_coord >> bit) & 1) << ((3 * bit) + 2)
+        end
+        code
+      end
+      private_class_method :morton3
+
+      def reorder(values, order)
+        indices = Numo::Int32.cast(order)
+        values.transform_values { |value| value[indices, *Array.new(value.ndim - 1, true)].dup }
+      end
+      private_class_method :reorder
+
+      def validate_lengths!(values)
+        required = %w[means scales quats opacities sh0 shN]
+        missing = required - values.keys
+        raise ArgumentError, "missing compression parameters: #{missing.join(', ')}" unless missing.empty?
+
+        count = values.fetch("means").shape[0]
+        invalid = values.find { |_name, value| !value.is_a?(Numo::NArray) || value.shape[0] != count }
+        raise ShapeError, "all compression parameters must share first dimension #{count}" if invalid
+        raise ShapeError, "expected means [N,3]" unless values.fetch("means").shape == [count, 3]
+      end
+      private_class_method :validate_lengths!
+
+      def data(value)
+        value.is_a?(Autograd::Variable) ? value.data : value
+      end
+      private_class_method :data
+    end
+  end
+end

data/lib/gsplat/compression/kmeans.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Compression
+    # Deterministic Manhattan K-means for SH coefficient vectors.
+    module KMeans
+      module_function
+
+      def compress(path, array, clusters:, iterations:, quantization:)
+        shape = array.shape
+        return empty_metadata(array, quantization) if array.empty?
+
+        vectors = array.reshape(shape[0], array.size / shape[0]).to_a
+        cluster_count = [clusters, vectors.length, 65_536].min
+        centroids, labels = cluster(vectors, cluster_count, iterations)
+        quantized, minimum, maximum = quantize_centroids(centroids, quantization)
+        IO::Npy.write_npz(
+          path,
+          {
+            centroids: Numo::UInt8.cast(quantized).reshape(cluster_count, vectors.first.length),
+            labels: Numo::UInt16.cast(labels)
+          }
+        )
+        {
+          "shape" => shape,
+          "dtype" => Quantizer.dtype_name(array.class),
+          "mins" => minimum,
+          "maxs" => maximum,
+          "quantization" => quantization,
+          "clusters" => cluster_count
+        }
+      end
+
+      def decompress(path, metadata)
+        shape = metadata.fetch("shape")
+        return Quantizer.numeric_type(metadata.fetch("dtype")).zeros(*shape) if shape.inject(1, :*).zero?
+
+        archive = IO::Npy.read_npz(path)
+        centroids = decode_centroids(archive.fetch("centroids"), shape, metadata)
+        values = archive.fetch("labels").to_a.map { |label| centroids.fetch(label) }.flatten
+        Quantizer.numeric_type(metadata.fetch("dtype")).cast(values).reshape(*shape)
+      end
+
+      def decode_centroids(values, shape, metadata)
+        limit = (1 << metadata.fetch("quantization")) - 1
+        minimum = metadata.fetch("mins")
+        span = metadata.fetch("maxs") - minimum
+        feature_count = shape.drop(1).inject(1, :*)
+        values.to_a.flatten.each_slice(feature_count).map do |row|
+          row.map { |value| minimum + ((value.to_f / limit) * span) }
+        end
+      end
+      private_class_method :decode_centroids
+
+      def cluster(vectors, count, iterations)
+        return [vectors.map(&:dup), (0...vectors.length).to_a] if count == vectors.length
+
+        centroids = Array.new(count) do |index|
+          vectors[(index * (vectors.length - 1).fdiv([count - 1, 1].max)).round].dup
+        end
+        labels = Array.new(vectors.length, 0)
+        iterations.times do
+          changed = assign!(vectors, centroids, labels)
+          update!(vectors, centroids, labels)
+          break unless changed
+        end
+        [centroids, labels]
+      end
+      private_class_method :cluster
+
+      def assign!(vectors, centroids, labels)
+        changed = false
+        vectors.each_with_index do |vector, index|
+          label = centroids.each_index.min_by do |centroid_index|
+            vector.zip(centroids[centroid_index]).sum { |left, right| (left - right).abs }
+          end
+          changed ||= labels[index] != label
+          labels[index] = label
+        end
+        changed
+      end
+      private_class_method :assign!
+
+      def update!(vectors, centroids, labels)
+        sums = centroids.map { |centroid| Array.new(centroid.length, 0.0) }
+        counts = Array.new(centroids.length, 0)
+        vectors.each_with_index do |vector, index|
+          label = labels[index]
+          counts[label] += 1
+          vector.each_with_index { |value, feature| sums[label][feature] += value }
+        end
+        centroids.each_index do |index|
+          next if counts[index].zero?
+
+          centroids[index] = sums[index].map { |value| value / counts[index] }
+        end
+      end
+      private_class_method :update!
+
+      def quantize_centroids(centroids, bits)
+        values = centroids.flatten
+        minimum, maximum = values.minmax
+        limit = (1 << bits) - 1
+        span = maximum - minimum
+        quantized = values.map { |value| span.zero? ? 0 : (((value - minimum) / span) * limit).round }
+        [quantized, minimum, maximum]
+      end
+      private_class_method :quantize_centroids
+
+      def empty_metadata(array, quantization)
+        {
+          "shape" => array.shape,
+          "dtype" => Quantizer.dtype_name(array.class),
+          "quantization" => quantization,
+          "empty" => true
+        }
+      end
+      private_class_method :empty_metadata
+    end
+  end
+end

data/lib/gsplat/compression/png.rb

@@ -0,0 +1,147 @@
+# frozen_string_literal: true
+
+require "fileutils"
+require "json"
+
+require_relative "grid_sort"
+require_relative "kmeans"
+require_relative "png_codec"
+require_relative "quantizer"
+
+module Gsplat
+  # Parameter codecs for compact Gaussian model storage.
+  module Compression
+    # Self-contained PNG and K-means Gaussian parameter compression.
+    class Png
+      # Parameter names encoded as 8-bit PNG images.
+      PNG8_PARAMETERS = %w[scales quats opacities sh0].freeze
+
+      attr_reader :use_sort, :kmeans_clusters, :kmeans_iterations, :sh_quantization
+
+      def initialize(use_sort: true, kmeans_clusters: 256, kmeans_iterations: 8, sh_quantization: 6)
+        @use_sort = use_sort
+        @kmeans_clusters = Integer(kmeans_clusters)
+        @kmeans_iterations = Integer(kmeans_iterations)
+        @sh_quantization = Integer(sh_quantization)
+        validate_options!
+      end
+
+      # Compresses Gaussian arrays into an upstream-compatible directory.
+      #
+      # @param directory [String] output directory
+      # @param parameters [Hash{String, Symbol=>Numo::NArray}] arrays with a shared first dimension
+      # @return [String] the output directory
+      def compress(directory, parameters)
+        FileUtils.mkdir_p(directory)
+        values, side = GridSort.prepare(parameters, use_sort: use_sort)
+        values["means"] = log_transform(values.fetch("means"))
+        values["quats"] = Math::Quaternion.normalize(values.fetch("quats"))
+        metadata = values.to_h do |name, value|
+          validate_name!(name)
+          [name, compress_parameter(directory, name, value, side)]
+        end
+        File.write(File.join(directory, "meta.json"), "#{JSON.pretty_generate(metadata)}\n")
+        directory
+      end
+
+      # Restores Gaussian arrays from a compressed directory.
+      #
+      # @param directory [String] directory produced by {#compress}
+      # @return [Hash{Symbol=>Numo::NArray}]
+      def decompress(directory)
+        metadata = JSON.parse(File.read(File.join(directory, "meta.json")))
+        values = metadata.to_h do |name, entry|
+          validate_name!(name)
+          [name.to_sym, decompress_parameter(directory, name, entry)]
+        end
+        values[:means] = inverse_log_transform(values.fetch(:means))
+        values
+      end
+
+      private
+
+      def compress_parameter(directory, name, value, side)
+        return compress_png16(directory, name, value, side) if name == "means"
+        return compress_png8(directory, name, value, side) if PNG8_PARAMETERS.include?(name)
+        if name == "shN"
+          return KMeans.compress(
+            File.join(directory, "shN.npz"),
+            value,
+            clusters: kmeans_clusters,
+            iterations: kmeans_iterations,
+            quantization: sh_quantization
+          )
+        end
+
+        IO::Npy.write_npz(File.join(directory, "#{name}.npz"), { arr: value })
+        { "shape" => value.shape, "dtype" => Quantizer.dtype_name(value.class), "codec" => "npz" }
+      end
+
+      def decompress_parameter(directory, name, metadata)
+        return decompress_png16(directory, name, metadata) if name == "means"
+        return decompress_png8(directory, name, metadata) if PNG8_PARAMETERS.include?(name)
+        return KMeans.decompress(File.join(directory, "shN.npz"), metadata) if name == "shN"
+
+        IO::Npy.read_npz(File.join(directory, "#{name}.npz")).fetch("arr")
+      end
+
+      def compress_png8(directory, name, value, side)
+        pixels, metadata = Quantizer.encode(value, side, bits: 8)
+        PngCodec.write(File.join(directory, "#{name}.png"), pixels)
+        metadata
+      end
+
+      def decompress_png8(directory, name, metadata)
+        Quantizer.decode(PngCodec.read(File.join(directory, "#{name}.png")), metadata)
+      end
+
+      def compress_png16(directory, name, value, side)
+        pixels, metadata = Quantizer.encode(value, side, bits: 16)
+        shape = pixels.shape
+        values = pixels.to_a.flatten
+        lower = Numo::UInt8.cast(values.map { |entry| entry & 0xff }).reshape(*shape)
+        upper = Numo::UInt8.cast(values.map { |entry| entry >> 8 }).reshape(*shape)
+        PngCodec.write(File.join(directory, "#{name}_l.png"), lower)
+        PngCodec.write(File.join(directory, "#{name}_u.png"), upper)
+        metadata
+      end
+
+      def decompress_png16(directory, name, metadata)
+        lower = PngCodec.read(File.join(directory, "#{name}_l.png")).to_a.flatten
+        upper = PngCodec.read(File.join(directory, "#{name}_u.png")).to_a.flatten
+        side = ::Math.sqrt(metadata.fetch("shape").first).to_i
+        shape = [side, side, metadata.fetch("shape").drop(1).inject(1, :*)]
+        pixels = Numo::UInt16.cast(
+          lower.zip(upper).map { |low, high| low | (high << 8) }
+        ).reshape(*shape)
+        Quantizer.decode(pixels, metadata)
+      end
+
+      def log_transform(values)
+        output = Numo::NMath.log(1 + values.abs)
+        negative = values.lt(0)
+        output[negative] *= -1 if negative.any?
+        output
+      end
+
+      def inverse_log_transform(values)
+        output = Numo::NMath.exp(values.abs) - 1
+        negative = values.lt(0)
+        output[negative] *= -1 if negative.any?
+        output
+      end
+
+      def validate_name!(name)
+        return if name.match?(/\A[a-zA-Z][a-zA-Z0-9_]*\z/)
+
+        raise ArgumentError, "invalid compression parameter name #{name.inspect}"
+      end
+
+      def validate_options!
+        raise ArgumentError, "kmeans_clusters must be between 1 and 65536" unless (1..65_536).cover?(kmeans_clusters)
+        raise ArgumentError, "kmeans_iterations must be positive" unless kmeans_iterations.positive?
+        raise ArgumentError, "sh_quantization must be between 1 and 8" unless (1..8).cover?(sh_quantization)
+      end
+    end
+  end
+end

data/lib/gsplat/compression/png_codec.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+require "zlib"
+
+module Gsplat
+  module Compression
+    # Dependency-free PNG codec for 8-bit parameter grids.
+    module PngCodec
+      SIGNATURE = "\x89PNG\r\n\x1A\n".b
+      COLOR_TYPES = { 1 => 0, 2 => 4, 3 => 2, 4 => 6 }.freeze
+      CHANNELS = COLOR_TYPES.invert.freeze
+
+      module_function
+
+      def write(path, pixels)
+        validate_pixels!(pixels)
+        height, width, channels = pixels.shape
+        rows = pixels.to_a.map { |row| "\0".b + row.flatten.pack("C*") }.join
+        header = [width, height, 8, COLOR_TYPES.fetch(channels), 0, 0, 0].pack("NNC5")
+        File.binwrite(
+          path,
+          SIGNATURE + chunk("IHDR", header) + chunk("IDAT", Zlib::Deflate.deflate(rows)) + chunk("IEND", "")
+        )
+        path
+      end
+
+      def read(path)
+        data = File.binread(path)
+        raise Gsplat::Error, "invalid PNG signature" unless data.start_with?(SIGNATURE)
+
+        width, height, channels, compressed = parse_chunks(data.byteslice(SIGNATURE.bytesize..))
+        bytes = unfilter(Zlib::Inflate.inflate(compressed), width, height, channels)
+        Numo::UInt8.from_binary(bytes).reshape(height, width, channels)
+      end
+
+      def validate_pixels!(pixels)
+        valid = pixels.is_a?(Numo::UInt8) && pixels.ndim == 3 && COLOR_TYPES.key?(pixels.shape[-1])
+        return if valid
+
+        actual = pixels.respond_to?(:shape) ? pixels.shape.inspect : pixels.class
+        raise ShapeError, "expected UInt8 PNG pixels [H,W,C] with C=1..4, got #{actual}"
+      end
+      private_class_method :validate_pixels!
+
+      def chunk(type, payload)
+        type = type.b
+        [payload.bytesize].pack("N") + type + payload + [Zlib.crc32(type + payload)].pack("N")
+      end
+      private_class_method :chunk
+
+      def parse_chunks(data)
+        width = height = channels = nil
+        compressed = +"".b
+        until data.empty?
+          length = data.byteslice(0, 4)&.unpack1("N")
+          raise Gsplat::Error, "truncated PNG chunk" unless length && data.bytesize >= length + 12
+
+          type = data.byteslice(4, 4)
+          payload = data.byteslice(8, length)
+          expected_crc = data.byteslice(8 + length, 4).unpack1("N")
+          raise Gsplat::Error, "invalid PNG CRC" unless Zlib.crc32(type + payload) == expected_crc
+
+          width, height, channels = parse_header(payload) if type == "IHDR"
+          compressed << payload if type == "IDAT"
+          data = data.byteslice((length + 12)..) || +"".b
+          break if type == "IEND"
+        end
+        raise Gsplat::Error, "PNG is missing IHDR or IDAT" unless width && !compressed.empty?
+
+        [width, height, channels, compressed]
+      end
+      private_class_method :parse_chunks
+
+      def parse_header(payload)
+        width, height, depth, color_type, compression, filter, interlace = payload.unpack("NNC5")
+        valid = depth == 8 && CHANNELS.key?(color_type) && compression.zero? && filter.zero? && interlace.zero?
+        raise NotSupportedError, "only non-interlaced 8-bit PNG is supported" unless valid
+
+        [width, height, CHANNELS.fetch(color_type)]
+      end
+      private_class_method :parse_header
+
+      def unfilter(data, width, height, channels)
+        stride = width * channels
+        output = String.new(capacity: stride * height, encoding: Encoding::BINARY)
+        previous = Array.new(stride, 0)
+        offset = 0
+        height.times do
+          filter = data.getbyte(offset)
+          encoded = data.byteslice(offset + 1, stride)&.bytes
+          raise Gsplat::Error, "truncated PNG scanline" unless encoded&.length == stride
+
+          row = reconstruct_row(encoded, previous, channels, filter)
+          output << row.pack("C*")
+          previous = row
+          offset += stride + 1
+        end
+        output
+      end
+      private_class_method :unfilter
+
+      def reconstruct_row(encoded, previous, channels, filter)
+        raise NotSupportedError, "unsupported PNG filter #{filter}" unless (0..4).cover?(filter)
+
+        encoded.each_with_index.map do |value, index|
+          left = index >= channels ? encoded[index - channels] : 0
+          upper = previous[index]
+          upper_left = index >= channels ? previous[index - channels] : 0
+          prediction = filter_prediction(filter, left, upper, upper_left)
+          encoded[index] = (value + prediction) & 0xff
+        end
+      end
+      private_class_method :reconstruct_row
+
+      def filter_prediction(filter, left, upper, upper_left)
+        return 0 if filter.zero?
+        return left if filter == 1
+        return upper if filter == 2
+        return (left + upper) / 2 if filter == 3
+
+        paeth(left, upper, upper_left)
+      end
+      private_class_method :filter_prediction
+
+      def paeth(left, upper, upper_left)
+        estimate = left + upper - upper_left
+        distances = [(estimate - left).abs, (estimate - upper).abs, (estimate - upper_left).abs]
+        [left, upper, upper_left][distances.each_with_index.min.last]
+      end
+      private_class_method :paeth
+    end
+  end
+end

data/lib/gsplat/compression/quantizer.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Compression
+    # Per-channel linear quantization used by PNG parameter grids.
+    module Quantizer
+      module_function
+
+      def encode(array, side, bits:)
+        feature_count = array.size / array.shape[0]
+        raise NotSupportedError, "PNG parameter grids support at most four channels" if feature_count > 4
+
+        flat = array.reshape(array.shape[0], feature_count)
+        mins, maxs = channel_extrema(flat, feature_count)
+        limit = (1 << bits) - 1
+        values = quantized_values(flat, mins, maxs, limit)
+        type = bits <= 8 ? Numo::UInt8 : Numo::UInt16
+        pixels = type.cast(values).reshape(side, side, feature_count)
+        [pixels, metadata(array, mins, maxs, bits)]
+      end
+
+      def decode(pixels, metadata)
+        shape = metadata.fetch("shape")
+        feature_count = shape.drop(1).inject(1, :*)
+        mins = metadata.fetch("mins")
+        maxs = metadata.fetch("maxs")
+        limit = (1 << metadata.fetch("bits")) - 1
+        values = pixels.reshape(pixels.shape[0] * pixels.shape[1], feature_count).to_a.flat_map do |row|
+          row.each_with_index.map do |value, channel|
+            mins[channel] + ((value.to_f / limit) * (maxs[channel] - mins[channel]))
+          end
+        end
+        numeric_type(metadata.fetch("dtype")).cast(values).reshape(*shape)
+      end
+
+      def channel_extrema(flat, feature_count)
+        [
+          feature_count.times.map { |channel| flat[true, channel].min.to_f },
+          feature_count.times.map { |channel| flat[true, channel].max.to_f }
+        ]
+      end
+      private_class_method :channel_extrema
+
+      def quantized_values(flat, mins, maxs, limit)
+        flat.to_a.flat_map do |row|
+          row.each_with_index.map do |value, channel|
+            span = maxs[channel] - mins[channel]
+            span.zero? ? 0 : (((value - mins[channel]) / span) * limit).round.clamp(0, limit)
+          end
+        end
+      end
+      private_class_method :quantized_values
+
+      def metadata(array, mins, maxs, bits)
+        {
+          "shape" => array.shape,
+          "dtype" => dtype_name(array.class),
+          "mins" => mins,
+          "maxs" => maxs,
+          "bits" => bits
+        }
+      end
+      private_class_method :metadata
+
+      def dtype_name(type)
+        return "float32" if type == Numo::SFloat
+        return "float64" if type == Numo::DFloat
+
+        raise NotSupportedError, "compression requires float32 or float64 parameters"
+      end
+
+      def numeric_type(name)
+        { "float32" => Numo::SFloat, "float64" => Numo::DFloat }.fetch(name) do
+          raise NotSupportedError, "unsupported compressed dtype #{name.inspect}"
+        end
+      end
+    end
+  end
+end