fix(png): fix multi-channel png decoding

Fixes google/tensorstore#209.

python/tests/n5_test.py

@@ -19,8 +19,6 @@
 import numpy as np
 import pytest
 
-TEST_DATA_DIR = pathlib.Path(__file__).parent.parent / "testdata"
-
 
 @pytest.mark.parametrize(
     "spec",

python/tests/precomputed_test.py

@@ -0,0 +1,93 @@
+# @license
+# Copyright 2023 Google Inc.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests the n5 datasource."""
+
+import pathlib
+
+import neuroglancer
+import numpy as np
+import pytest
+
+
+@pytest.mark.parametrize(
+    "spec",
+    [
+        {
+            "driver": "neuroglancer_precomputed",
+            "scale_metadata": {"encoding": "png", "chunk_size": [8, 9, 1]},
+            "dtype": "uint8",
+            "schema": {"domain": {"shape": [10, 20, 5, num_channels]}},
+        }
+        for num_channels in [1, 2, 3, 4]
+    ]
+    + [
+        # Currently TensorStore does not support uint16 with more than one channel.
+        {
+            "driver": "neuroglancer_precomputed",
+            "scale_metadata": {"encoding": "png"},
+            "dtype": "uint16",
+            "schema": {"domain": {"shape": [10, 20, 5, 1]}},
+        }
+    ]
+    + [
+        # Due to a tensorstore bug (https://github.com/google/neuroglancer/issues/677)
+        # the block shape must be square.
+        {
+            "driver": "neuroglancer_precomputed",
+            "scale_metadata": {"encoding": "jpeg", "chunk_size": [10, 10, 1]},
+            "dtype": "uint8",
+            "schema": {"domain": {"shape": [10, 20, 1, num_channels]}},
+        }
+        for num_channels in [1, 3]
+    ],
+    ids=str,
+)
+def test_precomputed(tempdir_server: tuple[pathlib.Path, str], webdriver, spec):
+    import tensorstore as ts
+
+    tmp_path, server_url = tempdir_server
+
+    full_spec = {
+        "kvstore": {
+            "driver": "file",
+            "path": str(tmp_path),
+        }
+    }
+    full_spec.update(spec)
+
+    store = ts.open(full_spec, create=True).result()
+
+    a = np.arange(np.prod(store.shape), dtype=store.dtype.numpy_dtype).reshape(
+        store.shape
+    )
+
+    store[...] = a
+
+    with webdriver.viewer.txn() as s:
+        s.layers.append(
+            name="a",
+            layer=neuroglancer.ImageLayer(source=f"precomputed://{server_url}"),
+        )
+
+    if store.shape[-1] == 1:
+        # Neuroglancer elides the channel dimension if there is only 1 channel
+        store = store[..., 0]
+
+    vol = webdriver.viewer.volume("a").result()
+    b = vol.read().result()
+
+    if spec["scale_metadata"]["encoding"] == "jpeg":
+        np.testing.assert_allclose(store.read().result(), b, atol=4, rtol=0)
+    else:
+        np.testing.assert_equal(store.read().result(), b)

src/sliceview/png/index.ts

@@ -15,6 +15,7 @@
  */
 
 import type { DecodedImage } from "#src/async_computation/decode_png_request.js";
+import { transposeArray2d } from "#src/util/array.js";
 
 const libraryEnv = {
   emscripten_notify_memory_growth: () => {},
@@ -160,7 +161,7 @@ function readHeader(buffer: Uint8Array): {
         `png: invalid bit depth for grayscale + alpha channel colorspace. Got: ${bitDepth}`,
       );
     }
-    numChannels = 4;
+    numChannels = 2;
   } else {
     throw new Error(`png: Invalid color space: ${colorSpace}`);
   }
@@ -232,11 +233,16 @@ export async function decompressPng(
     // Likewise, we reference memory.buffer instead of heap.buffer
     // because memory growth during decompress could have detached
     // the buffer.
-    const image = new Uint8Array(
+    let image = new Uint8Array(
       (m.exports.memory as WebAssembly.Memory).buffer,
       imagePtr,
       nbytes,
     );
+
+    if (numChannels !== 1) {
+      image = transposeArray2d(image, sx * sy, numChannels);
+    }
+
     // copy the array so it can be memory managed by JS
     // and we can free the emscripten buffer
     return {