weka spark classification

package de.hanslovsky.zspacing.spark.experiments.mask.weka;

import java.io.File;
import java.lang.invoke.MethodHandles;
import java.util.ArrayList;

import org.apache.log4j.Level;
import org.apache.log4j.LogManager;
import org.apache.log4j.Logger;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.broadcast.Broadcast;
import org.janelia.thickness.kryo.KryoSerialization;
import org.janelia.thickness.utility.Utility;

import ij.ImagePlus;
import ij.ImageStack;
import ij.io.FileSaver;
import ij.process.ImageProcessor;
import scala.Tuple2;
import trainableSegmentation.WekaSegmentation;

public class WekaClassifierSpark
{

    public static final Logger LOG = LogManager.getLogger( MethodHandles.lookup().lookupClass() );
    static
    {
        LOG.setLevel( Level.INFO );
    }

    public static void main( final String[] args )
    {

        final SparkConf conf = new SparkConf()
                .setAppName( MethodHandles.lookup().lookupClass().getSimpleName() )
                .set( "spark.network.timeout", "600" )
                .set( "spark.serializer", "org.apache.spark.serializer.KryoSerializer" )
                .set( "spark.kryoserializer.buffer.max", "1g" )
                .set( "spark.kryo.unsafe", "true" ) // supposed to give huge
                // performance boost for
                // primitive arrays
                .set( "spark.kryo.registrator", KryoSerialization.Registrator.class.getName() );
        final JavaSparkContext sc = new JavaSparkContext( conf );
        Logger.getRootLogger().setLevel( Level.ERROR );

        final String classifierPath = args[ 0 ];
        final String source = args[ 1 ];
        final String target = args[ 2 ];
        final int start = Integer.parseInt( args[ 3 ] );
        final int stop = Integer.parseInt( args[ 4 ] );
        final int step = Integer.parseInt( args[ 5 ] );
        final int classLabel = args.length > 6 ? Integer.parseInt( args[ 6 ] ) : -1;

        run( sc, classifierPath, source, target, start, stop, step, classLabel );

        sc.stop();

    }

    public static void run( final JavaSparkContext sc, final String classifierPath, final String source, final String target, final int start, final int stop, final int step, final int classLabel )
    {
        final ArrayList< Integer > indices = Utility.arange( start, stop, step );
        final JavaPairRDD< Integer, ImageProcessor > images = sc.parallelize( indices ).mapToPair( i -> {
            final String fp = String.format( source, i );
            LOG.debug( "Loading image locted at " + fp );
            return new Tuple2<>( i, new ImagePlus( fp ).getProcessor() );
        } );
        final JavaPairRDD< ?, ImageProcessor[] > classified = run( sc, images, classifierPath );

        new File( String.format( target, start ) ).getParentFile().mkdirs();
        if ( classLabel >= 0 )
        {
            final int count = classified.map( t -> {
                final ImageProcessor fp = t._2()[ classLabel ];
                final int length = fp.getWidth() * fp.getHeight();
                for ( int i = 0; i < length; ++i )
                    if ( fp.getf( i ) < 0.2f )
                        fp.setf( i, 0.0f );
                return new FileSaver( new ImagePlus( "", fp ) ).saveAsTiff( String.format( target, t._1() ) ) ? 1 : 0;
            } ).reduce( ( i1, i2 ) -> i1 + i2 );
            LOG.info( "Classified and wrote " + count + "/" + indices.size() + "images " );
        }
        else
        {
            final int count = classified.map( t -> new FileSaver( new ImagePlus( "", toStack( t._2() ) ) ).saveAsTiff( String.format( target, t._1() ) ) ? 1 : 0 ).reduce( ( i1, i2 ) -> i1 + i2 );
            LOG.info( "Classified and wrote " + count + "/" + indices.size() + "images " );
        }
    }

    public static < K > JavaPairRDD< K, ImageProcessor[] > run( final JavaSparkContext sc, final JavaPairRDD< K, ImageProcessor > imgs, final String classifierPath )
    {

        final Broadcast< String > classifierPathBC = sc.broadcast( classifierPath );

        final JavaPairRDD< K, ImageProcessor[] > classified = imgs.mapValues( img -> {
            final ImagePlus imp = new ImagePlus( "", img );
            final WekaSegmentation seg = new WekaSegmentation( imp );
            final String cp = classifierPathBC.getValue();
            LOG.debug( "Loading classifier from " + cp + " . File exists? " + new File( cp ).exists() );
            seg.loadClassifier( cp );
            final ImagePlus probs = seg.applyClassifier( imp, 1, true );
            final ImageProcessor[] result = new ImageProcessor[ probs.getStackSize() ];
            final ImageStack s = probs.getStack();
            for ( int i = 0; i < result.length; ++i )
                result[ i ] = s.getProcessor( i + 1 );
            return result;
        } );

        return classified;
    }

    public static ImageStack toStack( final ImageProcessor[] ips )
    {
        final int w = ips[ 0 ].getWidth();
        final int h = ips[ 0 ].getHeight();
        final ImageStack is = new ImageStack( w, h );
        for ( final ImageProcessor ip : ips )
            is.addSlice( ip );
        return is;
    }

}