js/ml-client.js

// Client for the server-side background removal endpoint. In-browser
// inference via transformers.js proved unreliable on iOS (see
// memory/project_in_browser_ml_dead_end.md), so the model lives on the
// HF Docker Space that also hosts this static frontend. Same-origin —
// no CORS, no absolute URL to keep in sync.

import { tagOp } from './diag.js';

// In production the frontend is served by the same FastAPI process as
// the API (HF Space), so a relative URL is same-origin. In local dev
// the frontend is served by a plain static server with no backend —
// fall back to the HF Space absolute URL so dev against prod Just Works.
const SPACE_ORIGIN = 'https://sdragly-background-removal.hf.space';
const API_BASE =
  typeof location !== 'undefined' &&
  (location.hostname === 'localhost' || location.hostname === '127.0.0.1')
    ? SPACE_ORIGIN
    : '';

// 90s covers the worst-case cold-start (~15-30s container boot) plus
// a slow image upload on a weak cellular connection. If we ever need
// more, something has gone wrong upstream and we should surface the
// error rather than hang the UI.
const REQUEST_TIMEOUT_MS = 90_000;

// Show a "waking up" hint if the response doesn't come back quickly.
// The typical warm response is <3s; anything slower is either a cold
// start or a slow network, and the user deserves to know.
const COLD_START_HINT_MS = 3_000;

/**
 * Remove background from an image blob via the HF Space.
 * @param {Blob} blob  JPEG/PNG image
 * @param {(p: {status: string, message: string, progress?: number}) => void} [onProgress]
 * @returns {Promise<Blob>}  PNG with transparent background
 */
export async function removeBackground(blob, onProgress) {
  const progress = (message) => {
    if (onProgress) onProgress({ status: 'processing', message });
  };

  tagOp('[remove-bg] uploading');
  progress('Sending to server...');

  // If the server is cold, the first POST will hang for ~15-30s. Flip
  // the message after a few seconds so the user knows nothing's broken.
  const coldTimer = setTimeout(() => {
    tagOp('[remove-bg] waking server');
    progress('Waking up server (first request is slow)...');
  }, COLD_START_HINT_MS);

  const abort = new AbortController();
  const timeoutTimer = setTimeout(() => abort.abort(), REQUEST_TIMEOUT_MS);

  try {
    const form = new FormData();
    // Use a stable filename — server only cares about content, but
    // FormData requires *something* here.
    form.append('file', blob, 'input.png');

    const res = await fetch(`${API_BASE}/remove-bg`, {
      method: 'POST',
      body: form,
      signal: abort.signal,
    });

    clearTimeout(coldTimer);

    if (!res.ok) {
      const text = await res.text().catch(() => '');
      throw new Error(
        `Server returned ${res.status}${text ? `: ${text.slice(0, 200)}` : ''}`,
      );
    }

    tagOp('[remove-bg] receiving');
    progress('Receiving result...');
    const result = await res.blob();
    tagOp('[remove-bg] done');
    if (onProgress) onProgress({ status: 'done', message: 'Background removed' });
    return result;
  } catch (err) {
    clearTimeout(coldTimer);
    if (err.name === 'AbortError') {
      throw new Error(
        `Server request timed out after ${REQUEST_TIMEOUT_MS / 1000}s`,
      );
    }
    throw err;
  } finally {
    clearTimeout(timeoutTimer);
  }
}

/**
 * Auto-segment a full drawing into distinct parts via the HF Space
 * SlimSAM endpoint.
 *
 * @param {Blob} blob  background-removed PNG
 * @param {(p: {status?: string, message: string, progress?: number}) => void} [onProgress]
 * @returns {Promise<Array>}  segments in the shape segment-review.js expects
 */
export async function autoSegment(blob, onProgress) {
  const progress = (message, pct) => {
    if (onProgress) onProgress({ status: 'processing', message, progress: pct });
  };

  tagOp('[auto-segment] uploading');
  progress('Looking for parts...', 5);

  // Auto-segment is heavier than bg removal (~5s warm, more on cold
  // start). Keep the same ceiling as remove-bg for now — if SAM ever
  // drifts past this we should investigate rather than paper over.
  const coldTimer = setTimeout(() => {
    tagOp('[auto-segment] waking server');
    progress('Waking up server (first request is slow)...', 10);
  }, COLD_START_HINT_MS);

  const abort = new AbortController();
  const timeoutTimer = setTimeout(() => abort.abort(), REQUEST_TIMEOUT_MS);

  try {
    const form = new FormData();
    form.append('file', blob, 'input.png');

    const res = await fetch(`${API_BASE}/auto-segment`, {
      method: 'POST',
      body: form,
      signal: abort.signal,
    });

    clearTimeout(coldTimer);

    if (!res.ok) {
      const text = await res.text().catch(() => '');
      throw new Error(
        `Server returned ${res.status}${text ? `: ${text.slice(0, 200)}` : ''}`,
      );
    }

    progress('Decoding parts...', 85);
    const json = await res.json();
    const rawSegments = Array.isArray(json.segments) ? json.segments : [];
    tagOp(`[auto-segment] ${rawSegments.length} segments`);

    // The server hands back base64 for both the mask bytes and the
    // cropped PNG. Decode once here so segment-review.js can work with
    // plain Uint8Array + Blob like the old in-browser path did.
    const segments = rawSegments.map((s) => ({
      id: s.id,
      score: s.score,
      area: s.area,
      bbox: s.bbox,
      maskW: s.maskW,
      maskH: s.maskH,
      mask: base64ToUint8Array(s.mask),
      croppedBlob: base64ToBlob(s.croppedPng, 'image/png'),
    }));

    if (onProgress) onProgress({ status: 'done', message: 'Parts found', progress: 100 });
    return segments;
  } catch (err) {
    clearTimeout(coldTimer);
    if (err.name === 'AbortError') {
      throw new Error(
        `Server request timed out after ${REQUEST_TIMEOUT_MS / 1000}s`,
      );
    }
    throw err;
  } finally {
    clearTimeout(timeoutTimer);
  }
}

function base64ToUint8Array(b64) {
  const bin = atob(b64);
  const out = new Uint8Array(bin.length);
  for (let i = 0; i < bin.length; i++) out[i] = bin.charCodeAt(i);
  return out;
}

function base64ToBlob(b64, type) {
  return new Blob([base64ToUint8Array(b64)], { type });
}