Module: Gsplat::Compression::KMeans

Defined in:
lib/gsplat/compression/kmeans.rb

Overview

Deterministic Manhattan K-means for SH coefficient vectors.

Class Method Summary collapse

Class Method Details

.compress(path, array, clusters:, iterations:, quantization:) ⇒ Object 



9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
 
# File 'lib/gsplat/compression/kmeans.rb', line 9

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

.decompress(path, metadata) ⇒ Object 



34
35
36
37
38
39
40
41
42
 
# File 'lib/gsplat/compression/kmeans.rb', line 34

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