Skip to content

Commit

Permalink
Merge pull request #12 from melissalinkert/z-downsample
Browse files Browse the repository at this point in the history 
Handle Zarr data that is downsampled in Z
  • Loading branch information
@sbesson
sbesson authored Dec 5, 2024
2 parents deabf58 + f5cb42d commit f3a00c5
Show file tree
Hide file tree
Showing 2 changed files with 124 additions and 40 deletions.
128 changes: 88 additions & 40 deletions src/main/java/com/glencoesoftware/omero/zarr/ZarrPixelBuffer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;
Expand Down Expand Up @@ -75,6 +76,12 @@ public class ZarrPixelBuffer implements PixelBuffer {
/** Zarr array corresponding to the current resolution level */
private ZarrArray array;

/**
* Mapping of Z plane indexes in full resolution to
* Z plane indexes in current resolution.
*/
private Map<Integer, Integer> zIndexMap;

/** { resolutionLevel, z, c, t, x, y, w, h } vs. tile byte array cache */
private final AsyncLoadingCache<List<Integer>, byte[]> tileCache;

Expand Down Expand Up @@ -191,49 +198,64 @@ private void read(byte[] buffer, int[] shape, int[] offset)
// Check planar read size (sizeX and sizeY only)
checkReadSize(Arrays.copyOfRange(shape, 3, 5));

// if reading from a resolution downsampled in Z,
// adjust the shape/offset for the Z coordinate only
// this ensures that the correct Zs are read from the correct offsets
// since the requested shape/offset may not match the underlying array
int planes = 1;
int originalZIndex = offset[2];
if (getSizeZ() != getTrueSizeZ()) {
offset[2] = zIndexMap.get(originalZIndex);
planes = shape[2];
shape[2] = 1;
}

try {
ByteBuffer asByteBuffer = ByteBuffer.wrap(buffer);
DataType dataType = array.getDataType();
switch (dataType) {
case u1:
case i1:
array.read(buffer, shape, offset);
break;
case u2:
case i2:
{
short[] data = (short[]) array.read(shape, offset);
asByteBuffer.asShortBuffer().put(data);
break;
}
case u4:
case i4:
{
int[] data = (int[]) array.read(shape, offset);
asByteBuffer.asIntBuffer().put(data);
break;
}
case i8:
{
long[] data = (long[]) array.read(shape, offset);
asByteBuffer.asLongBuffer().put(data);
break;
}
case f4:
{
float[] data = (float[]) array.read(shape, offset);
asByteBuffer.asFloatBuffer().put(data);
break;
}
case f8:
{
double[] data = (double[]) array.read(shape, offset);
asByteBuffer.asDoubleBuffer().put(data);
break;
}
default:
throw new IllegalArgumentException(
"Data type " + dataType + " not supported");
for (int z=0; z<planes; z++) {
offset[2] = zIndexMap.get(originalZIndex + z);
switch (dataType) {
case u1:
case i1:
array.read(buffer, shape, offset);
break;
case u2:
case i2:
{
short[] data = (short[]) array.read(shape, offset);
asByteBuffer.asShortBuffer().put(data);
break;
}
case u4:
case i4:
{
int[] data = (int[]) array.read(shape, offset);
asByteBuffer.asIntBuffer().put(data);
break;
}
case i8:
{
long[] data = (long[]) array.read(shape, offset);
asByteBuffer.asLongBuffer().put(data);
break;
}
case f4:
{
float[] data = (float[]) array.read(shape, offset);
asByteBuffer.asFloatBuffer().put(data);
break;
}
case f8:
{
double[] data = (double[]) array.read(shape, offset);
asByteBuffer.asDoubleBuffer().put(data);
break;
}
default:
throw new IllegalArgumentException(
"Data type " + dataType + " not supported");
}
}
} catch (InvalidRangeException e) {
log.error("Error reading Zarr data", e);
Expand Down Expand Up @@ -745,6 +767,14 @@ public int getSizeY() {

@Override
public int getSizeZ() {
// this is expected to be the Z size of the full resolution array
return zIndexMap.size();
}

/**
* @return Z size of the current underlying Zarr array
*/
private int getTrueSizeZ() {
return array.getShape()[2];
}

Expand Down Expand Up @@ -783,9 +813,27 @@ public void setResolutionLevel(int resolutionLevel) {
throw new IllegalArgumentException(
"This Zarr file has no pixel data");
}
if (zIndexMap == null) {
zIndexMap = new HashMap<Integer, Integer>();
}
else {
zIndexMap.clear();
}
try {
array = zarrArrayCache.get(
root.resolve(Integer.toString(this.resolutionLevel))).get();

ZarrArray fullResolutionArray = zarrArrayCache.get(
root.resolve("0")).get();

// map each Z index in the full resolution array
// to a Z index in the subresolution array
// if no Z downsampling, this is just an identity map
int fullResZ = fullResolutionArray.getShape()[2];
int arrayZ = array.getShape()[2];
for (int z=0; z<fullResZ; z++) {
zIndexMap.put(z, Math.round(z * arrayZ / fullResZ));
}
} catch (Exception e) {
// FIXME: Throw the right exception
throw new RuntimeException(e);
Expand Down
36 changes: 36 additions & 0 deletions src/test/java/com/glencoesoftware/omero/zarr/ZarrPixelBufferTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -750,4 +750,40 @@ public void testSetResolutionLevelOutOfBounds() throws IOException {
zpbuf.setResolutionLevel(3);
}
}

@Test
public void testDownsampledZ() throws IOException {
int sizeT = 1;
int sizeC = 1;
int sizeZ = 16;
int sizeY = 2048;
int sizeX = 2048;
int resolutions = 3;

Pixels pixels = new Pixels(
null, null, sizeX, sizeY, sizeZ, sizeC, sizeT, "", null);
Path output = writeTestZarr(
sizeT, sizeC, sizeZ, sizeY, sizeX, "uint8", resolutions);

// Hack the .zarray to hide Z sections in lower resolutions
for (int r=1; r<resolutions; r++) {
ObjectMapper mapper = new ObjectMapper();
HashMap<String, Object> zArray = mapper.readValue(
Files.readAllBytes(output.resolve("0/" + r + "/.zarray")),
HashMap.class);
List<Integer> shape = (List<Integer>) zArray.get("shape");
shape.set(2, sizeZ / (int) Math.pow(2, r));
mapper.writeValue(output.resolve("0/" + r + "/.zarray").toFile(), zArray);
}

try (ZarrPixelBuffer zpbuf =
createPixelBuffer(pixels, output.resolve("0"), sizeX, sizeY)) {
// get the last Z section, for each resolution level
for (int r=0; r<resolutions; r++) {
zpbuf.setResolutionLevel(r);

byte[] plane = zpbuf.getPlane(sizeZ - 1, 0, 0).getData().array();
}
}
}
}

0 comments on commit f3a00c5

Please sign in to comment.