chaos-tracker-online/src/lib/processImage.ts

// tsconfig tip: set "skipLibCheck": true if your bundler yells about opencv.js types.
import cv from "@techstark/opencv-js";

export type ProcessImageResult = {
  greenRectCenterX: number | null;
  greenRectCenterY: number | null;
  blueRectCenterX: number | null;
  blueRectCenterY: number | null;
  greenRectRotation: number | null; // degrees in [0,180)
  blueRectRotation: number | null;  // degrees in [0,180)
  greenRectLength: number | null;   // long side
  blueRectLength: number | null;    // long side (ray length)
  greenRectWidth: number | null;    // short side
  blueRectWidth: number | null;     // short side = red short side
};

// Geometry helpers use fixed-length tuples to satisfy TS with noUncheckedIndexedAccess
type Vec2 = [number, number];
type Box4 = [Vec2, Vec2, Vec2, Vec2];

type RectProps = {
  box: Box4;            // 4 points [x,y]
  center: Vec2;         // [x,y]
  longLen: number;
  shortLen: number;
  dirLong: Vec2;        // unit vector along long edge
};

export function processImage(imgData: ImageData): ProcessImageResult {
  const src = cv.matFromImageData(imgData); // CV_8UC4 RGBA
  const toFree: cv.Mat[] = [src];

  const rgb = new cv.Mat(); toFree.push(rgb);
  cv.cvtColor(src, rgb, cv.COLOR_RGBA2RGB);
  const hsv = new cv.Mat(); toFree.push(hsv);
  cv.cvtColor(rgb, hsv, cv.COLOR_RGB2HSV);
  const gray = new cv.Mat(); toFree.push(gray);
  cv.cvtColor(src, gray, cv.COLOR_RGBA2GRAY);

  // ---- color masks (HSV) ----
  const kernel3 = cv.getStructuringElement(cv.MORPH_RECT, new cv.Size(3, 3));
  // red = two ranges
  const redMask = (() => {
    const low1 = new cv.Mat(hsv.rows, hsv.cols, hsv.type(), [0, 80, 80, 0]);
    const up1  = new cv.Mat(hsv.rows, hsv.cols, hsv.type(), [5, 255, 255, 0]);
    const low2 = new cv.Mat(hsv.rows, hsv.cols, hsv.type(), [175, 80, 80, 0]);
    const up2  = new cv.Mat(hsv.rows, hsv.cols, hsv.type(), [179, 255, 255, 0]);
    const m1 = new cv.Mat(); const m2 = new cv.Mat();
    cv.inRange(hsv, low1, up1, m1);
    cv.inRange(hsv, low2, up2, m2);
    const rm = new cv.Mat();
    cv.bitwise_or(m1, m2, rm);
    // morphology open -> close
    cv.morphologyEx(rm, rm, cv.MORPH_OPEN, kernel3);
    cv.morphologyEx(rm, rm, cv.MORPH_CLOSE, kernel3);
    low1.delete(); up1.delete(); low2.delete(); up2.delete(); m1.delete(); m2.delete();
    toFree.push(rm);
    return rm;
  })();

  const grnMask = (() => {
    const low = new cv.Mat(hsv.rows, hsv.cols, hsv.type(), [35, 80, 80, 0]);
    const up  = new cv.Mat(hsv.rows, hsv.cols, hsv.type(), [85, 255, 255, 0]);
    const gm = new cv.Mat();
    cv.inRange(hsv, low, up, gm);
    cv.morphologyEx(gm, gm, cv.MORPH_OPEN, kernel3);
    cv.morphologyEx(gm, gm, cv.MORPH_CLOSE, kernel3);
    low.delete(); up.delete();
    toFree.push(gm);
    return gm;
  })();

  // ---- detect largest rect (by area) within mask ----
  const detectLargestBox = (graySrc: cv.Mat, mask: cv.Mat): RectProps | null => {
    const masked = new cv.Mat(); toFree.push(masked);
    cv.bitwise_and(graySrc, graySrc, masked, mask);

    const blur = new cv.Mat(); toFree.push(blur);
    cv.GaussianBlur(masked, blur, new cv.Size(0, 0), 3.0, 3.0, cv.BORDER_DEFAULT);

    const high = new cv.Mat(); toFree.push(high);
    // unsharp: (1+alpha)*I + (-alpha)*blur
    cv.addWeighted(masked, 3.0 /*1+2*/, blur, -2.0, 0, high);

    const edges = new cv.Mat(); toFree.push(edges);
    cv.Canny(high, edges, 40, 120, 3, true);

    const edgesD = new cv.Mat(); toFree.push(edgesD);
    cv.dilate(edges, edgesD, kernel3);

    const contours = new cv.MatVector(); const hierarchy = new cv.Mat();
    cv.findContours(edgesD, contours, hierarchy, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE);

    let bestArea = 0;
    let bestRect: any = null; // rotated rect (object)
    for (let i = 0; i < contours.size(); i++) {
      const cnt = contours.get(i);
      const area = cv.contourArea(cnt);
      if (area > 300 && area > bestArea) {
        bestArea = area;
        bestRect = cv.minAreaRect(cnt); // {center:{x,y}, size:{width,height}, angle:number}
      }
      cnt.delete();
    }
    contours.delete(); hierarchy.delete();

    if (!bestRect) return null;

    const box = rotatedRectToPoints(bestRect);          // Float32Array[2] * 4
    const props = rectPropsFromBox(box);
    return props;
  };

  const redProps = detectLargestBox(gray, redMask);
  const grnProps = detectLargestBox(gray, grnMask);

  // ---- build blue rectangle from rules ----
  let blueCenter: Vec2 | null = null;
  let blueDir: Vec2 | null = null;
  let blueLen: number | null = null;
  let blueW: number | null = null;

  if (redProps && grnProps) {
    const rBox = orderBoxClockwise(redProps.box);
    const gCenter = grnProps.center;

    // midpoints of 4 edges (typed as Box4 to avoid undefined indexes)
    const [r0, r1, r2, r3] = rBox;
    const mids: Box4 = [
      midpoint(r0, r1),
      midpoint(r1, r2),
      midpoint(r2, r3),
      midpoint(r3, r0)
    ];

    // farthest midpoint from green center
    let farIdx = 0, maxD2 = -1;
    mids.forEach((m, i) => {
      const dx = m[0] - gCenter[0];
      const dy = m[1] - gCenter[1];
      const d2 = dx * dx + dy * dy;
      if (d2 > maxD2) { maxD2 = d2; farIdx = i; }
    });
    const [m0, m1, m2, m3] = mids;
    const startPt: Vec2 = farIdx === 0 ? m0 : farIdx === 1 ? m1 : farIdx === 2 ? m2 : m3;

    // long direction of red; point toward projection to green center
  const baseDir = redProps.dirLong;
  let dir: Vec2 = [baseDir[0], baseDir[1]];
  const toG: Vec2 = v2(gCenter[0] - startPt[0], gCenter[1] - startPt[1]);
    if (dot(toG, dir) < 0) dir = scale(dir, -1);

    const proj = dot(toG, dir); // scalar projection
    if (proj > 1.0) {
      blueLen = proj;
      blueW = redProps.shortLen;
      const bb = blueBoxFromRay(startPt, dir, blueLen, blueW);
      blueCenter = bb.center;
      blueDir = dir;
      // const blueBox = bb.box; // (available if you later want corners)
    }
  }

  // ---- assemble metrics ----
  const res: ProcessImageResult = {
    greenRectCenterX: grnProps ? grnProps.center[0] : null,
    greenRectCenterY: grnProps ? grnProps.center[1] : null,
    blueRectCenterX:  blueCenter ? blueCenter[0] : null,
    blueRectCenterY:  blueCenter ? blueCenter[1] : null,
    greenRectRotation: grnProps ? angleDegFromDir(grnProps.dirLong) : null,
    blueRectRotation:  blueDir   ? angleDegFromDir(blueDir)         : null,
    greenRectLength:   grnProps ? grnProps.longLen  : null,
    blueRectLength:    blueLen,
    greenRectWidth:    grnProps ? grnProps.shortLen : null,
    blueRectWidth:     blueW
  };

  // ---- cleanup ----
  try {
    // nothing else
  } finally {
    kernel3.delete();
    toFree.forEach(m => { try { m.delete(); } catch { /* noop */ } });
  }

  return res;
}

/* ----------------------- helpers ----------------------- */

function v2(x: number, y: number): Vec2 { return [x, y]; }
function f32(arr: readonly number[]): Vec2 { return [arr[0] ?? 0, arr[1] ?? 0]; }

function dot(a: Vec2, b: Vec2): number { return a[0] * b[0] + a[1] * b[1]; }
function norm(a: Vec2): number { return Math.hypot(a[0], a[1]); }
function normalize(a: Vec2): Vec2 {
  const n = norm(a) + 1e-9;
  return f32([a[0]/n, a[1]/n]);
}
function scale(a: Vec2, k: number): Vec2 { return f32([a[0]*k, a[1]*k]); }
function add(a: Vec2, b: Vec2): Vec2 { return f32([a[0]+b[0], a[1]+b[1]]); }
function sub(a: Vec2, b: Vec2): Vec2 { return f32([a[0]-b[0], a[1]-b[1]]); }

function midpoint(a: Vec2, b: Vec2): Vec2 {
  return f32([(a[0] + b[0]) * 0.5, (a[1] + b[1]) * 0.5]);
}

function angleDegFromDir(d: Vec2): number {
  let ang = Math.atan2(d[1], d[0]) * 180 / Math.PI;
  ang = (ang + 360) % 180; // [0,180)
  return ang;
}

function orderBoxClockwise(box: Box4): Box4 {
  const c = centroid(box);
  const withAng = box.map(p => ({ p, a: Math.atan2(p[1] - c[1], p[0] - c[0]) }));
  withAng.sort((u, v) => u.a - v.a); // -pi..pi
  const sorted = withAng.map(o => o.p) as Vec2[];
  // ensure Box4 length 4 (non-null assertion since input is Box4)
  return [sorted[0]!, sorted[1]!, sorted[2]!, sorted[3]!];
}

function centroid(pts: ReadonlyArray<Vec2>): Vec2 {
  let x = 0, y = 0;
  for (const p of pts) { x += p[0]; y += p[1]; }
  return f32([x / pts.length, y / pts.length]);
}

function edgeLengths(box: Box4): number[] {
  const [p0, p1, p2, p3] = box;
  return [
    Math.hypot(p1[0] - p0[0], p1[1] - p0[1]),
    Math.hypot(p2[0] - p1[0], p2[1] - p1[1]),
    Math.hypot(p3[0] - p2[0], p3[1] - p2[1]),
    Math.hypot(p0[0] - p3[0], p0[1] - p3[1])
  ];
}

function rectPropsFromBox(boxIn: Box4): RectProps {
  const box = orderBoxClockwise(boxIn);
  const center = centroid(box);
  const L = edgeLengths(box);
  const longLen = Math.max(...L);
  const shortLen = Math.min(...L);
  const idx = L.indexOf(longLen) as 0 | 1 | 2 | 3;
  const [p0, p1, p2, p3] = box;
  let a: Vec2, b: Vec2;
  switch (idx) {
    case 0: a = p0; b = p1; break;
    case 1: a = p1; b = p2; break;
    case 2: a = p2; b = p3; break;
    case 3: a = p3; b = p0; break;
  }
  const dirLong = normalize(sub(b, a));
  return { box, center, longLen, shortLen, dirLong };
}

function rotatedRectToPoints(rr: { center: { x: number, y: number }, size: { width: number, height: number }, angle: number }): Box4 {
  const cx = rr.center.x, cy = rr.center.y;
  const w = rr.size.width, h = rr.size.height;
  const ang = rr.angle * Math.PI / 180.0;
  const cosA = Math.cos(ang), sinA = Math.sin(ang);
  const hw = w * 0.5, hh = h * 0.5;

  const local: [number, number][] = [
    [-hw, -hh], [ hw, -hh],
    [ hw,  hh], [-hw,  hh]
  ];

  const pts: Vec2[] = [];
  for (const [lx, ly] of local) {
    const x = cx + lx * cosA - ly * sinA;
    const y = cy + lx * sinA + ly * cosA;
    pts.push(f32([x, y]));
  }
  return [pts[0]!, pts[1]!, pts[2]!, pts[3]!];
}

function blueBoxFromRay(startPt: Vec2, dirUnitIn: Vec2, length: number, width: number): { box: Box4, center: Vec2 } {
  const dir = normalize(dirUnitIn);
  const perp = f32([-dir[1], dir[0]]);
  const c = add(startPt, scale(dir, length * 0.5));
  const halfL = scale(dir, length * 0.5);
  const halfW = scale(perp, width * 0.5);

  const p0 = sub(sub(c, halfL), halfW);
  const p1 = add(sub(c, halfW), halfL);
  const p2 = add(add(c, halfL), halfW);
  const p3 = add(sub(c, halfW), scale(halfL, -1));

  return { box: [p0, p1, p2, p3], center: c };
}