diff --git a/render-ws-spark-client/src/main/java/org/janelia/render/client/spark/multisem/Wafer6061Inpainter.java b/render-ws-spark-client/src/main/java/org/janelia/render/client/spark/multisem/Wafer6061Inpainter.java
index b05bc3090..938297519 100644
--- a/render-ws-spark-client/src/main/java/org/janelia/render/client/spark/multisem/Wafer6061Inpainter.java
+++ b/render-ws-spark-client/src/main/java/org/janelia/render/client/spark/multisem/Wafer6061Inpainter.java
@@ -2,23 +2,50 @@
import com.beust.jcommander.Parameter;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.function.Consumer;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import bdv.export.Downsample;
import net.imglib2.Cursor;
-import net.imglib2.Interval;
+import net.imglib2.IterableInterval;
+import net.imglib2.KDTree;
import net.imglib2.RandomAccess;
import net.imglib2.RandomAccessibleInterval;
+import net.imglib2.RealPoint;
import net.imglib2.img.Img;
import net.imglib2.img.array.ArrayImgs;
+import net.imglib2.img.basictypeaccess.AccessFlags;
+import net.imglib2.img.cell.CellGrid;
+import net.imglib2.neighborsearch.KNearestNeighborSearchOnKDTree;
import net.imglib2.type.numeric.integer.UnsignedByteType;
-import net.imglib2.util.Intervals;
+import net.imglib2.type.numeric.real.DoubleType;
import net.imglib2.view.Views;
+
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.broadcast.Broadcast;
+import org.janelia.alignment.multisem.MultiSemUtilities;
+import org.janelia.alignment.spec.Bounds;
+import org.janelia.alignment.spec.TileBounds;
import org.janelia.alignment.util.Grid;
import org.janelia.render.client.ClientRunner;
+import org.janelia.render.client.RenderDataClient;
import org.janelia.render.client.parameter.CommandLineParameters;
import org.janelia.render.client.spark.LogUtilities;
+import org.janelia.saalfeldlab.n5.DataBlock;
import org.janelia.saalfeldlab.n5.DatasetAttributes;
import org.janelia.saalfeldlab.n5.N5Reader;
import org.janelia.saalfeldlab.n5.N5Writer;
@@ -27,376 +54,957 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import java.io.Serializable;
-import java.util.Arrays;
-import java.util.List;
-import java.util.Objects;
-
/**
- * Class for inpainting small gaps between tiles in the wafer 60/61 dataset.
+ * Spark client for filling holes in the tissue of wafer 60/61 N5 volumes.
+ *
+ * Inputs are (a) an N5 tissue volume (only non-empty blocks stored), (b) an N5 mask marking where image data is
+ * present ({@code mask > 0}) vs. missing ({@code mask == 0}), and (c) the acquisition xlog zarr. The ROI point cloud
+ * is built in the tissue's s0 voxel frame from the ALIGNED render stack (so montage stitching is accounted for): the
+ * first layer's tile bounds are fetched from the render web service on the driver, and each tile center (render world
+ * pixels) is mapped to the voxel frame with {@code voxel = center - translate} (the neuroglancer group-level
+ * {@code translate}, i.e. the stack bounding-box min). Each tile carries the {@code distance_roi} of its SFOV, read
+ * from the xlog for the slab whose {@code id_serial} equals {@code serial}: the tile's mfov (0-based) maps to the xlog
+ * mfov axis as {@code render_mfov + 5} (the axis reserves 5 leading always-NaN rows), and its sfov number (the 1-based
+ * {@code _s##} field) maps to the 0-based xlog sfov axis as {@code _s## - 1}.
*
- * The regions to inpaint are determined by looking at mask pixels: if an unmasked pixel is encountered,
- * the algorithm looks at pairs of pixels `stepSize` away in the x and y directions to find non-masked pixels.
- * If such a pair is found, the pixel is most likely in a narrow gap between tiles and needs inpainting. The
- * inpainting is done by averaging the image data from the adjacent pixels in the z direction. If only one of the
- * image values in z is available, that value is used. If neither is available, the pixel is set to 0.
+ * The client processes the full-resolution ({@code s0}) blocks of the tissue in parallel. For each block it first
+ * applies the cheap ROI-distance filter (interpolating {@code distance_roi} at the block center via inverse-distance
+ * weighting; blocks whose interpolated distance is {@code >= maxRoiDistance} are skipped without any I/O), then reads
+ * the block (skipping absent/empty blocks) and fills every pixel where the mask is 0 with the average of the tissue in
+ * the sections above and below (copying the single neighbor at the volume's z-boundary). Before overwriting a modified
+ * block, the original block is copied verbatim into a sibling {@code _backup} N5 container. Finally, only the pyramid
+ * blocks affected by the modified {@code s0} blocks are re-downsampled (again backing up the originals first), using
+ * the same per-block averaging as {@code N5DownsamplerSpark} so the pyramid stays consistent.
*/
public class Wafer6061Inpainter {
+ // Inverse-distance-weighting knobs for the ROI-distance interpolation; fixed for wafer 60/61.
+ private static final int IDW_K = 8;
+ private static final double IDW_POWER = 2.0;
+
+ // Tissue containers/datasets are named w_s_r (e.g. w61_s109_r00); the serial is parsed from that.
+ private static final Pattern SERIAL_IN_NAME = Pattern.compile("_s(\\d+)");
+
+ // xlog field layout (consumed by loadPointCloud). The xlog is a zarr: its arrays are stored C-order, but the
+ // n5-zarr reader REVERSES the axis order, so the imglib2 axis indices below are the reverse of the C-order shape.
+ // field C-order shape imglib2 axes (what this code sees) meaning
+ // id_serial (slab,) [slab] serial label per slab position
+ // distance_roi (slab, mfov, sfov) [sfov, mfov, slab] distance to ROI, um (scan-independent)
+ // The slab axis index is HARDCODED (below) rather than discovered by matching sizes (sizes are not guaranteed to
+ // be unique). After slicing the slab axis, distance_roi reduces to a 2-D (sfov, mfov) plane.
+ private static final int DIST_AXIS_SLAB = 2;
+
+ // The distance_roi mfov axis reserves 5 leading (always-NaN) positions before the first real mfov: real mfov data
+ // starts at row 5 for every wafer-60/61 slab (verified constant + contiguous across all 399 finite slabs). Render
+ // numbers each section's mfovs from 0, so the render mfov maps to the xlog mfov axis position as render_mfov + 5.
+ private static final int MFOV_ROW_OFFSET = 5;
+
private static final Logger LOG = LoggerFactory.getLogger(Wafer6061Inpainter.class);
public static class Parameters extends CommandLineParameters {
@Parameter(
names = "--n5Path",
- description = "Path to the N5 container containing the data and the mask.",
+ description = "Path to the N5 container holding the tissue and mask (local path or gs://...).",
required = true)
public String n5Path;
@Parameter(
names = "--dataset",
- description = "Name of the dataset to inpaint; assumed to be a multiscale pyramid, only s0 is inpainted.",
+ description = "Name of the tissue multiscale group; the full-resolution data is at /s0.",
required = true)
public String dataset;
@Parameter(
names = "--mask",
- description = "Name of the mask dataset. This is supposed be a binary uint8 mask covering the whole dataset.",
+ description = "Name of the binary uint8 mask dataset (same grid as /s0). Pass /s0 if the mask is itself a pyramid.",
required = true)
public String mask;
@Parameter(
- names = "--output",
- description = "Name of the dataset to write the inpainted data to. Only blocks that are inpainted are written. "
- + "If omitted, input blocks are overwritten.")
- public String output;
+ names = "--xlogPath",
+ description = "Path to the acquisition xlog zarr, e.g. xlog_wafer_61.zarr (local path or gs://...).",
+ required = true)
+ public String xlogPath;
@Parameter(
- names = "--inpaintingSize",
- description = "Rough size of the inpainting region in pixels. This is used to determine the regions to inpaint, so better to be too large than too small.",
- required = true)
- public int stepSize;
+ names = "--serial",
+ description = "Serial label (id_serial) of the section stored in this N5, resolved to the reference " +
+ "arrays' slab position. Defaults to the serial parsed from the dataset name " +
+ "(w_s_r, e.g. w61_s109_r00 -> 109).")
+ public long serial = -1;
+
+ @Parameter(
+ names = "--maxRoiDistance",
+ description = "Keep (inpaint) a block only if its interpolated distance_roi (microns) is less than this.")
+ public double maxRoiDistance = 10.0;
+
+ @Parameter(
+ names = "--backupPath",
+ description = "N5 container where original (overwritten) blocks are backed up. Defaults to a sibling _backup.n5.")
+ public String backupPath;
+ @Parameter(
+ names = "--dryRun",
+ description = "Only compute and log candidate / would-be-modified counts; do not write or back up anything.")
+ public boolean dryRun = false;
+
+ public String fullDataset() {
+ return dataset + "/s0";
+ }
+
+ public String getBackupPath() {
+ if (backupPath != null) {
+ return backupPath;
+ }
+ String p = n5Path;
+ while (p.endsWith("/")) {
+ p = p.substring(0, p.length() - 1);
+ }
+ if (p.endsWith(".n5")) {
+ return p.substring(0, p.length() - 3) + "_backup.n5";
+ }
+ return p + "_backup.n5";
+ }
public void validate() {
- if (stepSize <= 0) {
- throw new IllegalArgumentException("Inpainting size must be positive");
+ if (maxRoiDistance <= 0) {
+ throw new IllegalArgumentException("--maxRoiDistance must be positive");
+ }
+ if (serial < 0) {
+ serial = inferSerial(dataset);
}
}
- public String fullDataset() {
- return dataset + "/s0";
+ /** Parses the serial label from a dataset/container name like {@code .../w61_s109_r00} (basename {@code _s}). */
+ static long inferSerial(final String path) {
+ final Matcher matcher = SERIAL_IN_NAME.matcher(basename(path));
+ if (matcher.find()) {
+ return Long.parseLong(matcher.group(1));
+ }
+ throw new IllegalArgumentException("could not infer the serial from '" + path + "'; pass --serial explicitly");
+ }
+
+ /** Last path segment of a path, stripped of any leading group/parent path and trailing slashes. */
+ static String basename(final String path) {
+ String p = path;
+ while (p.endsWith("/")) {
+ p = p.substring(0, p.length() - 1);
+ }
+ final int slash = p.lastIndexOf('/');
+ return slash >= 0 ? p.substring(slash + 1) : p;
}
}
- private final Parameters param;
+ private final Parameters params;
- private ExtendedAttributes tissueAttributes;
- private ExtendedAttributes maskAttributes;
+ public Wafer6061Inpainter(final Parameters params) {
+ this.params = params;
+ }
+ public static void main(final String[] args) {
+ final ClientRunner clientRunner = new ClientRunner(args) {
+ @Override
+ public void runClient(final String[] args) throws Exception {
+
+ final Parameters parameters = new Parameters();
+ parameters.parse(args);
+ parameters.validate();
+
+ LOG.info("runClient: entry, parameters={}", parameters);
- public Wafer6061Inpainter(final Parameters parameters) {
- this.param = parameters;
+ final Wafer6061Inpainter client = new Wafer6061Inpainter(parameters);
+ client.run();
+ }
+ };
+ clientRunner.run();
}
- public void run() {
- final String output = param.output == null ? "input dataset" : "'" + param.output + "'";
- LOG.info("Inpainting dataset '{}' in '{}' using mask '{}' and writing to {}",
- param.dataset, param.n5Path, param.mask, output);
-
- // Read and cache some metadata of the tissue and mask datasets
- // Assume that the tissue is a multiscale pyramid / mask is a standalone dataset
- try (final N5Reader n5 = new N5Factory().openReader(N5Factory.StorageFormat.N5, param.n5Path)) {
- LOG.info("Reading metadata from {}", param.n5Path);
- tissueAttributes = ExtendedAttributes.read(n5, param.fullDataset(), param.dataset);
- maskAttributes = ExtendedAttributes.read(n5, param.mask, param.mask);
-
- if (param.output == null) {
- param.output = param.fullDataset();
- LOG.info("Output dataset equals input dataset. Overwriting blocks in the input dataset '{}'", param.output);
- } else if (n5.exists(param.output)) {
- throw new IllegalArgumentException("Dataset '" + param.output + "' is different from the input dataset and already exists. Stopping.");
- } else {
- LOG.info("Output dataset is '{}'. Creating new dataset.", param.output);
- try (final N5Writer n5Writer = new N5Factory().openWriter(N5Factory.StorageFormat.N5, param.n5Path)) {
- n5Writer.createDataset(param.output, tissueAttributes.attrs);
+ public void run() throws IOException {
+
+ // Read the tissue s0 metadata, the render target, the world->voxel offset, and discover the pyramid levels.
+ final double[] worldToVoxel;
+ final RenderTarget renderTarget;
+ final int numDimensions;
+ final List levels = new ArrayList<>();
+ final Map levelAttributes = new LinkedHashMap<>();
+ int[] downsampleFactors = null;
+ try (final N5Reader n5 = openN5Reader(params.n5Path)) {
+
+ final DatasetAttributes s0Attributes = n5.getDatasetAttributes(params.fullDataset());
+ if (s0Attributes == null) {
+ throw new IllegalArgumentException("tissue dataset " + params.fullDataset() + " does not exist");
+ }
+ numDimensions = s0Attributes.getNumDimensions();
+ if (numDimensions != 3) {
+ throw new IllegalArgumentException("expected a 3D tissue volume but " + params.fullDataset() +
+ " has " + numDimensions + " dimensions");
+ }
+ final DatasetAttributes maskAttributes = n5.getDatasetAttributes(params.mask);
+ if (maskAttributes == null) {
+ throw new IllegalArgumentException("mask dataset " + params.mask + " does not exist");
+ }
+
+ // The inpainter reads and writes one chunk at a time (no whole-volume open, so no accumulating cell
+ // cache). That requires z to be a single chunk, so that every z-1 / z+1 section the z-average needs lives
+ // inside the block; and it reads the mask by the tissue block's grid position, so the two datasets must
+ // share the same block grid and origin.
+ if (s0Attributes.getBlockSize()[2] < s0Attributes.getDimensions()[2]) {
+ throw new IllegalArgumentException(
+ "block-local inpainting requires z to be a single chunk, but " + params.fullDataset() +
+ " has blockSize[2]=" + s0Attributes.getBlockSize()[2] +
+ " < dimensions[2]=" + s0Attributes.getDimensions()[2]);
+ }
+ if (! Arrays.equals(maskAttributes.getBlockSize(), s0Attributes.getBlockSize()) ||
+ ! Arrays.equals(maskAttributes.getDimensions(), s0Attributes.getDimensions())) {
+ throw new IllegalArgumentException(
+ "mask grid " + Arrays.toString(maskAttributes.getDimensions()) + " @ " +
+ Arrays.toString(maskAttributes.getBlockSize()) + " must match tissue grid " +
+ Arrays.toString(s0Attributes.getDimensions()) + " @ " +
+ Arrays.toString(s0Attributes.getBlockSize()));
+ }
+
+ // Render service parameters (baseDataUrl / owner / project / stack) come straight from the group's
+ // "renderExport" metadata written by render's N5 export — the same attributes.json that holds the pyramid
+ // scales and translate — so they never have to be passed on the command line.
+ renderTarget = readRenderTarget(n5, params.dataset);
+
+ // World->voxel offset for placing render tile centers: the neuroglancer 'translate' (the stack
+ // bounding-box min in world pixels) is written on the multiscale GROUP, not on s0 (s0 only carries a
+ // sub-pixel centering transform). May be null here; loadPointCloud falls back to the render stack bounds.
+ worldToVoxel = readGroupTranslate(n5, params.dataset);
+
+ levels.add(params.fullDataset());
+ levelAttributes.put(params.fullDataset(), s0Attributes);
+ for (int scale = 1; ; scale++) {
+ final String levelDataset = params.dataset + "/s" + scale;
+ if (! n5.datasetExists(levelDataset)) {
+ break;
}
+ levels.add(levelDataset);
+ levelAttributes.put(levelDataset, n5.getDatasetAttributes(levelDataset));
+ }
+
+ // The relative per-level downsampling factor is read from the pyramid itself rather than passed in: the
+ // group's neuroglancer "scales" attribute lists the cumulative factor per level relative to s0, and these
+ // pyramids are built with a constant factor at every step (see DownsampleHelper / N5DownsamplerSpark), so
+ // scales[1] applies to all levels.
+ if (levels.size() > 1) {
+ downsampleFactors = readDownsamplingFactors(n5, params.dataset, numDimensions);
}
}
+ LOG.info("run: world->voxel translate={}, pyramid levels={}, downsampleFactors={}",
+ Arrays.toString(worldToVoxel), levels, Arrays.toString(downsampleFactors));
+
+ // Load the (small) per-slab ROI point cloud on the driver: positions from the aligned render stack, distances
+ // from the xlog. A single render request keeps the server load light; the cloud is then broadcast to executors.
+ final RenderDataClient renderClient =
+ new RenderDataClient(renderTarget.baseDataUrl, renderTarget.owner, renderTarget.project);
+ final PointCloud cloud = loadPointCloud(params, renderClient, renderTarget.stack, worldToVoxel);
+ LOG.info("run: loaded {} ROI reference points for serial {} (render {} {}/{}/{})",
+ cloud.size(), params.serial, renderTarget.baseDataUrl, renderTarget.owner, renderTarget.project,
+ renderTarget.stack);
+
+ // Create the backup container and mirror all datasets that might receive backups.
+ final String backupPath = params.getBackupPath();
+ if (! params.dryRun) {
+ try (final N5Writer backup = openN5Writer(backupPath)) {
+ for (final String levelDataset : levels) {
+ if (! backup.datasetExists(levelDataset)) {
+ backup.createDataset(levelDataset, levelAttributes.get(levelDataset));
+ }
+ }
+ }
+ LOG.info("run: originals will be backed up to {}", backupPath);
+ }
- final SparkConf conf = new SparkConf().setAppName("Wafer6061Inpainter");
+ final SparkConf conf = new SparkConf().setAppName(getClass().getSimpleName());
try (final JavaSparkContext sparkContext = new JavaSparkContext(conf)) {
- runWithSparkContext(sparkContext);
+ LOG.info("run: appId is {}", sparkContext.getConf().getAppId());
+ runWithSparkContext(sparkContext, cloud, levels, levelAttributes, backupPath, downsampleFactors);
}
}
- private void runWithSparkContext(final JavaSparkContext sparkContext) {
- // Find out which blocks need inpainting (i.e., find blocks that are neither all mask nor all void)
- final List maskBlocks = Grid.create(maskAttributes.attrs.getDimensions(), maskAttributes.attrs.getBlockSize());
- final JavaRDD maskRDD = sparkContext.parallelize(maskBlocks);
- final Broadcast maskAttributesBroadcast = sparkContext.broadcast(maskAttributes);
- final Broadcast paramBroadcast = sparkContext.broadcast(param);
- LOG.info("Filtering empty mask blocks from {} blocks", maskBlocks.size());
-
- final List nonHomogeneousMaskBlocks = maskRDD
- .map(block -> translateAndCheckHomogeneity(block,
- maskAttributesBroadcast.value().min,
- paramBroadcast.value()))
- .filter(Objects::nonNull)
- .collect();
- LOG.info("Found {} non-homogeneous mask blocks", nonHomogeneousMaskBlocks.size());
-
- // Check which tissue blocks are covered by the potentially inpainted mask blocks determined above
- final List tissueBlocks = Grid.create(tissueAttributes.attrs.getDimensions(), tissueAttributes.attrs.getBlockSize());
- final JavaRDD tissueBlocksRDD = sparkContext.parallelize(tissueBlocks);
- final Broadcast tissueAttributesBroadcast = sparkContext.broadcast(tissueAttributes);
- final Broadcast> maskBlocksBroadcast = sparkContext.broadcast(nonHomogeneousMaskBlocks);
-
- final List tissueBlocksToInpaint = tissueBlocksRDD.map(
- block -> translateAndCheckIfOverlaps(block,
- tissueAttributesBroadcast.value().min,
- maskBlocksBroadcast.value()))
- .filter(Objects::nonNull)
- .collect();
- LOG.info("Found {} tissue blocks to inpaint", tissueBlocksToInpaint.size());
-
- // Inpaint the blocks
- final JavaRDD inpaintingBlocksRDD = sparkContext.parallelize(tissueBlocksToInpaint);
-
- inpaintingBlocksRDD.foreach(block -> inpaintBlock(block,
- maskAttributesBroadcast.value().min,
- tissueAttributesBroadcast.value().min,
- paramBroadcast.value(),
- tissueAttributesBroadcast.value().attrs));
+ private void runWithSparkContext(final JavaSparkContext sparkContext,
+ final PointCloud cloud,
+ final List levels,
+ final Map levelAttributes,
+ final String backupPath,
+ final int[] downsampleFactors) {
+
+ final DatasetAttributes s0Attributes = levelAttributes.get(params.fullDataset());
+ final List s0Blocks = new ArrayList<>(Grid.create(s0Attributes.getDimensions(),
+ s0Attributes.getBlockSize()));
+
+ // Grid.create returns blocks in raster order, and Spark's parallelize slices the list into contiguous
+ // partitions. The blocks that actually do work (present-check + inpaint) are the near-ROI ones, and the ROI is
+ // a small, spatially clustered region, so contiguous slicing piles all the expensive blocks into a few
+ // partitions while the rest only run the cheap no-I/O distance filter -> severe load skew. Shuffling first gives
+ // every partition a representative mix of near- and far-ROI blocks, so the pass is balanced. The seed (serial)
+ // keeps the partitioning reproducible, and the outputs (emitted grid positions, per-block decision logs) are
+ // order-independent, so this changes only the distribution of work, not the result.
+ Collections.shuffle(s0Blocks, new Random(params.serial));
+ LOG.info("runWithSparkContext: {} s0 grid blocks to consider (shuffled for load balance)", s0Blocks.size());
+
+ final Broadcast cloudBroadcast = sparkContext.broadcast(cloud);
+ final Broadcast paramsBroadcast = sparkContext.broadcast(params);
+
+ // A single driver line carries the grid extent and z-depth the visualization needs; the per-block
+ // 'blockDecision ...' lines are logged on the executors (see inpaintPartition / logDecision).
+ final long[] dims = s0Attributes.getDimensions();
+ final int[] blockSize = s0Attributes.getBlockSize();
+ LOG.info("runWithSparkContext: diagnostics metadata serial={} maxRoiDistance={} gridX={} gridY={} zLayers={}",
+ params.serial, params.maxRoiDistance,
+ (dims[0] + blockSize[0] - 1) / blockSize[0], (dims[1] + blockSize[1] - 1) / blockSize[1], dims[2]);
+
+ // 2. + 3. Filter (distance, then presence) and inpaint in one distributed pass.
+ final String backup = params.dryRun ? null : backupPath;
+ final JavaRDD modifiedRDD = sparkContext.parallelize(s0Blocks).mapPartitions(
+ blockIterator -> inpaintPartition(blockIterator,
+ cloudBroadcast.value(),
+ paramsBroadcast.value(),
+ backup));
+
+ final List modifiedS0 = modifiedRDD.collect();
+ LOG.info("runWithSparkContext: {} s0 block(s) were {}",
+ modifiedS0.size(), params.dryRun ? "identified for inpainting (dry run)" : "inpainted");
+
+ if (params.dryRun || modifiedS0.isEmpty()) {
+ return;
+ }
+
+ // 6. Selectively update the downsample pyramid, one level at a time.
+ updatePyramid(sparkContext, levels, levelAttributes, modifiedS0, backupPath, downsampleFactors);
}
- private static Grid.Block translateAndCheckHomogeneity(
- final Grid.Block block,
- final long[] shift,
- final Parameters param
- ) {
- LogUtilities.setupExecutorLog4j("");
-
- // Read the mask block and check if it is homogeneous
- boolean isHomogeneous = true;
- try (final N5Reader n5 = new N5Factory().openReader(N5Factory.StorageFormat.N5, param.n5Path)) {
- final Img mask = N5Utils.open(n5, param.mask);
- final Interval interval = Intervals.intersect(mask, block);
- final RandomAccessibleInterval maskPixels = Views.interval(mask, interval);
-
- final UnsignedByteType firstPixel = maskPixels.firstElement();
- for (final UnsignedByteType pixel : maskPixels) {
- if (! pixel.equals(firstPixel)) {
- isHomogeneous = false;
- break;
+ // ------------------------------------------------------------------------------------------------
+ // Step 1: load the per-slab ROI point cloud (positions from the aligned render stack, distances from the xlog).
+ // ------------------------------------------------------------------------------------------------
+
+ /**
+ * Builds the per-slab ROI point cloud in the tissue s0 voxel frame. Positions come from the ALIGNED render
+ * stack (so montage stitching is accounted for): the first layer's tile bounds are fetched from the render web
+ * service on the driver, and each tile center (render world pixels) is mapped to the voxel frame with
+ * {@code voxel = center - worldToVoxel} ({@code worldToVoxel} is the neuroglancer group-level {@code translate},
+ * i.e. the stack bounding-box min; if it is null we fall back to the render stack bounds). Each tile carries the
+ * {@code distance_roi} of its SFOV, read from the xlog for the slab whose {@code id_serial} equals
+ * {@code params.serial}: the tile's mfov (0-based, parsed from the tileId) maps to the xlog mfov axis as
+ * {@code render_mfov + 5} (the axis reserves 5 leading always-NaN rows; see {@link #MFOV_ROW_OFFSET}), and its sfov
+ * number (the 1-based {@code _s##} field) maps to the 0-based xlog sfov axis as {@code _s## - 1}.
+ * Tiles whose SFOV has no (NaN) {@code distance_roi}, or an out-of-range mfov/sfov, are dropped.
+ */
+ static PointCloud loadPointCloud(final Parameters params,
+ final RenderDataClient renderClient,
+ final String stack,
+ final double[] worldToVoxel)
+ throws IOException {
+
+ // (a) xlog: the slab's 2-D distance_roi grid, materialized into a plain array so it outlives the xlog reader.
+ final double[][] distBySfovMfov; // [rowMajorSfov0Based][mfov0Based]
+ try (final N5Reader xlog = new N5Factory().openReader(params.xlogPath)) {
+
+ final double[] idSerial = read1d(xlog);
+ final int slabPosition = findSlabPosition(idSerial, params.serial);
+ if (slabPosition < 0) {
+ final long[] available = new long[idSerial.length];
+ for (int i = 0; i < idSerial.length; i++) {
+ available[i] = Math.round(idSerial[i]);
}
+ throw new IllegalArgumentException("serial " + params.serial + " not found in id_serial; available serials are " +
+ Arrays.toString(available));
}
- }
- // Translate the block to physical coordinates
- final Interval blockInterval = Intervals.translate(block, shift);
- final Grid.Block translatedBlock = new Grid.Block(blockInterval, block.gridPosition);
+ final RandomAccessibleInterval distAll = openDoubles(xlog, "distance_roi");
+ requireSlabAxis(distAll, idSerial.length);
+
+ // distance_roi is scan-independent: slice the slab axis -> 2-D (sfov, mfov).
+ final RandomAccessibleInterval distSlab = Views.hyperSlice(distAll, DIST_AXIS_SLAB, slabPosition);
+ final int nSfov = (int) distSlab.dimension(0);
+ final int nMfov = (int) distSlab.dimension(1);
+ distBySfovMfov = new double[nSfov][nMfov];
+ final RandomAccess dra = distSlab.randomAccess();
+ final long[] pos = new long[2];
+ for (int s = 0; s < nSfov; s++) {
+ for (int m = 0; m < nMfov; m++) {
+ pos[0] = s;
+ pos[1] = m;
+ distBySfovMfov[s][m] = dra.setPositionAndGet(pos).get();
+ }
+ }
+ LOG.info("loadPointCloud: serial {} -> slab position {}, distance_roi grid is {} sfov x {} mfov",
+ params.serial, slabPosition, nSfov, nMfov);
+ }
+ final int nSfov = distBySfovMfov.length;
+ final int nMfov = nSfov > 0 ? distBySfovMfov[0].length : 0;
- final String blockType = isHomogeneous ? "homogeneous -> skip" : "non-homogeneous -> possibly inpaint";
- LOG.info("Mask block {} at {} is {}", translatedBlock.gridPosition, translatedBlock.offset, blockType);
- return isHomogeneous ? null : translatedBlock;
- }
+ // (b) render: the first layer's aligned tile centers, fetched once on the driver.
+ final List zValues = renderClient.getStackZValues(stack);
+ if (zValues.isEmpty()) {
+ throw new IllegalArgumentException("render stack " + stack + " has no z layers");
+ }
+ final double firstZ = zValues.get(0);
+ final List tiles = renderClient.getTileBounds(stack, firstZ);
+
+ final double offsetX;
+ final double offsetY;
+ if (worldToVoxel != null) {
+ offsetX = worldToVoxel[0];
+ offsetY = worldToVoxel[1];
+ } else {
+ final Bounds stackBounds = renderClient.getStackMetaData(stack).getStackBounds();
+ if (stackBounds == null || stackBounds.getMinX() == null || stackBounds.getMinY() == null) {
+ throw new IllegalArgumentException("N5 group " + params.dataset + " has no neuroglancer 'translate' and " +
+ "render stack " + stack + " has no bounds; cannot map world coordinates to voxels");
+ }
+ offsetX = stackBounds.getMinX();
+ offsetY = stackBounds.getMinY();
+ LOG.warn("loadPointCloud: N5 group {} has no 'translate'; using render stack bounds min ({}, {}) as the world->voxel offset",
+ params.dataset, offsetX, offsetY);
+ }
- private static Grid.Block translateAndCheckIfOverlaps(
- final Grid.Block block,
- final long[] shift,
- final List blocksToCheckAgainst
- ) {
- LogUtilities.setupExecutorLog4j("");
-
- // Translate the block to physical coordinates
- final Interval blockInterval = Intervals.translate(block, shift);
- final Grid.Block translatedBlock = new Grid.Block(blockInterval, block.gridPosition);
-
- // Check if the block overlaps with any of the mask blocks that might need inpainting
- for (final Interval maskBlock : blocksToCheckAgainst) {
- final boolean intervalsAreDisjoint = Intervals.isEmpty(Intervals.intersect(translatedBlock, maskBlock));
- if (! intervalsAreDisjoint) {
- LOG.info("Tissue block {} at {} is determined a candidate for inpainting",
- translatedBlock.gridPosition, translatedBlock.minAsLongArray());
- return translatedBlock;
+ // Attach each tile's distance_roi (by mfov + sfov) and place its center in the voxel frame.
+ final List xs = new ArrayList<>();
+ final List ys = new ArrayList<>();
+ final List ds = new ArrayList<>();
+ int noDistance = 0;
+ int outOfRange = 0;
+ for (final TileBounds tile : tiles) {
+ final String tileId = tile.getTileId();
+ // render mfov is 0-based per section; the xlog mfov axis reserves 5 leading NaN rows (see MFOV_ROW_OFFSET).
+ final int mfov = Integer.parseInt(MultiSemUtilities.getSimpleMfovForTileId(tileId).substring(1)) + MFOV_ROW_OFFSET; // m0013 -> 13 -> row 18
+ final int sfov = Integer.parseInt(MultiSemUtilities.getSFOVIndexForTileId(tileId)) - 1;
+ if (mfov < 0 || mfov >= nMfov || sfov < 0 || sfov >= nSfov) {
+ outOfRange++;
+ continue;
+ }
+ final double d = distBySfovMfov[sfov][mfov];
+ if (! Double.isFinite(d)) {
+ noDistance++;
+ continue;
}
+ xs.add(tile.getCenterX() - offsetX);
+ ys.add(tile.getCenterY() - offsetY);
+ ds.add(d);
+ }
+ if (xs.isEmpty()) {
+ throw new IllegalArgumentException("no render tiles in stack " + stack + " z " + firstZ +
+ " could be matched to a finite xlog distance_roi (fetched " + tiles.size() +
+ " tiles; " + noDistance + " had NaN distance, " + outOfRange + " had out-of-range mfov/sfov)");
}
+ LOG.info("loadPointCloud: matched {} of {} render tiles (z {}) to distance_roi; dropped {} NaN-distance, {} out-of-range; world->voxel offset ({}, {})",
+ xs.size(), tiles.size(), firstZ, noDistance, outOfRange, offsetX, offsetY);
+ return new PointCloud(toArray(xs), toArray(ys), toArray(ds));
+ }
- LOG.info("Tissue block {} at {} is not a candidate for inpainting",
- translatedBlock.gridPosition, translatedBlock.minAsLongArray());
- return null;
+ /** Fails fast if a hardcoded slab axis does not carry the expected slab count (guards the layout at the top). */
+ private static void requireSlabAxis(final RandomAccessibleInterval img, final long slabCount) {
+ final int distAxis = Wafer6061Inpainter.DIST_AXIS_SLAB;
+ if (img.numDimensions() <= distAxis || img.dimension(distAxis) != slabCount) {
+ throw new IllegalArgumentException("xlog field '" + "distance_roi" + "' has dimensions " +
+ Arrays.toString(img.dimensionsAsLongArray()) + " but the hardcoded layout " +
+ "expects " + slabCount + " slabs on axis " + Wafer6061Inpainter.DIST_AXIS_SLAB +
+ " (see the xlog axis layout at the top of the class)");
+ }
}
- private static void inpaintBlock(
- final Grid.Block block,
- final long[] maskMin,
- final long[] tissueMin,
- final Parameters param,
- final DatasetAttributes targetAttributes
- ) {
- LogUtilities.setupExecutorLog4j("Block " + Arrays.toString(block.gridPosition));
-
- // Preallocate the inpainted block
- final Img inpaintedBlock = ArrayImgs.unsignedBytes(block.dimensions);
-
- try (final N5Reader n5 = new N5Factory().openReader(N5Factory.StorageFormat.N5, param.n5Path)) {
- // Load and translate the tissue and mask data
- LOG.info("Loading data at {}", block.offset);
- final Img rawTissue = N5Utils.open(n5, param.fullDataset());
- final Img rawMask = N5Utils.open(n5, param.mask);
-
- final RandomAccessibleInterval tissue = Views.translate(rawTissue, tissueMin);
- final RandomAccessibleInterval mask = Views.translate(rawMask, maskMin);
-
- // For each pixel, determine if it should be inpainted and if so, inpaint it by interpolating in z
- LOG.info("Start inpainting");
- final long start = System.currentTimeMillis();
- final Cursor targetCursor = Views.translate(inpaintedBlock, block.offset).localizingCursor();
- final long[] location = new long[3];
- final PixelFiller interpolator = new PixelFiller(tissue, mask, param.stepSize);
-
- while (targetCursor.hasNext()) {
- final UnsignedByteType targetPixel = targetCursor.next();
- targetCursor.localize(location);
- final int value = interpolator.getPixel(location);
- targetPixel.set(value);
- }
- LOG.info("Finished inpainting in {} ms", System.currentTimeMillis() - start);
+ /** Copies a list of doubles into a primitive array. */
+ private static double[] toArray(final List list) {
+ final double[] array = new double[list.size()];
+ for (int i = 0; i < array.length; i++) {
+ array[i] = list.get(i);
}
+ return array;
+ }
- try (final N5Writer n5Writer = new N5Factory().openWriter(N5Factory.StorageFormat.N5, param.n5Path)) {
- N5Utils.saveBlock(inpaintedBlock, n5Writer, param.output, targetAttributes, block.gridPosition);
- LOG.info("Wrote tissue block to '{}'", param.output);
- } catch (final Exception e) {
- LOG.error("Failed to write inpainted block", e);
+ /** Reads a 1-D double dataset (blosc-safe: uses the block-not-found overload to avoid the getAttribute NPE). */
+ private static double[] read1d(final N5Reader n5) {
+ final RandomAccessibleInterval img = openDoubles(n5, "id_serial");
+ final long length = img.dimension(0);
+ final double[] values = new double[(int) length];
+ final RandomAccess ra = img.randomAccess();
+ for (int i = 0; i < length; i++) {
+ values[i] = ra.setPositionAndGet(new long[] {i}).get();
}
+ return values;
}
+ /**
+ * Opens a double-typed dataset using the {@code (blockNotFoundHandler, accessFlags)} overload. This avoids
+ * {@code N5Utils.open(reader, dataset)}'s {@code isLabelMultisetType -> getAttribute} path, which NPEs on blosc
+ * arrays in the shaded jar (the n5-blosc CompressionType service registration is filtered out). Missing blocks
+ * are filled with NaN so they are dropped downstream.
+ */
+ private static RandomAccessibleInterval openDoubles(final N5Reader n5, final String dataset) {
+ final Consumer> nanFill = it -> it.forEach(t -> t.set(Double.NaN));
+ return N5Utils.open(n5, dataset, nanFill, AccessFlags.setOf());
+ }
- public static void main(final String[] args) {
- final ClientRunner clientRunner = new ClientRunner(args) {
- @Override
- public void runClient(final String[] args) {
+ // ------------------------------------------------------------------------------------------------
+ // Steps 2 + 3: filter and inpaint one partition of s0 blocks.
+ // ------------------------------------------------------------------------------------------------
+
+ private static Iterator inpaintPartition(final Iterator blocks,
+ final PointCloud cloud,
+ final Parameters params,
+ final String backupPath) {
+
+ LogUtilities.setupExecutorLog4j("inpaint");
+
+ final List modified = new ArrayList<>();
+ final KNearestNeighborSearchOnKDTree search = cloud.buildSearch();
+
+ try (final N5Reader reader = openN5Reader(params.n5Path);
+ final N5Writer tissueWriter = params.dryRun ? null : openN5Writer(params.n5Path);
+ final N5Writer backupWriter = params.dryRun ? null : openN5Writer(backupPath)) {
+
+ final DatasetAttributes s0Attributes = reader.getDatasetAttributes(params.fullDataset());
+ final DatasetAttributes maskAttributes = reader.getDatasetAttributes(params.mask);
+
+ long considered = 0;
+ long nearRoi = 0;
+ long present = 0;
+ while (blocks.hasNext()) {
+ final Grid.Block block = blocks.next();
+ considered++;
+ final long gridX = block.gridPosition[0];
+ final long gridY = block.gridPosition[1];
+
+ // (1) ROI-distance filter, computed with no I/O. The cloud is already in the s0 voxel frame, so the
+ // block center (voxel index) queries it directly; distance_roi is interpolated by inverse-distance
+ // weighting (the interpolated value is in microns regardless of the voxel-space query units).
+ final double centerX = block.offset[0] + block.dimensions[0] / 2.0;
+ final double centerY = block.offset[1] + block.dimensions[1] / 2.0;
+ final double roiDistance = cloud.interpolate(search, centerX, centerY);
+ if (! (roiDistance < params.maxRoiDistance)) {
+ logDecision(gridX, gridY, "outside_roi", -1, roiDistance);
+ continue;
+ }
+ nearRoi++;
+
+ // (2) presence check: readBlock returns null for absent (empty) blocks. The block it returns is also
+ // the raw tissue we inpaint from and back up, so no whole-volume open (and no accumulating cell cache)
+ // is needed: because z is a single chunk (guarded on the driver), every z-1 / z+1 section the
+ // z-average reads lives inside this same block. Near-ROI empty blocks are not logged: the distance
+ // filter runs before this presence check, so block emptiness is not observable outside the ROI.
+ final DataBlock tissueBlock = reader.readBlock(params.fullDataset(), s0Attributes, block.gridPosition);
+ if (tissueBlock == null) {
+ continue;
+ }
+ present++;
+
+ // (3) inpaint: fill mask==0 pixels with the z-average of the sections above/below. The mask is read by
+ // the same grid position (the datasets share a grid, guarded on the driver); an absent mask block
+ // counts as all-background (all holes), matching a zero-filled whole-volume read of a missing chunk.
+ final DataBlock maskBlock = reader.readBlock(params.mask, maskAttributes, block.gridPosition);
+ final InpaintResult result = inpaintBlock(asByteImg(tissueBlock, block.dimensions),
+ asByteImg(maskBlock, block.dimensions),
+ block.dimensions);
+ if (! result.changed) {
+ logDecision(gridX, gridY, "inside_roi", -1, roiDistance);
+ continue;
+ }
- final Wafer6061Inpainter.Parameters parameters = new Wafer6061Inpainter.Parameters();
- parameters.parse(args);
- parameters.validate();
+ modified.add(block.gridPosition);
+ if (! params.dryRun) {
+ backupWriter.writeBlock(params.fullDataset(), s0Attributes, tissueBlock);
+ N5Utils.saveBlock(result.inpainted, tissueWriter, params.fullDataset(), s0Attributes, block.gridPosition);
+ }
+ logDecision(gridX, gridY, "inpainted", result.minChangedLayer, roiDistance);
+ }
+ LOG.info("inpaintPartition: partition summary: considered={}, nearRoi(<{}um)={}, present={}, inpainted={}",
+ considered, params.maxRoiDistance, nearRoi, present, modified.size());
+ }
- LOG.info("runClient: entry, parameters={}", parameters);
+ return modified.iterator();
+ }
- final Wafer6061Inpainter inpainter = new Wafer6061Inpainter(parameters);
- inpainter.run();
- }
- };
- clientRunner.run();
+ /**
+ * Logs one greppable per-block decision line consumed by the visualization (plot_inpainter_diagnostics.py). Keeping
+ * a fixed {@code blockDecision key=value ...} shape means the log is the single source of truth — no file is written
+ * and no decision logic is re-derived downstream.
+ */
+ private static void logDecision(final long gridX, final long gridY, final String decision,
+ final int minLayer, final double roiDistance) {
+ LOG.info("blockDecision gridX={} gridY={} decision={} minLayer={} roiDistance={}",
+ gridX, gridY, decision, minLayer, roiDistance);
}
+ /**
+ * Wraps a uint8 {@link DataBlock} as a block-local image. An absent (null) block becomes all-zeros, matching the
+ * zero fill an {@code N5Utils.open} whole-volume read would give for a missing chunk.
+ */
+ private static Img asByteImg(final DataBlock dataBlock, final long[] blockDimensions) {
+ if (dataBlock == null) {
+ return ArrayImgs.unsignedBytes(blockDimensions);
+ }
+ return ArrayImgs.unsignedBytes((byte[]) dataBlock.getData(), blockDimensions);
+ }
/**
- * Performs all the inpainting-logic, i.e., when and how to inpaint.
+ * Produces the inpainted version of a single block: every pixel where the mask is 0 is replaced with the average
+ * of the tissue in the sections above and below (the single neighbor is copied at the block's z-boundary). Both
+ * {@code tissueBlock} and {@code maskBlock} are block-local images with dimensions {@code blockDimensions}; because
+ * the volume is a single z-chunk, the block's z-boundary is the volume's z-boundary, so no neighbouring block is
+ * needed for the z-average.
*/
- private static class PixelFiller {
-
- private final RandomAccess tissueAccess;
- private final RandomAccess maskAccess;
- private final int posStep;
- private final int negStep;
-
- public PixelFiller(
- final RandomAccessibleInterval tissue,
- final RandomAccessibleInterval mask,
- final int stepSize
- ) {
- this.tissueAccess = tissue.randomAccess();
- this.maskAccess = Views.extendZero(mask).randomAccess();
- this.posStep = stepSize;
- this.negStep = -2 * stepSize;
- }
-
- public int getPixel(final long[] position) {
- if (shouldBeInpainted(position)) {
- return zAverage(position);
+ static InpaintResult inpaintBlock(final RandomAccessibleInterval tissueBlock,
+ final RandomAccessibleInterval maskBlock,
+ final long[] blockDimensions) {
+
+ final RandomAccess tissueAccess = tissueBlock.randomAccess();
+ final RandomAccess maskAccess = maskBlock.randomAccess();
+ final long zMax = blockDimensions[2] - 1;
+
+ final Img inpainted = ArrayImgs.unsignedBytes(blockDimensions);
+ final Cursor cursor = inpainted.localizingCursor();
+ final long[] local = new long[3];
+ boolean changed = false;
+ int minChangedLayer = Integer.MAX_VALUE;
+ while (cursor.hasNext()) {
+ final UnsignedByteType target = cursor.next();
+ cursor.localize(local);
+
+ final int original = tissueAccess.setPositionAndGet(local).get();
+ final int value;
+ if (maskAccess.setPositionAndGet(local).get() == 255) {
+ value = original;
} else {
- return tissueAccess.setPositionAndGet(position).get();
+ value = zAverage(tissueAccess, local, 0, zMax);
+ if (value != original) {
+ changed = true;
+ if (local[2] < minChangedLayer) {
+ minChangedLayer = (int) local[2];
+ }
+ }
}
+ target.set(value);
}
- /**
- * Average the z-values of the pixels above and below the current pixel.
- * If only one of the pixels is available, that value is used.
- * If neither is available, the pixel is set to 0.
- */
- private int zAverage(final long[] position) {
- maskAccess.setPosition(position);
- maskAccess.move(-1, 2);
- final boolean hasContentAbove = maskAccess.get().get() > 0;
- maskAccess.move(2, 2);
- final boolean hasContentBelow = maskAccess.get().get() > 0;
-
- tissueAccess.setPositionAndGet(position);
- if (hasContentAbove && hasContentBelow) {
- tissueAccess.move(-1, 2);
- final int above = tissueAccess.get().get();
-
- tissueAccess.move(2, 2);
- final int below = tissueAccess.get().get();
-
- return UnsignedByteType.getCodedSignedByteChecked((above + below) >>> 1);
- } else if (hasContentAbove) {
- tissueAccess.move(-1, 2);
- return tissueAccess.get().get();
- } else if (hasContentBelow) {
- tissueAccess.move(2, 2);
- return tissueAccess.get().get();
- } else {
- return 0;
+ return new InpaintResult(inpainted, changed, changed ? minChangedLayer : -1);
+ }
+
+ /** Average of the tissue in the z-1 and z+1 sections, copying the single neighbor at the volume z-boundary. */
+ private static int zAverage(final RandomAccess tissueAccess,
+ final long[] world,
+ final long zMin,
+ final long zMax) {
+ final long z = world[2];
+ int sum = 0;
+ int count = 0;
+ if (z - 1 >= zMin) {
+ world[2] = z - 1;
+ sum += tissueAccess.setPositionAndGet(world).get();
+ count++;
+ }
+ if (z + 1 <= zMax) {
+ world[2] = z + 1;
+ sum += tissueAccess.setPositionAndGet(world).get();
+ count++;
+ }
+ world[2] = z;
+ // both neighbors -> mean (== the old (above+below)>>>1 for byte values); one -> that neighbor; none -> unchanged.
+ return count > 0 ? sum / count : tissueAccess.setPositionAndGet(world).get();
+ }
+
+ // ------------------------------------------------------------------------------------------------
+ // Step 6: selectively re-downsample only the pyramid blocks affected by the modified s0 blocks.
+ // ------------------------------------------------------------------------------------------------
+
+ private void updatePyramid(final JavaSparkContext sparkContext,
+ final List levels,
+ final Map levelAttributes,
+ final List modifiedS0,
+ final String backupPath,
+ final int[] factors) {
+
+ List modifiedPrevious = modifiedS0;
+
+ for (int scale = 1; scale < levels.size(); scale++) {
+ final String fromDataset = levels.get(scale - 1);
+ final String toDataset = levels.get(scale);
+ final DatasetAttributes toAttributes = levelAttributes.get(toDataset);
+
+ // affected blocks: previous-level grid position p maps to this-level block p / factor.
+ final Map affected = new LinkedHashMap<>();
+ for (final long[] p : modifiedPrevious) {
+ final long[] g = affectedBlock(p, factors);
+ affected.putIfAbsent(Arrays.toString(g), g);
}
+ final List affectedBlocks = new ArrayList<>(affected.values());
+ // Every affected block does real work, so there is no near/far skew here, but the per-block cost still varies
+ // with how much of its source region is actually present (dense in the ROI interior, sparse at its edge).
+ // The affected blocks inherit a spatial order, so contiguous partitioning would group same-density
+ // neighbours together and leave some partitions all-dense and others all-sparse. Shuffle (as for s0) mixes
+ // densities across partitions; the seed varies per level but stays reproducible.
+ Collections.shuffle(affectedBlocks, new Random(params.serial * 31L + scale));
+ LOG.info("updatePyramid: re-downsampling {} block(s) for {}", affectedBlocks.size(), toDataset);
+
+ final String n5Path = params.n5Path;
+ sparkContext.parallelize(affectedBlocks).foreachPartition(
+ gridPositions -> downsamplePartition(gridPositions, n5Path, backupPath,
+ fromDataset, toDataset, toAttributes, factors));
+
+ modifiedPrevious = affectedBlocks;
}
+ }
+
+ /**
+ * Re-downsamples one partition's worth of pyramid blocks, opening the N5 handles and the source level once
+ * for the whole partition rather than per block (the source open is a lazy {@code N5Utils.open}, so each block
+ * still reads only the source chunks it needs). Each block is rebuilt from the (already updated) previous level
+ * with the same per-block math as {@code N5DownsamplerSpark} so the result matches the rest of the pyramid, backing
+ * up the original block before overwriting it.
+ */
+ static void downsamplePartition(final Iterator gridPositions,
+ final String n5Path,
+ final String backupPath,
+ final String fromDataset,
+ final String toDataset,
+ final DatasetAttributes toAttributes,
+ final int[] factors) {
+
+ LogUtilities.setupExecutorLog4j("downsample");
+
+ final CellGrid cellGrid = new CellGrid(toAttributes.getDimensions(), toAttributes.getBlockSize());
+
+ try (final N5Reader reader = openN5Reader(n5Path);
+ final N5Writer writer = openN5Writer(n5Path);
+ final N5Writer backupWriter = openN5Writer(backupPath)) {
+
+ final RandomAccessibleInterval source = N5Utils.open(reader, fromDataset);
+
+ int count = 0;
+ while (gridPositions.hasNext()) {
+ downsampleBlock(source, reader, writer, backupWriter, cellGrid,
+ gridPositions.next(), toDataset, toAttributes, factors);
+ count++;
+ }
+ LOG.info("downsamplePartition: re-downsampled {} {} block(s)", count, toDataset);
+ }
+ }
+
+ /** Re-downsamples a single block using the handles and source level already opened for the partition. */
+ private static void downsampleBlock(final RandomAccessibleInterval source,
+ final N5Reader reader,
+ final N5Writer writer,
+ final N5Writer backupWriter,
+ final CellGrid cellGrid,
+ final long[] gridPosition,
+ final String toDataset,
+ final DatasetAttributes toAttributes,
+ final int[] factors) {
+
+ final int n = toAttributes.getNumDimensions();
+ final long[] targetMin = new long[n];
+ final int[] cellDimensions = new int[n];
+ cellGrid.getCellDimensions(gridPosition, targetMin, cellDimensions);
+
+ final long[] sourceMin = new long[n];
+ final long[] sourceSize = new long[n];
+ final long[] targetSize = new long[n];
+ for (int d = 0; d < n; d++) {
+ sourceMin[d] = targetMin[d] * factors[d];
+ sourceSize[d] = (long) cellDimensions[d] * factors[d];
+ targetSize[d] = cellDimensions[d];
+ }
+
+ final RandomAccessibleInterval sourceBlock = Views.offsetInterval(source, sourceMin, sourceSize);
+ final Img targetBlock = ArrayImgs.unsignedBytes(targetSize);
+ Downsample.downsample(sourceBlock, targetBlock, factors);
+
+ // back up the original block (if present) before overwriting.
+ final DataBlock originalBlock = reader.readBlock(toDataset, toAttributes, gridPosition);
+ if (originalBlock != null) {
+ backupWriter.writeBlock(toDataset, toAttributes, originalBlock);
+ }
+
+ // delete first so a block that became empty does not leave a stale remnant.
+ N5Utils.deleteBlock(targetBlock, writer, toDataset, gridPosition);
+ N5Utils.saveNonEmptyBlock(targetBlock, writer, toDataset, gridPosition, new UnsignedByteType());
+ LOG.info("downsampleBlock: updated {} block {}", toDataset, Arrays.toString(gridPosition));
+ }
+
+ // ------------------------------------------------------------------------------------------------
+ // Helpers
+ // ------------------------------------------------------------------------------------------------
- /**
- * Determines if the pixel at the given position should be inpainted based on the local environment.
- */
- private boolean shouldBeInpainted(final long[] position) {
- final boolean hasContent = maskAccess.setPositionAndGet(position).get() > 0;
- if (hasContent) {
- return false;
+ /** Returns the slab-axis position whose id_serial label equals {@code serial}, or -1 if none matches. */
+ static int findSlabPosition(final double[] idSerial, final long serial) {
+ for (int i = 0; i < idSerial.length; i++) {
+ if (Math.round(idSerial[i]) == serial) {
+ return i;
}
+ }
+ return -1;
+ }
+
+ /** Maps a grid block position at pyramid level k-1 to the block it feeds at level k (per-dimension p / factor). */
+ static long[] affectedBlock(final long[] previousGridPosition, final int[] factors) {
+ final long[] g = new long[previousGridPosition.length];
+ for (int d = 0; d < g.length; d++) {
+ g[d] = previousGridPosition[d] / factors[d];
+ }
+ return g;
+ }
- // If the pixel has no content, check the pixels in +/- y direction
- // Only if both have content, the pixel should be inpainted (otherwise, it is a border pixel)
- maskAccess.move(posStep, 1);
- final boolean hasContentFront = maskAccess.get().get() > 0;
- maskAccess.move(negStep, 1);
- final boolean hasContentBack = maskAccess.get().get() > 0;
- if (hasContentFront && hasContentBack) {
- return true;
+ /**
+ * Reads the neuroglancer {@code translate} (the stack bounding-box min in world pixels) from the multiscale group,
+ * or returns null if it is absent. The render N5 export writes it on the group ({@code }), not on
+ * {@code s0} (s0 only carries a sub-pixel centering {@code transform}); callers fall back to the render stack bounds.
+ */
+ private static double[] readGroupTranslate(final N5Reader n5, final String group) {
+ try {
+ final double[] translate = n5.getAttribute(group, "translate", double[].class);
+ if (translate != null && translate.length >= 2) {
+ return translate;
}
+ } catch (final Exception e) {
+ LOG.warn("readGroupTranslate: could not read 'translate' from group {} ({})", group, e.getMessage());
+ }
+ return null;
+ }
+
+ /** Opens an N5 reader for a tissue/backup container (explicit N5 format; local path or gs://). */
+ private static N5Reader openN5Reader(final String path) {
+ return new N5Factory().openReader(N5Factory.StorageFormat.N5, path);
+ }
+
+ /** Opens an N5 writer for a tissue/backup container (explicit N5 format; local path or gs://). */
+ private static N5Writer openN5Writer(final String path) {
+ return new N5Factory().openWriter(N5Factory.StorageFormat.N5, path);
+ }
+
+ /**
+ * Reads the per-step pyramid downsampling factor from the multiscale group's neuroglancer {@code scales} attribute
+ * (as written by render's N5 export). {@code scales} is the cumulative factor per level relative to s0, e.g.
+ * {@code [[1,1,1],[2,2,1],[4,4,1],...]}, so {@code scales[1]} is the factor of s1 relative to s0 — and because these
+ * pyramids use a constant step at every level, it is the per-step factor for all levels. Fails fast when it is
+ * absent or has fewer than two levels, so the pyramid factor is never guessed.
+ */
+ private static int[] readDownsamplingFactors(final N5Reader n5, final String group, final int numDimensions) {
+ final int[][] scales = n5.getAttribute(group, "scales", int[][].class);
+ if (scales == null || scales.length < 2) {
+ throw new IllegalArgumentException("group " + group + " has no 'scales' attribute with at least two levels " +
+ "to derive the pyramid downsampling factor from");
+ }
+ final int[] factors = scales[1]; // s1 relative to s0 == the per-step factor (constant-step pyramids)
+ if (factors.length != numDimensions) {
+ throw new IllegalArgumentException("group " + group + " has scales[1] " + Arrays.toString(factors) +
+ " but the volume is " + numDimensions + "-dimensional");
+ }
+ return factors;
+ }
- // If that is inconclusive, check the pixels in +/- x direction
- maskAccess.move(posStep, 1);
- maskAccess.move(posStep, 0);
- final boolean hasContentRight = maskAccess.get().get() > 0;
- maskAccess.move(negStep, 0);
- final boolean hasContentLeft = maskAccess.get().get() > 0;
- return hasContentRight && hasContentLeft;
+ /**
+ * Reads the render service coordinates (baseDataUrl / owner / project / stack) straight from the multiscale group's
+ * {@code renderExport} metadata, written by render's N5 export (the same attributes.json that holds the pyramid
+ * {@code scales} and {@code translate}). Fails fast when it is absent so the render target is never guessed.
+ */
+ private static RenderTarget readRenderTarget(final N5Reader n5, final String group) {
+ final RenderExport export = n5.getAttribute(group, "renderExport", RenderExport.class);
+ if (export == null || export.runParameters == null || export.runParameters.renderWeb == null ||
+ export.runParameters.renderWeb.baseDataUrl == null || export.runParameters.stack == null) {
+ throw new IllegalArgumentException(
+ "N5 group " + group + " has no usable 'renderExport' metadata (need runParameters.renderWeb." +
+ "baseDataUrl/owner/project and runParameters.stack); this client reads the render service parameters from there");
}
+ final RenderExport.RenderWeb web = export.runParameters.renderWeb;
+ return new RenderTarget(web.baseDataUrl, web.owner, web.project, export.runParameters.stack);
+ }
+ /** Render service coordinates resolved from the group's {@code renderExport} metadata. */
+ static class RenderTarget {
+ final String baseDataUrl;
+ final String owner;
+ final String project;
+ final String stack;
+
+ RenderTarget(final String baseDataUrl, final String owner, final String project, final String stack) {
+ this.baseDataUrl = baseDataUrl;
+ this.owner = owner;
+ this.project = project;
+ this.stack = stack;
+ }
}
- private static class ExtendedAttributes implements Serializable {
- public final DatasetAttributes attrs;
- public final long[] min;
+ /**
+ * Minimal GSON view of the group's {@code renderExport} attribute (only the render coordinates this client needs;
+ * all other fields written by the export are ignored).
+ */
+ private static class RenderExport {
+ RunParameters runParameters;
+
+ private static class RunParameters {
+ RenderWeb renderWeb;
+ String stack;
+ }
- public ExtendedAttributes(final DatasetAttributes attrs, final long[] min) {
- this.attrs = attrs;
- this.min = min;
+ private static class RenderWeb {
+ String baseDataUrl;
+ String owner;
+ String project;
}
+ }
- public static ExtendedAttributes read(final N5Reader n5, final String attrsPath, final String minPath) {
- final DatasetAttributes attrs = n5.getDatasetAttributes(attrsPath);
- final long[] min = n5.getAttribute(minPath, "translate", long[].class);
- return new ExtendedAttributes(attrs, min);
+ /** Result of inpainting one block: the (block-sized) inpainted image and whether any pixel changed. */
+ static class InpaintResult {
+ final Img inpainted;
+ final boolean changed;
+ final int minChangedLayer; // minimal z-layer with an inpainted (changed) pixel, or -1 when nothing changed
+
+ InpaintResult(final Img inpainted, final boolean changed, final int minChangedLayer) {
+ this.inpainted = inpainted;
+ this.changed = changed;
+ this.minChangedLayer = minChangedLayer;
+ }
+ }
+
+ /** Serializable 2-D point cloud of ROI reference points with a distance value per point. */
+ static class PointCloud implements Serializable {
+ private final double[] xs;
+ private final double[] ys;
+ private final double[] dists;
+
+ PointCloud(final double[] xs, final double[] ys, final double[] dists) {
+ this.xs = xs;
+ this.ys = ys;
+ this.dists = dists;
+ }
+
+ int size() {
+ return xs.length;
+ }
+
+ KNearestNeighborSearchOnKDTree buildSearch() {
+ final List points = new ArrayList<>(xs.length);
+ final List values = new ArrayList<>(xs.length);
+ for (int i = 0; i < xs.length; i++) {
+ points.add(new RealPoint(xs[i], ys[i]));
+ values.add(new DoubleType(dists[i]));
+ }
+ final KDTree tree = new KDTree<>(values, points);
+ return new KNearestNeighborSearchOnKDTree<>(tree, Math.min(Wafer6061Inpainter.IDW_K, xs.length));
+ }
+
+ /** Inverse-distance-weighted interpolation of the distance value at (x, y). */
+ double interpolate(final KNearestNeighborSearchOnKDTree search,
+ final double x,
+ final double y) {
+ search.search(new RealPoint(x, y));
+ final int numNeighbors = search.getK();
+ double numerator = 0;
+ double denominator = 0;
+ for (int i = 0; i < numNeighbors; i++) {
+ final double r = search.getDistance(i);
+ final double v = search.getSampler(i).get().get();
+ if (r == 0.0) {
+ return v;
+ }
+ final double w = 1.0 / Math.pow(r, Wafer6061Inpainter.IDW_POWER);
+ numerator += w * v;
+ denominator += w;
+ }
+ return numerator / denominator;
}
}
}