// Background removal with finegrain-box-segmenter-ONNX, via onnxruntime-node + sharp.
//
//   npm i onnxruntime-node sharp
//   # download the model to onnx/model.onnx (from this repo's Files tab)
//   node remove_bg.mjs --image photo.jpg --model onnx/model.onnx
//   node remove_bg.mjs --image photo.jpg --model onnx/model.onnx --ep webgpu
//
// Writes <name>_mask.png (grayscale alpha) and <name>_cutout.png (RGBA, original colours + mask alpha).
import * as ort from 'onnxruntime-node';
import sharp from 'sharp';
import path from 'node:path';
import fs from 'node:fs/promises';

const SIZE = 1024;
const MEAN = [0.485, 0.456, 0.406];
const STD = [0.229, 0.224, 0.225];
const sigmoid = (x) => 1 / (1 + Math.exp(-x));

async function imageToTensor(inputPath) {
  const meta = await sharp(inputPath).metadata();
  const { data: rgb } = await sharp(inputPath)
    .removeAlpha().resize(SIZE, SIZE, { fit: 'fill', kernel: 'cubic' })
    .raw().toBuffer({ resolveWithObject: true });
  const plane = SIZE * SIZE, chw = new Float32Array(3 * plane);
  for (let p = 0; p < plane; p++) {
    chw[p] = (rgb[p * 3] / 255 - MEAN[0]) / STD[0];
    chw[plane + p] = (rgb[p * 3 + 1] / 255 - MEAN[1]) / STD[1];
    chw[2 * plane + p] = (rgb[p * 3 + 2] / 255 - MEAN[2]) / STD[2];
  }
  return { tensor: new ort.Tensor('float32', chw, [1, 3, SIZE, SIZE]), W: meta.width, H: meta.height };
}

function logitsToMask1024(out) {
  const data = out.data, plane = SIZE * SIZE, u8 = Buffer.allocUnsafe(plane);
  let needSig = false;
  for (let i = 0; i < plane; i++) { const v = data[i]; if (v < 0 || v > 1) { needSig = true; break; } }
  for (let i = 0; i < plane; i++) u8[i] = Math.max(0, Math.min(255, Math.round((needSig ? sigmoid(data[i]) : data[i]) * 255)));
  return u8;
}

async function main() {
  const args = process.argv.slice(2); const o = { ep: 'cpu', model: 'onnx/model.onnx' };
  for (let i = 0; i < args.length; i++) {
    if (args[i] === '--image') o.image = args[++i];
    else if (args[i] === '--model') o.model = args[++i];
    else if (args[i] === '--ep') o.ep = args[++i];
  }
  if (!o.image) { console.error('usage: node remove_bg.mjs --image <path> [--model onnx/model.onnx] [--ep cpu|webgpu|dml]'); process.exit(1); }

  const providers = o.ep === 'webgpu' ? ['webgpu', 'cpu'] : o.ep === 'dml' ? ['dml', 'cpu'] : ['cpu'];
  const session = await ort.InferenceSession.create(o.model, { executionProviders: providers });

  const { tensor, W, H } = await imageToTensor(o.image);
  const res = await session.run({ [session.inputNames[0]]: tensor });
  const mask1024 = logitsToMask1024(res[session.outputNames[0]]);
  // resize mask back to original size; force b-w so the raw buffer is exactly W*H bytes (1 channel)
  const maskFull = await sharp(mask1024, { raw: { width: SIZE, height: SIZE, channels: 1 } })
    .resize(W, H, { fit: 'fill', kernel: 'cubic' }).toColourspace('b-w').raw().toBuffer();

  const base = path.basename(o.image).replace(/\.[^.]+$/, '');
  await sharp(maskFull, { raw: { width: W, height: H, channels: 1 } }).png().toFile(`${base}_mask.png`);

  const { data: orig } = await sharp(o.image).removeAlpha().raw().toBuffer({ resolveWithObject: true });
  const rgba = Buffer.allocUnsafe(W * H * 4);
  for (let i = 0; i < W * H; i++) {
    rgba[i * 4] = orig[i * 3]; rgba[i * 4 + 1] = orig[i * 3 + 1]; rgba[i * 4 + 2] = orig[i * 3 + 2]; rgba[i * 4 + 3] = maskFull[i];
  }
  await fs.writeFile(`${base}_cutout.png`, await sharp(rgba, { raw: { width: W, height: H, channels: 4 } }).png().toBuffer());
  console.log(`wrote ${base}_mask.png and ${base}_cutout.png (${W}x${H}, ep=${o.ep})`);
  await session.release?.();
}
main().catch((e) => { console.error(e); process.exit(1); });