diff --git a/src/main/java/fiji/plugin/trackmate/omnipose/advanced/AdvancedOmniposeDetectorConfigurationPanel.java b/src/main/java/fiji/plugin/trackmate/omnipose/advanced/AdvancedOmniposeDetectorConfigurationPanel.java new file mode 100644 index 00000000..e3ee7ef4 --- /dev/null +++ b/src/main/java/fiji/plugin/trackmate/omnipose/advanced/AdvancedOmniposeDetectorConfigurationPanel.java @@ -0,0 +1,146 @@ +package fiji.plugin.trackmate.omnipose.advanced; + +import static fiji.plugin.trackmate.cellpose.advanced.AdvancedCellposeDetectorFactory.KEY_CELL_PROB_THRESHOLD; +import static fiji.plugin.trackmate.cellpose.advanced.AdvancedCellposeDetectorFactory.KEY_FLOW_THRESHOLD; +import static fiji.plugin.trackmate.gui.Fonts.SMALL_FONT; + +import java.awt.GridBagConstraints; +import java.awt.Insets; +import java.util.Map; + +import javax.swing.JLabel; + +import fiji.plugin.trackmate.Model; +import fiji.plugin.trackmate.Settings; +import fiji.plugin.trackmate.cellpose.advanced.AdvancedCellposeDetectorConfigurationPanel; +import fiji.plugin.trackmate.detection.SpotDetectorFactoryBase; +import fiji.plugin.trackmate.gui.displaysettings.SliderPanelDouble; +import fiji.plugin.trackmate.gui.displaysettings.StyleElements; +import fiji.plugin.trackmate.omnipose.OmniposeDetectorConfigurationPanel; +import fiji.plugin.trackmate.omnipose.OmniposeSettings.PretrainedModelOmnipose; + +public class AdvancedOmniposeDetectorConfigurationPanel extends OmniposeDetectorConfigurationPanel +{ + + private static final long serialVersionUID = 1L; + + private static final String TITLE = AdvancedOmniposeDetectorFactory.NAME;; + + private final StyleElements.BoundedDoubleElement flowThresholdEl = new StyleElements.BoundedDoubleElement( "Flow threshold", 0.0, 3.0 ) + { + + private double flowThreshold = 0.; + + @Override + public double get() + { + return flowThreshold; + } + + @Override + public void set( final double v ) + { + flowThreshold = v; + } + }; + + private final StyleElements.BoundedDoubleElement cellProbThresholdEl = new StyleElements.BoundedDoubleElement("Cell prob", -6.0, 6.0) + { + + private double cellProbThreshold = 0.; + + @Override + public double get() + { + return cellProbThreshold; + } + + @Override + public void set( final double v ) + { + cellProbThreshold = v; + } + }; + + public AdvancedOmniposeDetectorConfigurationPanel( final Settings settings, final Model model ) + { + super( settings, model, TITLE, ICON, DOC1_URL, "omnipose", PretrainedModelOmnipose.values() ); + + /* + * Add flow threshold. + */ + + int gridy = 12; + + final JLabel lblFlowThreshold = new JLabel( "Flow threshold:" ); + lblFlowThreshold.setFont( SMALL_FONT ); + final GridBagConstraints gbcLblFlowThreshold = new GridBagConstraints(); + gbcLblFlowThreshold.anchor = GridBagConstraints.EAST; + gbcLblFlowThreshold.insets = new Insets( 0, 5, 5, 5 ); + gbcLblFlowThreshold.gridx = 0; + gbcLblFlowThreshold.gridy = gridy; + add( lblFlowThreshold, gbcLblFlowThreshold ); + + final SliderPanelDouble sliderPanelFlowThreshold = StyleElements.linkedSliderPanel( flowThresholdEl, 3, 0.1 ); + AdvancedCellposeDetectorConfigurationPanel.setFont( sliderPanelFlowThreshold, SMALL_FONT ); + final GridBagConstraints gbcFlowThresholdSlider = new GridBagConstraints(); + gbcFlowThresholdSlider.anchor = GridBagConstraints.EAST; + gbcFlowThresholdSlider.insets = new Insets( 0, 5, 5, 5 ); + gbcFlowThresholdSlider.fill = GridBagConstraints.HORIZONTAL; + gbcFlowThresholdSlider.gridx = 1; + gbcFlowThresholdSlider.gridwidth = 2; + gbcFlowThresholdSlider.gridy = gridy; + add( sliderPanelFlowThreshold, gbcFlowThresholdSlider ); + + /* + * Add cell probability threshold. + */ + + gridy++; + + final JLabel lblCellProb = new JLabel( "Mask threshold:" ); + lblCellProb.setFont( SMALL_FONT ); + final GridBagConstraints gbcLblCellProb = new GridBagConstraints(); + gbcLblCellProb.anchor = GridBagConstraints.EAST; + gbcLblCellProb.insets = new Insets( 0, 5, 5, 5 ); + gbcLblCellProb.gridx = 0; + gbcLblCellProb.gridy = gridy; + add( lblCellProb, gbcLblCellProb ); + + final SliderPanelDouble sliderPanelCellProbThreshold = StyleElements.linkedSliderPanel( cellProbThresholdEl, 3, 0.4 ); + AdvancedCellposeDetectorConfigurationPanel.setFont( sliderPanelCellProbThreshold, SMALL_FONT ); + final GridBagConstraints gbcCellProbThresholdSlider = new GridBagConstraints(); + gbcCellProbThresholdSlider.anchor = GridBagConstraints.EAST; + gbcCellProbThresholdSlider.insets = new Insets( 0, 5, 5, 5 ); + gbcCellProbThresholdSlider.fill = GridBagConstraints.HORIZONTAL; + gbcCellProbThresholdSlider.gridx = 1; + gbcCellProbThresholdSlider.gridwidth = 2; + gbcCellProbThresholdSlider.gridy = gridy; + add( sliderPanelCellProbThreshold, gbcCellProbThresholdSlider ); + } + + @Override + protected SpotDetectorFactoryBase< ? > getDetectorFactory() + { + return new AdvancedOmniposeDetectorFactory<>(); + } + + @Override + public void setSettings( final Map< String, Object > settings ) + { + super.setSettings( settings ); + flowThresholdEl.set( ( double ) settings.get( KEY_FLOW_THRESHOLD ) ); + flowThresholdEl.update(); + cellProbThresholdEl.set( ( double ) settings.get( KEY_CELL_PROB_THRESHOLD ) ); + cellProbThresholdEl.update(); + } + + @Override + public Map< String, Object > getSettings() + { + final Map< String, Object > settings = super.getSettings(); + settings.put( KEY_FLOW_THRESHOLD, flowThresholdEl.get() ); + settings.put( KEY_CELL_PROB_THRESHOLD, cellProbThresholdEl.get() ); + return settings; + } +} diff --git a/src/main/java/fiji/plugin/trackmate/omnipose/advanced/AdvancedOmniposeDetectorFactory.java b/src/main/java/fiji/plugin/trackmate/omnipose/advanced/AdvancedOmniposeDetectorFactory.java new file mode 100644 index 00000000..15225234 --- /dev/null +++ b/src/main/java/fiji/plugin/trackmate/omnipose/advanced/AdvancedOmniposeDetectorFactory.java @@ -0,0 +1,289 @@ +/*- + * #%L + * TrackMate: your buddy for everyday tracking. + * %% + * Copyright (C) 2021 - 2023 TrackMate developers. + * %% + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU General Public License as + * published by the Free Software Foundation, either version 3 of the + * License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU General Public License for more details. + * + * You should have received a copy of the GNU General Public + * License along with this program. If not, see + * . + * #L% + */ +package fiji.plugin.trackmate.omnipose.advanced; + +import static fiji.plugin.trackmate.cellpose.advanced.AdvancedCellposeDetectorFactory.DEFAULT_CELL_PROB_THRESHOLD; +import static fiji.plugin.trackmate.cellpose.advanced.AdvancedCellposeDetectorFactory.DEFAULT_FLOW_THRESHOLD; +import static fiji.plugin.trackmate.cellpose.advanced.AdvancedCellposeDetectorFactory.KEY_CELL_PROB_THRESHOLD; +import static fiji.plugin.trackmate.cellpose.advanced.AdvancedCellposeDetectorFactory.KEY_FLOW_THRESHOLD; +import static fiji.plugin.trackmate.detection.DetectorKeys.KEY_TARGET_CHANNEL; +import static fiji.plugin.trackmate.detection.ThresholdDetectorFactory.KEY_SIMPLIFY_CONTOURS; +import static fiji.plugin.trackmate.io.IOUtils.readBooleanAttribute; +import static fiji.plugin.trackmate.io.IOUtils.readDoubleAttribute; +import static fiji.plugin.trackmate.io.IOUtils.readIntegerAttribute; +import static fiji.plugin.trackmate.io.IOUtils.readStringAttribute; +import static fiji.plugin.trackmate.io.IOUtils.writeAttribute; +import static fiji.plugin.trackmate.util.TMUtils.checkMapKeys; +import static fiji.plugin.trackmate.util.TMUtils.checkParameter; + +import java.util.Arrays; +import java.util.List; +import java.util.Map; + +import javax.swing.ImageIcon; + +import org.jdom2.Element; +import org.scijava.Priority; +import org.scijava.plugin.Plugin; + +import fiji.plugin.trackmate.Logger; +import fiji.plugin.trackmate.Model; +import fiji.plugin.trackmate.Settings; +import fiji.plugin.trackmate.cellpose.CellposeDetector; +import fiji.plugin.trackmate.detection.SpotDetectorFactory; +import fiji.plugin.trackmate.detection.SpotGlobalDetector; +import fiji.plugin.trackmate.io.IOUtils; +import fiji.plugin.trackmate.omnipose.OmniposeDetectorFactory; +import fiji.plugin.trackmate.omnipose.OmniposeSettings.PretrainedModelOmnipose; +import fiji.plugin.trackmate.util.TMUtils; +import net.imglib2.Interval; +import net.imglib2.type.NativeType; +import net.imglib2.type.numeric.RealType; + +@Plugin( type = SpotDetectorFactory.class, priority = Priority.LOW - 1. ) +public class AdvancedOmniposeDetectorFactory< T extends RealType< T > & NativeType< T > > extends OmniposeDetectorFactory< T > +{ + + /** A string key identifying this factory. */ + public static final String DETECTOR_KEY = "OMNIPOSE_ADVANCED_DETECTOR"; + + /** The pretty name of the target detector. */ + public static final String NAME = "Omnipose advanced detector"; + + /** An html information text. */ + public static final String INFO_TEXT = "" + + "This detector relies on omnipose to detect objects." + + "

" + "It is identical to the Omnipose detector, except that it allows to " + + "tweak the 'flow threshold' and 'cell probability threshold' parameters of the " + + "algorithm." + + "

" + + "If you use this detector for your work, please be so kind as to " + + "also cite the omnipose paper: Cutler, Kevin J., et al., " + + "'Omnipose: A High-Precision Morphology-Independent Solution for Bacterial Cell Segmentation.' " + + "Nature Methods 19, no. 11 (November 2022): 1438–48." + + "

" + + "Documentation for this module " + + "on the ImageJ Wiki." + + ""; + + /* + * METHODS + */ + + @Override + public SpotGlobalDetector< T > getDetector( final Interval interval ) + { + // Base settings. + + final String omniposePythonPath = ( String ) settings.get( KEY_OMNIPOSE_PYTHON_FILEPATH ); + final PretrainedModelOmnipose model = ( PretrainedModelOmnipose ) settings.get( KEY_OMNIPOSE_MODEL ); + final String customModelPath = ( String ) settings.get( KEY_OMNIPOSE_CUSTOM_MODEL_FILEPATH ); + final boolean simplifyContours = ( boolean ) settings.get( KEY_SIMPLIFY_CONTOURS ); + final boolean useGPU = ( boolean ) settings.get( KEY_USE_GPU ); + + // Channels are 0-based (0: grayscale, then R & G & B). + final int channel = ( Integer ) settings.get( KEY_TARGET_CHANNEL ); + final int channel2 = ( Integer ) settings.get( KEY_OPTIONAL_CHANNEL_2 ); + + // Convert to diameter in pixels. + final double[] calibration = TMUtils.getSpatialCalibration( img ); + final double diameter = ( double ) settings.get( KEY_CELL_DIAMETER ) / calibration[ 0 ]; + + // Advanced settings. + + final double flowThreshold = ( Double ) settings.get( KEY_FLOW_THRESHOLD ); + final double cellProbThreshold = ( Double ) settings.get( KEY_CELL_PROB_THRESHOLD ); + + final AdvancedOmniposeSettings cellposeSettings = AdvancedOmniposeSettings + .create() + .omniposePythonPath( omniposePythonPath ) + .customModel( customModelPath ) + .model( model ) + .channel1( channel ) + .channel2( channel2 ) + .diameter( diameter ) + .useGPU( useGPU ) + .simplifyContours( simplifyContours ) + .flowThreshold( flowThreshold ) + .cellProbThreshold( cellProbThreshold ) + .get(); + + // Logger. + final Logger logger = ( Logger ) settings.get( KEY_LOGGER ); + final CellposeDetector< T > detector = new CellposeDetector<>( img, interval, cellposeSettings, logger ); + return detector; + } + + @Override + public boolean marshall( final Map< String, Object > settings, final Element element ) + { + if ( !super.marshall( settings, element ) ) + return false; + + final StringBuilder errorHolder = new StringBuilder(); + boolean ok = writeAttribute( settings, element, KEY_FLOW_THRESHOLD, Double.class, errorHolder ); + ok = ok && writeAttribute( settings, element, KEY_CELL_PROB_THRESHOLD, Double.class, errorHolder ); + if ( !ok ) + errorMessage = errorHolder.toString(); + return ok; + } + + @Override + public boolean unmarshall( final Element element, final Map< String, Object > settings ) + { + settings.clear(); + final StringBuilder errorHolder = new StringBuilder(); + boolean ok = true; + ok = ok && readStringAttribute( element, settings, KEY_OMNIPOSE_PYTHON_FILEPATH, errorHolder ); + ok = ok && readStringAttribute( element, settings, KEY_OMNIPOSE_CUSTOM_MODEL_FILEPATH, errorHolder ); + ok = ok && readIntegerAttribute( element, settings, KEY_TARGET_CHANNEL, errorHolder ); + ok = ok && readIntegerAttribute( element, settings, KEY_OPTIONAL_CHANNEL_2, errorHolder ); + ok = ok && readDoubleAttribute( element, settings, KEY_CELL_DIAMETER, errorHolder ); + ok = ok && readBooleanAttribute( element, settings, KEY_USE_GPU, errorHolder ); + ok = ok && readBooleanAttribute( element, settings, KEY_SIMPLIFY_CONTOURS, errorHolder ); + ok = ok && readDoubleAttribute( element, settings, KEY_FLOW_THRESHOLD, errorHolder ); + ok = ok && readDoubleAttribute( element, settings, KEY_CELL_PROB_THRESHOLD, errorHolder ); + + // Read model. + final String str = element.getAttributeValue( KEY_OMNIPOSE_MODEL ); + if ( null == str ) + { + errorHolder.append( "Attribute " + KEY_OMNIPOSE_MODEL + " could not be found in XML element.\n" ); + ok = false; + } + settings.put( KEY_OMNIPOSE_MODEL, PretrainedModelOmnipose.valueOf( str ) ); + + return checkSettings( settings ); + } + + @Override + public AdvancedOmniposeDetectorConfigurationPanel getDetectorConfigurationPanel( final Settings settings, final Model model ) + { + return new AdvancedOmniposeDetectorConfigurationPanel( settings, model ); + } + + @Override + public Map< String, Object > getDefaultSettings() + { + final Map< String, Object > settings = super.getDefaultSettings(); + settings.put( KEY_FLOW_THRESHOLD, DEFAULT_FLOW_THRESHOLD ); + settings.put( KEY_CELL_PROB_THRESHOLD, DEFAULT_CELL_PROB_THRESHOLD ); + return settings; + } + + @Override + public boolean checkSettings( final Map< String, Object > settings ) + { + boolean ok = true; + final StringBuilder errorHolder = new StringBuilder(); + ok = ok & checkParameter( settings, KEY_OMNIPOSE_PYTHON_FILEPATH, String.class, errorHolder ); + ok = ok & checkParameter( settings, KEY_OMNIPOSE_CUSTOM_MODEL_FILEPATH, String.class, errorHolder ); + ok = ok & checkParameter( settings, KEY_OMNIPOSE_MODEL, PretrainedModelOmnipose.class, errorHolder ); + ok = ok & checkParameter( settings, KEY_TARGET_CHANNEL, Integer.class, errorHolder ); + ok = ok & checkParameter( settings, KEY_OPTIONAL_CHANNEL_2, Integer.class, errorHolder ); + ok = ok & checkParameter( settings, KEY_CELL_DIAMETER, Double.class, errorHolder ); + ok = ok & checkParameter( settings, KEY_USE_GPU, Boolean.class, errorHolder ); + ok = ok & checkParameter( settings, KEY_SIMPLIFY_CONTOURS, Boolean.class, errorHolder ); + ok = ok & checkParameter( settings, KEY_FLOW_THRESHOLD, Double.class, errorHolder ); + ok = ok & checkParameter( settings, KEY_CELL_PROB_THRESHOLD, Double.class, errorHolder ); + + // If we have a logger, test it is of the right class. + final Object loggerObj = settings.get( KEY_LOGGER ); + if ( loggerObj != null && !Logger.class.isInstance( loggerObj ) ) + { + errorHolder.append( "Value for parameter " + KEY_LOGGER + " is not of the right class. " + + "Expected " + Logger.class.getName() + ", got " + loggerObj.getClass().getName() + ".\n" ); + ok = false; + } + + final List< String > mandatoryKeys = Arrays.asList( + KEY_OMNIPOSE_PYTHON_FILEPATH, + KEY_OMNIPOSE_MODEL, + KEY_TARGET_CHANNEL, + KEY_OPTIONAL_CHANNEL_2, + KEY_CELL_DIAMETER, + KEY_USE_GPU, + KEY_SIMPLIFY_CONTOURS ); + final List< String > optionalKeys = Arrays.asList( + KEY_OMNIPOSE_CUSTOM_MODEL_FILEPATH, + KEY_LOGGER, + KEY_FLOW_THRESHOLD, + KEY_CELL_PROB_THRESHOLD ); + ok = ok & checkMapKeys( settings, mandatoryKeys, optionalKeys, errorHolder ); + if ( !ok ) + errorMessage = errorHolder.toString(); + + // Extra test to make sure we can read the classifier file. + if ( ok ) + { + final Object obj = settings.get( KEY_OMNIPOSE_PYTHON_FILEPATH ); + if ( obj == null ) + { + errorMessage = "The path to the Omnipose python executable is not set."; + return false; + } + + if ( !IOUtils.canReadFile( ( String ) obj, errorHolder ) ) + { + errorMessage = "Problem with Omnipose python executable: " + errorHolder.toString(); + return false; + } + } + return ok; + } + + @Override + public String getInfoText() + { + return INFO_TEXT; + } + + @Override + public ImageIcon getIcon() + { + return null; + } + + @Override + public String getKey() + { + return DETECTOR_KEY; + } + + @Override + public String getName() + { + return NAME; + } + + @Override + public boolean has2Dsegmentation() + { + return true; + } + + @Override + public AdvancedOmniposeDetectorFactory< T > copy() + { + return new AdvancedOmniposeDetectorFactory<>(); + } +} diff --git a/src/main/java/fiji/plugin/trackmate/omnipose/advanced/AdvancedOmniposeSettings.java b/src/main/java/fiji/plugin/trackmate/omnipose/advanced/AdvancedOmniposeSettings.java new file mode 100644 index 00000000..bfe05105 --- /dev/null +++ b/src/main/java/fiji/plugin/trackmate/omnipose/advanced/AdvancedOmniposeSettings.java @@ -0,0 +1,143 @@ +package fiji.plugin.trackmate.omnipose.advanced; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import fiji.plugin.trackmate.omnipose.OmniposeSettings; + +public class AdvancedOmniposeSettings extends OmniposeSettings +{ + private final double flowThreshold; + + private final double cellProbThreshold; + + public AdvancedOmniposeSettings( + final String omniposePythonPath, + final PretrainedModelOmnipose model, + final String customModelPath, + final int chan, + final int chan2, + final double diameter, + final boolean useGPU, + final boolean simplifyContours, + final double flowThreshold, + final double cellProbThreshold ) + { + super( omniposePythonPath, model, customModelPath, chan, chan2, diameter, useGPU, simplifyContours ); + this.flowThreshold = flowThreshold; + this.cellProbThreshold = cellProbThreshold; + } + + @Override + public List< String > toCmdLine( final String imagesDir ) + { + final List< String > cmd = new ArrayList<>( super.toCmdLine( imagesDir ) ); + cmd.add( "--flow_threshold" ); + cmd.add( String.valueOf( flowThreshold ) ); + /* + * Careful! Because omnipose is still based on cellpose 1, the + * cellprob_threshold parameter is still called mask_threshold. + */ + cmd.add( "--mask_threshold" ); + cmd.add( String.valueOf( cellProbThreshold ) ); + return Collections.unmodifiableList( cmd ); + } + + public static Builder create() + { + return new Builder(); + } + + public static final class Builder extends OmniposeSettings.Builder + { + + private double flowThreshold = 0.4; + + private double cellProbThreshold = 0.0; + + public Builder flowThreshold( final double flowThreshold ) + { + this.flowThreshold = flowThreshold; + return this; + } + + public Builder cellProbThreshold( final double cellProbThreshold ) + { + this.cellProbThreshold = cellProbThreshold; + return this; + } + + @Override + public Builder channel1( final int ch ) + { + super.channel1( ch ); + return this; + } + + @Override + public Builder channel2( final int ch ) + { + super.channel2( ch ); + return this; + } + + @Override + public Builder omniposePythonPath( final String cellposePythonPath ) + { + super.omniposePythonPath( cellposePythonPath ); + return this; + } + + @Override + public Builder model( final PretrainedModelOmnipose model ) + { + super.model( model ); + return this; + } + + @Override + public Builder diameter( final double diameter ) + { + super.diameter( diameter ); + return this; + } + + @Override + public Builder useGPU( final boolean useGPU ) + { + super.useGPU( useGPU ); + return this; + } + + @Override + public Builder simplifyContours( final boolean simplifyContours ) + { + super.simplifyContours( simplifyContours ); + return this; + } + + @Override + public Builder customModel( final String customModelPath ) + { + super.customModel( customModelPath ); + return this; + } + + @Override + public AdvancedOmniposeSettings get() + { + return new AdvancedOmniposeSettings( + omniposePythonPath, + model, + customModelPath, + chan, + chan2, + diameter, + useGPU, + simplifyContours, + flowThreshold, + cellProbThreshold ); + } + } +}