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,242 @@
+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 java.util.Map;
+
+import net.imglib2.util.Cast;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import io.bioimage.modelrunner.bioimageio.BioimageioRepo;
+import io.bioimage.modelrunner.bioimageio.description.ModelDescriptor;
+import io.bioimage.modelrunner.bioimageio.description.exceptions.ModelSpecsException;
+import io.bioimage.modelrunner.utils.JSONUtils;
+
+public class StarDist extends Segmentation3D
+{
+	private static final Logger logger = LoggerFactory.getLogger( MethodHandles.lookup().lookupClass() );
+
+	private final ModelType modelType;
+
+	private String installationFolderName;
+
+	private boolean dataIs2D;
+
+	public StarDist( final ModelType model ) throws IOException, InterruptedException
+	{
+		super();
+		logger.info( "Initializing StarDist, model: {}", model );
+		this.modelType = model;
+		Path starDistModelRoot = getStarDistModelRoot();
+		if ( starDistModelRoot == null )
+			logger.debug( "StarDist model path is null. This is normal for the built-in demo models" );
+		else
+		{
+			File directory = starDistModelRoot.toFile();
+			if ( !directory.exists() )
+			{
+				if ( !directory.mkdirs() )
+				{
+					throw new RuntimeException( "Failed to create environment directory: " + directory );
+				}
+			}
+			if ( directory.isDirectory() )
+			{
+				File[] files = directory.listFiles();
+				if ( files != null && files.length > 0 )
+				{
+					logger.info( "Found {} files/directories in {}", files.length, directory.getAbsolutePath() );
+					for ( File file : files )
+						logger.debug( "File/Directory: {}", file.getName() );
+					installationFolderName = Paths.get( files[ 0 ].getAbsolutePath() ).getFileName().toString();
+					logger.info( "Reusing model in folder: {}", installationFolderName );
+				}
+				else
+				{
+					try
+					{
+						logger.info( "Downloading model to {}", directory.getAbsolutePath() );
+						BioimageioRepo repo = BioimageioRepo.connect();
+						ModelDescriptor descriptor = repo.selectByName( modelType.getModelName() );
+						String installationFolder = repo.downloadByName( modelType.getModelName(), directory.getAbsolutePath() );
+						installationFolderName = Paths.get( installationFolder ).getFileName().toString();
+						createConfigFromBioimageio( descriptor, directory.getAbsolutePath() + File.separator + installationFolderName );
+						logger.info( "Downloading finished. Installation folder: {}", installationFolderName );
+					}
+					catch ( IllegalArgumentException e )
+					{
+						logger.info( "Exception while downloading model: {}", e.getMessage() );
+					}
+					catch ( ModelSpecsException e )
+					{
+						logger.info( "Exception while creating config.json: {}", e.getMessage() );
+					}
+				}
+			}
+			else
+				logger.error( "The specified path is not a directory: {}", directory.getAbsolutePath() );
+		}
+	}
+
+	private Path getStarDistModelRoot()
+	{
+		if ( modelType.getModelPath() == null )
+			return null;
+		return Paths.get( System.getProperty( "user.home" ), ".local", "share", "appose", "stardist", "models", modelType.getModelPath() );
+	}
+
+	private static void createConfigFromBioimageio( final ModelDescriptor descriptor, final String modelDir )
+			throws IOException, ModelSpecsException
+	{
+		Map< String, Object > stardistMap = Cast.unchecked( descriptor.getConfig().getSpecMap().get( "stardist" ) );
+		Map< String, Object > stardistConfig = Cast.unchecked( stardistMap.get( "config" ) );
+		File jsonFile = new File( modelDir, "config.json" );
+		logger.info( "Creating config.json file: {}", jsonFile.getAbsolutePath() );
+		JSONUtils.writeJSONFile( jsonFile.getAbsolutePath(), stardistConfig );
+	}
+
+	@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.getModelPath() == 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.getModelPath() == 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='" + installationFolderName + "', basedir=r\"models" + File.separator
+				+ modelType.getModelPath() + "\")"
+				+ "\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( "StarDist Demo", null, null );
+
+		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;
+		}
+
+		public String getDisplayName()
+		{
+			String dimensionality;
+			if ( is2D == null )
+				dimensionality = " (2D/3D)";
+			else if ( is2D )
+				dimensionality = " (2D)";
+			else
+				dimensionality = " (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 );
+		}
+	}
+}