diff --git a/render-ws-java-client/src/main/java/org/janelia/render/client/StackStraighteningClient.java b/render-ws-java-client/src/main/java/org/janelia/render/client/StackStraighteningClient.java new file mode 100644 index 000000000..f9e183b5c --- /dev/null +++ b/render-ws-java-client/src/main/java/org/janelia/render/client/StackStraighteningClient.java @@ -0,0 +1,218 @@ +package org.janelia.render.client; + +import com.beust.jcommander.Parameter; +import com.beust.jcommander.ParametersDelegate; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import mpicbg.trakem2.transform.AffineModel2D; + +import org.janelia.alignment.spec.Bounds; +import org.janelia.alignment.spec.LeafTransformSpec; +import org.janelia.alignment.spec.ResolvedTileSpecCollection; +import org.janelia.alignment.spec.TileSpec; +import org.janelia.alignment.spec.stack.StackMetaData; +import org.janelia.render.client.parameter.CommandLineParameters; +import org.janelia.render.client.parameter.RenderWebServiceParameters; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Java client for straightening a stack by gradually applying an offset + * computed from the midpoints of the first and last layer bounding boxes. + *
+ * The offset is linearly interpolated across layers: + * layer z is moved by (z - zmin) / (zmax - zmin) * offset + */ +public class StackStraighteningClient { + + public static class Parameters extends CommandLineParameters { + + @ParametersDelegate + public RenderWebServiceParameters renderWeb = new RenderWebServiceParameters(); + + @Parameter( + names = "--stack", + description = "Name of source stack", + required = true) + public String stack; + + @Parameter( + names = "--targetStack", + description = "Name of target stack", + required = true) + public String targetStack; + + @Parameter( + names = "--numberOfZLayersPerChunk", + description = "Number of Z layers to process per chunk") + public int numberOfZLayersPerChunk = 1000; + } + + 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); + + LOG.info("runClient: entry, parameters={}", parameters); + + final StackStraighteningClient client = new StackStraighteningClient(parameters); + client.straightenStack(); + } + }; + clientRunner.run(); + } + + private final Parameters parameters; + private final RenderDataClient renderDataClient; + private final List zValues; + private final double minZ; + private final double maxZ; + private final double zRange; + + private StackStraighteningClient(final Parameters parameters) + throws IOException { + + this.parameters = parameters; + this.renderDataClient = parameters.renderWeb.getDataClient(); + this.zValues = renderDataClient.getStackZValues(parameters.stack); + + if (zValues.size() < 2) { + throw new IllegalArgumentException("Stack must have at least 2 layers for straightening"); + } + + this.minZ = zValues.get(0); + this.maxZ = zValues.get(zValues.size() - 1); + this.zRange = maxZ - minZ; + + if (parameters.numberOfZLayersPerChunk < 1) { + throw new IllegalArgumentException("numberOfZLayersPerChunk must be at least 1"); + } + } + + private void straightenStack() throws Exception { + + final StackMetaData sourceStackMetaData = renderDataClient.getStackMetaData(parameters.stack); + renderDataClient.setupDerivedStack(sourceStackMetaData, parameters.targetStack); + + LOG.info("straightenStack: minZ={}, maxZ={}", minZ, maxZ); + + // Get bounding box midpoints for first and last layers + final double[] firstLayerMidpoint = getLayerBoundingBoxMidpoint(minZ); + final double[] lastLayerMidpoint = getLayerBoundingBoxMidpoint(maxZ); + + // Compute the total offset between first and last layer midpoints + final double totalOffsetX = lastLayerMidpoint[0] - firstLayerMidpoint[0]; + final double totalOffsetY = lastLayerMidpoint[1] - firstLayerMidpoint[1]; + + LOG.info("straightenStack: firstLayerMidpoint=({}, {}), lastLayerMidpoint=({}, {}), totalOffset=({}, {})", + firstLayerMidpoint[0], firstLayerMidpoint[1], + lastLayerMidpoint[0], lastLayerMidpoint[1], + totalOffsetX, totalOffsetY); + + for (int start = 0; start < zValues.size(); start += parameters.numberOfZLayersPerChunk) { + final int end = Math.min(start + parameters.numberOfZLayersPerChunk, zValues.size()); + final List zLayersChunk = zValues.subList(start, end); + straightenLayers(zLayersChunk, totalOffsetX, totalOffsetY); + } + + renderDataClient.setStackState(parameters.targetStack, StackMetaData.StackState.COMPLETE); + + LOG.info("straightenStack: exit, processed {} layers", zValues.size()); + } + + private double[] getLayerBoundingBoxMidpoint(final double z) throws Exception { + final Bounds layerBounds = renderDataClient.getLayerBounds(parameters.stack, z); + + if (layerBounds == null || layerBounds.getMinX() == null) { + throw new IllegalArgumentException("Cannot compute bounding box for layer z=" + z); + } + + final double midX = (layerBounds.getMinX() + layerBounds.getMaxX()) / 2.0; + final double midY = (layerBounds.getMinY() + layerBounds.getMaxY()) / 2.0; + + LOG.info("getLayerBoundingBoxMidpoint: z={}, bounds={}, midpoint=({}, {})", + z, layerBounds, midX, midY); + + return new double[] { midX, midY }; + } + + private void straightenLayers(final List zLayersChunk, + final double totalOffsetX, + final double totalOffsetY) + throws Exception { + + final Double firstZ = zLayersChunk.get(0); + final Double lastZ = zLayersChunk.get(zLayersChunk.size() - 1); + + LOG.info("straightenLayers: entry, firstZ={}, lastZ={}", firstZ, lastZ); + + final ResolvedTileSpecCollection resolvedTiles = renderDataClient.getResolvedTiles(parameters.stack, + firstZ, + lastZ, + null, + null, + null, + null, + null, + null); + + final Map> zToTileSpecs = new HashMap<>(); + final Map zToTranslation = new HashMap<>(); + + for (final TileSpec tileSpec : resolvedTiles.getTileSpecs()) { + + final List tileSpecsForZ = zToTileSpecs.computeIfAbsent(tileSpec.getZ(), + k -> new ArrayList<>()); + tileSpecsForZ.add(tileSpec); + + if (! zToTranslation.containsKey(tileSpec.getZ())) { + + final Double z = tileSpec.getZ(); + + // Compute the interpolation factor for this layer + // At zMin, factor = 0 (no movement) + // At zMax, factor = 1 (full offset applied, but negated to bring it back to first layer position) + final double factor = (z - minZ) / zRange; + + // The translation needed to straighten this layer + final double translateX = -factor * totalOffsetX; + final double translateY = -factor * totalOffsetY; + + final AffineModel2D translationModel = new AffineModel2D(); + translationModel.set(1, 0, 0, 1, translateX, translateY); + + zToTranslation.put(z, new LeafTransformSpec(translationModel.getClass().getName(), + translationModel.toDataString())); + } + } + + for (final Double z : zLayersChunk) { + + final List layerTileSpecs = zToTileSpecs.get(z); + final LeafTransformSpec transformSpec = zToTranslation.get(z); + + for (final TileSpec tileSpec : layerTileSpecs) { + resolvedTiles.addTransformSpecToTile(tileSpec.getTileId(), + transformSpec, + ResolvedTileSpecCollection.TransformApplicationMethod.PRE_CONCATENATE_LAST); + } + + LOG.debug("straightenLayers: applied affine transform {} to {} tiles for z={}", + transformSpec.getDataString(), layerTileSpecs.size(), z); + } + + renderDataClient.saveResolvedTiles(resolvedTiles, parameters.targetStack, null); + + LOG.info("straightenLayers: exit, saved {} tiles for z {} to {}", resolvedTiles.getTileCount(), firstZ, lastZ); + } + + private static final Logger LOG = LoggerFactory.getLogger(StackStraighteningClient.class); +}