Add stardist model selection

Author: Stefan Hahmann

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

@@ -0,0 +1,180 @@
+package org.mastodon.mamut.detection;
+
+import java.io.File;
+import java.io.IOException;
+import java.lang.invoke.MethodHandles;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import io.bioimage.modelrunner.bioimageio.BioimageioRepo;
+
+public class StarDist extends Segmentation3D
+{
+	private static final Logger logger = LoggerFactory.getLogger( MethodHandles.lookup().lookupClass() );
+
+	private final ModelType modelType;
+
+	private boolean dataIs2D;
+
+	private static final String PATH = "stardist";
+
+	public StarDist( final ModelType model ) throws IOException, InterruptedException
+	{
+		super();
+		this.modelType = model;
+		if ( modelType == null )
+			logger.info( "No star dist model path specified. Using pretrained demo model." );
+		else
+		{
+			Path starDistModelRoot =
+					Paths.get( System.getProperty( "user.home" ), ".local", "share", "stardist", "models", modelType.getModelPath() );
+			File directory = starDistModelRoot.toFile();
+			if ( directory.isDirectory() )
+			{
+				String[] files = directory.list();
+				if ( files != null && files.length > 0 )
+					logger.debug( "Found {} files in {}", files.length, directory.getAbsolutePath() );
+				else
+					BioimageioRepo.connect().downloadByName( "StarDist Plant Nuclei 3D ResNet", directory.getAbsolutePath() );
+			}
+			else
+				logger.error( "The specified path is not a directory: {}", directory.getAbsolutePath() );
+		}
+	}
+
+	@Override
+	String generateEnvFileContent()
+	{
+		return "name: stardist\n"
+				+ "channels:\n"
+				+ "  - conda-forge\n"
+				+ "dependencies:\n"
+				+ "  - python=3.10\n"
+				+ "  - cudatoolkit=11.2\n"
+				+ "  - cudnn=8.1.0\n"
+				+ "  - numpy<1.24\n"
+				+ "  - pip\n"
+				+ "  - pip:\n"
+				+ "    - numpy<1.24\n"
+				+ "    - tensorflow==2.10\n"
+				+ "    - stardist==0.8.5\n"
+				+ "    - appose\n";
+	}
+
+	@Override
+	String generateScript()
+	{
+		return "import numpy as np" + "\n"
+				+ "import appose" + "\n"
+				+ "from csbdeep.utils import normalize" + "\n"
+				+ getImportStarDistCommand()
+				+ "\n"
+				+ "task.update(message=\"Imports completed\")" + "\n"
+				+ "np.random.seed(6)" + "\n"
+				+ "axes_normalize = (0, 1, 2)" + "\n"
+				+ "\n"
+				+ "task.update(message=\"Loading StarDist pretrained 3D model\")" + "\n"
+				+ getLoadModelCommand()
+				+ "image_ndarray = image.ndarray()" + "\n"
+				+ "image_normalized = normalize(image_ndarray, 1, 99.8, axis=axes_normalize)" + "\n"
+				+ "task.update(message=\"Image shape:\" + str(image_normalized.shape))" + "\n"
+				+ "\n"
+				+ "guessed_tiles = model._guess_n_tiles(image_normalized)" + "\n"
+				+ "task.update(message=\"Guessed tiles: \" + str(guessed_tiles))" + "\n"
+				+ "\n"
+				+ "label_image, details = model.predict_instances(image_normalized, axes='ZYX', n_tiles=guessed_tiles)" + "\n"
+				+ "shared = appose.NDArray(image.dtype, image.shape)" + "\n"
+				+ "shared.ndarray()[:] = label_image" + "\n"
+				+ "task.outputs['label_image'] = shared" + "\n";
+	}
+
+	public ModelType getModelType()
+	{
+		return modelType;
+	}
+
+	private String getImportStarDistCommand()
+	{
+		if ( modelType == null )
+		{
+			if ( dataIs2D )
+				return "from stardist.models import StarDist2D" + "\n ";
+			return "from stardist.models import StarDist3D" + "\n ";
+		}
+		if ( modelType.is2D() )
+			return "from stardist.models import StarDist2D" + "\n ";
+		return "from stardist.models import StarDist3D" + "\n ";
+	}
+
+	private String getLoadModelCommand()
+	{
+		if ( modelType == null )
+		{
+			if ( dataIs2D )
+				return "model = StarDist2D.from_pretrained('2D_demo')" + "\n";
+			else
+				return "model = StarDist3D.from_pretrained('3D_demo')" + "\n";
+		}
+		String starDistModel = modelType.is2D() ? "StarDist2D" : "StarDist3D";
+		return "model = " + starDistModel + "(None, name='" + modelType.getModelPath() + "', basedir=r\"" + PATH + "\")" + "\n";
+	}
+
+	public enum ModelType
+	{
+		PLANT_NUCLEI_3D( "StarDist Plant Nuclei 3D ResNet", "stardist-plant-nuclei-3d", false ),
+		FLUO_2D( "StarDist Fluorescence Nuclei Segmentation", "stardist-fluo-2d", true ),
+		H_E( "StarDist H&E Nuclei Segmentation", "stardist-h-e-nuclei", true ),
+		DEMO_2D( "StarDist Demo", "stardist-demo-2d", true ),
+		DEMO_3D( "StarDist Demo", "stardist-demo-3d", false );
+
+		private final String modelName;
+
+		private final String modelPath;
+
+		private final boolean is2D;
+
+		ModelType( final String modelName, final String modelPath, final boolean is2D )
+		{
+			this.modelName = modelName;
+			this.modelPath = modelPath;
+			this.is2D = is2D;
+		}
+
+		public String getModelName()
+		{
+			return modelName;
+		}
+
+		public String getModelPath()
+		{
+			return modelPath;
+		}
+
+		public boolean is2D()
+		{
+			return is2D;
+		}
+
+		@Override
+		public String toString()
+		{
+			String dimensionality = is2D ? " (2D)" : " (3D)";
+			return modelName + dimensionality;
+		}
+
+		public static ModelType fromString( final String modelName )
+		{
+			for ( ModelType type : ModelType.values() )
+			{
+				if ( type.modelName.equalsIgnoreCase( modelName ) )
+				{
+					return type;
+				}
+			}
+			throw new IllegalArgumentException( "No enum constant for model name: " + modelName );
+		}
+	}
+}

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

@@ -6,6 +6,7 @@
 import static org.mastodon.tracking.detection.DetectorKeys.KEY_MAX_TIMEPOINT;
 import static org.mastodon.tracking.detection.DetectorKeys.KEY_MIN_TIMEPOINT;
 import static org.mastodon.tracking.detection.DetectorKeys.KEY_SETUP_ID;
+import static org.mastodon.tracking.mamut.trackmate.wizard.descriptors.StarDistDetectorDescriptor.KEY_MODEL_TYPE;
 import static org.mastodon.tracking.linking.LinkingUtils.checkParameter;
 
 import java.lang.invoke.MethodHandles;
@@ -39,7 +40,7 @@
 		+ "A cell probability threshold can be set to allow more/less sensitive detection.<br><br>"
 		+ "<strong>When this detection method is used for the first time, internet connection is needed, since an internal installation process is started. The installation consumes ~7GB hard disk space.</strong><br>"
 		+ "</html>" )
-public class StarDist3DDetector extends AbstractSpotDetectorOp
+public class StarDistDetector extends AbstractSpotDetectorOp
 {
 
 	private static final Logger logger = LoggerFactory.getLogger( MethodHandles.lookup().lookupClass() );
@@ -55,7 +56,8 @@ public void compute( final List< SourceAndConverter< ? > > sources, final ModelG
 		// A. Read the settings map, and check the validity.
 		final StringBuilder errorHolder = new StringBuilder();
 		boolean good;
-		good = checkParameter( settings, KEY_SETUP_ID, Integer.class, errorHolder );
+		good = checkParameter( settings, KEY_MODEL_TYPE, StarDist.ModelType.class, errorHolder );
+		good = good & checkParameter( settings, KEY_SETUP_ID, Integer.class, errorHolder );
 		good = good & checkParameter( settings, KEY_MIN_TIMEPOINT, Integer.class, errorHolder );
 		good = good & checkParameter( settings, KEY_MAX_TIMEPOINT, Integer.class, errorHolder );
 		if ( !good )
@@ -68,6 +70,7 @@ public void compute( final List< SourceAndConverter< ? > > sources, final ModelG
 		final int minTimepoint = ( int ) settings.get( KEY_MIN_TIMEPOINT );
 		final int maxTimepoint = ( int ) settings.get( KEY_MAX_TIMEPOINT );
 		final int setup = ( int ) settings.get( KEY_SETUP_ID );
+		final StarDist.ModelType modelType = ( StarDist.ModelType ) settings.get( KEY_MODEL_TYPE );
 
 		if ( setup < 0 || setup >= sources.size() )
 		{
@@ -86,7 +89,7 @@ public void compute( final List< SourceAndConverter< ? > > sources, final ModelG
 		// The `statusService` can be used to show short messages.
 		statusService.showStatus( "Detecting spots using Cellpose." );
 
-		try (StarDist3D starDist3D = new StarDist3D( null ))
+		try (StarDist starDist = new StarDist( modelType ))
 		{
 			for ( int timepoint = minTimepoint; timepoint <= maxTimepoint; timepoint++ )
 			{
@@ -115,13 +118,14 @@ public void compute( final List< SourceAndConverter< ? > > sources, final ModelG
 				 */
 				final int level = 0;
 				final RandomAccessibleInterval< ? > image = source.getSource( timepoint, level );
+				boolean is3D = is3D( image );
 
 				/*
 				 * This is the 3D image of the current time-point, specified
 				 * channel. It is always 3D. If the source is 2D, the 3rd dimension
 				 * will have a size of 1.
 				 */
-				Img< ? > segmentation = starDist3D.segmentImage( Cast.unchecked( image ) );
+				Img< ? > segmentation = starDist.segmentImage( Cast.unchecked( image ) );
 
 				final AffineTransform3D transform = DetectionUtil.getTransform( sources, timepoint, setup, level );
 				LabelImageUtils.createSpotsForFrame( graph, Cast.unchecked( segmentation ), timepoint, transform, 1d );