Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Handle Zarr data that is downsampled in Z #12

Merged
merged 3 commits into from
Dec 5, 2024
Merged

Handle Zarr data that is downsampled in Z #12

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
09670b0
Add basic test for downsampled Z data
melissalinkert Oct 4, 2024
ae60b7f
Handle different Z sizes in different resolutions
melissalinkert Oct 4, 2024
f5cb42d
Fix Z index mapping
melissalinkert Dec 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 89 additions & 40 deletions src/main/java/com/glencoesoftware/omero/zarr/ZarrPixelBuffer.java
View file
Open in desktop
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;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Went back and forth on whether this should become a Map<Integer, Map<Integer, Integer>> and store all the mappings for all the resolutions once rather than invalidating them and recomputing every time setResolutionLevel is called.

At least in the context of the image-region endpoints invoked by a viewer like PathViewer, my understanding is that each resolution call will initialize its own ZarrPixelBuffer in a separate thread. So the value of such computation would be limited.


/** { 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++) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why the loop over planes?

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,28 @@ 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];
int zStep = fullResZ / arrayZ;
for (int z=0; z<fullResZ; z++) {
zIndexMap.put(z, z * zStep);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this the right calculation? Using the example of a resolution 0 with [1, 4, 218, 632, 370] and a resolution 1 with [1, 4, 109, 316, 185], this gives

jshell> int fullResZ = 218;
fullResZ ==> 218

jshell> int arrayZ=109;
arrayZ ==> 109

jshell> int zStep = fullResZ / arrayZ;
zStep ==> 2

jshell> Map<Integer, Integer> zIndexMap  = new HashMap<Integer, Integer>();
zIndexMap ==> {}

jshell> for (int z=0; z<fullResZ; z++) {
   ...> zIndexMap.put(z, z * zStep);
   ...> }

jshell> zIndexMap
zIndexMap ==> {0=0, 1=2, 2=4, 3=6, 4=8, 5=10, 6=12, 7=14, 8=16, 9=18, 10=20, 11=22, 12=24, 13=26, 14=28, 15=30, 16=32, 17=34, 18=36, 19=38, 20=40, 21=42, 22=44, 23=46, 24=48, 25=50, 26=52, 27=54, 28=56, 29=58, 30=60, 31=62, 32=64, 33=66, 34=68, 35=70, 36=72, 37=74, 38=76, 39=78, 40=80, 41=82, 42=84, 43=86, 44=88, 45=90, 46=92, 47=94, 48=96, 49=98, 50=100, 51=102, 52=104, 53=106, 54=108, 55=110, 56=112, 57=114, 58=116, 59=118, 60=120, 61=122, 62=124, 63=126, 64=128, 65=130, 66=132, 67=134, 68=136, 69=138, 70=140, 71=142, 72=144, 73=146, 74=148, 75=150, 76=152, 77=154, 78=156, 79=158, 80=160, 81=162, 82=164, 83=166, 84=168, 85=170, 86=172, 87=174, 88=176, 89=178, 90=180 ...  181=362, 182=364, 183=366, 184=368, 185=370, 186=372, 187=374, 188=376, 189=378, 190=380, 191=382, 192=384, 193=386, 194=388, 195=390, 196=392, 197=394, 198=396, 199=398, 200=400, 201=402, 202=404, 203=406, 204=408, 205=410, 206=412, 207=414, 208=416, 209=418, 210=420, 211=422, 212=424, 213=426, 214=428, 215=430, 216=432, 217=434}

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was able to get some decent results looking at a tile using the same z,c,t,x,y,w,h value and different resolutions with the following changes

-            int zStep = fullResZ / arrayZ;
             for (int z=0; z<fullResZ; z++) {
-                zIndexMap.put(z, z * zStep);
+                zIndexMap.put(z, Math.round(z * arrayZ / fullResZ));
             }

The exact calculation including whether round is the correct rounding operation should obviously be reviewed

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, that makes sense. Proposed arithmetic fix is in f5cb42d.

}
} 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
Open in desktop
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();
}
}
}
}