Big Volume Viewer for large multi-scale N5 dataset

Example8.java

package bvv.examples;

import static net.imglib2.img.basictypeaccess.AccessFlags.VOLATILE;
import static net.imglib2.type.PrimitiveType.BYTE;
import static net.imglib2.type.PrimitiveType.DOUBLE;
import static net.imglib2.type.PrimitiveType.FLOAT;
import static net.imglib2.type.PrimitiveType.INT;
import static net.imglib2.type.PrimitiveType.LONG;
import static net.imglib2.type.PrimitiveType.SHORT;

import java.io.IOException;
import java.util.Arrays;

import org.janelia.saalfeldlab.n5.N5FSReader;
import org.janelia.saalfeldlab.n5.N5Reader;
import org.janelia.saalfeldlab.n5.imglib2.N5Utils;
import org.janelia.saalfeldlab.n5.imglib2.RandomAccessibleLoader;

import bdv.util.RandomAccessibleIntervalMipmapSource;
import bdv.util.volatiles.SharedQueue;
import bdv.util.volatiles.VolatileViews;
import bvv.util.Bvv;
import bvv.util.BvvFunctions;
import bvv.util.BvvOptions;
import bvv.util.BvvStackSource;
import mpicbg.spim.data.sequence.FinalVoxelDimensions;
import net.imglib2.RandomAccessibleInterval;
import net.imglib2.Volatile;
import net.imglib2.cache.Cache;
import net.imglib2.cache.img.CachedCellImg;
import net.imglib2.cache.img.LoadedCellCacheLoader;
import net.imglib2.cache.ref.SoftRefLoaderCache;
import net.imglib2.cache.volatiles.CacheHints;
import net.imglib2.cache.volatiles.LoadingStrategy;
import net.imglib2.converter.Converter;
import net.imglib2.converter.Converters;
import net.imglib2.img.basictypeaccess.AccessFlags;
import net.imglib2.img.basictypeaccess.ArrayDataAccessFactory;
import net.imglib2.img.cell.Cell;
import net.imglib2.img.cell.CellGrid;
import net.imglib2.realtransform.AffineTransform3D;
import net.imglib2.type.NativeType;
import net.imglib2.type.numeric.ARGBType;
import net.imglib2.type.numeric.IntegerType;
import net.imglib2.type.numeric.RealType;
import net.imglib2.type.numeric.integer.GenericByteType;
import net.imglib2.type.numeric.integer.GenericIntType;
import net.imglib2.type.numeric.integer.GenericLongType;
import net.imglib2.type.numeric.integer.GenericShortType;
import net.imglib2.type.numeric.integer.UnsignedShortType;
import net.imglib2.type.numeric.real.DoubleType;
import net.imglib2.type.numeric.real.FloatType;
import net.imglib2.type.volatiles.AbstractVolatileNativeRealType;
import net.imglib2.type.volatiles.VolatileUnsignedShortType;
import net.imglib2.util.Intervals;
import net.imglib2.util.Pair;
import net.imglib2.util.Util;
import net.imglib2.util.ValuePair;
import net.imglib2.view.Views;

public class Example8 {

	private static <T extends NativeType<T>> Pair<RandomAccessibleInterval<T>[], double[][]> open(final N5Reader n5, final String groupName, final double[] res) throws IOException {

		Pair<RandomAccessibleInterval<T>[], double[][]> n5Sources;
		int n;
		if (n5.datasetExists(groupName)) {
			final RandomAccessibleInterval<T> source = N5Utils.openVolatile(n5, groupName);
			n = source.numDimensions();
			final double[] scale = new double[n];
			Arrays.fill(scale, 1.0);
			n5Sources = new ValuePair<>(new RandomAccessibleInterval[]{source}, new double[][]{scale});
		} else {
			n5Sources = N5Utils.openMipmaps(n5, groupName, true);
			n = n5Sources.getA()[0].numDimensions();
		}
		for (int k = 0; k < n5Sources.getA().length; ++k) {
			final double[] scale = n5Sources.getB()[k];
			scale[0] = scale[0] * res[0];
			scale[1] = scale[1] * res[1];
			scale[2] = scale[2] * res[2];
		}
		return n5Sources;
	}

	private static <T> RandomAccessibleInterval<UnsignedShortType>[] convertToCachedUnsignedShort(final RandomAccessibleInterval<T>[] sources, final Converter<T, UnsignedShortType> converter, final int[] blockSize) {

		final RandomAccessibleInterval<UnsignedShortType>[] uShortSources = (RandomAccessibleInterval<UnsignedShortType>[])new RandomAccessibleInterval[sources.length];

		for (int k = 0; k < sources.length; ++k) {
			uShortSources[k] = cache(
					Converters.convert(sources[k], converter, new UnsignedShortType()),
					blockSize);
		}
		return uShortSources;
	}

	private static <T extends NativeType<T>, V extends Volatile<T>> RandomAccessibleInterval<V>[] wrapAsVolatile(final RandomAccessibleInterval<T>[] sources, final SharedQueue queue) {

		final RandomAccessibleInterval<V>[] volatileSources = (RandomAccessibleInterval<V>[])new RandomAccessibleInterval[sources.length];

		for (int k = 0; k < sources.length; ++k) {
			volatileSources[k] = VolatileViews.wrapAsVolatile(
				sources[k],
				queue,
				new CacheHints(LoadingStrategy.VOLATILE, 0, true));
		}
		return volatileSources;
	}

	private static <T extends RealType<T> & NativeType<T>, V extends AbstractVolatileNativeRealType<T, V>> RandomAccessibleIntervalMipmapSource<V> openSource(final N5Reader n5, final String groupName, final double[] res, final double[] offset, final SharedQueue queue, final V type) throws IOException {

		final Pair<RandomAccessibleInterval<T>[], double[][]> n5Sources = open(n5, groupName, res);
		final RandomAccessibleInterval<V>[] volatileSources = wrapAsVolatile(n5Sources.getA(), queue);

		final AffineTransform3D sourceTransform = new AffineTransform3D();
		sourceTransform.setTranslation(offset);

		final RandomAccessibleIntervalMipmapSource<V> mipmapSource =
			new RandomAccessibleIntervalMipmapSource<>(
				volatileSources,
				type,
				n5Sources.getB(),
				new FinalVoxelDimensions("px", res),
				sourceTransform,
				groupName);

		return mipmapSource;
	}

	private static <T extends IntegerType<T> & NativeType<T>> RandomAccessibleIntervalMipmapSource<VolatileUnsignedShortType> openIntegerSourceAsUnsignedShortSource(
			final N5Reader n5,
			final String groupName,
			final double[] res,
			final double[] offset,
			final SharedQueue queue,
			final Converter<T, UnsignedShortType> converter,
			final int[] blockSize) throws IOException {

		final Pair<RandomAccessibleInterval<T>[], double[][]> n5Sources = open(n5, groupName, res);
		final RandomAccessibleInterval<UnsignedShortType>[] sources = convertToCachedUnsignedShort(
				n5Sources.getA(),
				converter,
				blockSize);
		final RandomAccessibleInterval<VolatileUnsignedShortType>[] volatileSources = wrapAsVolatile(
				sources,
				queue);

		final AffineTransform3D sourceTransform = new AffineTransform3D();
		sourceTransform.setTranslation(offset);

		final RandomAccessibleIntervalMipmapSource<VolatileUnsignedShortType> mipmapSource =
			new RandomAccessibleIntervalMipmapSource<>(
				volatileSources,
				new VolatileUnsignedShortType(),
				n5Sources.getB(),
				new FinalVoxelDimensions("px", res),
				sourceTransform,
				groupName);

		return mipmapSource;
	}

	private static <T extends RealType<T> & NativeType<T>> RandomAccessibleIntervalMipmapSource<VolatileUnsignedShortType> openRealSourceAsUnsignedShortSource(
			final N5Reader n5,
			final String groupName,
			final double[] res,
			final double[] offset,
			final SharedQueue queue,
			final Converter<T, UnsignedShortType> converter,
			final int[] blockSize) throws IOException {

		final Pair<RandomAccessibleInterval<T>[], double[][]> n5Sources = open(n5, groupName, res);
		final RandomAccessibleInterval<VolatileUnsignedShortType>[] volatileSources = wrapAsVolatile(
				convertToCachedUnsignedShort(
						n5Sources.getA(),
						converter,
						blockSize),
				queue);

		final AffineTransform3D sourceTransform = new AffineTransform3D();
		sourceTransform.setTranslation(offset);

		final RandomAccessibleIntervalMipmapSource<VolatileUnsignedShortType> mipmapSource =
			new RandomAccessibleIntervalMipmapSource<>(
				volatileSources,
				new VolatileUnsignedShortType(),
				n5Sources.getB(),
				new FinalVoxelDimensions("px", res),
				sourceTransform,
				groupName);

		return mipmapSource;
	}

	@SuppressWarnings( { "unchecked", "rawtypes" } )
	public static final <T extends NativeType<T>> RandomAccessibleInterval<T> cache(
			final RandomAccessibleInterval<T> source,
			final int[] blockSize) {

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

		final RandomAccessibleLoader<T> loader = new RandomAccessibleLoader<T>(Views.zeroMin(source));

		final T type = Util.getTypeFromInterval(source);

		final CachedCellImg<T, ?> img;
		final Cache<Long, Cell<?>> cache =
				new SoftRefLoaderCache().withLoader(LoadedCellCacheLoader.get(grid, loader, type, AccessFlags.setOf(VOLATILE)));

		if (GenericByteType.class.isInstance(type)) {
			img = new CachedCellImg(grid, type, cache, ArrayDataAccessFactory.get(BYTE, AccessFlags.setOf(VOLATILE)));
		} else if (GenericShortType.class.isInstance(type)) {
			img = new CachedCellImg(grid, type, cache, ArrayDataAccessFactory.get(SHORT, AccessFlags.setOf(VOLATILE)));
		} else if (GenericIntType.class.isInstance(type)) {
			img = new CachedCellImg(grid, type, cache, ArrayDataAccessFactory.get(INT, AccessFlags.setOf(VOLATILE)));
		} else if (GenericLongType.class.isInstance(type)) {
			img = new CachedCellImg(grid, type, cache, ArrayDataAccessFactory.get(LONG, AccessFlags.setOf(VOLATILE)));
		} else if (FloatType.class.isInstance(type)) {
			img = new CachedCellImg(grid, type, cache, ArrayDataAccessFactory.get(FLOAT, AccessFlags.setOf(VOLATILE)));
		} else if (DoubleType.class.isInstance(type)) {
			img = new CachedCellImg(grid, type, cache, ArrayDataAccessFactory.get(DOUBLE, AccessFlags.setOf(VOLATILE)));
		} else {
			img = null;
		}

		return img;
	}


	/**
	 * ImgLib2 :-)
	 * @throws IOException
	 */
	public static void main( final String[] args ) throws IOException
	{
		final BvvOptions bvvOptions = Bvv.options()
				.maxAllowedStepInVoxels( 4 )
				.renderWidth( 512 )
				.renderHeight( 512 )
				.preferredSize( 1024, 1024 )
				.numDitherSamples(20)
				.cacheBlockSize(new int[]{64, 64, 64})
				.maxCacheSizeInMB(5000);

		final int numProc = Runtime.getRuntime().availableProcessors();
		final SharedQueue queue = new SharedQueue(Math.min(8, Math.max(1, numProc / 2)));

		final N5Reader n5 = new N5FSReader("/nrs/saalfeld/FAFB00/v14_align_tps_20170818_dmg.n5");

		final double[] resolution = new double[]{4, 4, 40};

//		 BdvFunctions.show(
//				openIntegerSourceAsUnsignedShortSource(
//						n5,
//						"/volumes/predictions/synapses_dt_reblocked",
//						resolution,
//						new double[]{0, 0, 0},
//						queue,
//						(a, b) -> b.set(a.getInteger()),
//						new int[] {64, 64, 8}));

		final BvvStackSource<VolatileUnsignedShortType> source = BvvFunctions.show(
				openIntegerSourceAsUnsignedShortSource(
						n5,
						"/volumes/predictions/synapses_dt_reblocked",
						resolution,
						new double[]{0, 0, 0},
						queue,
						(a, b) -> b.set(a.getInteger()),
						new int[] {128, 128, 128}),
				bvvOptions);
		source.setDisplayRange( 0, 10 );
		source.setColor( new ARGBType( 0xffff0088 ) );
	}
}