Merge pull request #290 from trackmate-sc/export-csv-headless

Export to CSV file in headless mode

src/main/java/fiji/plugin/trackmate/io/CSVExporter.java

@@ -0,0 +1,309 @@
+package fiji.plugin.trackmate.io;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.FileReader;
+import java.io.FileWriter;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+
+import com.opencsv.CSVWriter;
+
+import fiji.plugin.trackmate.Dimension;
+import fiji.plugin.trackmate.Model;
+import fiji.plugin.trackmate.Spot;
+import fiji.plugin.trackmate.TrackModel;
+import fiji.plugin.trackmate.util.TMUtils;
+
+public class CSVExporter
+{
+
+ /**
+ * The symbol that will be appended when a value is missing.
+ */
+ public static final String MISSING_STR = "ø";
+
+ /**
+ * Exports the spots in the specified model to a CSV file.
+ * <p>
+ * The table contains the spots, one spot per line. All the declared spot
+ * feature values are exported, one feature per column. Additionally, the
+ * following columns are added:
+ * <ol start="0">
+ * <li>The spot ID.
+ * <li>The spot name.
+ * <li>The ID of the track the spot belongs to, if any.
+ * <li>The name of the track the spot belongs to, if any.
+ * <li>Whether the track the spot belongs to is visible, if any. 1 if
+ * visible, 0 otherwise.
+ * </ol>
+ * 
+ * The table will have a header on 4 lines containing:
+ * <ol start="0">
+ * <li>The feature key.
+ * <li>The feature name.
+ * <li>The feature short name.
+ * <li>The feature unit if any.
+ * </ol>
+ * 
+ * @param csvFile
+ * the path of the file to export to.
+ * @param model
+ * the model to export.
+ * @param visibleOnly
+ * if <code>true</code>, only the spots in visible tracks will be
+ * exported. Otherwise all spots are exported.
+ * @throws IOException
+ * if something wrong happens while writing to the CSV file.
+ */
+ public static final void exportSpots( final String csvFile, final Model model, final boolean visibleOnly ) throws IOException
+ {
+ final TrackModel tm = model.getTrackModel();
+ final Set< Integer > trackIDs = tm.trackIDs( visibleOnly );
+ final List< String > features = new ArrayList<>( model.getFeatureModel().getSpotFeatures() );
+
+ // Function to return a feature value.
+ final Map< String, Boolean > isInt = model.getFeatureModel().getSpotFeatureIsInt();
+ final BiFunction< Spot, String, Number > featureVal = ( s, f ) -> {
+ final Double obj = s.getFeature( f );
+ if ( obj == null )
+ return null;
+
+ if ( isInt.getOrDefault( f, Boolean.FALSE ) )
+ return obj.intValue();
+ else
+ return obj;
+ };
+
+ /*
+ * Extra columns.
+ */
+
+ // Spot name function.
+ final Function< Spot, String > labelFun = ( s ) -> s.getName();
+ // Spot id function/
+ final Function< Spot, String > idFun = ( s ) -> Integer.toString( s.ID() );
+ // Track id function.
+ final Function< Spot, String > trackIdFun = ( s ) -> {
+ final Integer trackID = tm.trackIDOf( s );
+ if ( trackID == null )
+ return MISSING_STR;
+ else
+ return Integer.toString( trackID );
+ };
+ // Track name function.
+ final Function< Spot, String > trackNameFun = ( s ) -> {
+ final Integer trackID = tm.trackIDOf( s );
+ if ( trackID == null )
+ return MISSING_STR;
+ else
+ return tm.name( trackID );
+ };
+ // Track visibility function.
+ final Function< Spot, String > trackVisiblityFun = ( s ) -> {
+ final Integer trackID = tm.trackIDOf( s );
+ if ( trackID == null || !tm.isVisible( trackID ) )
+ return "0";
+ else
+ return "1";
+ };
+
+ final List< Function< Spot, String > > extraFuns = Arrays.asList( idFun, labelFun, trackIdFun, trackNameFun, trackVisiblityFun );
+ final String[][] extraHeaders = new String[][] {
+ { "ID", "ID", "ID", "" },
+ { "LABEL", "Label", "Label", "" },
+ { "TRACK_ID", "Track ID", "Track ID", "" },
+ { "TRACK_NAME", "Track name", "Track name", "" },
+ { "TRACK_VISIBLE", "Track visibility", "Track visibility", "" }
+ };
+
+ try (CSVWriter writer = new CSVWriter( new FileWriter( new File( csvFile ) ),
+ CSVWriter.DEFAULT_SEPARATOR,
+ CSVWriter.NO_QUOTE_CHARACTER,
+ CSVWriter.DEFAULT_ESCAPE_CHARACTER,
+ CSVWriter.DEFAULT_LINE_END ))
+ {
+ // Header.
+ writeHeaderSpots( writer, model, extraHeaders );
+
+ final SpotWriter spotWriter = new SpotWriter( writer, features, featureVal, extraFuns );
+
+ // Spots in tracks.
+ for ( final Integer trackID : trackIDs )
+ {
+ final List< Spot > spots = new ArrayList<>( tm.trackSpots( trackID ) );
+ spots.sort( Spot.frameComparator );
+ for ( final Spot spot : spots )
+ spotWriter.write( spot );
+ }
+
+ // Possibly spots not in tracks.
+ if ( !visibleOnly )
+ {
+ for ( final Spot spot : model.getSpots().iterable( false ) )
+ {
+ if ( tm.trackIDOf( spot ) != null )
+ continue; // Already done abobe.
+
+ spotWriter.write( spot );
+ }
+ }
+ }
+ }
+
+ private static void writeHeaderSpots( final CSVWriter writer, final Model model, final String[][] extra )
+ {
+ final List< String > features = new ArrayList<>( model.getFeatureModel().getSpotFeatures() );
+ final Map< String, String > featureNames = model.getFeatureModel().getSpotFeatureNames();
+ final Map< String, String > featureShortNames = model.getFeatureModel().getSpotFeatureShortNames();
+ final Map< String, String > featureUnits = new HashMap<>();
+ for ( final String feature : features )
+ {
+ final Dimension dimension = model.getFeatureModel().getSpotFeatureDimensions().get( feature );
+ final String units = TMUtils.getUnitsFor( dimension, model.getSpaceUnits(), model.getTimeUnits() );
+ featureUnits.put( feature, units );
+ }
+ writeHeader( writer, features, featureNames, featureShortNames, featureUnits, extra );
+ }
+
+ private static void writeHeader(
+ final CSVWriter writer,
+ final List< String > features,
+ final Map< String, String > featureNames,
+ final Map< String, String > featureShortNames,
+ final Map< String, String > featureUnits,
+ final String[][] extra )
+ {
+
+ final int columnShift = extra.length;
+ final int nCols = columnShift + features.size();
+ final String[] content = new String[ nCols ];
+
+ /*
+ * Determine whether we can skip 2nd or 3d line, if it's identical to
+ * the 2nd one (happens when the names are repeated).
+ */
+
+ boolean skipThirdLine = true;
+ boolean skipSecondLine = true;
+ for ( int i = columnShift; i < content.length; i++ )
+ {
+ final String feature = features.get( i - columnShift );
+ final String name = featureNames.get( features.get( i - columnShift ) );
+ final String shortName = featureShortNames.get( features.get( i - columnShift ) );
+ if ( !feature.equals( name ) )
+ skipSecondLine = false;
+ if ( !name.equals( shortName ) )
+ skipThirdLine = false;
+ }
+
+ // Header 1st line.
+ for ( int i = 0; i < extra.length; i++ )
+ content[ i ] = extra[ i ][ 0 ];
+ for ( int i = columnShift; i < content.length; i++ )
+ content[ i ] = features.get( i - columnShift );
+ writer.writeNext( content );
+
+ // Header 2nd line.
+ if ( !skipSecondLine )
+ {
+ for ( int i = 0; i < extra.length; i++ )
+ content[ i ] = extra[ i ][ 1 ];
+ for ( int i = columnShift; i < content.length; i++ )
+ content[ i ] = featureNames.get( features.get( i - columnShift ) );
+ writer.writeNext( content );
+ }
+
+ // Header 3rd line.
+ if ( !skipThirdLine )
+ {
+ for ( int i = 0; i < extra.length; i++ )
+ content[ i ] = extra[ i ][ 2 ];
+ for ( int i = columnShift; i < content.length; i++ )
+ content[ i ] = featureShortNames.get( features.get( i - columnShift ) );
+ writer.writeNext( content );
+ }
+
+ // Header 4th line.
+ for ( int i = 0; i < extra.length; i++ )
+ content[ i ] = extra[ i ][ 3 ];
+ for ( int i = columnShift; i < content.length; i++ )
+ {
+ final String feature = features.get( i - columnShift );
+ final String units = featureUnits.get( feature );
+ final String unitsStr = ( units == null || units.isEmpty() ) ? "" : "(" + units + ")";
+ content[ i ] = unitsStr;
+ }
+ writer.writeNext( content );
+ }
+
+ private static class SpotWriter
+ {
+
+ private final CSVWriter writer;
+
+ private final List< String > features;
+
+ private final List< Function< Spot, String > > extraFuns;
+
+ private final String[] content;
+
+ private final BiFunction< Spot, String, Number > featureVal;
+
+ private final int columnShift;
+
+ public SpotWriter( final CSVWriter writer, final List< String > features, final BiFunction< Spot, String, Number > featureVal, final List< Function< Spot, String > > extraFuns )
+ {
+ this.writer = writer;
+ this.features = features;
+ this.featureVal = featureVal;
+ this.extraFuns = extraFuns;
+ this.columnShift = extraFuns.size();
+ final int nCols = columnShift + features.size();
+ this.content = new String[ nCols ];
+ }
+
+ public void write( final Spot spot )
+ {
+ for ( int i = 0; i < columnShift; i++ )
+ content[ i ] = extraFuns.get( i ).apply( spot );
+ for ( int i = columnShift; i < content.length; i++ )
+ {
+ final Object val = featureVal.apply( spot, features.get( i - columnShift ) );
+ content[ i ] = ( val == null ) ? MISSING_STR : val.toString();
+ }
+
+ writer.writeNext( content );
+ }
+ }
+
+ public static void main( final String[] args ) throws IOException
+ {
+ final String filename = "samples/FakeTracks.xml";
+ final TmXmlReader reader = new TmXmlReader( new File( filename ) );
+ if ( !reader.isReadingOk() )
+ {
+ System.out.println( reader.getErrorMessage() );
+ return;
+ }
+
+ final String csvFile = filename.replace( ".xml", "-spots_in_tracks.csv" );
+ CSVExporter.exportSpots( csvFile, reader.getModel(), false );
+
+ System.out.println( "Done." );
+ System.out.println( "__________________" );
+ try (BufferedReader br = new BufferedReader( new FileReader( csvFile ) ))
+ {
+ String line;
+ while ( ( line = br.readLine() ) != null )
+ System.out.println( line );
+ }
+ }
+}