Merge pull request #144 from mastodon-sc/cancel-deep-learning-detectors

Add Cancelation support for deep learning detectors

Author: stefanhahmann

src/main/java/org/mastodon/mamut/detection/DeepLearningDetector.java

@@ -41,11 +41,14 @@
 import net.imglib2.view.IntervalView;
 import net.imglib2.view.Views;
 
+import org.mastodon.mamut.util.ApposeProcess;
 import org.mastodon.mamut.util.ByteFormatter;
 import org.mastodon.mamut.detection.util.SpimImageProperties;
 import org.mastodon.mamut.model.ModelGraph;
 import org.mastodon.mamut.util.ImgUtils;
 import org.mastodon.mamut.util.LabelImageUtils;
+import org.mastodon.tracking.detection.AbstractDetectorOp;
+import org.mastodon.tracking.detection.DetectionCreatorFactory;
 import org.mastodon.tracking.detection.DetectionUtil;
 import org.mastodon.tracking.detection.DetectorKeys;
 import org.mastodon.tracking.mamut.detection.AbstractSpotDetectorOp;
@@ -80,6 +83,8 @@ public abstract class DeepLearningDetector extends AbstractSpotDetectorOp
 	 */
 	private static final int MAX_SIZE_IN_BYTES = Integer.MAX_VALUE - 1;
 
+	protected ApposeProcess apposeProcess;
+
 	@Override
 	public void compute( final List< SourceAndConverter< ? > > sources, final ModelGraph graph )
 	{
@@ -103,6 +108,9 @@ public void compute( final List< SourceAndConverter< ? > > sources, final ModelG
 		statusService.showStatus( "Detecting spots..." );
 		try
 		{
+			log.info( "Initialize python environment." );
+			log.info( "On first time use this requires internet connection and may take a couple of minutes." );
+			log.info( "Progress can be observed using FIJI console: FIJI > Window > Console." );
 			for ( int timepoint = minTimepoint; timepoint <= maxTimepoint; timepoint++ )
 			{
 				// We use the `statusService to show progress.
@@ -203,7 +211,6 @@ private void detectAndAddSpots( final List< SourceAndConverter< ? > > sources, f
 			image = Views.interval( image, roi );
 		}
 
-		System.err.println(); // show the FIJI console
 		final Img< ? > segmentation =
 				performSegmentation( Views.dropSingletonDimensions( image ), source.getVoxelDimensions().dimensionsAsDoubleArray() );
 
@@ -299,4 +306,31 @@ public Map< String, Object > getDefaultSettings()
 	protected abstract void addSpecificDefaultSettings( final Map< String, Object > defaultSettings );
 
 	protected abstract String getDetectorName();
+
+	/** Cancels the command execution, with the given reason for doing so. */
+	@Override
+	public void cancel( final String reason )
+	{
+		// this is a workaround to avoid a null pointer exception during the cancel operation
+		detector = new DummyDetectorOp();
+		try
+		{
+			if ( apposeProcess != null )
+				apposeProcess.cancel();
+		}
+		catch ( InterruptedException e )
+		{
+			logger.info( "Interrupted while cancelling detector: {}", e.getMessage() );
+		}
+		super.cancel( reason );
+	}
+
+	private class DummyDetectorOp extends AbstractDetectorOp
+	{
+		@Override
+		public void mutate1( final DetectionCreatorFactory arg, final List< SourceAndConverter< ? > > in )
+		{
+			// Do nothing
+		}
+	}
 }

src/main/java/org/mastodon/mamut/detection/cellpose/Cellpose3Detector.java

@@ -103,6 +103,7 @@ && checkParameter( settings, KEY_GPU_ID, Integer.class, errorHolder )
 	{
 		try (Cellpose3 cellpose = new Cellpose3( ( Cellpose3.ModelType ) settings.get( KEY_MODEL_TYPE ) ))
 		{
+			this.apposeProcess = cellpose;
 			boolean is3D = is3D( image );
 			cellpose.set3D( is3D );
 			cellpose.setCellProbThreshold( ( double ) settings.get( KEY_CELL_PROBABILITY_THRESHOLD ) );

src/main/java/org/mastodon/mamut/detection/cellpose/Cellpose4Detector.java

@@ -94,6 +94,7 @@ && checkParameter( settings, KEY_GPU_ID, Integer.class, errorHolder )
 	{
 		try (Cellpose4 cellpose = new Cellpose4())
 		{
+			this.apposeProcess = cellpose;
 			cellpose.set3D( is3D( image ) );
 			cellpose.setCellProbThreshold( ( double ) settings.get( KEY_CELL_PROBABILITY_THRESHOLD ) );
 			cellpose.setFlowThreshold( ( double ) settings.get( KEY_FLOW_THRESHOLD ) );

src/main/java/org/mastodon/mamut/detection/stardist/StarDistDetector.java

@@ -105,6 +105,7 @@ protected boolean validateSettings( StringBuilder errorHolder )
 	{
 		try (StarDist starDist = new StarDist( ( StarDist.ModelType ) settings.get( KEY_MODEL_TYPE ) ))
 		{
+			this.apposeProcess = starDist;
 			boolean isData3D = is3D( image );
 			Boolean isModelType2D = starDist.getModelType().is2D();
 			if ( isModelType2D != null )

src/main/java/org/mastodon/mamut/util/ApposeProcess.java

@@ -15,6 +15,7 @@
 import java.util.function.Consumer;
 
 import org.apache.commons.lang3.time.StopWatch;
+import org.apposed.appose.Builder;
 import org.apposed.appose.Environment;
 import org.apposed.appose.Service;
 import org.apposed.appose.TaskEvent;
@@ -46,6 +47,8 @@ public abstract class ApposeProcess implements AutoCloseable
 
 	protected final Map< String, Object > inputs;
 
+	protected Service.Task currentTask;
+
 	protected ApposeProcess() throws IOException
 	{
 		this.stopWatch = StopWatch.createStarted();
@@ -101,8 +104,9 @@ private Environment setUpEnv()
 				writer.write( content );
 			}
 			logEnvFile( envFile );
-			environment = org.apposed.appose.Appose.file( envFile, "environment.yml" )
-					.subscribeProgress( ( title, cur, max ) -> logger.info( "{}: {}/{}", title, cur, max ) )
+			Builder envBuilder = org.apposed.appose.Appose.file( envFile, "environment.yml" );
+			environment =
+					envBuilder.subscribeProgress( ( title, cur, max ) -> logger.info( "{}: {}/{}", title, cur, max ) )
 					.subscribeOutput( msg -> {
 						if ( !msg.isEmpty() )
 							logger.info( msg );
@@ -139,6 +143,9 @@ protected Consumer< TaskEvent > getTaskListener( final StopWatch stopWatch, fina
 			case FAILURE:
 				logger.error( "Task failed with error: {}. Time elapsed: {}", task.error, stopWatch.formatSplitTime() );
 				break;
+			case CRASH:
+				logger.error( "Task crashed with error. Error: {}. Time elapsed: {}", task.error, stopWatch.formatSplitTime() );
+				break;
 			default:
 				logger.warn( "Unhandled task event: {}.", taskEvent.responseType );
 				break;
@@ -201,21 +208,26 @@ protected Service.Task runScript() throws IOException
 		}
 
 		String script = generateScript();
-		Service.Task task = pythonWorker.task( script, inputs, null );
+		currentTask = pythonWorker.task( script, inputs, null );
 		stopWatch.split();
 		if ( logger.isInfoEnabled() )
 			logger.info( "Created python task. Time elapsed: {}", stopWatch.formatSplitTime() );
-		task.listen( getTaskListener( stopWatch, task ) );
-		if ( isPythonTaskInterrupted( task ) )
+		currentTask.listen( getTaskListener( stopWatch, currentTask ) );
+		if ( isPythonTaskInterrupted( currentTask ) )
 			return null;
 
 		// Verify that it worked.
-		if ( task.status != Service.TaskStatus.COMPLETE )
-			throw new org.mastodon.mamut.detection.PythonRuntimeException( "Python task failed with error: " + task.error );
+		if ( currentTask.status != Service.TaskStatus.COMPLETE )
+			throw new PythonRuntimeException( "Python task failed with error: " + currentTask.error );
 
 		stopWatch.split();
 		if ( logger.isInfoEnabled() )
 			logger.info( "Python task completed. Total time {}", stopWatch.formatSplitTime() );
-		return task;
+		return currentTask;
+	}
+
+	public void cancel() throws InterruptedException
+	{
+		pythonWorker.kill();
 	}
 }