js/output-parser.js

/**
 * Output parser: extracts structured predictions from model text output.
 * Ported from hf_space/output_parser.py.
 */
import { COORD_SCALE } from "./constants.js";

export function parseModelOutput(text) {
  const result = { detections: [], segmentations: [], classifications: [], keypoints: [], rawText: text, parseErrors: [] };
  const jsonData = extractJson(text);
  if (jsonData === null) {
    result.parseErrors.push("No valid JSON found in output");
    return result;
  }
  const items = Array.isArray(jsonData) ? jsonData : [jsonData];
  for (const item of items) {
    if (typeof item === "object" && item !== null) parseItem(item, result);
  }
  return result;
}

function extractJson(text) {
  text = text.trim();
  try { return JSON.parse(text); } catch {}
  let m = text.match(/\[[\s\S]*\]/);
  if (m) { try { return JSON.parse(m[0]); } catch {} }
  m = text.match(/\{[\s\S]*\}/);
  if (m) { try { return JSON.parse(m[0]); } catch {} }
  const cleaned = text.replace(/,\s*([}\]])/g, "$1");
  try { return JSON.parse(cleaned); } catch {}
  return null;
}

function parseItem(item, result) {
  const label = item.label || "";
  const bbox2d = item.bbox_2d || item.box_2d;
  const mask = item.mask;
  const kpData = item.keypoints;
  if (!label) { result.parseErrors.push("Item missing label"); return; }
  if (!bbox2d) { result.classifications.push({ label }); return; }
  if (!Array.isArray(bbox2d) || bbox2d.length !== 4) { result.parseErrors.push("Invalid bbox_2d"); return; }
  const bbox = bbox2d.map(v => Math.round(Number(v)));
  if (bbox.some(isNaN)) { result.parseErrors.push("Non-numeric bbox_2d"); return; }

  if (kpData != null) {
    try {
      const kps = kpData.map(k => [Math.round(k[0]), Math.round(k[1]), Math.round(k[2])]);
      result.keypoints.push({ label, bbox, keypoints: kps });
    } catch (e) { result.parseErrors.push("Invalid keypoints: " + e); }
    return;
  }
  if (mask != null) {
    try {
      const polygon = mask.map(pt => [Math.round(pt[0]), Math.round(pt[1])]);
      if (polygon.length >= 3) result.segmentations.push({ label, bbox, polygon });
      else result.parseErrors.push("Mask polygon too few points");
    } catch (e) { result.parseErrors.push("Invalid mask: " + e); }
    return;
  }
  result.detections.push({ label, bbox });
}

export function rescalePredictions(parsed, origW, origH) {
  const sx = origW / COORD_SCALE, sy = origH / COORD_SCALE;
  const scaleBox = b => [Math.round(b[0]*sx), Math.round(b[1]*sy), Math.round(b[2]*sx), Math.round(b[3]*sy)];
  const scalePt = (x, y) => [Math.round(x*sx), Math.round(y*sy)];

  return {
    detections: parsed.detections.map(d => ({ label: d.label, bbox: scaleBox(d.bbox) })),
    segmentations: parsed.segmentations.map(s => ({
      label: s.label, bbox: scaleBox(s.bbox),
      polygon: s.polygon.map(([x,y]) => scalePt(x,y)),
    })),
    classifications: [...parsed.classifications],
    keypoints: parsed.keypoints.map(k => ({
      label: k.label, bbox: scaleBox(k.bbox),
      keypoints: k.keypoints.map(([x,y,v]) => [...scalePt(x,y), v]),
    })),
    rawText: parsed.rawText,
    parseErrors: [...parsed.parseErrors],
  };
}