From 0c3bb74e1af6ef7f6df7d42f6ee26e8981ad4145 Mon Sep 17 00:00:00 2001
From: lucasoptml <47575527+lucasoptml@users.noreply.github.com>
Date: Fri, 11 Jul 2025 13:30:57 -0500
Subject: [PATCH] Add sh1, sh2 support to sogs (#109)

* Add sh1, sh2 support to sogs

* Fix lint issues

---------

Co-authored-by: Diego Marcos Segura <diego.marcos@gmail.com>
---
 src/pcsogs.ts | 61 +++++++++++++++++++++++++++++++--------------------
 1 file changed, 37 insertions(+), 24 deletions(-)

diff --git a/src/pcsogs.ts b/src/pcsogs.ts
index 572deee..d86b924 100644
--- a/src/pcsogs.ts
+++ b/src/pcsogs.ts
@@ -97,9 +97,14 @@ export async function unpackPcSogs(
   }
 
   if (json.shN) {
-    extra.sh1 = new Uint32Array(numSplats * 2);
-    extra.sh2 = new Uint32Array(numSplats * 4);
-    extra.sh3 = new Uint32Array(numSplats * 4);
+    const useSH3 = json.shN.shape[1] >= 48 - 3;
+    const useSH2 = json.shN.shape[1] >= 27 - 3;
+    const useSH1 = json.shN.shape[1] >= 12 - 3;
+
+    if (useSH1) extra.sh1 = new Uint32Array(numSplats * 2);
+    if (useSH2) extra.sh2 = new Uint32Array(numSplats * 4);
+    if (useSH3) extra.sh3 = new Uint32Array(numSplats * 4);
+
     const sh1 = new Float32Array(9);
     const sh2 = new Float32Array(15);
     const sh3 = new Float32Array(21);
@@ -116,32 +121,40 @@ export async function unpackPcSogs(
       const offset = row * centroids.width + col;
 
       for (let d = 0; d < 3; ++d) {
-        for (let k = 0; k < 3; ++k) {
-          sh1[k * 3 + d] =
-            json.shN.mins +
-            ((json.shN.maxs - json.shN.mins) *
-              centroids.rgba[(offset + k) * 4 + d]) /
-              255;
+        if (useSH1) {
+          for (let k = 0; k < 3; ++k) {
+            sh1[k * 3 + d] =
+              json.shN.mins +
+              ((json.shN.maxs - json.shN.mins) *
+                centroids.rgba[(offset + k) * 4 + d]) /
+                255;
+          }
         }
-        for (let k = 0; k < 5; ++k) {
-          sh2[k * 3 + d] =
-            json.shN.mins +
-            ((json.shN.maxs - json.shN.mins) *
-              centroids.rgba[(offset + 3 + k) * 4 + d]) /
-              255;
+
+        if (useSH2) {
+          for (let k = 0; k < 5; ++k) {
+            sh2[k * 3 + d] =
+              json.shN.mins +
+              ((json.shN.maxs - json.shN.mins) *
+                centroids.rgba[(offset + 3 + k) * 4 + d]) /
+                255;
+          }
         }
-        for (let k = 0; k < 7; ++k) {
-          sh3[k * 3 + d] =
-            json.shN.mins +
-            ((json.shN.maxs - json.shN.mins) *
-              centroids.rgba[(offset + 8 + k) * 4 + d]) /
-              255;
+
+        if (useSH3) {
+          for (let k = 0; k < 7; ++k) {
+            sh3[k * 3 + d] =
+              json.shN.mins +
+              ((json.shN.maxs - json.shN.mins) *
+                centroids.rgba[(offset + 8 + k) * 4 + d]) /
+                255;
+          }
         }
       }
 
-      encodeSh1Rgb(extra.sh1 as Uint32Array, i, sh1);
-      encodeSh2Rgb(extra.sh2 as Uint32Array, i, sh2);
-      encodeSh3Rgb(extra.sh3 as Uint32Array, i, sh3);
+      if (useSH1) encodeSh1Rgb(extra.sh1 as Uint32Array, i, sh1);
+      if (useSH2) encodeSh2Rgb(extra.sh2 as Uint32Array, i, sh2);
+      if (useSH3) encodeSh3Rgb(extra.sh3 as Uint32Array, i, sh3);
     }
   }