lasercut-scrap-modeller/src/lib/detect.ts

import type { Mat } from '@techstark/opencv-js';
import type { CV } from './opencv-loader';
import type { DetectParams, Point, RasterImage, Shape } from './types';

/** Blurred, lighting-corrected grayscale; shared by thresholding and Otsu so both see the same pixels. */
function toGray(cv: CV, image: RasterImage): { gray: Mat; background: number } {
  const src = cv.matFromImageData(image as unknown as ImageData);
  const gray = new cv.Mat();
  try {
    cv.cvtColor(src, gray, cv.COLOR_RGBA2GRAY);
    cv.GaussianBlur(gray, gray, new cv.Size(5, 5), 0);
    return { gray, background: flattenLighting(gray) };
  } finally {
    src.delete();
  }
}

/**
 * Removes smooth lighting gradients (side light, vignetting) by fitting a quadratic surface to the
 * photo's border band, which is almost always background. Returns the background brightness.
 */
function flattenLighting(gray: Mat): number {
  const w = gray.cols;
  const h = gray.rows;
  const g = gray.data;
  const band = Math.max(1, Math.round(Math.min(w, h) * 0.08));
  const step = Math.max(1, Math.round(Math.min(w, h) / 400));
  const terms = (x: number, y: number) => {
    const u = x / w - 0.5;
    const v = y / h - 0.5;
    return [1, u, v, u * u, u * v, v * v];
  };
  const samples: { t: number[]; v: number }[] = [];
  for (let y = 0; y < h; y += step) {
    for (let x = 0; x < w; x += step) {
      if (x >= band && x < w - band && y >= band && y < h - band) continue;
      samples.push({ t: terms(x, y), v: g[y * w + x] });
    }
  }
  // Keep samples near the median first, so a piece touching the border barely affects the fit.
  const sorted = samples.map((s) => s.v).sort((a, b) => a - b);
  const median = sorted[sorted.length >> 1];
  const mad = samples.map((s) => Math.abs(s.v - median)).sort((a, b) => a - b)[samples.length >> 1];
  const nearMedian = samples.filter((s) => Math.abs(s.v - median) <= 3.7 * mad + 2);
  let coef = fitSurface(nearMedian);
  // Second pass drops outliers (a piece touching the border) so they don't skew the fit.
  const residuals = nearMedian.map((s) => s.v - dot(coef, s.t));
  const sd = Math.sqrt(residuals.reduce((a, r) => a + r * r, 0) / residuals.length);
  const inliers = nearMedian.filter((_, i) => Math.abs(residuals[i]) <= 2 * sd + 1);
  if (inliers.length >= 6) coef = fitSurface(inliers);
  const level = inliers.reduce((a, s) => a + s.v, 0) / inliers.length;

  for (let y = 0; y < h; y++) {
    for (let x = 0; x < w; x++) {
      const i = y * w + x;
      g[i] = Math.max(0, Math.min(255, Math.round(g[i] - dot(coef, terms(x, y)) + level)));
    }
  }
  return level;
}

const dot = (a: number[], b: number[]) => a.reduce((sum, ai, i) => sum + ai * b[i], 0);

/** Least-squares fit via the normal equations, solved with Gaussian elimination. */
function fitSurface(samples: { t: number[]; v: number }[]): number[] {
  const n = 6;
  const m = Array.from({ length: n }, () => new Array<number>(n + 1).fill(0));
  for (const { t, v } of samples) {
    for (let i = 0; i < n; i++) {
      for (let j = 0; j < n; j++) m[i][j] += t[i] * t[j];
      m[i][n] += t[i] * v;
    }
  }
  for (let c = 0; c < n; c++) {
    let pivot = c;
    for (let r = c + 1; r < n; r++) if (Math.abs(m[r][c]) > Math.abs(m[pivot][c])) pivot = r;
    [m[c], m[pivot]] = [m[pivot], m[c]];
    if (Math.abs(m[c][c]) < 1e-12) continue;
    for (let r = 0; r < n; r++) {
      if (r === c) continue;
      const f = m[r][c] / m[c][c];
      for (let k = c; k <= n; k++) m[r][k] -= f * m[c][k];
    }
  }
  return m.map((row, i) => (Math.abs(row[i]) < 1e-12 ? 0 : row[n] / row[i]));
}

export function otsuThreshold(cv: CV, image: RasterImage): number {
  const { gray } = toGray(cv, image);
  const out = new cv.Mat();
  try {
    return Math.round(cv.threshold(gray, out, 0, 255, cv.THRESH_BINARY | cv.THRESH_OTSU));
  } finally {
    gray.delete();
    out.delete();
  }
}

export function traceMask(cv: CV, mask: Mat, marginPx: number): Shape | null {
  const radius = Math.round(Math.abs(marginPx));
  // Padding keeps dilation from being clipped and contours off the image edge.
  const pad = radius + 2;
  const padded = new cv.Mat();
  const contours = new cv.MatVector();
  const hierarchy = new cv.Mat();
  try {
    cv.copyMakeBorder(mask, padded, pad, pad, pad, pad, cv.BORDER_CONSTANT, new cv.Scalar(0));
    if (radius > 0) offsetMask(cv, padded, marginPx < 0 ? -radius : radius);
    cv.findContours(padded, contours, hierarchy, cv.RETR_CCOMP, cv.CHAIN_APPROX_SIMPLE);

    const count = contours.size();
    const h = hierarchy.data32S; // per contour: [next, prev, firstChild, parent]
    let best = -1;
    let bestArea = 0;
    for (let i = 0; i < count; i++) {
      if (h[i * 4 + 3] !== -1) continue;
      const c = contours.get(i);
      const area = cv.contourArea(c);
      c.delete();
      if (area > bestArea) {
        bestArea = area;
        best = i;
      }
    }
    if (best === -1) return null;

    const ring = (i: number): Point[] => {
      const c = contours.get(i);
      const d = c.data32S;
      const pts: Point[] = [];
      for (let k = 0; k < d.length; k += 2) pts.push({ x: d[k] - pad, y: d[k + 1] - pad });
      c.delete();
      return pts;
    };
    const holes: Point[][] = [];
    for (let i = 0; i < count; i++) {
      if (h[i * 4 + 3] === best) holes.push(ring(i));
    }
    return { outer: ring(best), holes };
  } finally {
    padded.delete();
    contours.delete();
    hierarchy.delete();
  }
}

/**
 * Erodes (negative) or dilates (positive) a mask by a radius, in place. Uses a distance transform:
 * an elliptical kernel costs radius² per pixel and froze the UI for large margins.
 */
function offsetMask(cv: CV, mask: Mat, radius: number) {
  const source = new cv.Mat();
  const dist = new cv.Mat();
  try {
    // Dilating the piece is eroding the background, so measure from the other side.
    if (radius > 0) cv.bitwise_not(mask, source);
    else mask.copyTo(source);
    cv.distanceTransform(source, dist, cv.DIST_L2, cv.DIST_MASK_PRECISE);
    const r = Math.abs(radius);
    cv.threshold(dist, dist, r, 255, radius > 0 ? cv.THRESH_BINARY_INV : cv.THRESH_BINARY);
    dist.convertTo(mask, cv.CV_8U);
  } finally {
    source.delete();
    dist.delete();
  }
}

export function detect(cv: CV, image: RasterImage, params: DetectParams): Shape | null {
  const { gray, background } = toGray(cv, image);
  const mask = new cv.Mat();
  const kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, new cv.Size(5, 5));
  const labels = new cv.Mat();
  const stats = new cv.Mat();
  const centroids = new cv.Mat();
  try {
    const segment = (dark: boolean, seed: Point | null) => {
      cv.threshold(gray, mask, params.threshold, 255, dark ? cv.THRESH_BINARY_INV : cv.THRESH_BINARY);
      cv.morphologyEx(mask, mask, cv.MORPH_OPEN, kernel);
      cv.morphologyEx(mask, mask, cv.MORPH_CLOSE, kernel);
      const count = cv.connectedComponentsWithStats(mask, labels, stats, centroids, 8, cv.CV_32S);
      return pickComponent(cv, labels, stats, count, image.width, image.height, seed);
    };
    const seedDark = seedIsDark(gray, params.threshold, params.seed);
    let label = seedDark === null ? null : segment(seedDark, params.seed);
    // A click on the background makes it "the piece", which is rejected; fall back to the automatic choice.
    if (label === null) label = segment(background > params.threshold, null);
    if (label === null) return null;

    const lab = labels.data32S;
    const m = mask.data;
    for (let i = 0; i < lab.length; i++) m[i] = lab[i] === label ? 255 : 0;
    fillSmallHoles(cv, mask, params.minHolePx);
    return traceMask(cv, mask, params.marginPx);
  } finally {
    gray.delete();
    mask.delete();
    kernel.delete();
    labels.delete();
    stats.delete();
    centroids.delete();
  }
}

function pickComponent(
  cv: CV, labels: Mat, stats: Mat, count: number, width: number, height: number, seed: Point | null,
): number | null {
  const stat = (label: number, field: number) => stats.intAt(label, field);
  const box = (label: number) => {
    const left = stat(label, cv.CC_STAT_LEFT);
    const top = stat(label, cv.CC_STAT_TOP);
    return { left, top, right: left + stat(label, cv.CC_STAT_WIDTH), bottom: top + stat(label, cv.CC_STAT_HEIGHT) };
  };

  if (seed) {
    const x = Math.round(seed.x);
    const y = Math.round(seed.y);
    if (x >= 0 && y >= 0 && x < width && y < height) {
      const label = labels.intAt(y, x);
      if (label !== 0) {
        const b = box(label);
        const coversImage = b.left === 0 && b.top === 0 && b.right === width && b.bottom === height;
        return coversImage ? null : label;
      }
    }
  }

  // Without a click, specks (dust, grain) must never pass as the piece.
  const minArea = 0.01 * width * height;
  let best: number | null = null;
  let bestArea = 0;
  for (let label = 1; label < count; label++) {
    const b = box(label);
    if (b.left === 0 || b.top === 0 || b.right === width || b.bottom === height) continue;
    const area = stat(label, cv.CC_STAT_AREA);
    if (area >= minArea && area > bestArea) {
      bestArea = area;
      best = label;
    }
  }
  return best;
}

/** Fill holes under minHolePx in both directions, before the margin can grow them. */
function fillSmallHoles(cv: CV, mask: Mat, minHolePx: number) {
  if (minHolePx <= 0) return;
  const contours = new cv.MatVector();
  const hierarchy = new cv.Mat();
  try {
    cv.findContours(mask, contours, hierarchy, cv.RETR_CCOMP, cv.CHAIN_APPROX_SIMPLE);
    const h = hierarchy.data32S;
    for (let i = 0; i < contours.size(); i++) {
      if (h[i * 4 + 3] === -1) continue;
      const c = contours.get(i);
      const rect = cv.boundingRect(c);
      c.delete();
      if (rect.width < minHolePx && rect.height < minHolePx) {
        cv.drawContours(mask, contours, i, new cv.Scalar(255), cv.FILLED);
      }
    }
  } finally {
    contours.delete();
    hierarchy.delete();
  }
}

/** The clicked pixel's side of the threshold is the piece's; null without a usable click. */
function seedIsDark(gray: Mat, threshold: number, seed: Point | null): boolean | null {
  if (!seed) return null;
  const x = Math.round(seed.x);
  const y = Math.round(seed.y);
  if (x < 0 || y < 0 || x >= gray.cols || y >= gray.rows) return null;
  return gray.ucharAt(y, x) <= threshold;
}