[TrackMate with Custom StarDist Model] Run TrackMate from a script and save the results. #fiji #stardist #trackmate

TrackMate-StarDist.groovy

/**
 * Use StarDist in TrackMate from the latest version and export results
 * Using a custom StarDist Model
 * Author: Olivier Burri, EPFL SV PTECH BIOP
 * Last Modification: March 2022
 * 
 * Due to the simple nature of this code, no copyright is applicable
 * 
 */

#@ File model (label="StarDist Model File")
#@ Double scoreThr (label="StarDist Score Threshold")
#@ Double overlapThr (label="StarDist Overlap Threshold")
#@ Integer frameGap (label="Max Frame Gap [frames]")
#@ Double linkDistance (label="Linking Max Distance [calibrated]")
#@ Double gapDistance (label="Gap Closing Max Distance [calibrated]")
#@ Boolean allowSplit (label="Allow Track Splitting")
#@ Double splitDistance (label="Split Distance [calibrated]")
#@ Double areaThr (label="Area Threshold Filter [calibrated^2]")

#@Integer nuc_ch (label="Channel to Detect", value=2)
#@String measure_ch (label="Channels to Measure", value="1")
#@File image_path

IJ.run("Close All", "")

// Convert string of numbers into a list
def measure_channels = measure_ch.split(",").collect { it as int }


// Prepare save folder
def saveDir = new File(image_path.getParent(), "csv_results")
saveDir.mkdirs()


// Open Image
def imp = IJ.openImage(image_path.toString())
imp.show()

def tracker = Tracker.builder().rawImage(imp)
        .starDistModel(model)
        .scoreThreshold(scoreThr)
        .overlapThreshold(overlapThr)
        .nuclearChannel(nuc_ch)
        .measureChannels(measure_channels)
        .linkDistance(linkDistance)
        .frameGap(frameGap)
        .gapDistance(gapDistance)
        .allowSplit(allowSplit)
        .splitDistance(splitDistance)
        .areaThreshold(areaThr)
        .build().initialize()


def OK = tracker.process()


def selectionModel = new SelectionModel(tracker.trackmate.getModel())

def ds = DisplaySettingsIO.readUserDefault()

def displayer = new HyperStackDisplayer(tracker.trackmate.getModel(), selectionModel, imp, ds)
displayer.render()
displayer.refresh()

// Export results
def results = tracker.obtainTrackResults()
results.saveColumnHeaders(true)
results.show("Tracking Results")
results.save(new File(saveDir, imp.getTitle()).getAbsolutePath() + ".csv")


// Export XML file
tracker.export(new File(saveDir, imp.getTitle() + ".xml"))

println "END"

@Builder(excludes = 'trackmate')
class Tracker {

    TrackMate trackmate
    ImagePlus rawImage

    File starDistModel
    double scoreThreshold
    double overlapThreshold
    int nuclearChannel
    List<Integer> measureChannels

    double linkDistance
    int frameGap
    double gapDistance
    boolean allowSplit
    double splitDistance

    double areaThreshold

    def initialize() {

        def model = new Model()
        def cal = rawImage.getCalibration()
        model.setLogger(Logger.IJ_LOGGER)
        model.setPhysicalUnits(cal.getUnit(), cal.getTimeUnit())

        //Prepare settings
        def settings = new Settings(rawImage)

        //Configure detector - We use the Strings for the keys
        settings.detectorFactory = new StarDistCustomDetectorFactory()
        settings.detectorSettings["TARGET_CHANNEL"] = nuclearChannel
        settings.detectorSettings["SCORE_THRESHOLD"] = scoreThreshold
        settings.detectorSettings["OVERLAP_THRESHOLD"] = overlapThreshold
        settings.detectorSettings["MODEL_FILEPATH"] = starDistModel.getAbsolutePath()


        // Add ALL the feature analyzers known to TrackMate, via
        // providers.
        settings.addAllAnalyzers()

        //Configure tracker
        settings.trackerFactory = new SparseLAPTrackerFactory()
        settings.trackerSettings = LAPUtils.getDefaultLAPSettingsMap()
        settings.trackerSettings["LINKING_MAX_DISTANCE"] = this.linkDistance
        settings.trackerSettings["GAP_CLOSING_MAX_DISTANCE"] = this.gapDistance
        settings.trackerSettings["MAX_FRAME_GAP"] = this.frameGap
        settings.trackerSettings["ALLOW_TRACK_SPLITTING"] = this.allowSplit
        settings.trackerSettings["SPLITTING_MAX_DISTANCE"] = this.splitDistance

        settings.initialSpotFilterValue = -1.0

        def trackmate = new TrackMate(model, settings)

        trackmate.getModel().getLogger().log(settings.toStringFeatureAnalyzersInfo())
        trackmate.computeSpotFeatures(true)
        trackmate.computeEdgeFeatures(true)
        trackmate.computeTrackFeatures(true)

        // Add size filter
        def areaFilter = new FeatureFilter("AREA", this.areaThreshold, true)
        settings.addSpotFilter(areaFilter)


        this.trackmate = trackmate

        return this
    }

    def process() {
        //Execute the full process EXCEPT for the detection step.

        // Check settings.
        def ok = this.trackmate.checkInput()
        // Initial filtering
        ok = trackmate.process()

        return ok
    }
    
    /*
     * Return the results of this tracking as a simple results table
     */

    ResultsTable obtainTrackResults() {

        // TODO Export Trackmate results directly
        def rt = new ResultsTable()
        def model = this.trackmate.getModel()
        def trackIDs = model.getTrackModel().trackIDs(false)
        trackIDs.each { id ->
            def spots = model.getTrackModel().trackSpots(id).toSorted { it.getFeature("POSITION_T") }
            spots.each { spot ->
                rt.incrementCounter()
                rt.addLabel(this.rawImage.getTitle())
                rt.addValue("Track ID", id)
                rt.addValue("ROI name", spot.getName())
                rt.addValue("Frame", spot.getFeature("FRAME"))
                rt.addValue("Position T", spot.getFeature("POSITION_T"))
                rt.addValue("Area", spot.getFeature("AREA"))

                this.measureChannels.each { channel ->
                    rt.addValue("Mean Ch" + channel, spot.getFeature("MEAN_INTENSITY_CH" + channel))
                    rt.addValue("StDev Ch" + channel, spot.getFeature("STD_INTENSITY_CH" + channel))
                }

            }
        }
        return rt
    }

    /*
     * Export this trackmate result to an XML file to reuse later
     * This also records the location of the ImagePlus that was opened. So it is important to save it **BEFORE** running trackmate
     */

    void export(File outFile) {
        def fname = outFile.getName()
        def pos = fname.lastIndexOf(".")
        fname = fname.substring(0, pos)

        // Export the tracks
        def tracksFile = new File(outFile.getParent(), fname + "_Tracks.xml")
        ExportTracksToXML.export(this.trackmate.getModel(), this.trackmate.getSettings(), tracksFile)

        // Export the Trackmate file
        def writer = new TmXmlWriter(outFile)
        writer.appendModel(trackmate.getModel())
        writer.appendSettings(trackmate.getSettings())
        writer.writeToFile()
    }
}


// All Imports
import fiji.plugin.trackmate.*
import fiji.plugin.trackmate.action.ExportTracksToXML
import fiji.plugin.trackmate.features.FeatureFilter
import fiji.plugin.trackmate.gui.displaysettings.DisplaySettingsIO
import fiji.plugin.trackmate.io.TmXmlWriter
import fiji.plugin.trackmate.stardist.StarDistCustomDetectorFactory
import fiji.plugin.trackmate.tracking.LAPUtils
import fiji.plugin.trackmate.tracking.sparselap.SparseLAPTrackerFactory
import fiji.plugin.trackmate.visualization.hyperstack.HyperStackDisplayer
import groovy.transform.builder.Builder
import ij.IJ
import ij.ImagePlus
import ij.measure.ResultsTable