Apply_Bigwarp_Xfm_XML.groovy

Apply_Bigwarp_Xfm_XML.groovy

#@ File (label="Landmark file", style="extensions:csv") landmarksPath
#@ File (label="Moving image (BDV/BigStitcher xml)", style="extensions:xml") movingXml
#@ File (label="Target image (BDV/BigStitcher xml)", style="extensions:xml") targetXml
#@ File (label="Output zarr path", style="save") outputPath
#@ String (label="Transform type", choices={"Thin Plate Spline", "Affine", "Similarity", "Rotation", "Translation" }) transformType
#@ String (label="Interpolation", choices={"Linear", "Nearest Neighbor"}) interpType
#@ Integer (label="Number of threads", min=1, max=64, value=4) nThreads
#@ String (label="Compression", choices={"gzip", "raw", "lz4", "xz", "blosc"}) compressionType
#@ Integer (label="Block size override (0 = auto-detect from moving image)", value=0) blockSizeOverride

// Applies a bigwarp landmark transform to a moving BDV/BigStitcher xml and writes the warped
// result to an OME-Zarr container, one chunk at a time, and a block size matched to the
// moving image whenever possible
//
// NB:
// - Block size is read from the moving image's own cell/chunk grid (via imglib2's CachedCellImg)
// - This writes a single full-resolution scale ("s0") per channel, with minimal OME-Zarr
//   (v0.4) multiscale metadata written by hand as plain JSON via N5Writer.setAttribute rather
//   than through n5-universe's OmeNgffMetadata/OmeNgffMetadataParser classes: I am getting a
//   bunch of conflicts between what I have in IntelliJ (scijava pom 45.0.0) and what is available
//   in regular fiji


import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

import bdv.ij.ApplyBigwarpPlugin;
import bdv.img.WarpedSource;
import bdv.viewer.Interpolation;
import bdv.viewer.Source;
import bdv.viewer.SourceAndConverter;
import bigwarp.BigWarp;
import bigwarp.BigWarpInit;
import bigwarp.landmarks.LandmarkTableModel;
import bigwarp.transforms.BigWarpTransform;

import net.imglib2.Interval;
import net.imglib2.RandomAccessibleInterval;
import net.imglib2.RealRandomAccessible;
import net.imglib2.img.array.ArrayImgFactory;
import net.imglib2.img.cell.AbstractCellImg;
import net.imglib2.img.cell.CellGrid;
import net.imglib2.realtransform.AffineGet;
import net.imglib2.realtransform.AffineRandomAccessible;
import net.imglib2.realtransform.AffineTransform3D;
import net.imglib2.realtransform.BoundingBoxEstimation;
import net.imglib2.realtransform.InvertibleRealTransform;
import net.imglib2.realtransform.RealViews;
import net.imglib2.util.Intervals;
import net.imglib2.util.Util;
import net.imglib2.view.Views;

import org.janelia.saalfeldlab.n5.Compression;
import org.janelia.saalfeldlab.n5.DatasetAttributes;
import org.janelia.saalfeldlab.n5.N5Writer;
import org.janelia.saalfeldlab.n5.imglib2.N5Utils;
import org.janelia.saalfeldlab.n5.universe.N5Factory;

import sc.fiji.snt.util.ImgUtils;

println( "[1/6] Loading moving/target BDV-XML..." );

bigwarpdata = BigWarpInit.initData()

// the dataset argument is ignored whenever the root path ends in "xml": only used for n5/zarr/h5 roots
BigWarpInit.addToData( bigwarpdata, true, 0, movingXml.getAbsolutePath(), "" )

// the moving xml may have contributed more than one source (one per channel).
// We'll read the count back instead of assuming +1, otherwise a multichannel
// moving xml would collide setup ids with the target
int nextId = bigwarpdata.sources.size();
BigWarpInit.addToData( bigwarpdata, false, nextId, targetXml.getAbsolutePath(), "" );

println( "  " + bigwarpdata.numMovingSources() + " moving channel(s), " + bigwarpdata.numTargetSources() + " target channel(s)" );

// grab a handle to channel 0's raw (pre-warp) moving source now, so we can inspect
// its native chunking before wrapMovingSources() wraps it in a WarpedSource
movingRawSource = bigwarpdata.sources.get( 0 ).getSpimSource()

int nd = BigWarp.detectNumDims( bigwarpdata.sources );
BigWarp.wrapMovingSources( nd, bigwarpdata );

println( "[2/6] Detecting chunk/block size..." )
int[] blockSize;
if ( blockSizeOverride > 0 ) {
	blockSize = new int[ nd ];
	Arrays.fill( blockSize, blockSizeOverride );
	println( "  Using block size override: " + Arrays.toString( blockSize ) );

} else {
	movingRai0 = movingRawSource.getSource( 0, 0 );
	// AbstractCellImg covers CachedCellImg, VolatileCachedCellImg, and other cell-backed imgs
	if ( movingRai0 instanceof AbstractCellImg ) {
		CellGrid grid = ( (AbstractCellImg) movingRai0 ).getCellGrid();
		blockSize = new int[ nd ];
		grid.cellDimensions( blockSize );
		println( "  Detected chunk size from moving image (channel 0): " + Arrays.toString( blockSize ) );
	
	} else {
		blockSize = new int[ nd ];
		Arrays.fill( blockSize, 64 );
		println( "  Could not read a chunk grid from the moving source (" + movingRai0.getClass().getSimpleName() +
				"); falling back to default block size: " + Arrays.toString( blockSize ) );
	}
}

println( "[3/6] Loading landmarks and building transform..." );
ltm = new LandmarkTableModel( nd );
ltm.load( landmarksPath );

bwTransform = new BigWarpTransform( ltm, transformType );
InvertibleRealTransform invXfm = bwTransform.getTransformation();

int numMovingChannels = bigwarpdata.numMovingSources();
for (int i = 0; i < numMovingChannels; i++) {
	SourceAndConverter movingSourceSAC = bigwarpdata.getMovingSource( i );
	WarpedSource ws = (WarpedSource) movingSourceSAC.getSpimSource();
	if ( ws.getTransform() == null && invXfm != null ) {
		ws.updateTransform( invXfm );
		ws.setIsTransformed( true );
	}
}

Interpolation interp = Interpolation.NLINEAR;
if ( interpType.equals( "Nearest Neighbor" ) )
	interp = Interpolation.NEARESTNEIGHBOR;

bboxEst = new BoundingBoxEstimation( BoundingBoxEstimation.Method.CORNERS );

// field of view / resolution taken from the target image
resolutionSpec = [ 1.0, 1.0, 1.0 ] as double[];
fovSpec        = [ -1.0, -1.0, -1.0 ] as double[];
offsetSpec     = [ 0.0, 0.0, 0.0 ] as double[];

println( "[4/6] Computing output field of view..." );
Source tgtSrc = bigwarpdata.getTargetSource( 0 ).getSpimSource();

// computed once and shared across all moving channels
double[] res = ApplyBigwarpPlugin.getResolution( bigwarpdata, "Target", resolutionSpec );
List<Interval> outputIntervalList = ApplyBigwarpPlugin.getPixelInterval( bigwarpdata, ltm, invXfm, "Target",
		"", bboxEst, fovSpec, offsetSpec, res );
double[] offset = ApplyBigwarpPlugin.getPhysicalOffset( "Target", offsetSpec, ltm, invXfm,
		"", bboxEst, res, bigwarpdata.getMovingSource( 0 ).getSpimSource(), tgtSrc );
String unit = ApplyBigwarpPlugin.getUnit( bigwarpdata, "Target" );

Interval outputInterval = outputIntervalList.get( 0 );

// same rendering transform construction as ApplyBigwarpPlugin.runN5Export()
AffineTransform3D resolutionTransform = new AffineTransform3D();
resolutionTransform.set( res[ 0 ], 0, 0 );
resolutionTransform.set( res[ 1 ], 1, 1 );
if ( res.length > 2 )
	resolutionTransform.set( res[ 2 ], 2, 2 );

AffineTransform3D offsetTransform = new AffineTransform3D();
offsetTransform.set( offset[ 0 ], 0, 3 );
offsetTransform.set( offset[ 1 ], 1, 3 );
if ( offset.length > 2 )
	offsetTransform.set( offset[ 2 ], 2, 3 );

AffineTransform3D pixelRenderToPhysical = new AffineTransform3D();
pixelRenderToPhysical.concatenate( resolutionTransform );
pixelRenderToPhysical.concatenate( offsetTransform );

println( "[5/6] Opening zarr output..." );
outStr = outputPath.getAbsolutePath();
if ( !outStr.toLowerCase().endsWith( ".zarr" ) ) {
	outStr = outStr + ".zarr";
	println( "  Output path did not end in .zarr; using: " + outStr );
}

N5Writer n5 = new N5Factory().openWriter( outStr );
Compression compression = ApplyBigwarpPlugin.getCompression( compressionType );

long[] blockDimsLong = new long[ blockSize.length ];
for ( int d = 0; d < blockSize.length; d++ )
	blockDimsLong[ d ] = (long) blockSize[ d ];

println( "[6/6] Writing " + numMovingChannels + " channel(s) to zarr..." );
for ( int ch = 0; ch < numMovingChannels; ch++ ) {

	SourceAndConverter movingSourceForExport = bigwarpdata.getMovingSource( ch );
	String chName = movingSourceForExport.getSpimSource().getName();
	String safeChName = chName.replaceAll( "[^a-zA-Z0-9_.-]", "_" );
	String baseDataset = numMovingChannels > 1 ? ( "warped/" + safeChName ) : "warped";
	String dataPath = baseDataset + "/s0";

	println( "  Channel " + ( ch + 1 ) + "/" + numMovingChannels + ": \"" + chName + "\" -> " + dataPath );

	RealRandomAccessible raiRaw = movingSourceForExport.getSpimSource().getInterpolatedSource( 0, 0, interp );
	AffineRandomAccessible< ?, AffineGet > rai = RealViews.affine( raiRaw, pixelRenderToPhysical.inverse() );
	RandomAccessibleInterval img = Views.interval( Views.raster( rai ), Intervals.zeroMin( outputInterval ) );
	long[] dims = img.dimensionsAsLongArray();

	if ( ch == 0 )
		println( "  Output volume size (per channel): " + Arrays.toString( dims ) + " " + unit );

	DatasetAttributes requestedAttributes = new DatasetAttributes(
			dims, blockSize, N5Utils.dataType( Util.getTypeFromInterval( img ) ), compression );
	// don't rely on createDataset()'s return value: some n5 core versions return the written
	// DatasetAttributes, others return void!? Read the attributes back with getDatasetAttributes()
	// instead, which has been stable across versions.
	n5.createDataset( dataPath, requestedAttributes );
	DatasetAttributes attributes = n5.getDatasetAttributes( dataPath );
	if ( attributes == null )
		throw new RuntimeException( "Failed to create or read back dataset attributes at " + dataPath );

	List<Interval> blocks = ImgUtils.createIntervals( dims, blockDimsLong );
	int totalBlocks = blocks.size();
	println( "    " + totalBlocks + " blocks (block size " + Arrays.toString( blockSize ) + "), " + nThreads + " threads" );

	AtomicInteger done = new AtomicInteger( 0 );
	AtomicInteger failed = new AtomicInteger( 0 );
	int progressEvery = Math.max( 1, (int) ( totalBlocks / 20 ) );
	long startTime = System.currentTimeMillis();

	// The moving image is backed by a VolatileCachedCellImg. When using nThreads
	// at once, each sampling random distant regions, and it can hand back a
	// stale/empty cell instead of blocking for the real data!?. How to avoid that?
	// For now read each block interpolated + copied into a plain in-memory array)
	// synchronously under a lock
	final Object readLock = new Object();

	ExecutorService exec = Executors.newFixedThreadPool( nThreads );
	List<Future> futures = new ArrayList<>();
	for ( Interval blk : blocks )
	{
		final Interval blkFinal = blk;
		final RandomAccessibleInterval imgFinal = img;
		final DatasetAttributes attributesFinal = attributes;
		final String dataPathFinal = dataPath;
		futures.add( exec.submit( new Runnable() {
			void run() {
				try
				{
					long[] gridOffset = new long[ blkFinal.numDimensions() ];
					for ( int d = 0; d < gridOffset.length; d++ )
						gridOffset[ d ] = blkFinal.min( d ) / blockDimsLong[ d ];

					RandomAccessibleInterval materialized;
					synchronized ( readLock )
					{
						def factory = new ArrayImgFactory( Util.getTypeFromInterval( imgFinal ) );
						def buffer = factory.create( blkFinal );
						def srcView = Views.zeroMin( Views.interval( imgFinal, blkFinal ) );
						def cursor = buffer.cursor();
						def ra = srcView.randomAccess();
						while ( cursor.hasNext() )
						{
							cursor.fwd();
							ra.setPosition( cursor );
							cursor.get().set( ra.get() );
						}
						materialized = buffer;
					}

					// N5Utils.saveBlock() has several overloads sharing the same arg count
					// with different argument types (DatasetAttributes vs long[]); Groovy
					// gets confused about which one to use!?. Passing the grid offset
					// explicitly picks the 5-arg (source, n5, dataset, attributes, gridOffset)
					// form, a signature no other overload shares
					N5Utils.saveBlock( (RandomAccessibleInterval) materialized, n5, dataPathFinal, attributesFinal, gridOffset );
				}
				catch ( Exception e ) {
					int f = failed.incrementAndGet();
					// cap full stack traces to avoid flooding/freezing the console over
					// potentially thousands of blocks if something is systematically wrong
					if ( f <= 3 )
					{
						System.err.println( "Failed to write block " + blkFinal + " of " + dataPathFinal + ": " + e );
						e.printStackTrace();
					}
					else if ( f == 4 ) {
						System.err.println( "  (suppressing further per-block error details after 3 failures)" );
					}
				}
				int d = done.incrementAndGet();
				if ( d == totalBlocks || d % progressEvery == 0 ) {
					double pct = 100.0 * d / totalBlocks;
					double elapsedSec = ( System.currentTimeMillis() - startTime ) / 1000.0;
					double etaSec = ( elapsedSec / d ) * ( totalBlocks - d );
					println( String.format( "    %d/%d blocks (%.1f%%) - elapsed %.0fs, ETA %.0fs",
							d, totalBlocks, pct, elapsedSec, etaSec ) );
				}
			}
		} ) );
	}
	futures.each { it.get() };
	exec.shutdown();

	if ( failed.get() > 0 )
		println( "  WARNING: " + failed.get() + "/" + totalBlocks + " blocks failed to write for this channel - output is incomplete, see errors above" );

	// minimal OME-Zarr (v0.4) multiscale metadata, written manully N5Writer.setAttribute
	// Using n5-universe's OmeNgffMetadata/OmeNgffMetadataParser works on the IDE but not
	// on vanilla Fiji-latest!? //TODO: FIGURE OUT WHY
	//
	// From the web/claude: zarr stores array shape in the opposite dimension order from N5's
	// native (x-fastest) convention, so axes/scale/translation are written here as (z, y, x)
	// rather than (x, y, z) to match how N5Utils actually wrote this array's shape into the
	// zarr container. This mirrors what the "reverse for zarr" handling n5-universe's does!?
	List<String> axesLabelsRev = nd == 2 ? [ "y", "x" ] : [ "z", "y", "x" ];
	List<Double> resRev = new ArrayList<>();
	List<Double> offsetRev = new ArrayList<>();
	for ( int i = nd - 1; i >= 0; i-- ) {
		resRev.add( res[ i ] );
		offsetRev.add( offset[ i ] );
	}

	List<Map> axesList = new ArrayList<>();
	for ( String label : axesLabelsRev )
		axesList.add( [ "name": label, "type": "space", "unit": unit ] );

	Map dataset0 = [
		"path": "s0",
		"coordinateTransformations": [
			[ "type": "scale", "scale": resRev ],
			[ "type": "translation", "translation": offsetRev ]
		]
	];

	Map multiscale = [
		"version": "0.4",
		"name": chName,
		"axes": axesList,
		"datasets": [ dataset0 ]
	];

	n5.setAttribute( baseDataset, "multiscales", [ multiscale ] );
}

n5.close();

println( "Done. Wrote " + numMovingChannels + " channel(s) to " + outStr );