Lazy and cached border computation of a Label Image with BigDataViewer

CachedBorders.groovy

/**
 *  GROOVY SCRIPT COMPUTING ON THE FLY EDGES OF A LABEL IMAGE
 *  -----
 * 	 - Computes a sample label image
 *   - Computes lazily with a disk cached cell img an image which finds the border of the labels in 3D
 *   - Displays both images in BigDataViewer
 *
 *  Code to create the label image mainly copied from:
 *  - https://imagej.net/ImgLib2_Examples#Example_8_-_Working_with_sparse_data
 *  - And https://github.com/imglib/imglib2-cache-examples
 *
 *  Improvements
 *  - How to know when the computation is done ? Would be nice to benchmark options
 *
 *  See linked forum thread:
 *   - https://forum.image.sc/t/cached-imglib2-lazy-processing-challenge/25119
 */

import bdv.util.BdvFunctions
import bdv.util.BdvOptions
import bdv.util.BdvSource
import bdv.util.volatiles.SharedQueue
import bdv.util.volatiles.VolatileViews
import net.imglib2.*
import net.imglib2.algorithm.util.Grids
import net.imglib2.cache.img.*
import net.imglib2.img.Img
import net.imglib2.img.array.ArrayImgFactory
import net.imglib2.img.basictypeaccess.AccessFlags
import net.imglib2.img.basictypeaccess.array.ArrayDataAccess
import net.imglib2.img.cell.CellGrid
import net.imglib2.interpolation.neighborsearch.NearestNeighborSearchInterpolatorFactory
import net.imglib2.neighborsearch.NearestNeighborSearch
import net.imglib2.neighborsearch.NearestNeighborSearchOnKDTree
import net.imglib2.type.NativeType
import net.imglib2.type.Type
import net.imglib2.type.numeric.ARGBType
import net.imglib2.type.numeric.integer.ByteType
import net.imglib2.type.numeric.real.FloatType
import net.imglib2.util.Intervals
import net.imglib2.util.Util
import net.imglib2.view.Views

import java.io.IOException
import java.util.Random

import static java.lang.Math.abs
import static net.imglib2.cache.img.DiskCachedCellImgOptions.options


// Get test label image
RandomAccessibleInterval< FloatType > labelImage = getTestLabelImage([ 1600, 1600, 25 ] as long[]);

BdvOptions opt = new BdvOptions();

final BdvSource source = BdvFunctions.show(
        labelImage,
        "Label Image" );

source.setColor(new ARGBType((int)0xFF00FF00));
source.setDisplayRange(0,1);

final SharedQueue queue = new SharedQueue( 7 );
BdvFunctions.show(
        VolatileViews.wrapAsVolatile( get3DBorderLabelImage(labelImage), queue ),
        "Borders", opt.addTo(source));



public static <T extends Type<T> & Comparable<T> >  Img<ByteType>  get3DBorderLabelImage(RandomAccessibleInterval<T> lblImg) {
        // Make edge display on demand
        final int[] cellDimensions = [ 32, 32, 32 ] as int[];

        // Cached Image Factory Options
        final DiskCachedCellImgOptions factoryOptions = options()
                .cellDimensions( cellDimensions )
                .cacheType( DiskCachedCellImgOptions.CacheType.BOUNDED )
                .maxCacheSize( 100 );

        // Expand label image by one pixel to avoid out of bounds exception
        final RandomAccessibleInterval<T> lblImgWithBorder =  Views.expandBorder(lblImg,[1,1,1] as long[]);

        // Creates cached image factory of Type Byte
        final DiskCachedCellImgFactory< ByteType > factory = new DiskCachedCellImgFactory<>( new ByteType(), factoryOptions );

        // Creates shifted views by one pixel in each dimension
        RandomAccessibleInterval<T> lblImgXShift = Views.translate(lblImgWithBorder,[1,0,0] as long[]);
        RandomAccessibleInterval<T> lblImgYShift = Views.translate(lblImgWithBorder,[0,1,0] as long[]);
        RandomAccessibleInterval<T> lblImgZShift = Views.translate(lblImgWithBorder,[0,0,1] as long[]);

        // Creates border image, with cell Consumer method, which creates the image
        final Img<ByteType> borderLabel = factory.create( lblImg, { cell -> 

            // Cursor on the source image
            final Cursor<T> inNS = Views.flatIterable( Views.interval( lblImg, cell ) ).cursor();

            // Cursor on shifted source image
            final Cursor<T> inXS = Views.flatIterable( Views.interval( lblImgXShift, cell ) ).cursor();
            final Cursor<T> inYS = Views.flatIterable( Views.interval( lblImgYShift, cell ) ).cursor();
            final Cursor<T> inZS = Views.flatIterable( Views.interval( lblImgZShift, cell ) ).cursor();

            // Cursor on output image
            final Cursor<ByteType> out = Views.flatIterable( cell ).cursor();

            // Loops through voxels
            while ( out.hasNext() ) {
            	T v = inNS.next();
                if (v.compareTo(inXS.next())!=0) {
                	out.next().set( (byte) 126 ) 
                	inYS.next()
                	inZS.next()
                } else {
                	if (v.compareTo(inYS.next())!=0) {
	                	out.next().set( (byte) 126 ) 
	                	inZS.next()
	                } else {
	                	if (v.compareTo(inZS.next())!=0) {
		                	out.next().set( (byte) 126 ) 
		                } else {
		                	out.next() 
		                }
	                }
                }
            }
        }, 
        options().initializeCellsAsDirty( true ) );

        return borderLabel;
    }

    //------------------------------------- METHODS TO CREATE TEST LABEL IMAGE

    public static RandomAccessibleInterval< FloatType > getTestLabelImage(final long[] imgTestSize) {

        // the interval in which to create random points
        FinalInterval interval = new FinalInterval( imgTestSize );

        // create an IterableRealInterval
        IterableRealInterval< FloatType > realInterval = createRandomPoints( interval, 800 );

        // using nearest neighbor search we will be able to return a value an any position in space
        NearestNeighborSearch< FloatType > search =
                new NearestNeighborSearchOnKDTree<>(
                        new KDTree<>( realInterval ) );

        // make it into RealRandomAccessible using nearest neighbor search
        RealRandomAccessible< FloatType > realRandomAccessible =
                Views.interpolate( search, new NearestNeighborSearchInterpolatorFactory< FloatType >() );

        // convert it into a RandomAccessible which can be displayed
        RandomAccessible< FloatType > randomAccessible = Views.raster( realRandomAccessible );

        // set the initial interval as area to view
        RandomAccessibleInterval< FloatType > labelImage = Views.interval( randomAccessible, interval );

        final RandomAccessibleInterval< FloatType > labelImageCopy = new ArrayImgFactory( Util.getTypeFromInterval( labelImage ) ).create( labelImage );

        // Image copied to avoid computing it on the fly
        // https://github.com/imglib/imglib2-algorithm/blob/47cd6ed5c97cca4b316c92d4d3260086a335544d/src/main/java/net/imglib2/algorithm/util/Grids.java#L221 used for parallel copy
        Grids.collectAllContainedIntervals(imgTestSize, [128,128,32] as int[]).parallelStream().forEach( {blockinterval->
            copy(labelImage, Views.interval(labelImageCopy, blockinterval));
        });

        // Returning it
        return labelImageCopy;
    }

    /**
     * Copy from a source that is just RandomAccessible to an IterableInterval. Latter one defines
     * size and location of the copy operation. It will query the same pixel locations of the
     * IterableInterval in the RandomAccessible. It is up to the developer to ensure that these
     * coordinates match.
     *
     * Note that both, input and output could be Views, Img or anything that implements
     * those interfaces.
     *
     * @param source - a RandomAccess as source that can be infinite
     * @param target - an IterableInterval as target
     */
    public static < T extends Type< T > > void copy( final RandomAccessible< T > source,
                                              final IterableInterval< T > target )
    {
        // create a cursor that automatically localizes itself on every move
        Cursor< T > targetCursor = target.localizingCursor();
        RandomAccess< T > sourceRandomAccess = source.randomAccess();

        // iterate over the input cursor
        while ( targetCursor.hasNext())
        {
            // move input cursor forward
            targetCursor.fwd();

            // set the output cursor to the position of the input cursor
            sourceRandomAccess.setPosition( targetCursor );

            // set the value of this pixel of the output image, every Type supports T.set( T type )
            targetCursor.get().set( sourceRandomAccess.get() );
        }

    }


    /**
     * Create a number of n-dimensional random points in a certain interval
     * having a random intensity 0...1
     *
     * @param interval - the interval in which points are created
     * @param numPoints - the amount of points
     *
     * @return a RealPointSampleList (which is an IterableRealInterval)
     */
    public static RealPointSampleList< FloatType > createRandomPoints(
            RealInterval interval, int numPoints )
    {
        // the number of dimensions
        int numDimensions = interval.numDimensions();

        // a random number generator
        Random rnd = new Random( 2001);//System.currentTimeMillis() );

        // a list of Samples with coordinates
        RealPointSampleList< FloatType > elements =
                new RealPointSampleList<>( numDimensions );

        for ( int i = 0; i < numPoints; ++i )
        {
            RealPoint point = new RealPoint( numDimensions );

            for ( int d = 0; d < numDimensions; ++d )
                point.setPosition( rnd.nextDouble() *
                        ( interval.realMax( d ) - interval.realMin( d ) ) + interval.realMin( d ), d );

            // add a new element with a random intensity in the range 0...1
            elements.add( point, new FloatType( rnd.nextFloat() ) );
        }

        return elements;
    }


    public static final <T extends NativeType<T>, A extends ArrayDataAccess<A>, CA extends ArrayDataAccess<CA>> CachedCellImg<T, CA> wrapAsCachedCellImg(
            final RandomAccessibleInterval<T> source,
            final int[] blockSize) throws IOException {

        final long[] dimensions = Intervals.dimensionsAsLongArray(source);
        final CellGrid grid = new CellGrid(dimensions, blockSize);

        final RandomAccessibleCacheLoader<T, A, CA> loader = RandomAccessibleCacheLoader.get(grid, Views.zeroMin(source), AccessFlags.setOf());
        return new ReadOnlyCachedCellImgFactory().createWithCacheLoader(dimensions, source.randomAccess().get(), loader);
    }