wip: start function to test aligning stacks

opatiny · opatiny · commit ed97aef3ad3f · 2024-01-26T16:25:32.000+01:00
diff --git a/scripts/alignStack/testAlignStack.ts b/scripts/alignStack/testAlignStack.ts
@@ -0,0 +1,41 @@
+import { readdirSync } from 'node:fs';
+import {
+  alignStack,
+  createStackFromPaths,
+  cropCommonArea,
+} from '../../src/stack/index';
+import { findCommonArea } from '../../src/stack/align/utils/findCommonArea';
+import { crop, writeSync } from '../../src';
+
+const folder = `${__dirname}/BloodMoon`;
+let imageNames = readdirSync(folder);
+
+imageNames = imageNames.slice(0, 2);
+
+const paths = imageNames.map((name) => `${folder}/${name}`);
+
+console.log(paths);
+
+console.log('Number of images in stack:', paths.length);
+
+const stack = createStackFromPaths(paths);
+
+console.log(
+  `Images dimensions: ${stack.getImage(0).width}x${stack.getImage(0).height}`,
+);
+
+console.log('Compute absolute translations');
+const translations = alignStack(stack, {
+  minNbPixels: 1000,
+  debug: true,
+  scalingFactor: 100,
+});
+
+console.log('Cropping all images');
+const crops = cropCommonArea(stack, translations.absolute);
+
+console.log('Writing all images');
+for (let i = 0; i < crops.size; i++) {
+  const image = crops.getImage(i);
+  writeSync(`${__dirname}/result/${imageNames[i]}`, image);
+}
diff --git a/src/stack/align/__tests__/alignDifferentSize.test.ts b/src/stack/align/__tests__/alignDifferentSize.test.ts
@@ -59,6 +59,7 @@ test('different image sizes, no rounding problem', () => {
     maxNbOperations: 1e8,
     xFactor: 0.5,
     yFactor: 0.5,
+    debug: true,
   });
 
   const overlap = overlapImages(source, destination, { origin: result });
diff --git a/src/stack/align/alignDifferentSize.ts b/src/stack/align/alignDifferentSize.ts
@@ -1,26 +1,41 @@
-import { Image, Point, ThresholdAlgorithm } from '../..';
+import {
+  Image,
+  Point,
+  ThresholdAlgorithm,
+  overlapImages,
+  writeSync,
+} from '../..';
 import {
   LevelingAlgorithm,
   getAlignMask,
   prepareForAlign,
 } from '../../align/align';
+import { max } from '../../operations/greyAlgorithms';
 
 import {
-  ComputeNbOperationsOptions,
+  ComputeXYMarginsOptions,
   computeNbOperations,
   computeXYMargins,
 } from './utils/computeNbOperations';
 import { findOverlap } from './utils/findOverlap';
 import { getMinDiffTranslation } from './utils/getMinDiffTranslation';
 
-export interface AlignDifferentSizeOptions extends ComputeNbOperationsOptions {
+export interface AlignDifferentSizeOptions extends ComputeXYMarginsOptions {
   /**
    * Maximal number of operations that the algorithm can perform.
    * This number is used to rescale the images if they are too big so that
    * the algorithm takes roughly always the same time to compute.
+   * You can use scalingFactor instead to specify the scaling factor to apply,
+   * in that case, maxNbOperations should be undefined.
    * @default 1e8
    */
   maxNbOperations?: number;
+  /**
+   * Scaling factor to apply to the images before the rough alignment phase.
+   * Can only be used if maxnbOperations is undefined.
+   * @default undefined
+   */
+  scalingFactor?: number;
   /**
    * Threshold algorithm to use for the alignment masks.
    * @default 'otsu'
@@ -47,10 +62,20 @@ export interface AlignDifferentSizeOptions extends ComputeNbOperationsOptions {
    * @default 0.1
    */
   minFractionPixels?: number;
+  /**
+   * Minimal number of overlapping pixels to apply the algorithm.
+   * @default undefined
+   */
+  minNbPixels?: number;
+  /**
+   * Display debug information?
+   * @default false
+   */
+  debug?: boolean;
 }
 
 /**
- * Align two different size images by finding the position wchich minimises the difference.
+ * Align two different size images by finding the position which minimises the difference.
  * @param source - Source image.
  * @param destination - Destination image.
  * @param options - Align different size options.
@@ -64,34 +89,82 @@ export function alignDifferentSize(
   const {
     xFactor = 0.5,
     yFactor = 0.5,
-    maxNbOperations = 1e8,
     precisionFactor = 1.5,
     thresholdAlgoritm = 'otsu',
     level = 'minMax',
     blurKernelSize,
     minFractionPixels,
+    debug = false,
   } = options;
 
-  const margins = computeXYMargins(source, destination, { xFactor, yFactor });
+  let maxNbOperations;
+  let scalingFactor;
+  if (
+    options.maxNbOperations === undefined &&
+    options.scalingFactor === undefined
+  ) {
+    maxNbOperations = 1e8;
+  } else if (
+    options.maxNbOperations !== undefined &&
+    options.scalingFactor === undefined
+  ) {
+    maxNbOperations = options.maxNbOperations;
+  } else if (
+    options.maxNbOperations === undefined &&
+    options.scalingFactor !== undefined
+  ) {
+    scalingFactor = options.scalingFactor;
+  } else {
+    throw new Error('You cannot define both maxNbOperations and scalingFactor');
+  }
 
-  const nbOperations = computeNbOperations(source, destination, margins);
+  const margins = computeXYMargins(source, destination, { xFactor, yFactor });
 
-  let scalingFactor = 1;
-  if (nbOperations > maxNbOperations) {
-    scalingFactor = Math.sqrt(nbOperations / maxNbOperations);
+  if (maxNbOperations !== undefined) {
+    const initialSrcMask = getAlignMask(source, thresholdAlgoritm);
+    const initialDstMask = getAlignMask(destination, thresholdAlgoritm);
+    const nbOperations = computeNbOperations(source, destination, {
+      sourceMask: initialSrcMask,
+      destinationMask: initialDstMask,
+      margins,
+    });
+    if (debug) {
+      console.log({ nbOperations });
+    }
+    scalingFactor = 1;
+    if (nbOperations > maxNbOperations) {
+      scalingFactor = Math.sqrt(nbOperations / maxNbOperations);
+    }
+  }
+  if (debug) {
+    console.log({ scalingFactor });
   }
+
   // Rough alignment
   const smallSource = prepareForAlign(source, {
     scalingFactor,
     level,
     blurKernelSize,
   });
-  const smallMask = getAlignMask(smallSource, thresholdAlgoritm);
+  const smallSrcMask = getAlignMask(smallSource, thresholdAlgoritm);
   const smallDestination = prepareForAlign(destination, {
     scalingFactor,
     level,
     blurKernelSize,
   });
+  const smallDstMask = getAlignMask(smallDestination, thresholdAlgoritm);
+
+  const smallNbPixels =
+    (smallSrcMask.getNbNonZeroPixels() + smallDstMask.getNbNonZeroPixels()) / 2;
+
+  if (debug) {
+    console.log({ smallNbPixels });
+    writeSync(`${__dirname}/smallSource.png`, smallSource);
+    writeSync(`${__dirname}/smallDestination.png`, smallDestination);
+    writeSync(`${__dirname}/smallSrcMask.png`, smallSrcMask);
+    writeSync(`${__dirname}/smallDstMask.png`, smallDstMask);
+  }
+
   const smallMargins = computeXYMargins(smallSource, smallDestination, {
     xFactor,
     yFactor,
@@ -101,14 +174,26 @@ export function alignDifferentSize(
     smallSource,
     smallDestination,
     {
-      sourceMask: smallMask,
+      sourceMask: smallSrcMask,
+      destinationMask: smallDstMask,
       leftRightMargin: smallMargins.xMargin,
       topBottomMargin: smallMargins.yMargin,
-      minFractionPixels,
+      minNbPixels: smallNbPixels,
     },
   );
 
+  if (debug) {
+    console.log({ roughTranslation });
+    const overlap = overlapImages(smallSource, smallDestination, {
+      origin: roughTranslation,
+    });
+    writeSync(`${__dirname}/roughOverlap.png`, overlap);
+  }
+
   // Find overlapping surface and source and destination origins
+  if (debug) {
+    console.log('Find overlap');
+  }
   const scaledTranslation = {
     column: Math.round(roughTranslation.column * scalingFactor),
     row: Math.round(roughTranslation.row * scalingFactor),
@@ -142,23 +227,38 @@ export function alignDifferentSize(
     blurKernelSize,
     scalingFactor: 1,
   });
-  const mask = getAlignMask(preciseSource, thresholdAlgoritm);
+  const srcMask = getAlignMask(preciseSource, thresholdAlgoritm);
   const preciseDestination = prepareForAlign(destinationCrop, {
     level,
     blurKernelSize,
     scalingFactor: 1,
   });
+  const dstMask = getAlignMask(preciseDestination, thresholdAlgoritm);
+  if (debug) {
+    writeSync(`${__dirname}/preciseSource.png`, preciseSource);
+    writeSync(`${__dirname}/preciseDestination.png`, preciseDestination);
+    writeSync(`${__dirname}/srcMask.png`, srcMask);
+    writeSync(`${__dirname}/dstMask.png`, dstMask);
+  }
+
+  const nbPixels =
+    (srcMask.getNbNonZeroPixels() + dstMask.getNbNonZeroPixels()) / 2;
 
   const preciseMargins = Math.round(precisionFactor * scalingFactor);
 
+  if (debug) {
+    console.log('Compute precise translation');
+  }
   const preciseTranslation = getMinDiffTranslation(
     preciseSource,
     preciseDestination,
     {
       leftRightMargin: preciseMargins,
       topBottomMargin: preciseMargins,
       minFractionPixels,
-      sourceMask: mask,
+      sourceMask: srcMask,
+      destinationMask: dstMask,
+      minNbPixels: nbPixels,
     },
   );
 
diff --git a/src/stack/align/alignStack.ts b/src/stack/align/alignStack.ts
@@ -1,7 +1,10 @@
 import { Stack } from '../../Stack';
 import { Point, sum } from '../../utils/geometry/points';
 
-import { alignDifferentSize } from './alignDifferentSize';
+import {
+  AlignDifferentSizeOptions,
+  alignDifferentSize,
+} from './alignDifferentSize';
 
 export interface Translations {
   absolute: Point[];
@@ -15,15 +18,21 @@ export interface Translations {
  * Internally, the images are aligned between two consecutive images
  * (we imagine that there is less variation in consecutive images).
  * @param stack - Stack to align.
+ * @param options - Options to align different size images.
  * @returns The relative and absolute translations.
  */
-export function alignStack(stack: Stack): Translations {
+export function alignStack(
+  stack: Stack,
+  options: AlignDifferentSizeOptions = {},
+): Translations {
   const relativeTranslations: Point[] = [{ column: 0, row: 0 }];
   const absoluteTranslations: Point[] = [{ column: 0, row: 0 }];
   for (let i = 1; i < stack.size; i++) {
+    console.log(`Aligning image ${i}`);
     const currentRelativeTranslation = alignDifferentSize(
       stack.getImage(i),
       stack.getImage(i - 1),
+      options,
     );
     relativeTranslations.push(currentRelativeTranslation);
     absoluteTranslations.push(
diff --git a/src/stack/align/utils/computeNbOperations.ts b/src/stack/align/utils/computeNbOperations.ts
@@ -85,6 +85,11 @@ export interface ComputeNbOperationsOptions {
    * @default new Mask(source.width, source.height).fill(1)
    */
   sourceMask?: Mask;
+  /**
+   * The mask of the destination image.
+   * @default new Mask(destination.width, destination.height).fill(1)
+   */
+  destinationMask?: Mask;
   /**
    * The margins around the destination image in which the source can be translated.
    * @default { xMargin: 0, yMargin: 0 }
@@ -106,13 +111,17 @@ export function computeNbOperations(
 ): number {
   const {
     sourceMask = new Mask(source.width, source.height).fill(1),
+    destinationMask = new Mask(destination.width, destination.height).fill(1),
     margins = defaultMargins,
   } = options;
   const nbTranslations = computeNbTranslations(source, destination, margins);
   const minHeight = Math.min(source.height, destination.height);
   const minWidth = Math.min(source.width, destination.width);
   // take mask into account
-  const fractionWhitePixels = sourceMask.getNbNonZeroPixels() / source.size;
+  const fractionWhitePixels =
+    (sourceMask.getNbNonZeroPixels() / source.size +
+      destinationMask.getNbNonZeroPixels() / destination.size) /
+    2;
   console.log({ fractionWhitePixels });
 
   return nbTranslations * minHeight * minWidth * fractionWhitePixels;
diff --git a/src/stack/index.ts b/src/stack/index.ts
@@ -3,5 +3,9 @@ export * from './compute/maxImage';
 export * from './compute/meanImage';
 export * from './compute/medianImage';
 export * from './load/decodeStack';
+export * from './load/createStackFromPaths';
 export * from './compute/minImage';
 export * from './compute/sum';
+export * from './align/alignStack';
+export * from './align/alignDifferentSize';
+export * from './align/cropCommonArea';
