Proper way to make a volatile source for lazy loading and bdv display ?

[NicoKiaru]

Hello,

Excited to see this project going forward! I've already played quite a bit with it and it's really nice.

I have a question because for a better integration of Warpy and ABBA into QuPath I need to use BigDataViewer (support of transformed images, BigWarp), and BDV works better with data implementing the Source interface. Together with VolatileSource, this brings async loading to BDV and it's really life-changing for responsiveness.

However, for VolatileSource to work, there's this line which is called:

https://github.com/bigdataviewer/bigdataviewer-core/blob/f5226a680e61fcbccc329b1e92bb7d75edb320b7/src/main/java/bdv/util/VolatileSource.java#L71

This wrapping requires that the RAI implements cell-based access:

https://github.com/bigdataviewer/bigdataviewer-core/blob/f5226a680e61fcbccc329b1e92bb7d75edb320b7/src/main/java/bdv/util/volatiles/VolatileViews.java#L118

So it's not an issue, you could say, because you're already providing a LazyCellImg, so that should work. Unfortunately, the Source interface requires a 3D RAI, which is an XYZ hyperslice for a certain c and t of the 5D LazyCellImg provided in this repo. And hyperslicing the LazyCellImg causes it to lose its cell-based structure (in the sense that WrapAsVolatile throws an IllegalArgumentException).

So I wonder if you have any idea how this could be fixed?

Sorry it's long; I hope it's clear. In brief, here's my ideal solution (and I don't know if it's doable):

---

Is it possible to add a method that creates a LazyCellImg, but for a specific c and t index, and thus return a LazyCell 3D image? Something like

public Img<T> createForLevelCT(int level, int c, int t)

---

This method needs to preserve the cell-based structure while providing the 3D RAI that BDV expects. Maybe that's not easily doable, I don't know, but I'd love to have your input on this. Also pinging @tpietzsch in case he has time to chime in!

[petebankhead]

Hi @NicoKiaru I've spent a day looking into this and made a bit of progress.

I think the main problem is that CachedCellImg is a subclass of LazyCellImg - and so VotatileViews won't like that we're using LazyCellImg directly anyway.

However I'm not entirely sure that we need it.

I've written a Java class and included it in a Groovy script to show how an ImageServer from QuPath can provide the pixels to an image that is then directly viewed in BigDataViewer.

https://gist.github.com/petebankhead/d387639dc9ff3f437450634ea5d5813d

I've tried this with large RGB and multiplexed images, and to me it looks very smooth indeed.
It should also work with 5D images (mitosis.tif was fine... although it's pretty tiny).

The script doesn't require any changes to this repo. It doesn't use imglib2-cache and (presumably) won't do any disk caching, but it will use QuPath's existing tile cache and I think @Rylern has optimized pixel access to avoid unnecessary copies to make that enough.

I haven't used it for long enough to see any memory leaks emerge, and I don't know how it would cope with extra transforms, but I think it's worth a try. If it's promising enough, we're thinking of including qupath-imglib2 in QuPath v0.7.0 and then the script should work for anyone who runs from source with ./gradlew run -Pfiji.

[NicoKiaru]

Thanks @Rylern and @petebankhead for the time spent on it. I'll give it a go next week.

Cheers

---

In fact I tried now ;-). Yes it's fairly smooth - but I still have some lag for 3D images; and for Warpy, streaming an image server will also slow things down. I think it's great and is very useful in its current state, but I'm still in love with async loading. So I'll keep exploring ways to achieve that.

I think @Rylern has optimized pixel access to avoid unnecessary copies to make that enough.

Yes, yes that's really great! Indeed I'm not sure it's possible to have async loading without an extra other cache layer, but maybe the memory cost is worth. Let's see

[NicoKiaru]

Ok, quick and dirty (not type agnostic, and single channel for the moment) but:

import qupath.ext.imglib2.QuPathBdvSource
import net.imglib2.type.numeric.ARGBType
import net.imglib2.realtransform.AffineTransform3D
import net.imglib2.type.numeric.integer.UnsignedShortType
import net.imglib2.type.volatiles.VolatileUnsignedShortType

def source = new QuPathBdvSource(getCurrentServer(), new UnsignedShortType(), "")
def ans = bdv.util.BdvFunctions.show(new bdv.util.VolatileSource(source, new VolatileUnsignedShortType(), new bdv.cache.SharedQueue(10)))
ans.getBdvHandle().getViewerPanel().state().setViewerTransform(new AffineTransform3D());

With the class QuPathBdvSource like that:

package qupath.ext.imglib2;

import bdv.util.DefaultInterpolators;
import bdv.viewer.Interpolation;
import bdv.viewer.Source;
import mpicbg.spim.data.sequence.VoxelDimensions;
import net.imglib2.Cursor;
import net.imglib2.RandomAccessibleInterval;
import net.imglib2.RealRandomAccessible;
import net.imglib2.cache.img.CellLoader;
import net.imglib2.cache.img.ReadOnlyCachedCellImgFactory;
import net.imglib2.cache.img.ReadOnlyCachedCellImgOptions;
import net.imglib2.realtransform.AffineTransform3D;
import net.imglib2.type.NativeType;
import net.imglib2.type.numeric.NumericType;
import net.imglib2.type.volatiles.VolatileUnsignedShortType;
import net.imglib2.util.Intervals;
import net.imglib2.view.ExtendedRandomAccessibleInterval;
import net.imglib2.view.Views;
import qupath.lib.images.servers.ImageServer;

import java.awt.image.BufferedImage;
import java.util.HashMap;
import java.util.Map;

/**
 * Implementation of BigDataViewer's Source interface for QuPath ImageServers.
 * This allows QuPath images to be displayed in BigDataViewer.
 *
 * @param <T> the pixel type (e.g., ARGBType, UnsignedByteType, etc.)
 */
public class QuPathBdvSource<T extends NumericType<T> & NativeType<T>> implements Source<T> {
    VolatileUnsignedShortType
    private final ImageServer<BufferedImage> server;
    private final ImgCreator imgCreator;
    private final String name;
    private final T type;

    /**
     * Create a new QuPath BDV Source
     *
     * @param server the QuPath ImageServer
     * @param type the ImgLib2 type instance
     * @param name the name for this source
     */
    public QuPathBdvSource(ImageServer<BufferedImage> server, T type, String name) {
        this.server = server;
        this.type = type;
        this.name = server.getMetadata().getName();
        this.imgCreator = ImgCreator.builder(server).build();
    }

    @Override
    public boolean isPresent(int t) {
        // Check if timepoint exists
        return t >= 0 && t < server.nTimepoints();
    }

    final Map<Integer, Map<Integer, RandomAccessibleInterval<T>>> tLevelToRAI = new HashMap<>();

    @Override
    public RandomAccessibleInterval<T> getSource(int t, int level) {
        if (!isPresent(t)) {
            return null;
        }
        synchronized (tLevelToRAI) {
            if (!tLevelToRAI.containsKey(t)) {
                tLevelToRAI.put(t, new HashMap<>());
            }
            if (!tLevelToRAI.get(t).containsKey(level)) {
                // Create the image for the specified level
                RandomAccessibleInterval<T> source = imgCreator.createForLevel(level);
                //BdvStackSource<T> ans = BdvFunctions.show(this);
                //ans.getBdvHandle().getViewerPanel().state().setViewerTransform(new AffineTransform3D());

                // BDV expects 3D (XYZ), so we need to remove C and T dimensions via hyperslicing
                // Hyperslice at the given timepoint t and channel 0 (or make channel configurable)
                int cIndex = ImgCreator.getIndexOfDimension(qupath.ext.imglib2.Dimension.CHANNEL);
                int tIndex = ImgCreator.getIndexOfDimension(qupath.ext.imglib2.Dimension.TIME);

                // Hyperslice to remove C dimension (use channel 0)
                //if (cIndex >= 0 && source.numDimensions() > cIndex) {
                source = Views.hyperSlice(source, cIndex, 0);
                //}

                // Hyperslice to remove T dimension (use the provided timepoint)
                //if (tIndex >= 0 && source.numDimensions() > tIndex) {
                // Adjust index if C was already removed
                int adjustedTIndex = tIndex;
                if (cIndex >= 0 && cIndex < tIndex) {
                    adjustedTIndex--;
                }
                source = Views.hyperSlice(source, adjustedTIndex, t);

                tLevelToRAI.get(t).put(level, cacheWithCells(source,512,512,1));
            }
        }
        return tLevelToRAI.get(t).get(level);
    }

    transient protected final DefaultInterpolators interpolators = new DefaultInterpolators<>();


    @Override
    public RealRandomAccessible<T> getInterpolatedSource(int t, int level, Interpolation method) {
                ExtendedRandomAccessibleInterval
                eView = Views.extendZero(getSource( t, level ));
        RealRandomAccessible realRandomAccessible = Views.interpolate( eView, interpolators.get(method) );
        return realRandomAccessible;
    }


    @Override
    public void getSourceTransform(int t, int level, AffineTransform3D transform) {
        transform.identity();

        if (!isPresent(t)) {
            return;
        }

        // Get the downsample factor for this level
        double downsample = server.getDownsampleForResolution(level);

        // Get pixel calibration
        var cal = server.getPixelCalibration();
        double pixelWidth = cal.getPixelWidthMicrons();
        double pixelHeight = cal.getPixelHeightMicrons();
        double zSpacing = cal.getZSpacingMicrons();

        transform.identity();
        // Apply scaling for downsample and physical units
        transform.scale( downsample, downsample, 1);
    }

    @Override
    public T getType() {
        return type;
    }

    @Override
    public String getName() {
        return name;
    }

    @Override
    public VoxelDimensions getVoxelDimensions() {
        var cal = server.getPixelCalibration();
        return new QuPathVoxelDimensions(
                cal.getPixelWidthMicrons(),
                cal.getPixelHeightMicrons(),
                cal.getZSpacingMicrons()
        );
    }

    @Override
    public int getNumMipmapLevels() {
        return server.nResolutions()-1;
    }

    /**
     * Simple VoxelDimensions implementation for QuPath calibration
     */
    private static class QuPathVoxelDimensions implements VoxelDimensions {
        private final double[] dimensions;
        private final String unit;

        public QuPathVoxelDimensions(double pixelWidth, double pixelHeight, double zSpacing) {
            this.dimensions = new double[] { pixelWidth, pixelHeight, zSpacing };
            this.unit = "µm";
        }

        @Override
        public String unit() {
            return unit;
        }

        @Override
        public void dimensions(double[] target) {
            System.arraycopy(dimensions, 0, target, 0, Math.min(dimensions.length, target.length));
        }

        @Override
        public double dimension(int d) {
            return dimensions[d];
        }

        @Override
        public int numDimensions() {
            return dimensions.length;
        }
    }


    /**
     * Wraps a 3D RandomAccessibleInterval in a cell-cached image
     *
     * @param source The source 3D RandomAccessibleInterval to cache
     * @param cellDimensions The dimensions of each cache cell (e.g., {512, 512, 1})
     * @return A cached version of the input
     */
    public static <T extends NativeType<T>> RandomAccessibleInterval<T> cacheWithCells(
            final RandomAccessibleInterval<T> source,
            final int... cellDimensions) {

        // Validate dimensions
        if (source.numDimensions() != 3 || cellDimensions.length != 3) {
            throw new IllegalArgumentException("This function expects 3D data");
        }

        // Get the dimensions of the source
        final long[] dimensions = Intervals.dimensionsAsLongArray(source);

        // Get the type from the source
        final T type = source.randomAccess().get().createVariable();

        // Create a CellLoader that reads from the source
        final CellLoader<T> cellLoader = cell -> {
            // Copy data from source to the cell
            final Cursor<T> sourceCursor = Views.flatIterable(
                    Views.interval(source, cell)
            ).cursor();
            final Cursor<T> cellCursor = cell.cursor();

            while (sourceCursor.hasNext()) {
                cellCursor.next().set(sourceCursor.next());
            }
        };

        // Create cached cell image options
        final ReadOnlyCachedCellImgOptions options = ReadOnlyCachedCellImgOptions
                .options()
                .cellDimensions(cellDimensions)
                .maxCacheSize(100);

        // Create the cached cell image factory
        final ReadOnlyCachedCellImgFactory factory = new ReadOnlyCachedCellImgFactory(options);

        // Create and return the cached image
        return factory.create(
                dimensions,
                type,
                cellLoader
        );
    }
}

I've got lazy loading:

demo-qupath-bdv-async.mp4

There's an extra cache layer, but it can be bounded so the memory increase can be small.

[tpietzsch]

Hi @NicoKiaru

I didn’t read your code in detail, but I think it might be worth trying to just wrap the 5D LazyCellImg in a 5D CachedCellImg directly:

LazyCellImg.getCells() gives you a I extends RandomAccessible<C> & IterableInterval<C>>, with C extends Cell<A>.
You can just put these Cells directly into a ReadOnlyCachedCellImg without touching or copying the pixel data at all, using a CacheLoader: https://github.com/imglib/imglib2-cache/blob/5dc5b59ff9862a63796e2e5cc214f614b2bac9c5/src/main/java/net/imglib2/cache/img/ReadOnlyCachedCellImgFactory.java#L98-L102
Then you just go on from that, slicing into it etc.

VolatileViews didn’t fail because of the hyperslicing (it is clever enough to see through that), but because of the LazyCellImg instead of CachedCellImg, as @petebankhead noted.
So it should totally work on hyperslices of the 5D CachedCellImg.

There is a catch, and that is, that the DataAccess behind the pixel Type must be a VolatileAccess. (This is what is checked here: https://github.com/bigdataviewer/bigdataviewer-core/blob/f5226a680e61fcbccc329b1e92bb7d75edb320b7/src/main/java/bdv/util/volatiles/VolatileViews.java#L173-L175). So that must be taken care of when constructing the LazyCellImg.

If you point me to where that happens, I can probably make a suggestion how to make that happen (if it isn’t already the case).

(sorry for formatting, this started as an email...)

[tpietzsch]

I looked into the qupath-imglib2 code a bit.

Looking at the implementation of qupath.ext.imglib2.bufferedimageaccesses.IntRasterAccess etc, they do some work (index math etc) in the int getValue(int index) etc methods. So it might actually make sense to copy into a flat array for the 5D CachedCellImg, to avoid doing that work more than once (for the copying).

In any case, it may be worthwhile to explore the "5D CachedCellImg" option (with copying).

[NicoKiaru]

Thanks for your answer @tpietzsch !

In any case, it may be worthwhile to explore the "5D CachedCellImg" option (with copying).

Sorry, I'm not fully understanding the details, but would this mean that there will be an extra copy of the data - thus using 2 times more RAM ?

[petebankhead]

Thanks @tpietzsch !

I'm curious about the answer to @NicoKiaru's question and I had also run into this

There is a catch, and that is, that the DataAccess behind the pixel Type must be a VolatileAccess. (This is what is checked here: https://github.com/bigdataviewer/bigdataviewer-core/blob/f5226a680e61fcbccc329b1e92bb7d75edb320b7/src/main/java/bdv/util/volatiles/VolatileViews.java#L173-L175). So that must be taken care of when constructing the LazyCellImg.

Is there any reason we shouldn't simply add implements VolatileAccess to our various DataAccess implementations, and return true for isValid()?

[tpietzsch]

@NicoKiaru

In any case, it may be worthwhile to explore the "5D CachedCellImg" option (with copying).

Sorry, I'm not fully understanding the details, but would this mean that there will be an extra copy of the data - thus using 2 times more RAM ?

Not a full extra copy, but, as you said before:

There's an extra cache layer, but it can be bounded so the memory increase can be small.

@petebankhead

Is there any reason we shouldn't simply add implements VolatileAccess to our various DataAccess implementations, and return true for isValid()?

No, not really. I think that should work!

I was just seeing that your DataAccess implementations do a bit of work in the getValue methods (not much, but more than a simple array access.) So I'm wondering what the trade-offs are in terms of copying the data once (a whole cell, so probably copying more data than will be immediately needed) and then having simple array accesses afterwards. Vs always going through your DataAccess implementation, therefore doing a bit more work but only when a pixel is really accessed.

I have no good intuition of what the tradeoff looks like. Probably both should be relatively easy to implement.