Merge pull request #143 from mastodon-sc/cancel-trackastra-linking

Cancel trackastra linking on user request

Author: stefanhahmann

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

@@ -46,11 +46,15 @@ public class TrackastraLinker< V extends Vertex< E > & HasTimepoint & RealLocali
 	public void mutate1( final ReadOnlyGraph< V, E > graph, final SpatioTemporalIndex< V > index )
 	{
 		long start = System.currentTimeMillis();
+		statusService.clearStatus();
+		statusService.showStatus( "Trackastra linking." );
 		try
 		{
 			List< SingleTimepointRegionProps > regionProps = computeRegionProps( index );
-			runLinkPrediction( index, regionProps );
+			if ( !isCanceled() )
+				runLinkPrediction( index, regionProps );
 			ok = true;
+			statusService.clearStatus();
 		}
 		catch ( TrackastraLinkingException e )
 		{
@@ -86,7 +90,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, this.statusService ))
 		{
 			return computation.computeRegionPropsForSource( source, level, Cast.unchecked( index ), minTimepoint, maxTimepoint );
 		}
@@ -102,7 +106,8 @@ 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,
+								statusService ))
 		{
 			prediction.predictAndCreateLinks();
 		}

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

@@ -33,6 +33,8 @@
 import org.mastodon.mamut.util.ResourceUtils;
 import org.mastodon.spatial.SpatioTemporalIndex;
 import org.mastodon.tracking.linking.EdgeCreator;
+import org.scijava.Cancelable;
+import org.scijava.app.StatusService;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -49,8 +51,13 @@ public class LinkPrediction extends ApposeProcess
 
 	private final RegionProps regionProps;
 
+	private final Cancelable cancelable;
+
+	private final StatusService statusService;
+
 	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, final StatusService statusService
 	) throws IOException
 	{
 		super();
@@ -59,6 +66,8 @@ public LinkPrediction( final Map< String, Object > settings, final SpatioTempora
 		this.edgeCreator = edgeCreator;
 		this.regionProps = regionProps;
 		this.uiLogger = uiLogger;
+		this.cancelable = cancelable;
+		this.statusService = statusService;
 	}
 
 	public void predictAndCreateLinks() throws IOException
@@ -80,7 +89,15 @@ 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;
+			}
+			uiLogger.info( "Link prediction finished. Now writing edges.\n" );
+			statusService.showProgress( 9, 10 ); // 90% after prediction
 			writeEdges( result );
+			statusService.showProgress( 1, 1 ); // 100% after writing edges
 		}
 		finally
 		{
@@ -107,6 +124,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,8 @@
 import org.mastodon.mamut.util.ImgUtils;
 import org.mastodon.mamut.util.ResourceUtils;
 import org.mastodon.spatial.SpatioTemporalIndex;
+import org.scijava.Cancelable;
+import org.scijava.app.StatusService;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -51,11 +52,19 @@ 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;
+
+	private final StatusService statusService;
+
+	public RegionPropsComputation( final org.scijava.log.Logger uiLogger, final String model, final Cancelable cancelable,
+			final StatusService statusService )
+			throws IOException
 	{
 		super();
 		this.uiLogger = uiLogger;
 		this.model = model;
+		this.cancelable = cancelable;
+		this.statusService = statusService;
 	}
 
 	public List< SingleTimepointRegionProps > computeRegionPropsForSource( final Source< ? > source, final int level,
@@ -70,15 +79,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;
@@ -134,6 +151,7 @@ private SingleTimepointRegionProps computeSingleTimepointProps( final Source< ?
 	private void formatProgress( final int done, final int todo )
 	{
 		double progress = ( double ) done / todo;
+		statusService.showProgress( ( int ) ( 0.75d * done ), todo ); // reserve 25% for link prediction
 		NumberFormat percentFormatter = NumberFormat.getPercentInstance();
 		percentFormatter.setMinimumFractionDigits( 0 );
 		percentFormatter.setMaximumFractionDigits( 0 );