destripe.java

package com.cplab.jstripe;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import io.scif.ImageMetadata;
import io.scif.config.SCIFIOConfig;
import io.scif.services.DatasetIOService;
import net.imagej.axis.Axes;
import org.scijava.app.StatusService;
import org.scijava.command.Command;
import org.scijava.command.Previewable;
import org.scijava.Initializable;
import org.scijava.log.LogService;
import org.scijava.plugin.Parameter;
import org.scijava.plugin.Plugin;
import org.scijava.command.DynamicCommand;
import org.scijava.ui.UIService;
import org.scijava.module.MutableModuleItem;

import io.scif.Metadata;
import io.scif.img.SCIFIOImgPlus;
import io.scif.ome.OMEMetadata;
import io.scif.services.TranslatorService;

import loci.formats.ome.OMEXMLMetadata;

import net.imagej.Dataset;
import net.imagej.DatasetService;
import net.imagej.ImageJ;
import net.imagej.display.ImageDisplayService;
import net.imagej.axis.AxisType;
import net.imagej.ops.OpService;
import net.imglib2.Cursor;
import net.imglib2.IterableInterval;
import net.imglib2.RandomAccess;
import net.imglib2.type.NativeType;
import net.imglib2.type.numeric.RealType;
import net.imagej.ImgPlus;
import net.imglib2.view.Views;
import net.imglib2.type.numeric.real.FloatType;
import net.imglib2.type.numeric.integer.UnsignedShortType;



@Plugin(type = Command.class, menuPath = "Plugins>SPIM Destriping")
public class DeStripe extends DynamicCommand implements Command, Previewable, Initializable {

    @Parameter
    private UIService ui;

    @Parameter
    private LogService log;

    @Parameter
    private StatusService statusService;

    @Parameter(required = false)
    private ImageDisplayService idService;

    @Parameter
    private DatasetIOService datasetIOService;

    @Parameter
    private DatasetService datasetService;

    @Parameter
    TranslatorService translatorService;

    @Parameter
    private OpService opService;

    final private String defaultStringFile = "<Use filename below>";

    @Parameter(
            label = "Use open image (optional)",
            choices = { defaultStringFile },
            persist = false,
            description = "If you want to use an open dataset select it here, otherwise " +
                    "select '" + defaultStringFile + "' and select a file below."
    )
    private String inputImage = defaultStringFile;

    @Parameter(label = "Open from disk (optional)", required = false, persist = false)
    private File inputFile;

    @Parameter(label = "Load whole file into memory")
    private boolean bufferInput = true;

    @Parameter(label = "Display file (if opened from disk)")
    private boolean displayInput = false;

    @Parameter(choices = {"X", "Y"}, label = "Axis parallel to the lightsheet laser (and stripes)")
    private String laserAxis;

    @Parameter(label = "Process only current slice (for stack)")
    private boolean currentSlice = true;

    @Parameter(label = "Process slice range (or all slices - for stack)")
    private boolean doSliceRange = false;

    @Parameter(label = "Start slice (if range) - inclusive", min = "0")
    private int startSlice = 0;

    @Parameter(label = "Last slice (if range) - inclusive", min = "0")
    private int lastSlice = 0;

    @Parameter(label = "Display result")
    private boolean displayResult = true;

    @Parameter(label = "Save to disk (optional - if provided)", required = false, persist = false)
    private File outputFile;

    @Override
    public void initialize() {
        // update the dataset field with list of open datasets
        ArrayList<String> openDatasets = new ArrayList<>();

        List<Dataset> datasets = datasetService.getDatasets();
        final MutableModuleItem<String> inputImages = getInfo().getMutableInput(
                "inputImage", String.class);

        openDatasets.add(defaultStringFile);
        for (Dataset dataset : datasets) {
            openDatasets.add(dataset.getName());
        }

        inputImages.setChoices(openDatasets);
    }

    @Override
    public void run() {
        Dataset resultF;
        Dataset resultUInt;
        Dataset dataset;
        SCIFIOImgPlus<?> imgIn;
        final long[] zRange = {0, 0};

        // input dataset
        dataset = getUserDataset();
        if (dataset == null)
            return;
        try {
            // ensure we opened it with scifio
            imgIn = getImgPlus(dataset);
        } catch (final ClassCastException e) {
            log.error(e);
            return;
        }

        // requested range
        computeZExtent(dataset, zRange);

        // dataset for result
        resultF = create(
                dataset, new FloatType(),
                zRange[1] - zRange[0] + 1, imgIn, datasetService
        );

        processDataset(dataset, resultF, zRange);
//        resultUInt.setImgPlus(new ImgPlus<>(opService.convert().uint16((IterableInterval) resultF.getImgPlus())));

        if (displayResult) {
            ui.show(resultF);
        }

        saveDataset(resultF);
    }

    private Dataset getUserDataset() {
        Dataset dataset = null;

        // if no open dataset specified, open file
        if (inputImage.equals(defaultStringFile)) {
            if (inputFile == null) {
                log.error("No input file provided");
                return null;
            }

            SCIFIOConfig config = new SCIFIOConfig();
            config.enableBufferedReading(bufferInput);
            try {
                dataset = datasetIOService.open(inputFile.getAbsolutePath(), config);
            } catch (final IOException e) {
                log.error(e);
                return null;
            }
            if (displayInput)
                ui.show(dataset);
        } else {
            // otherwise use open dataset
            List<Dataset> datasets = datasetService.getDatasets();
            for (Dataset datasetItem : datasets) {
                if (datasetItem.getName().equals(inputImage)) {
                    dataset = datasetItem;
                    break;
                }
            }
            if (dataset == null)
            {
                log.error(new IllegalArgumentException("Can't find dataset by name"));
                return null;
            }
        }

        return dataset;
    }

    /** Get the z stack range requested by user, returning a valid range
     * for the dataset.
     */
    private void computeZExtent(Dataset dataset, long[] zRange) {
        if (currentSlice) {
            if (idService == null) {
                log.error("No display, but current slice requested");
                return;
            }
            if (dataset.dimensionIndex(Axes.Z) >= 0)
                zRange[0] = zRange[1] = idService.getActivePosition().getLongPosition(dataset.dimensionIndex(Axes.Z));
        } else if (doSliceRange) {
            zRange[0] = startSlice;
            zRange[1] = lastSlice;
        } else {
            zRange[1] = Math.max(dataset.getDepth() - 1, 0);
        }

        zRange[0] = Math.min(Math.max(zRange[0], 0), dataset.getDepth() - 1);
        zRange[1] = Math.min(Math.max(zRange[0], zRange[1]), dataset.getDepth() - 1);
    }

    /** Save the dataset to the file, if requested by user
     */
    private void saveDataset(Dataset dataset) {
        if (outputFile != null && !outputFile.getAbsolutePath().isEmpty()) {
            SCIFIOConfig config = new SCIFIOConfig();
            try {
                datasetIOService.save(dataset, outputFile.getAbsolutePath(), config);
            } catch (final IOException e) {
                log.error(e);
            }
        }
    }

    @SuppressWarnings("rawtypes")
    private void processDataset(final Dataset src, Dataset dst, long[] zRange) {
        final ImgPlus<? extends RealType<?>> srcImg = src.getImgPlus();
        final ImgPlus<? extends RealType<?>> dstImg = dst.getImgPlus();

        final long[] min = new long[src.numDimensions()];
        final long[] max = new long[src.numDimensions()];

        final int indexT, indexC, indexX, indexY, indexZ;
        indexT = src.dimensionIndex(Axes.TIME);
        indexC = src.dimensionIndex(Axes.CHANNEL);
        indexX = src.dimensionIndex(Axes.X);
        indexY = src.dimensionIndex(Axes.Y);
        indexZ = src.dimensionIndex(Axes.Z);

        final long sizeT = src.dimension(Axes.TIME);
        final long sizeC = src.dimension(Axes.CHANNEL);

        if (indexX >= 0){
            min[indexX] = 0;
            max[indexX] = src.dimension(Axes.X) - 1;
        }
        if (indexY >= 0){
            min[indexY] = 0;
            max[indexY] = src.dimension(Axes.Y) - 1;
        }

        for (long t = 0; t < sizeT; t++) {
            if (indexT >= 0)
                min[indexT] = max[indexT] = t;
            for (long c = 0; c < sizeC; c++) {
                if (indexC >= 0)
                    min[indexC] = max[indexC] = c;
                for (long z = zRange[0]; z <= zRange[1]; z++) {
                    if (indexZ >= 0)
                        min[indexZ] = max[indexZ] = z;

                    logSlice(src, dst, zRange, min, max);
                    computeParallel(srcImg, dstImg, min, max);
                    // computeSerial(src, srcImg, dstImg, min, max);
                }
            }
        }
    }

    private void logSlice(Dataset src, Dataset dst, long[] zRange, long[] min, long[] max) {
        Object[] items = {
                src.numDimensions(), dst.numDimensions(),

                src.dimension(Axes.TIME), src.dimension(Axes.CHANNEL),
                src.dimension(Axes.Z), src.dimension(Axes.X),
                src.dimension(Axes.Y),

                src.dimensionIndex(Axes.TIME), src.dimensionIndex(Axes.CHANNEL),
                src.dimensionIndex(Axes.Z), src.dimensionIndex(Axes.X),
                src.dimensionIndex(Axes.Y),

                String.join(", ", Arrays.toString(zRange)),
                String.join(", ", Arrays.toString(min)),
                String.join(", ", Arrays.toString(max))
        };
        String[] strings = Arrays.stream(items).map(Object::toString).toArray(String[]::new);
        log.info(String.join(", ", strings));
    }

    private void computeParallel(
            ImgPlus<? extends RealType<?>> srcImg,
            ImgPlus<? extends RealType<?>> dstImg, long[] min, long[] max) {
        final IterableInterval<FloatType> srcImgF = opService.convert().float32(
                (IterableInterval) Views.interval(srcImg, min, max));
        final IterableInterval<FloatType> dstImgF = (IterableInterval) Views.interval(dstImg, min, max);
        opService.math().multiply(dstImgF, srcImgF, srcImgF);
    }

    private void computeSerial(
            Dataset src, ImgPlus<? extends RealType<?>> srcImg,
            ImgPlus<? extends RealType<?>> dstImg, long[] min, long[] max
    ) {
        final long[] pos = new long[src.numDimensions()];
        RandomAccess<? extends RealType<?>> ra = Views.interval(srcImg, min, max).randomAccess();
        Cursor<? extends RealType<?>> cursor = Views.interval(dstImg, min, max).localizingCursor();

        while (cursor.hasNext()) {
            cursor.fwd();
            cursor.localize(pos);
            ra.setPosition(pos);
            final double sum = Math.pow(ra.get().getRealDouble(), 2);
            cursor.get().setReal(sum);
        }
        log.info(String.join(", ", Arrays.toString(pos)));

    }

    /** Creates a new dataset of the specified type, based on the input dataset
     * and metadata, if provided.
     */
    private static <T extends RealType<T> & NativeType<T>> Dataset create(
            final Dataset d, final T type, long zExtent, SCIFIOImgPlus<?> imgIn,
            DatasetService datasetService)
    {
        final int dimCount = d.numDimensions();
        final long[] dims = new long[d.numDimensions()];
        final AxisType[] axes = new AxisType[dimCount];

        d.dimensions(dims);
        int index = d.dimensionIndex(Axes.Z);
        if (index >= 0)
            dims[index] = zExtent;

        for (int i = 0; i < dimCount; i++) {
            axes[i] = d.axis(i).type();
        }

        Dataset result = datasetService.create(type, dims, "result", axes);
        SCIFIOImgPlus<?> img = getImgPlus(result);
        if (imgIn != null) {
            img.setMetadata(imgIn.getMetadata());
            img.setImageMetadata(imgIn.getImageMetadata());
        }

        return result;
    }

    /** Gets a SCIFIOImgPlus image object representing the dataset.
     */
    private static SCIFIOImgPlus<?> getImgPlus(Dataset data) {
        ImgPlus<?> imp = data.getImgPlus();

        if (imp instanceof SCIFIOImgPlus) {
            return  (SCIFIOImgPlus<?>) imp;
        }

        return new SCIFIOImgPlus<>(imp);
    }

    public static void main(final String... args) throws Exception {
        // Create the ImageJ application context with all available services
        final ImageJ ij = new ImageJ();
        ij.ui().showUI();
//        ij.launch(args);

        // Launch the "Add Two Datasets" command right away.
        ij.command().run(DeStripe.class, true);
    }

}