CompressedSplatPlyFixture.kt
src/jvmMain/kotlin/afield/graphics/scene/CompressedSplatPlyFixture.kt · 225 lines
package afield.graphics.scene
import afield.graphics.render.before
import java.io.File
import java.nio.ByteBuffer
import java.nio.ByteOrder
import kotlin.math.exp
import kotlin.math.floor
import kotlin.math.sqrt
internal object CompressedSplatPlyFixture
Writes PlayCanvas compressed PLY files for PlyLoaderUnitTest.
This is a fixture, not a feature: it is a line-by-line transcription of splat-transform's
CompressedChunk.pack and writeCompressedPly, so that CompressedSplatPly is checked against
the format's own encoder instead of against a second copy of whatever the reader believes. Field
has no reason to write these files, so this deliberately has no options, no streaming and no
Morton sort — the reader does not care what order the splats arrive in.
@see CompressedSplatPly
internal object CompressedSplatPlyFixture {
data class Splat
One splat in the uncompressed 3DGS conventions: log scales, logit opacity, SH-DC colour.
data class Splat(
val x: Float,
val y: Float,
val z: Float,
val fDc: FloatArray,
val opacity: Float,
/** `rot_0` first, the real part, as a 3DGS PLY spells it. */
val rot: FloatArray,
val scale: FloatArray,
/** `f_rest_0..` in channel-major order; 0, 9, 24 or 45 of them. */
val shRest: FloatArray = FloatArray(0)
)
fun write
@param withChunkColors emit the six min_r..max_b chunk properties. Off exercises the older
12-property chunk layout, where colour is a plain unorm instead of a per-chunk lerp.
@param withSh emit the sh element from Splat.shRest.
body — 105 lines, from line 44
fun write(
file: File,
splats: List<Splat>,
withChunkColors: Boolean = true,
withSh: Boolean = false
) {
require(splats.isNotEmpty()) { "need at least one splat" }
val shCount = if (withSh) splats.first().shRest.size else 0
require(splats.all { it.shRest.size == shCount || !withSh }) {
"every splat needs the same number of f_rest_* coefficients"
}
val chunkProps =
requiredChunkNames + if (withChunkColors) CompressedSplatPly.colorChunkProperties else emptyList()
val chunkCount = (splats.size + CompressedSplatPly.CHUNK_SIZE - 1) / CompressedSplatPly.CHUNK_SIZE
val header = buildString {
append("ply\n")
append("format binary_little_endian 1.0\n")
append("comment Generated by CompressedSplatPlyFixture\n")
append("element chunk $chunkCount\n")
chunkProps.forEach { append("property float $it\n") }
append("element vertex ${splats.size}\n")
CompressedSplatPly.vertexProperties.forEach { append("property uint $it\n") }
if (shCount > 0) {
append("element sh ${splats.size}\n")
repeat(shCount) { append("property uchar f_rest_$it\n") }
}
append("end_header\n")
}
val chunkData = ByteBuffer.allocate(chunkCount * chunkProps.size * 4).order(ByteOrder.LITTLE_ENDIAN)
val vertexData = ByteBuffer.allocate(splats.size * 16).order(ByteOrder.LITTLE_ENDIAN)
val shData = ByteBuffer.allocate(splats.size * shCount)
for (chunkIndex in 0 until chunkCount) {
val start = chunkIndex * CompressedSplatPly.CHUNK_SIZE
val end = minOf(splats.size, start + CompressedSplatPly.CHUNK_SIZE)
// The reference repeats the last real gaussian to fill a partial chunk, which cannot move
// the min/max because the repeated value is itself real data.
val members = (start until end).map { splats[it] }
val padded = members + List(CompressedSplatPly.CHUNK_SIZE - members.size) { members.last() }
val px = minMax(padded.map { it.x })
val py = minMax(padded.map { it.y })
val pz = minMax(padded.map { it.z })
val sx = minMax(padded.map { it.scale[0] }).clampScale()
val sy = minMax(padded.map { it.scale[1] }).clampScale()
val sz = minMax(padded.map { it.scale[2] }).clampScale()
// f_dc becomes a colour before the range is measured.
val colors = padded.map { splat ->
FloatArray(3) { splat.fDc[it] * CompressedSplatPly.SH_C0 + 0.5f }
}
val cr = minMax(colors.map { it[0] })
val cg = minMax(colors.map { it[1] })
val cb = minMax(colors.map { it[2] })
val ranges = mapOf(
"min_x" to px.min, "min_y" to py.min, "min_z" to pz.min,
"max_x" to px.max, "max_y" to py.max, "max_z" to pz.max,
"min_scale_x" to sx.min, "min_scale_y" to sy.min, "min_scale_z" to sz.min,
"max_scale_x" to sx.max, "max_scale_y" to sy.max, "max_scale_z" to sz.max,
"min_r" to cr.min, "min_g" to cg.min, "min_b" to cb.min,
"max_r" to cr.max, "max_g" to cg.max, "max_b" to cb.max
)
chunkProps.forEach { chunkData.putFloat(ranges.getValue(it)) }
for ((offset, splat) in members.withIndex()) {
val color = colors[offset]
vertexData.putInt(
pack111011(
normalize(splat.x, px), normalize(splat.y, py), normalize(splat.z, pz)
)
)
vertexData.putInt(packRot(splat.rot[0], splat.rot[1], splat.rot[2], splat.rot[3]))
vertexData.putInt(
pack111011(
normalize(splat.scale[0], sx), normalize(splat.scale[1], sy), normalize(splat.scale[2], sz)
)
)
vertexData.putInt(
if (chunkProps.size == 18) pack8888(
normalize(color[0], cr), normalize(color[1], cg), normalize(color[2], cb),
sigmoid(splat.opacity)
) else pack8888(color[0], color[1], color[2], sigmoid(splat.opacity))
)
}
}
if (shCount > 0) {
for (splat in splats) {
for (coefficient in splat.shRest) {
val normalized = coefficient / 8f + 0.5f
shData.put(((normalized * 256f).toInt()).coerceIn(0, 255).toByte())
}
}
}
file.outputStream().buffered().use { stream ->
stream.write(header.toByteArray())
stream.write(chunkData.array())
stream.write(vertexData.array())
stream.write(shData.array())
}
}
fun writeUncompressed(file: File, splats: List<Splat>)
Writes an uncompressed 3DGS PLY carrying the same splats, so a test can read both and compare.
fun writeUncompressed(file: File, splats: List<Splat>) {
val shCount = splats.first().shRest.size
val columns = CompressedSplatPly.decompressedColumns + (0 until shCount).map { "f_rest_$it" }
val header = buildString {
append("ply\n")
append("format binary_little_endian 1.0\n")
append("element vertex ${splats.size}\n")
columns.forEach { append("property float $it\n") }
append("end_header\n")
}
val body = ByteBuffer.allocate(splats.size * columns.size * 4).order(ByteOrder.LITTLE_ENDIAN)
for (splat in splats) {
body.putFloat(splat.x).putFloat(splat.y).putFloat(splat.z)
body.putFloat(splat.fDc[0]).putFloat(splat.fDc[1]).putFloat(splat.fDc[2])
body.putFloat(splat.opacity)
splat.rot.forEach { body.putFloat(it) }
splat.scale.forEach { body.putFloat(it) }
splat.shRest.forEach { body.putFloat(it) }
}
file.outputStream().buffered().use { stream ->
stream.write(header.toByteArray())
stream.write(body.array())
}
}
splat-transform's CompressedChunk.pack, transcribed
private val requiredChunkNames = CompressedSplatPly.requiredChunkProperties
data class Range(val min: Float, val max: Float)
private fun minMax(values: List<Float>) = Range(values.min(), values.max())
private fun Range.clampScale() = Range(min.coerceIn(-20f, 20f), max.coerceIn(-20f, 20f))
private fun normalize(x: Float, range: Range): Float = when {
x <= range.min -> 0f
x >= range.max -> 1f
range.max - range.min < 0.00001f -> 0f
else -> (x - range.min) / (range.max - range.min)
}
private fun packUnorm(value: Float, bits: Int): Int {
val t = (1 shl bits) - 1
return floor(value * t + 0.5f).toInt().coerceIn(0, t)
}
private fun pack111011(x: Float, y: Float, z: Float): Int =
(packUnorm(x, 11) shl 21) or (packUnorm(y, 10) shl 11) or packUnorm(z, 11)
private fun pack8888(x: Float, y: Float, z: Float, w: Float): Int =
(packUnorm(x, 8) shl 24) or (packUnorm(y, 8) shl 16) or (packUnorm(z, 8) shl 8) or packUnorm(w, 8)
private fun packRot(r0: Float, r1: Float, r2: Float, r3: Float): Int {
val length = sqrt(r0 * r0 + r1 * r1 + r2 * r2 + r3 * r3)
val a = floatArrayOf(r0 / length, r1 / length, r2 / length, r3 / length)
var largest = 0
for (i in 1..3) if (kotlin.math.abs(a[i]) > kotlin.math.abs(a[largest])) largest = i
if (a[largest] < 0f) for (i in 0..3) a[i] = -a[i]
val norm = sqrt(2f) * 0.5f
var result = largest
for (i in 0..3) {
if (i != largest) result = (result shl 10) or packUnorm(a[i] * norm + 0.5f, 10)
}
return result
}
private fun sigmoid(value: Float): Float = (1.0 / (1.0 + exp(-value.toDouble()))).toFloat()
}