math-tasks/src/scripts/remove-bg.mjs

import sharp from 'sharp';
import { resolve, basename, extname, join } from 'path';
import { existsSync, readdirSync, statSync } from 'fs';
import { fileURLToPath } from 'url';

const PROJECT_ROOT = resolve(fileURLToPath(import.meta.url), '../../..');
const DEFAULT_INPUT = resolve(PROJECT_ROOT, 'assets/icons/pack2');
const DEFAULT_THRESHOLD = 230;
const DEFAULT_FUZZ = 0;
const IMAGE_EXTENSION = /\.(png|jpe?g|webp)$/i;

function parseArgs(args) {
  const result = { input: DEFAULT_INPUT, threshold: DEFAULT_THRESHOLD, fuzz: DEFAULT_FUZZ };
  for (let i = 0; i < args.length; i++) {
    if (args[i] === '--threshold') result.threshold = parseInt(args[++i], 10);
    else if (args[i] === '--fuzz') result.fuzz = parseInt(args[++i], 10);
    else if (!args[i].startsWith('--')) result.input = resolve(args[i]);
  }
  return result;
}

function collectFiles(input) {
  if (statSync(input).isFile()) return [input];
  return readdirSync(input)
    .filter(f => IMAGE_EXTENSION.test(f))
    .sort()
    .map(f => join(input, f));
}

function isWhitish(data, offset, threshold) {
  return data[offset] >= threshold && data[offset + 1] >= threshold && data[offset + 2] >= threshold;
}

function floodFillFromEdges(data, width, height, threshold) {
  const total = width * height;
  const visited = new Uint8Array(total);
  const transparent = new Uint8Array(total);
  const queue = [];
  let head = 0;
  let count = 0;

  const seed = (x, y) => {
    const idx = y * width + x;
    if (visited[idx]) return;
    visited[idx] = 1;
    const off = idx * 4;
    if (data[off + 3] === 0 || isWhitish(data, off, threshold)) {
      queue.push(idx);
    }
  };

  for (let x = 0; x < width; x++) {
    seed(x, 0);
    seed(x, height - 1);
  }
  for (let y = 1; y < height - 1; y++) {
    seed(0, y);
    seed(width - 1, y);
  }

  const dx = [-1, 1, 0, 0];
  const dy = [0, 0, -1, 1];

  while (head < queue.length) {
    const idx = queue[head++];
    transparent[idx] = 1;
    data[idx * 4 + 3] = 0;
    count++;

    const x = idx % width;
    const y = (idx - x) / width;

    for (let d = 0; d < 4; d++) {
      const nx = x + dx[d];
      const ny = y + dy[d];
      if (nx < 0 || nx >= width || ny < 0 || ny >= height) continue;
      const nIdx = ny * width + nx;
      if (visited[nIdx]) continue;
      visited[nIdx] = 1;
      const off = nIdx * 4;
      if (data[off + 3] === 0 || isWhitish(data, off, threshold)) {
        queue.push(nIdx);
      }
    }
  }

  return { transparent, count };
}

function applyAntiAlias(data, width, height, transparent, fuzz) {
  if (fuzz <= 0) return;

  const total = width * height;
  const dist = new Float32Array(total);
  const visited = new Uint8Array(total);
  const queue = [];
  let head = 0;
  const dx = [-1, 1, 0, 0];
  const dy = [0, 0, -1, 1];

  for (let idx = 0; idx < total; idx++) {
    if (transparent[idx]) continue;
    const x = idx % width;
    const y = (idx - x) / width;
    for (let d = 0; d < 4; d++) {
      const nx = x + dx[d];
      const ny = y + dy[d];
      if (nx < 0 || nx >= width || ny < 0 || ny >= height) continue;
      if (transparent[ny * width + nx]) {
        dist[idx] = 1;
        visited[idx] = 1;
        queue.push(idx);
        break;
      }
    }
  }

  while (head < queue.length) {
    const idx = queue[head++];
    data[idx * 4 + 3] = Math.round(255 * dist[idx] / (fuzz + 1));

    if (dist[idx] >= fuzz) continue;
    const x = idx % width;
    const y = (idx - x) / width;
    for (let d = 0; d < 4; d++) {
      const nx = x + dx[d];
      const ny = y + dy[d];
      if (nx < 0 || nx >= width || ny < 0 || ny >= height) continue;
      const nIdx = ny * width + nx;
      if (visited[nIdx] || transparent[nIdx]) continue;
      visited[nIdx] = 1;
      dist[nIdx] = dist[idx] + 1;
      queue.push(nIdx);
    }
  }
}

function outputPath(filePath) {
  const ext = extname(filePath).toLowerCase();
  if (ext === '.png') return filePath;
  return join(resolve(filePath, '..'), basename(filePath, extname(filePath)) + '.png');
}

async function removeBackground(filePath, threshold, fuzz) {
  const { data, info } = await sharp(filePath)
    .ensureAlpha()
    .raw()
    .toBuffer({ resolveWithObject: true });

  const { width, height } = info;
  const { transparent, count } = floodFillFromEdges(data, width, height, threshold);
  applyAntiAlias(data, width, height, transparent, fuzz);

  const outFile = outputPath(filePath);
  await sharp(data, { raw: { width, height, channels: 4 } })
    .png()
    .toFile(outFile);

  return { count, outFile };
}

async function main() {
  const { input, threshold, fuzz } = parseArgs(process.argv.slice(2));

  if (!existsSync(input)) {
    console.error(`Not found: ${input}`);
    process.exit(1);
  }

  const files = collectFiles(input);

  if (files.length === 0) {
    console.error(`No image files found in: ${input}`);
    process.exit(1);
  }

  console.log(`Removing backgrounds from ${files.length} image(s) (threshold: ${threshold}, fuzz: ${fuzz})...`);

  const results = await Promise.all(
    files.map(async (filePath) => {
      const { count, outFile } = await removeBackground(filePath, threshold, fuzz);
      const suffix = outFile !== filePath ? ` (saved as ${basename(outFile)})` : '';
      console.log(`  ${basename(filePath)} -> ${count} pixels made transparent${suffix}`);
      return count;
    })
  );

  const totalPixels = results.reduce((sum, r) => sum + r, 0);
  console.log(`Done. ${totalPixels} total pixels made transparent across ${files.length} file(s).`);
}

main();