Merge pull request #140 from mastodon-sc/appose

Deep Learning Detectors Improvements

Author: stefanhahmann

pom.xml

@@ -147,7 +147,7 @@
         <dependency>
             <groupId>org.apposed</groupId>
             <artifactId>appose</artifactId>
-            <version>0.5.1</version>
+            <version>0.7.0</version>
         </dependency>
 
         <!-- mastodon tracking as base for custom detectors (cellpose, stardist) -->
@@ -164,6 +164,13 @@
             <version>0.5.10</version>
         </dependency>
 
+        <!-- oshi for system / GPU information (used to limit gpu usage of cellpose) -->
+        <dependency>
+            <groupId>com.github.oshi</groupId>
+            <artifactId>oshi-core</artifactId>
+            <version>6.8.3</version>
+        </dependency>
+
         <!-- Test dependencies -->
         <dependency>
             <groupId>org.junit.jupiter</groupId>
@@ -210,6 +217,10 @@
             <type>test-jar</type>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-api</artifactId>
+        </dependency>
     </dependencies>
 
     <mailingLists>
@@ -259,6 +270,11 @@
             <id>scijava.public</id>
             <url>https://maven.scijava.org/content/repositories/public/</url>
         </repository>
+        <repository>
+            <!-- mvn central -->
+            <id>central</id>
+            <url>https://repo.maven.apache.org/maven2</url>
+        </repository>
     </repositories>
 
     <profiles>

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

@@ -41,13 +41,16 @@
 import net.imglib2.view.IntervalView;
 import net.imglib2.view.Views;
 
-import org.apache.commons.lang3.tuple.ImmutableTriple;
-import org.apache.commons.lang3.tuple.Triple;
+import org.mastodon.mamut.util.ByteFormatter;
+import org.mastodon.mamut.util.ImgSizeUtils;
+import org.mastodon.mamut.detection.util.SpimImageProperties;
 import org.mastodon.mamut.model.ModelGraph;
 import org.mastodon.mamut.util.LabelImageUtils;
 import org.mastodon.tracking.detection.DetectionUtil;
 import org.mastodon.tracking.detection.DetectorKeys;
 import org.mastodon.tracking.mamut.detection.AbstractSpotDetectorOp;
+import org.scijava.Context;
+import org.scijava.plugin.Parameter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -69,6 +72,14 @@ public abstract class DeepLearningDetector extends AbstractSpotDetectorOp
 
 	private static final Logger logger = LoggerFactory.getLogger( MethodHandles.lookup().lookupClass() );
 
+	@Parameter
+	protected Context context;
+
+	/**
+	 * 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;
+
 	@Override
 	public void compute( final List< SourceAndConverter< ? > > sources, final ModelGraph graph )
 	{
@@ -80,14 +91,13 @@ public void compute( final List< SourceAndConverter< ? > > sources, final ModelG
 		if ( !validateAndInitializeSettings() )
 			return;
 
-		Triple< Integer, Integer, Integer > settings = extractSettings( sources );
+		SpimImageProperties settings = extractSettings( sources );
 		if ( settings == null )
 			return;
 		// Now we are sure the settings are valid.
 
-		int minTimepoint = settings.getLeft();
-		int maxTimepoint = settings.getMiddle();
-		int setup = settings.getRight();
+		int minTimepoint = settings.getMinTimepoint();
+		int maxTimepoint = settings.getMaxTimepoint();
 
 		// Perform detection.
 		statusService.showStatus( "Detecting spots..." );
@@ -101,8 +111,8 @@ public void compute( final List< SourceAndConverter< ? > > sources, final ModelG
 				if ( isCanceled() )
 					break; // Exit but don't fail.
 
-				if ( DetectionUtil.isPresent( sources, setup, timepoint ) )
-					detectAndAddSpots( sources, graph, setup, timepoint );
+				if ( DetectionUtil.isPresent( sources, settings.getSetupId(), timepoint ) )
+					detectAndAddSpots( sources, graph, settings.getSetupId(), timepoint, settings.getResolutionLevel() );
 			}
 		}
 		catch ( Exception e )
@@ -119,14 +129,16 @@ public void compute( final List< SourceAndConverter< ? > > sources, final ModelG
 		ok = true;
 	}
 
-	private Triple< Integer, Integer, Integer > extractSettings( final List< SourceAndConverter< ? > > sources )
+	private SpimImageProperties extractSettings( final List< SourceAndConverter< ? > > sources )
 	{
 		// Extract settings.
 		final int minTimepoint = ( int ) settings.get( DetectorKeys.KEY_MIN_TIMEPOINT );
 		final int maxTimepoint = ( int ) settings.get( DetectorKeys.KEY_MAX_TIMEPOINT );
 		final int setup = ( int ) settings.get( DetectorKeys.KEY_SETUP_ID );
+		final int level = ( int ) settings.get( DeepLearningDetectorKeys.KEY_LEVEL );
 
-		logger.info( "Settings contain, minTimepoint: {}, maxTimepoint: {} and setup {}", minTimepoint, maxTimepoint, setup );
+		logger.info( "Settings contain, minTimepoint: {}, maxTimepoint: {}, setup: {} and level: {}", minTimepoint, maxTimepoint, setup,
+				level );
 
 		if ( setup < 0 || setup >= sources.size() )
 		{
@@ -142,27 +154,58 @@ private Triple< Integer, Integer, Integer > extractSettings( final List< SourceA
 			logger.error( "Invalid time-point range: {}", errorMessage );
 			return null;
 		}
-		return new ImmutableTriple<>( minTimepoint, maxTimepoint, setup );
+		int availableLevels = sources.get( setup ).getSpimSource().getNumMipmapLevels();
+		if ( level < 0 || level >= availableLevels )
+		{
+			errorMessage = "The parameter " + DeepLearningDetectorKeys.KEY_LEVEL
+					+ " is not in the range of available resolution levels ( 0 to " + ( availableLevels - 1 ) + "): " + level;
+			logger.error( "Invalid level: {}.", errorMessage );
+			return null;
+		}
+		return new SpimImageProperties( minTimepoint, maxTimepoint, setup, level );
 	}
 
 	private void detectAndAddSpots( final List< SourceAndConverter< ? > > sources, final ModelGraph graph, final int setup,
-			final int timepoint )
+			final int timepoint, final int level )
 	{
 		// 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();
 
-		// level 0 is always the highest resolution.
-		final int level = 0;
-
 		// This is the 3D image of the current time-point and specified channel. It is always 3D. If the source is 2D, the 3rd dimension has a size of 1.
 		RandomAccessibleInterval< ? > image = source.getSource( timepoint, level );
 
+		long theoreticalImageSize = 0;
+		try
+		{
+			theoreticalImageSize = ImgSizeUtils.getSizeInBytes( image );
+		}
+		catch ( IllegalArgumentException e )
+		{
+			logger.info( "Could not estimate image size." );
+		}
+		if ( theoreticalImageSize > MAX_SIZE_IN_BYTES )
+		{
+			String actualSize = ByteFormatter.humanReadableByteCount( theoreticalImageSize );
+			String maxSize = ByteFormatter.humanReadableByteCount( MAX_SIZE_IN_BYTES );
+			logger.warn( "Size of image at timepoint {}, setup {} and level {} is {}, which is larger than the maximum size of {}.",
+					timepoint, setup, level, actualSize, maxSize );
+			String message = "Size of image at timepoint " + timepoint + ", setup " + setup + " and level " + level + " is " + actualSize
+					+ ", which is larger than the supported maximum size of " + maxSize
+					+ " per timepoint. Consider using a different resolution level or process only a region of interest.\n";
+			throw new InferenceException( message );
+		}
+
 		// Crop the image to the region of interest (ROI) if specified in the settings.
 		final Interval roi = ( Interval ) settings.get( DetectorKeys.KEY_ROI );
 		if ( roi != null )
+		{
+			logger.info( "Settings contained a roi: {}", roi );
 			image = Views.interval( image, roi );
+		}
 
-		final Img< ? > segmentation = performSegmentation( image, source.getVoxelDimensions().dimensionsAsDoubleArray() );
+		System.err.println(); // show the FIJI console
+		final Img< ? > segmentation =
+				performSegmentation( Views.dropSingletonDimensions( image ), source.getVoxelDimensions().dimensionsAsDoubleArray() );
 
 		if ( segmentation != null )
 		{
@@ -192,8 +235,11 @@ private boolean validateAndInitializeSettings()
 		return true;
 	}
 
-	protected double getAnisotropy( double[] voxelSizes )
+	protected double getAnisotropy( double[] voxelSizes, boolean is3D )
 	{
+		if ( !is3D )
+			return 1.0;
+
 		if ( voxelSizes == null || voxelSizes.length == 0 )
 		{
 			throw new IllegalArgumentException( "Array must not be empty" );
@@ -241,6 +287,7 @@ public Map< String, Object > getDefaultSettings()
 		defaultSettings.put( DetectorKeys.KEY_MIN_TIMEPOINT, DetectorKeys.DEFAULT_MIN_TIMEPOINT );
 		defaultSettings.put( DetectorKeys.KEY_MAX_TIMEPOINT, DetectorKeys.DEFAULT_MAX_TIMEPOINT );
 		defaultSettings.put( DetectorKeys.KEY_ROI, null ); // No ROI by default.
+		defaultSettings.put( DeepLearningDetectorKeys.KEY_LEVEL, DeepLearningDetectorKeys.DEFAULT_LEVEL );
 		addSpecificDefaultSettings( defaultSettings );
 		return defaultSettings;
 	}

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

@@ -0,0 +1,48 @@
+package org.mastodon.mamut.detection;
+
+public class DeepLearningDetectorKeys
+{
+	private DeepLearningDetectorKeys()
+	{
+		// Prevent instantiation
+	}
+
+	/**
+	 * Key for the parameter specifying what resolution level to perform on. Expected values must be
+	 * {@link Integer} larger or equal to 0.
+	 */
+	public static final String KEY_LEVEL = "LEVEL";
+
+	/**
+	 * Key for the parameter specifying what GPU to use. Expected values must be
+	 * {@link Integer} representing the GPU ID, or {@code null} to use the CPU.
+	 */
+	public static final String KEY_GPU_ID = "GPU_ID";
+
+	/**
+	 * Key for the parameter specifying what fraction of the GPU memory to allocate. Expected values must be
+	 * {@link Double} between 0 and 1.
+	 */
+	public static final String KEY_GPU_MEMORY_FRACTION = "GPU_MEMORY_FRACTION";
+
+	/**
+	 * Default value for the {@link #KEY_LEVEL} parameter.
+	 * <br>
+	 * The highest resolution level (0).
+	 */
+	public static final int DEFAULT_LEVEL = 0;
+
+	/**
+	 * Default value for the {@link #KEY_GPU_ID} parameter.
+	 * <br>
+	 * The first gpu (0).
+	 */
+	public static final int DEFAULT_GPU_ID = 0;
+
+	/**
+	 * Default value for the {@link #KEY_GPU_MEMORY_FRACTION} parameter.
+	 * <br>
+	 * Half of the GPU memory (0.5).
+	 */
+	public static final double DEFAULT_GPU_MEMORY_FRACTION = 0.5d;
+}

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

@@ -0,0 +1,24 @@
+package org.mastodon.mamut.detection;
+
+/**
+ * Represents an exception that occurs during an inference process.
+ * This exception extends RuntimeException and provides constructors
+ * for initializing with a message, a cause, or both.
+ */
+public class InferenceException extends RuntimeException
+{
+	public InferenceException( final String message )
+	{
+		super( message );
+	}
+
+	public InferenceException( final String message, final Throwable cause )
+	{
+		super( message, cause );
+	}
+
+	public InferenceException( final Throwable cause )
+	{
+		super( cause );
+	}
+}