Demo_Pixel_Replication.java

import ij.*;
import ij.plugin.filter.PlugInFilter;
import ij.process.*;
import ij.gui.*;
import java.awt.*;
import java.awt.image.*;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Random;
import java.util.Collections;
import java.util.List;

import org.opencv.core.Core;
import org.opencv.core.Mat;
import org.opencv.core.CvType;
import org.opencv.ml.EM;
import org.opencv.imgproc.Imgproc;

public class Demo_Pixel_Replication implements PlugInFilter {

    int k;
    Random rnd;
    int Replicates = 5;
    int MaxIter = 100;

    // foreground points in ip
    ArrayList<Point> UniqPts;
    // pixel-replicated points in ip
    ArrayList<Point> PrPts;

    static{ System.loadLibrary(Core.NATIVE_LIBRARY_NAME); }

    public Demo_Pixel_Replication() {
        UniqPts = new ArrayList<>();
        PrPts = new ArrayList<>();
        rnd = new Random();
        Core.setRNGSeed(rnd.nextInt());
    }

    public int setup(String arg, ImagePlus imp) {
        if (IJ.versionLessThan("1.37j"))
            return DONE;
        else    
            return DOES_ALL+DOES_STACKS+SUPPORTS_MASKING;
    }

    public boolean showDialog() {
        String[] kChoices = {"2", "3", "4", "5", "6"};

        GenericDialog gd = new GenericDialog("Number of ellipses");
        gd.addChoice("Number of ellipses:", kChoices, kChoices[2]);
        gd.addNumericField("Number of replicates", Replicates, 0);
        gd.addNumericField("Maximum number of iterations", MaxIter, 0);

        gd.showDialog();
        if (gd.wasCanceled()) {
            return false;
        } else {
            k = gd.getNextChoiceIndex() + 2;
            Replicates = (int) gd.getNextNumber();
            MaxIter = (int) gd.getNextNumber();
            return true;
        }
    }

    public void run(ImageProcessor ip) {
        if (!showDialog())
            return;

        Mat ipMat = ip2mat(ip);
        Rectangle r = ip.getRoi();

        // use distance transform to populate UniqPts and PrPts
        Mat distMat = DistTransform(ipMat);

        for (int y=r.y; y<(r.y+r.height); y++) {
            for (int x=r.x; x<(r.x+r.width); x++) {
                double[] val = distMat.get(y, x);
                long n = (int) Math.round(val[0]);
                if (n > 0) {
                    Point p = new Point(x, y);
                    UniqPts.add(p);
                    long rep = Math.round((double)n);
                    for (int i = 0; i < rep; i++) {
                        PrPts.add(p);
                    }
                }
            }
        }

        // IJ.log(String.format("%d unique points, %d pixel rep points", UniqPts.size(), PrPts.size()));

        // compute the gmm
        EM em = computeGMM(k);

        // create a new image showing the partition
        ImageProcessor ipOutPr = ip.duplicate();
        Mat PrMat = new Mat(1, 2, CvType.CV_32FC1);
        int clrOffset = 256 / k;
        for (Point p : UniqPts) {
            PrMat.put(0, 0, p.x);
            PrMat.put(0, 1, p.y);
            double[] pr = em.predict(PrMat);
            int val = ((int) pr[1] + 1) * clrOffset - 1;
            ipOutPr.set((int) p.x, (int) p.y, val);
        }

        ImagePlus segIP = new ImagePlus("Segmentation", ipOutPr);
        segIP.show();

        // draw ellipses around the regions found by the gmms
        Mat means = em.getMat("means");
        List<Mat> covs = em.getMatVector("covs");
        Overlay ellipses = new Overlay();
        for (int i = 0; i < k; i++) {
            // compute eigenvalues and eigenvectors of the covariance matrix
            Mat eigVal = new Mat();
            Mat eigVec = new Mat();
            int lo, hi;
            Core.eigen(covs.get(i), true, eigVal, eigVec);
            double[] valIdx = new double[2];
            eigVal.get(0, 0, valIdx);
            if (valIdx[0] > valIdx[1]) {
                hi = 0;
                lo = 1;
            } else {
                hi = 1;
                lo = 0;
            }

            double A = Math.sqrt(valIdx[hi] * 20 / 3);
            double B = Math.sqrt(valIdx[lo] * 20 / 3);

            EllipseRoi elRoi = makeEllipseRoi(means.row(i), eigVec.row(hi), A, B);
            elRoi.setStrokeWidth(2);
            elRoi.setStrokeColor(Color.red);
            ellipses.add(elRoi);
        }
        segIP.setOverlay(ellipses);
    }

    private Mat ip2mat(ImageProcessor src) {
        BufferedImage bi = src.getBufferedImage();
        return bufferedImageToMat(bi);
    }

    public Mat bufferedImageToMat(BufferedImage bi) {
        Mat mat = new Mat(bi.getHeight(), bi.getWidth(), CvType.CV_8U);
        byte[] data = ((DataBufferByte) bi.getRaster().getDataBuffer()).getData();
        mat.put(0, 0, data);
        return mat;
    }

    private Mat DistTransform(Mat matIn) {
        Mat matOut = new Mat(matIn.size(), CvType.CV_32F);
        Imgproc.distanceTransform(matIn, matOut, Imgproc.CV_DIST_L2, 3);
        return matOut;
    }

    private Mat PrPts2Mat(ArrayList<Point> PrPts, double pct) {
        ArrayList<Point> ShufPts = new ArrayList<Point>(PrPts);
        Collections.shuffle(ShufPts);
        return PrPts2Mat(new ArrayList<Point>(ShufPts.subList(0, (int) (ShufPts.size() * pct))));
    }

    private Mat PrPts2Mat(ArrayList<Point> PrPts) {
        Mat emMat = new Mat(PrPts.size(), 2, CvType.CV_32FC1);
        for (int i = 0; i < PrPts.size(); i++) {
            Point p = PrPts.get(i);
            float[] vals = { p.x, p.y };
            emMat.put(i, 0, vals);
        }

        return emMat;
    }

    private EM computeGMM(int k) {
        EM bestEm = new EM();
        double bestScore = -1e300;
        int bestI = 0;

        for (int i = 1; i <= Replicates; i++) {
            IJ.log(String.format("replicate %d", i));
            EM em = new EM();
            Mat logLikelihoods = new Mat(PrPts.size(), 1, CvType.CV_64FC1);
            Mat labels = new Mat(PrPts.size(), 1, CvType.CV_32FC1);
            Mat covs0 = makeInitCovs(k);
            Mat weights0 = makeInitWeights(k);
            em.setInt("nclusters", k);
            em.setInt("covMatType", EM.COV_MAT_GENERIC);
            em.setInt("maxIters", MaxIter);
            em.setDouble("epsilon", 1e-6);

            // Mat emMat = PrPts2Mat(PrPts, 0.10);
            Mat emMat = PrPts2Mat(PrPts);
            Mat initMeans = randomInitMeans(k);
            if (!em.trainE(emMat, initMeans, covs0, weights0, logLikelihoods, labels, new Mat()))
                IJ.log("trainE() returned false!");
            double score = addUp(logLikelihoods.col(0));
            IJ.log(String.format("score = %f", score));
            if (score > bestScore) {
                bestEm = em;
                bestScore = score;
                bestI = i;
            }
        }

        // IJ.log(String.format("best replicate was %d", bestI));
        return bestEm;
    }

    private double addUp(Mat m) {
        double total = 0;
        double[] data = new double[(int) m.total()];
        m.get(0, 0, data);

        for (double val : data)
            total += val;

        return total;
    }

    private Mat makeInitCovs(int k) {
        return new Mat();
    }

    private Mat makeInitWeights(int k) {
        return new Mat();
    }

    private EllipseRoi makeEllipseRoi(Mat center, Mat unitVec, double A, double B) {
        Mat end1 = new Mat();
        Mat end2 = new Mat();

        double[] data1 = new double[2];
        double[] data2 = new double[2];

        Core.scaleAdd(unitVec, A, center, end1);
        Core.scaleAdd(unitVec, -A, center, end2);
        double aspectRatio = B / A;

        end1.get(0, 0, data1);
        end2.get(0, 0, data2);

        return new EllipseRoi(data1[0], data1[1], data2[0], data2[1], aspectRatio);
    }

    // Initialize the means by picking 2 distinct points at random from UniqPts.
    // We're using a simple naive algorithm since we assume k << n.
    private Mat randomInitMeans(int k) {
        HashSet<Integer> used = new HashSet<Integer>();
        Mat means = new Mat(k, 2, CvType.CV_64F);
        double[] meansData = new double[k * 2];

        int i = 0;
        int n = 0;
        while (i < k) {
            int j = rnd.nextInt(UniqPts.size());
            if (!used.contains(j)) {
                Point p = UniqPts.get(j);
                meansData[n++] = p.x;
                meansData[n++] = p.y;
                used.add(j);
                // IJ.log(String.format("random pt %d: (%d,%d)", i, p.x, p.y));
                i++;
            }
        }

        means.put(0, 0, meansData);

        return means;
    }
}