Add StarDist3D detector

Author: Stefan Hahmann

pom.xml

@@ -150,13 +150,20 @@
             <version>0.4.0-SNAPSHOT</version>
         </dependency>
 
-        <!-- mastodon tracking as base for python based detection -->
+        <!-- mastodon tracking as base for Python-based detection -->
         <dependency>
             <groupId>${mastodon.group}</groupId>
             <artifactId>mastodon-tracking</artifactId>
             <version>${mastodon-tracking.version}</version>
         </dependency>
 
+        <!-- dl-modelrunner -->
+        <dependency>
+            <groupId>io.bioimage</groupId>
+            <artifactId>dl-modelrunner</artifactId>
+            <version>0.5.10</version>
+        </dependency>
+
         <!-- Test dependencies -->
         <dependency>
             <groupId>org.junit.jupiter</groupId>

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

@@ -0,0 +1,202 @@
+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 = getStarDistModelRoot();
+			File directory = starDistModelRoot.toFile();
+			if ( !directory.exists() )
+			{
+				if ( !directory.mkdirs() )
+				{
+					throw new RuntimeException( "Failed to create environment directory: " + directory );
+				}
+			}
+			if ( directory.isDirectory() )
+			{
+				String[] files = directory.list();
+				if ( files != null && files.length > 0 )
+					logger.info( "Found {} files/directories in {}", files.length, directory.getAbsolutePath() );
+				else
+				{
+					try
+					{
+						logger.info( "Downloading model to {}", directory.getAbsolutePath() );
+						BioimageioRepo.connect().downloadByName( modelType.getModelName(), directory.getAbsolutePath() );
+					}
+					catch ( IllegalArgumentException e )
+					{
+						logger.info( "Exception while downloading model: {}", e.getMessage() );
+					}
+				}
+			}
+			else
+				logger.error( "The specified path is not a directory: {}", directory.getAbsolutePath() );
+		}
+	}
+
+	private Path getStarDistModelRoot()
+	{
+		return Paths.get( System.getProperty( "user.home" ), ".local", "share", "stardist", "models", modelType.getModelPath() );
+	}
+
+	@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.getModelName() + "', basedir=r\"" + getStarDistModelRoot() + "\")"
+				+ "\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 );
+		}
+	}
+}