Merge pull request #149 from mastodon-sc/re-use-envs

Re use already loaded python environments

Author: stefanhahmann

pom.xml

@@ -147,7 +147,7 @@
         <dependency>
             <groupId>org.apposed</groupId>
             <artifactId>appose</artifactId>
-            <version>0.8.0</version>
+            <version>0.8.1</version>
         </dependency>
 
         <!-- mastodon tracking as base for custom detectors (cellpose, stardist) -->

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

@@ -41,10 +41,12 @@
 import net.imglib2.view.IntervalView;
 import net.imglib2.view.Views;
 
-import org.mastodon.mamut.util.appose.ApposeProcess;
-import org.mastodon.mamut.util.ByteFormatter;
+import org.apposed.appose.Appose;
+import org.apposed.appose.Environment;
+import org.apposed.appose.Service;
 import org.mastodon.mamut.detection.util.SpimImageProperties;
 import org.mastodon.mamut.model.ModelGraph;
+import org.mastodon.mamut.util.ByteFormatter;
 import org.mastodon.mamut.util.ImgUtils;
 import org.mastodon.mamut.util.LabelImageUtils;
 import org.mastodon.tracking.detection.AbstractDetectorOp;
@@ -78,13 +80,13 @@ public abstract class DeepLearningDetector extends AbstractSpotDetectorOp
 	@Parameter
 	protected Context context;
 
+	protected Service pythonService;
+
 	/**
 	 * Represents the maximum allowable size in bytes for datasets handle via the appose java-python bridge.
 	 */
 	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 )
 	{
@@ -111,17 +113,39 @@ public void compute( final List< SourceAndConverter< ? > > sources, final ModelG
 			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++ )
+			Environment environment = Appose.mamba().scheme( "environment.yml" ).content( getPythonEnvContent() ).logDebug()
+					.subscribeProgress( ( title, cur, max ) -> logger.info( "{}: {}/{}", title, cur, max ) )
+					.subscribeOutput( logger::info )
+					.subscribeError( logger::error ).build();
+			if ( logger.isInfoEnabled() )
+				logger.info( "Set up environment. Path: {}", environment.base() );
+			try (Service python = environment.python())
 			{
-				// We use the `statusService to show progress.
-				statusService.showProgress( timepoint - minTimepoint + 1, maxTimepoint - minTimepoint + 1 );
-
-				if ( isCanceled() )
-					break; // Exit but don't fail.
-
-				if ( DetectionUtil.isPresent( sources, settings.getSetupId(), timepoint ) )
-					detectAndAddSpots( sources, graph, settings.getSetupId(), timepoint, settings.getResolutionLevel() );
+				if ( isWindows() )
+					python.init( getPythonEnvInit() );
+
+				// First, get the source for the current channel (or setup) at the desired time-point. In BDV jargon, this is a source.
+				final Source< ? > source = sources.get( settings.getSetupId() ).getSpimSource();
+				RandomAccessibleInterval< ? > image = source.getSource( 0, 0 );
+
+				Service.Task importTask = python.task( getImportScript( !is3D( image ) ), "main" );
+				importTask.waitFor();
+				this.pythonService = python;
+				for ( int timepoint = minTimepoint; timepoint <= maxTimepoint; timepoint++ )
+				{
+					// We use the `statusService to show progress.
+					statusService.showProgress( timepoint - minTimepoint + 1, maxTimepoint - minTimepoint + 1 );
+
+					if ( isCanceled() )
+						break; // Exit but don't fail.
+
+					if ( DetectionUtil.isPresent( sources, settings.getSetupId(), timepoint ) )
+						detectAndAddSpots( sources, graph, settings.getSetupId(), timepoint, settings.getResolutionLevel(), python );
+				}
+				if ( logger.isInfoEnabled() )
+					logger.info( "Finished python process." );
 			}
+
 		}
 		catch ( Exception e )
 		{
@@ -174,7 +198,7 @@ private SpimImageProperties extractSettings( final List< SourceAndConverter< ? >
 	}
 
 	private void detectAndAddSpots( final List< SourceAndConverter< ? > > sources, final ModelGraph graph, final int setup,
-			final int timepoint, final int level )
+			final int timepoint, final int level, final Service python )
 	{
 		// First, get the source for the current channel (or setup) at the desired time-point. In BDV jargon, this is a source.
 		final Source< ? > source = sources.get( setup ).getSpimSource();
@@ -212,7 +236,8 @@ private void detectAndAddSpots( final List< SourceAndConverter< ? > > sources, f
 		}
 
 		final Img< ? > segmentation =
-				performSegmentation( Views.dropSingletonDimensions( image ), source.getVoxelDimensions().dimensionsAsDoubleArray() );
+				performSegmentation( Views.dropSingletonDimensions( image ), source.getVoxelDimensions().dimensionsAsDoubleArray(),
+						python );
 
 		if ( segmentation != null )
 		{
@@ -301,20 +326,36 @@ public Map< String, Object > getDefaultSettings()
 
 	protected abstract boolean validateSettings( final StringBuilder errorHolder );
 
-	protected abstract Img< ? > performSegmentation( final RandomAccessibleInterval< ? > image, final double[] voxelDimensions );
+	protected abstract Img< ? > performSegmentation( final RandomAccessibleInterval< ? > image, final double[] voxelDimensions,
+			final Service python );
 
 	protected abstract void addSpecificDefaultSettings( final Map< String, Object > defaultSettings );
 
 	protected abstract String getDetectorName();
 
+	protected abstract String getPythonEnvContent();
+
+	protected String getPythonEnvInit()
+	{
+		return "import numpy\n";
+	}
+
+	protected abstract String getImportScript( final boolean dataIs2D );
+
+	private static boolean isWindows()
+	{
+		String os = System.getProperty( "os.name" ).toLowerCase();
+		return os.contains( "win" );
+	}
+
 	/** 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();
-		if ( apposeProcess != null )
-			apposeProcess.cancel();
+		if ( pythonService != null )
+			pythonService.kill();
 		super.cancel( reason );
 	}
 

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

@@ -54,13 +54,13 @@ public abstract class Segmentation extends ApposeProcess
 {
 	private static final Logger logger = LoggerFactory.getLogger( MethodHandles.lookup().lookupClass() );
 
-	protected Segmentation() throws IOException
+	protected Segmentation( final Service pythonService )
 	{
-		super();
+		super( pythonService );
 	}
 
 	/**
-	 * Segments the input image using the configured Python environment and
+	 * Segments the input image using the configured Python runtime environment and
 	 * returns the segmented image as an {@link Img}.
 	 *
 	 * @param inputImage the input image to be segmented.

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

@@ -30,6 +30,7 @@
 
 import java.io.IOException;
 
+import org.apposed.appose.Service;
 import org.mastodon.mamut.detection.Segmentation;
 
 /**
@@ -58,9 +59,9 @@ public abstract class Cellpose extends Segmentation
 
 	public static final double DEFAULT_DIAMETER = 0d;
 
-	protected Cellpose() throws IOException
+	protected Cellpose( final Service python ) throws IOException
 	{
-		super();
+		super( python );
 	}
 
 	public void setCellProbThreshold( final double cellProbThreshold )
@@ -151,8 +152,7 @@ protected String generateScript()
 				+ "task.outputs['label_image'] = shared" + "\n";
 	}
 
-	@Override
-	protected String generateImportStatements()
+	public static String generateImportStatements()
 	{
 		return "import numpy as np" + "\n"
 				+ "from cellpose import models" + "\n"

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

@@ -30,6 +30,8 @@
 
 import java.io.IOException;
 
+import org.apposed.appose.Service;
+
 /**
  * Cellpose3 is a specialized implementation of the {@link Cellpose} class, specifically
  * designed to use Cellpose version 3 model for cell segmentation tasks.<br>
@@ -56,11 +58,12 @@ public class Cellpose3 extends Cellpose
 			+ "  - numpy\n";
 
 	private final ModelType modelType;
+
 	private double anisotropy = 1;
 
-	public Cellpose3( final ModelType modelType ) throws IOException
+	public Cellpose3( final ModelType modelType, final Service python ) throws IOException
 	{
-		super();
+		super( python );
 		this.modelType = modelType;
 	}
 
@@ -102,12 +105,6 @@ protected String getEvaluateModelCommand()
 				+ "cellprob_threshold=" + cellProbThreshold + ")" + "\n";
 	}
 
-	@Override
-	protected String generateEnvFileContent()
-	{
-		return ENV_FILE_CONTENT;
-	}
-
 	public enum ModelType
 	{
 		CYTO3( "cyto3", true ),

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

@@ -51,8 +51,8 @@
 import net.imglib2.RandomAccessibleInterval;
 import net.imglib2.img.Img;
 import net.imglib2.util.Cast;
-import net.imglib2.view.Views;
 
+import org.apposed.appose.Service;
 import org.mastodon.mamut.detection.DeepLearningDetector;
 import org.mastodon.tracking.mamut.detection.SpotDetectorOp;
 import org.scijava.Priority;
@@ -99,11 +99,13 @@ && checkParameter( settings, KEY_GPU_ID, Integer.class, errorHolder )
 	}
 
 	@Override
-	protected Img< ? > performSegmentation( final RandomAccessibleInterval< ? > image, final double[] voxelDimensions )
+	protected Img< ? > performSegmentation( final RandomAccessibleInterval< ? > image, final double[] voxelDimensions,
+			final Service python )
 	{
-		try (Cellpose3 cellpose = new Cellpose3( ( Cellpose3.ModelType ) settings.get( KEY_MODEL_TYPE ) ))
+
+		try
 		{
-			this.apposeProcess = cellpose;
+			Cellpose3 cellpose = new Cellpose3( ( Cellpose3.ModelType ) settings.get( KEY_MODEL_TYPE ), python );
 			boolean is3D = is3D( image );
 			cellpose.set3D( is3D );
 			cellpose.setCellProbThreshold( ( double ) settings.get( KEY_CELL_PROBABILITY_THRESHOLD ) );
@@ -143,4 +145,16 @@ protected String getDetectorName()
 	{
 		return "Cellpose3";
 	}
+
+	@Override
+	protected String getPythonEnvContent()
+	{
+		return Cellpose3.ENV_FILE_CONTENT;
+	}
+
+	@Override
+	protected String getImportScript( final boolean dataIs2D )
+	{
+		return Cellpose3.generateImportStatements();
+	}
 }

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

@@ -30,6 +30,8 @@
 
 import java.io.IOException;
 
+import org.apposed.appose.Service;
+
 /**
  * Cellpose3 is a specialized implementation of the {@link Cellpose} class, specifically
  * designed to use Cellpose version 4 model for cell segmentation tasks.<br>
@@ -54,9 +56,9 @@ public class Cellpose4 extends Cellpose
 			+ "  - pip:\n"
 			+ "    - appose==" + APPOSE_PYTHON_VERSION + "\n";
 
-	public Cellpose4() throws IOException
+	public Cellpose4( final Service python ) throws IOException
 	{
-		super();
+		super( python );
 	}
 
 	@Override
@@ -80,10 +82,4 @@ protected String getEvaluateModelCommand()
 				+ "flow_threshold=" + flowThreshold + ", "
 				+ "cellprob_threshold=" + cellProbThreshold + ")" + "\n";
 	}
-
-	@Override
-	protected String generateEnvFileContent()
-	{
-		return ENV_FILE_CONTENT;
-	}
 }

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

@@ -48,6 +48,7 @@
 import net.imglib2.img.Img;
 import net.imglib2.util.Cast;
 
+import org.apposed.appose.Service;
 import org.mastodon.mamut.detection.DeepLearningDetector;
 import org.mastodon.tracking.mamut.detection.SpotDetectorOp;
 import org.mastodon.tracking.mamut.trackmate.wizard.descriptors.cellpose.Cellpose4DetectorDescriptor;
@@ -90,11 +91,12 @@ && checkParameter( settings, KEY_GPU_ID, Integer.class, errorHolder )
 	}
 
 	@Override
-	protected Img< ? > performSegmentation( final RandomAccessibleInterval< ? > image, final double[] voxelDimensions )
+	protected Img< ? > performSegmentation( final RandomAccessibleInterval< ? > image, final double[] voxelDimensions,
+			final Service python )
 	{
-		try (Cellpose4 cellpose = new Cellpose4())
+		try
 		{
-			this.apposeProcess = cellpose;
+			Cellpose4 cellpose = new Cellpose4( python );
 			cellpose.set3D( is3D( image ) );
 			cellpose.setCellProbThreshold( ( double ) settings.get( KEY_CELL_PROBABILITY_THRESHOLD ) );
 			cellpose.setFlowThreshold( ( double ) settings.get( KEY_FLOW_THRESHOLD ) );
@@ -128,4 +130,16 @@ protected String getDetectorName()
 	{
 		return "Cellpose4";
 	}
+
+	@Override
+	protected String getPythonEnvContent()
+	{
+		return Cellpose4.ENV_FILE_CONTENT;
+	}
+
+	@Override
+	protected String getImportScript( final boolean dataIs2D )
+	{
+		return Cellpose.generateImportStatements();
+	}
 }