Add (interactive) background correction

render-app/src/main/java/org/janelia/alignment/filter/ShadingCorrectionFilter.java

@@ -0,0 +1,102 @@
+package org.janelia.alignment.filter;
+
+import ij.process.ImageProcessor;
+import org.janelia.alignment.filter.emshading.FourthOrderShading;
+import org.janelia.alignment.filter.emshading.QuadraticShading;
+import org.janelia.alignment.filter.emshading.ShadingModel;
+
+import java.util.Arrays;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+public class ShadingCorrectionFilter implements Filter {
+
+    public enum ShadingCorrectionMethod {
+        QUADRATIC(QuadraticShading::new),
+        FOURTH_ORDER(FourthOrderShading::new);
+
+        private final Function<double[], ShadingModel> modelFactory;
+
+        ShadingCorrectionMethod(final Function<double[], ShadingModel> modelFactory) {
+            this.modelFactory = modelFactory;
+        }
+
+        public ShadingModel create(final double[] coefficients) {
+            return modelFactory.apply(coefficients);
+        }
+
+        public static ShadingCorrectionMethod fromString(final String method) {
+            for (final ShadingCorrectionMethod shadingCorrectionMethod : values()) {
+                if (shadingCorrectionMethod.name().equalsIgnoreCase(method)) {
+                    return shadingCorrectionMethod;
+                }
+            }
+            throw new IllegalArgumentException("Unknown shading correction method: " + method);
+        }
+
+        public static ShadingCorrectionMethod forModel(final ShadingModel model) {
+            if (model instanceof QuadraticShading) {
+                return QUADRATIC;
+            } else if (model instanceof FourthOrderShading) {
+                return FOURTH_ORDER;
+            } else {
+                throw new IllegalArgumentException("Unknown shading model class: " + model.getClass().getName());
+            }
+        }
+    }
+
+    private ShadingCorrectionMethod correctionMethod;
+    private double[] coefficients;
+
+    // empty constructor required to create instances from specifications
+    @SuppressWarnings("unused")
+    public ShadingCorrectionFilter() {
+    }
+
+    public ShadingCorrectionFilter(final ShadingModel model) {
+        this(ShadingCorrectionMethod.forModel(model), model.getCoefficients());
+    }
+
+    public ShadingCorrectionFilter(final ShadingCorrectionMethod correctionMethod, final double[] coefficients) {
+        this.correctionMethod = correctionMethod;
+        this.coefficients = coefficients;
+    }
+
+    @Override
+    public void init(final Map<String, String> params) {
+        this.correctionMethod = ShadingCorrectionMethod.fromString(Filter.getStringParameter("correctionMethod", params));
+        final String[] rawCoefficients = Filter.getCommaSeparatedStringParameter("coefficients", params);
+        this.coefficients = Arrays.stream(rawCoefficients).mapToDouble(Double::parseDouble).toArray();
+    }
+
+    @Override
+    public Map<String, String> toParametersMap() {
+        final Map<String, String> map = new LinkedHashMap<>();
+        map.put("correctionMethod", correctionMethod.name());
+        map.put("coefficients", Arrays.stream(coefficients).mapToObj(String::valueOf).collect(Collectors.joining(",")));
+        return map;
+    }
+
+    @Override
+    public void process(final ImageProcessor ip, final double scale) {
+        // transform pixel coordinates into [-1, 1] x [-1, 1]
+        final double scaleX = ip.getWidth() / 2.0;
+        final double scaleY = ip.getHeight() / 2.0;
+        final ShadingModel shadingModel = correctionMethod.create(coefficients);
+
+        // subtract shading model from image
+        final double[] location = new double[2];
+        for (int i = 0; i < ip.getWidth(); i++) {
+            for (int j = 0; j < ip.getHeight(); j++) {
+                location[0] = ShadingModel.scaleCoordinate(i, scaleX);
+                location[1] = ShadingModel.scaleCoordinate(j, scaleY);
+                shadingModel.applyInPlace(location);
+
+                final double value = ip.getPixelValue(i, j) - location[0];
+                ip.putPixelValue(i, j, value);
+            }
+        }
+    }
+}

render-app/src/main/java/org/janelia/alignment/filter/emshading/FourthOrderShading.java

@@ -0,0 +1,74 @@
+package org.janelia.alignment.filter.emshading;
+
+import java.util.List;
+
+
+/**
+ * A 4th order model for shading correction in 2D slices of EM data.
+ * To ensure that the model is concave, no 3rd order terms are included.
+ */
+public class FourthOrderShading extends ShadingModel {
+
+	public FourthOrderShading() {
+		super();
+	}
+
+	public FourthOrderShading(final double[] coefficients) {
+		super(coefficients);
+	}
+
+	public FourthOrderShading(final FourthOrderShading shading) {
+		super(shading.getCoefficients());
+	}
+
+
+	@Override
+	protected int nCoefficients() {
+		return 9;
+	}
+
+	@Override
+	protected void fillRowA(final double[] rowA, final double x, final double y) {
+		final double xx = x * x;
+		final double yy = y * y;
+		rowA[0] = 1;
+		rowA[1] = y;
+		rowA[2] = x;
+		rowA[3] = yy;
+		rowA[4] = x * y;
+		rowA[5] = xx;
+		rowA[6] = yy * yy;
+		rowA[7] = xx * yy;
+		rowA[8] = xx * xx;
+	}
+
+	@Override
+	protected List<String> coefficientNames() {
+		return List.of("1", "y", "x", "y^2", "xy", "x^2", "y^4", "x^2 y^2", "x^4");
+	}
+
+	@Override
+	public void applyInPlace(final double[] location) {
+		final double x = location[0];
+		final double y = location[1];
+		final double xx = x * x;
+		final double yy = y * y;
+		final double[] coefficients = getCoefficients();
+		final double result = coefficients[8] * xx * xx
+				+ coefficients[7] * xx * yy
+				+ coefficients[6] * yy * yy
+				+ coefficients[5] * xx
+				+ coefficients[4] * x * y
+				+ coefficients[3] * yy
+				+ coefficients[2] * x
+				+ coefficients[1] * y
+				+ coefficients[0];
+		location[0] = result;
+		location[1] = 0.0;
+	}
+
+	@Override
+	public FourthOrderShading copy() {
+		return new FourthOrderShading(this);
+	}
+}

render-app/src/main/java/org/janelia/alignment/filter/emshading/QuadraticShading.java

@@ -0,0 +1,63 @@
+package org.janelia.alignment.filter.emshading;
+
+import java.util.List;
+
+
+/**
+ * A quadratic model for shading correction in 2D slices of EM data.
+ */
+public class QuadraticShading extends ShadingModel {
+
+	public QuadraticShading() {
+		super();
+	}
+
+	public QuadraticShading(final double[] coefficients) {
+		super(coefficients);
+	}
+
+	public QuadraticShading(final QuadraticShading shading) {
+		super(shading.getCoefficients());
+	}
+
+
+	@Override
+	protected int nCoefficients() {
+		return 6;
+	}
+
+	@Override
+	protected void fillRowA(final double[] rowA, final double x, final double y) {
+		rowA[0] = 1;
+		rowA[1] = y;
+		rowA[2] = x;
+		rowA[3] = y * y;
+		rowA[4] = x * y;
+		rowA[5] = x * x;
+	}
+
+	@Override
+	protected List<String> coefficientNames() {
+		return List.of("1", "y", "x", "y^2", "xy", "x^2");
+	}
+
+	@Override
+	public void applyInPlace(final double[] location) {
+		final double x = location[0];
+		final double y = location[1];
+		final double[] coefficients = getCoefficients();
+		final double result = coefficients[5] * x * x
+				+ coefficients[4] * x * y
+				+ coefficients[3] * y * y
+				+ coefficients[2] * x
+				+ coefficients[1] * y
+				+ coefficients[0];
+		location[0] = result;
+		location[1] = 0.0;
+	}
+
+	@Override
+	public QuadraticShading copy() {
+		return new QuadraticShading(this);
+	}
+}

render-app/src/main/java/org/janelia/alignment/filter/emshading/ShadingModel.java

@@ -0,0 +1,152 @@
+package org.janelia.alignment.filter.emshading;
+
+import mpicbg.models.AbstractModel;
+import mpicbg.models.IllDefinedDataPointsException;
+import mpicbg.models.NotEnoughDataPointsException;
+import mpicbg.models.PointMatch;
+import org.ejml.data.DMatrixRMaj;
+import org.ejml.dense.row.factory.LinearSolverFactory_DDRM;
+import org.ejml.interfaces.linsol.LinearSolverDense;
+
+import java.io.Serializable;
+import java.util.Collection;
+import java.util.List;
+
+
+/**
+ * An abstract base model for shading correction in 2D slices of EM data.
+ */
+public abstract class ShadingModel extends AbstractModel<ShadingModel> implements Serializable {
+
+	private final double[] coefficients;
+
+	protected abstract int nCoefficients();
+	protected abstract void fillRowA(final double[] rowA, final double x, final double y);
+	protected abstract List<String> coefficientNames();
+
+
+	public ShadingModel() {
+		coefficients = new double[nCoefficients()];
+	}
+
+	public ShadingModel(final double[] coefficients) {
+		this();
+		System.arraycopy(coefficients, 0, this.coefficients, 0, nCoefficients());
+	}
+
+	public double[] getCoefficients() {
+		return coefficients;
+	}
+
+	/**
+	 * Normalize the model so that it approximately integrates to zero over [-1, 1] x [-1, 1].
+	 * This assumes that the constant term is stored in the first coefficient.
+	 */
+	public void normalize() {
+		final int N = 10 * nCoefficients();
+		final double[] location = new double[2];
+		double avg = 0;
+
+		// approximate integral over [-1, 1] x [-1, 1] by averaging the function values at a grid of points
+		for (int i = 0; i < N; i++) {
+			for (int j = 0; j < N; j++) {
+				location[0] = 2.0 * i / N - 1.0;
+				location[1] = 2.0 * j / N - 1.0;
+				applyInPlace(location);
+				avg += location[0];
+			}
+		}
+
+		avg /= (N * N);
+		coefficients[0] -= avg;
+	}
+
+	/**
+	 * Scale the given coordinate to the model coordinates (in [-1, 1]).
+	 * @param coordinate the coordinate to scale
+	 * @param scale the scale factor (usually half the image size in the respective dimension)
+	 * @return the scaled coordinate in the range [-1, 1]
+	 */
+	public static double scaleCoordinate(final double coordinate, final double scale) {
+		return (coordinate - scale) / scale;
+	}
+
+	@Override
+	public int getMinNumMatches() {
+		return nCoefficients();
+	}
+
+	@Override
+	public <P extends PointMatch> void fit(final Collection<P> matches) throws NotEnoughDataPointsException, IllDefinedDataPointsException {
+		final DMatrixRMaj ATA = new DMatrixRMaj(nCoefficients(), nCoefficients(), true, new double[nCoefficients() * nCoefficients()]);
+		final DMatrixRMaj ATb = new DMatrixRMaj(nCoefficients(), 1);
+
+		final double[] rowA = new double[nCoefficients()];
+
+		for (final P match : matches) {
+			final double x = match.getP1().getL()[0];
+			final double y = match.getP1().getL()[1];
+			final double z = match.getP2().getL()[0];
+
+			// compute one row of the least-squares matrix A
+			fillRowA(rowA, x, y);
+
+			// update upper triangle of A^T * A
+			for (int i = 0; i < nCoefficients(); i++) {
+				for (int j = i; j < nCoefficients(); j++) {
+					ATA.data[i * nCoefficients() + j] += rowA[i] * rowA[j];
+				}
+			}
+
+			// update right-hand side A^T * b
+			for (int i = 0; i < nCoefficients(); i++) {
+				ATb.data[i] += rowA[i] * z;
+			}
+		}
+
+		// set up Cholesky decomposition for A^T * A x = A^T * b (only upper triangle of A^T * A is used)
+		final LinearSolverDense<DMatrixRMaj> solver = LinearSolverFactory_DDRM.chol(nCoefficients());
+		solver.setA(ATA);
+
+		// coefficients are modified in place
+		final DMatrixRMaj x = new DMatrixRMaj(nCoefficients(), 1);
+		x.setData(coefficients);
+		solver.solve(ATb, x);
+	}
+
+	@Override
+	public double[] apply(final double[] location) {
+		final double[] result = location.clone();
+		applyInPlace(result);
+		return result;
+	}
+
+	@Override
+	public void set(final ShadingModel model) {
+		System.arraycopy(model.getCoefficients(), 0, coefficients, 0, nCoefficients());
+	}
+
+	@Override
+	public String toString() {
+		final StringBuilder sb = new StringBuilder(this.getClass().getSimpleName()).append("{");
+
+		if (coefficients[nCoefficients() - 1] < 0) {
+			sb.append("-");
+		}
+
+		for (int i = nCoefficients() - 1; i > 0; i--) {
+			sb.append(String.format("%.2f", Math.abs(coefficients[i])))
+					.append(" ")
+					.append(coefficientNames().get(i));
+
+			if (coefficients[i - 1] >= 0) {
+				sb.append(" + ");
+			} else {
+				sb.append(" - ");
+			}
+		}
+
+		sb.append(String.format("%.2f", Math.abs(coefficients[0])));
+		return sb.append("}").toString();
+	}
+}