Add cancellation support to LinkPrediction, RegionPropsComputation, and TrackastraLinker.

Author: Stefan Hahmann

src/main/java/org/mastodon/mamut/linking/trackastra/TrackastraLinker.java

@@ -49,7 +49,8 @@ public void mutate1( final ReadOnlyGraph< V, E > graph, final SpatioTemporalInde
 		try
 		{
 			List< SingleTimepointRegionProps > regionProps = computeRegionProps( index );
-			runLinkPrediction( index, regionProps );
+			if ( !isCanceled() )
+				runLinkPrediction( index, regionProps );
 			ok = true;
 		}
 		catch ( TrackastraLinkingException e )
@@ -86,7 +87,7 @@ private List< SingleTimepointRegionProps > computeRegionProps( final SpatioTempo
 			throw new IllegalArgumentException(
 					String.format( "Window size (%d) exceeds time range (%d). Adjust window size or time range.", windowSize, timeRange ) );
 
-		try (RegionPropsComputation computation = new RegionPropsComputation( logger, model ))
+		try (RegionPropsComputation computation = new RegionPropsComputation( logger, model, this ))
 		{
 			return computation.computeRegionPropsForSource( source, level, Cast.unchecked( index ), minTimepoint, maxTimepoint );
 		}
@@ -102,7 +103,7 @@ private void runLinkPrediction( final SpatioTemporalIndex< V > index, final List
 		log.info( "Performing Trackastra link prediction" );
 		try (RegionProps props = new RegionProps( regionProps );
 				LinkPrediction prediction =
-						new LinkPrediction( settings, Cast.unchecked( index ), Cast.unchecked( edgeCreator ), props, logger ))
+						new LinkPrediction( settings, Cast.unchecked( index ), Cast.unchecked( edgeCreator ), props, logger, this ))
 		{
 			prediction.predictAndCreateLinks();
 		}

src/main/java/org/mastodon/mamut/linking/trackastra/appose/computation/LinkPrediction.java

@@ -33,6 +33,7 @@
 import org.mastodon.mamut.util.ResourceUtils;
 import org.mastodon.spatial.SpatioTemporalIndex;
 import org.mastodon.tracking.linking.EdgeCreator;
+import org.scijava.Cancelable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -49,8 +50,11 @@ public class LinkPrediction extends ApposeProcess
 
 	private final RegionProps regionProps;
 
+	private final Cancelable cancelable;
+
 	public LinkPrediction( final Map< String, Object > settings, final SpatioTemporalIndex< Spot > index,
-			final EdgeCreator< Spot > edgeCreator, final RegionProps regionProps, final org.scijava.log.Logger uiLogger
+			final EdgeCreator< Spot > edgeCreator, final RegionProps regionProps, final org.scijava.log.Logger uiLogger,
+			final Cancelable cancelable
 	) throws IOException
 	{
 		super();
@@ -59,6 +63,7 @@ public LinkPrediction( final Map< String, Object > settings, final SpatioTempora
 		this.edgeCreator = edgeCreator;
 		this.regionProps = regionProps;
 		this.uiLogger = uiLogger;
+		this.cancelable = cancelable;
 	}
 
 	public void predictAndCreateLinks() throws IOException
@@ -80,6 +85,11 @@ public void predictAndCreateLinks() throws IOException
 			inputs.put( BORDER_DIST, NDArrays.asNDArray( regionProps.borderDists ) );
 
 			Service.Task result = runScript();
+			if ( cancelable.isCanceled() )
+			{
+				uiLogger.info( "Link prediction canceled by user.\n" );
+				return;
+			}
 			writeEdges( result );
 		}
 		finally
@@ -107,6 +117,11 @@ private void writeEdges( final Service.Task task )
 
 		for ( int row = 0; row < rows; row++ )
 		{
+			if ( cancelable.isCanceled() )
+			{
+				uiLogger.info( "Link creation canceled by user.\n" );
+				return;
+			}
 			int startFrame = getInt( randomAccess, 0, row );
 			int startId = getInt( randomAccess, 1, row ) - 1; // Trackastra uses 1-based labels
 			int endFrame = getInt( randomAccess, 2, row );

src/main/java/org/mastodon/mamut/linking/trackastra/appose/computation/RegionPropsComputation.java

@@ -16,7 +16,6 @@
 import java.text.NumberFormat;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.Consumer;
 
 import net.imglib2.RandomAccessibleInterval;
 import net.imglib2.appose.NDArrays;
@@ -38,6 +37,7 @@
 import org.mastodon.mamut.util.ImgUtils;
 import org.mastodon.mamut.util.ResourceUtils;
 import org.mastodon.spatial.SpatioTemporalIndex;
+import org.scijava.Cancelable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -51,11 +51,15 @@ public class RegionPropsComputation extends ApposeProcess
 
 	private final String model;
 
-	public RegionPropsComputation( final org.scijava.log.Logger uiLogger, final String model ) throws IOException
+	private final Cancelable cancelable;
+
+	public RegionPropsComputation( final org.scijava.log.Logger uiLogger, final String model, final Cancelable cancelable )
+			throws IOException
 	{
 		super();
 		this.uiLogger = uiLogger;
 		this.model = model;
+		this.cancelable = cancelable;
 	}
 
 	public List< SingleTimepointRegionProps > computeRegionPropsForSource( final Source< ? > source, final int level,
@@ -70,15 +74,23 @@ public List< SingleTimepointRegionProps > computeRegionPropsForSource( final Sou
 		List< SingleTimepointRegionProps > singleTimepointRegionProps = new ArrayList<>();
 		for ( int timepoint = minTimepoint; timepoint <= maxTimepoint; timepoint++, done++ )
 		{
+			if ( cancelable.isCanceled() )
+			{
+				log.info( "Region props computation canceled at timepoint: {}", timepoint );
+				uiLogger.info( "Region props computation canceled at timepoint: " + timepoint + "\n" );
+				break;
+			}
 			if ( spatioTemporalIndex.getSpatialIndex( timepoint ).isEmpty() )
 			{
 				log.info( "No spots. Adding empty region props for timepoint: {}", timepoint );
 				uiLogger.info( "No spots. Adding empty region props for timepoint: " + timepoint + "\n" );
 				singleTimepointRegionProps.add( null );
-				continue;
 			}
-			SingleTimepointRegionProps regionProps = computeSingleTimepointProps( source, timepoint, level, spatioTemporalIndex );
-			singleTimepointRegionProps.add( regionProps );
+			else
+			{
+				SingleTimepointRegionProps regionProps = computeSingleTimepointProps( source, timepoint, level, spatioTemporalIndex );
+				singleTimepointRegionProps.add( regionProps );
+			}
 			formatProgress( done + 1, todo );
 		}
 		return singleTimepointRegionProps;