js/class-mapper.js

/**
 * Class mapper: interpretation-to-class cascade mapping.
 * Ported from hf_space/class_mapper.py.
 */
import {
  CLASS_KEYWORDS, CLASSIFY_TO_DETECT, CLASSIFY_TO_SEGMENT,
  CO_OCCURRENCE_GROUPS, DEFAULT_DETECT_CLASSES, FIELD_WEIGHTS,
  GROUP_ALIASES, NEGATIVE_MEASUREMENT_PATTERNS, PLANE_TO_GROUP,
  SPECIFIC_CLS_LABELS, keywordLookup,
} from "./constants.js";

export function parseInterpretationJson(rawText) {
  const text = rawText.trim();
  try {
    const data = JSON.parse(text);
    if (typeof data === "object" && !Array.isArray(data)) { data._parseSuccess = true; return data; }
  } catch {}
  const bs = text.indexOf("{"), be = text.lastIndexOf("}");
  if (bs !== -1 && be > bs) {
    try {
      const data = JSON.parse(text.slice(bs, be + 1));
      if (typeof data === "object") { data._parseSuccess = true; return data; }
    } catch {}
  }
  return { _parseSuccess: false, _rawText: text };
}

function flattenFieldText(value) {
  if (typeof value === "string") return value;
  if (Array.isArray(value)) return value.map(String).join(" ");
  if (typeof value === "object" && value !== null) return Object.values(value).map(String).join(" ");
  return String(value);
}

function scoreClassesPerField(interp, weights = FIELD_WEIGHTS) {
  const scores = {};
  for (const cls of Object.keys(CLASS_KEYWORDS)) scores[cls] = 0;
  for (const [fieldName, weight] of Object.entries(weights)) {
    if (weight === 0) continue;
    const raw = interp[fieldName];
    if (!raw) continue;
    const text = flattenFieldText(raw).toLowerCase();
    if (!text) continue;
    for (const [cls, keywords] of Object.entries(CLASS_KEYWORDS)) {
      for (const kw of keywords) {
        if (text.includes(kw)) { scores[cls] += weight; break; }
      }
    }
  }
  return scores;
}

function aggregateGroupScores(classScores, interp) {
  let fullText = "";
  for (const [k, v] of Object.entries(interp)) {
    if (!k.startsWith("_") && v) fullText += " " + flattenFieldText(v);
  }
  fullText = fullText.toLowerCase();

  const results = CO_OCCURRENCE_GROUPS.map(([name, detCls, segCls, compat]) => {
    let gscore = detCls.reduce((s, c) => s + (classScores[c] || 0), 0);
    const aliases = GROUP_ALIASES[name] || [];
    for (const alias of aliases) {
      if (fullText.includes(alias)) { gscore += 2.0; break; }
    }
    return { name, gscore, detCls, segCls, compat };
  });
  results.sort((a, b) => b.gscore - a.gscore);
  return results;
}

export function extractClsLabel(clsResult) {
  if (!clsResult) return null;
  const parsed = clsResult.parsed || clsResult;
  const label = parsed.label;
  return (typeof label === "string" && label.trim()) ? label.trim() : null;
}

export function mapInterpretationToClasses(interpParsed, clsLabel = null) {
  // P1: SPECIFIC cls_label exact match
  if (clsLabel && SPECIFIC_CLS_LABELS.has(clsLabel) && CLASSIFY_TO_DETECT[clsLabel]) {
    return { det: CLASSIFY_TO_DETECT[clsLabel], seg: CLASSIFY_TO_SEGMENT[clsLabel] || null, tier: `cls_specific_${clsLabel}` };
  }

  // P2: imaging plane / orientation -> group
  const planeText = flattenFieldText(interpParsed.imaging_plane || "").toLowerCase();
  const orientText = flattenFieldText(interpParsed.fetal_orientation || "").toLowerCase();
  const combined = `${planeText} ${orientText}`;
  for (const [planeKey, groupName] of Object.entries(PLANE_TO_GROUP)) {
    if (combined.includes(planeKey)) {
      for (const [gname, detCls, segCls] of CO_OCCURRENCE_GROUPS) {
        if (gname === groupName) {
          return { det: detCls.join(", "), seg: segCls ? segCls.join(", ") : null, tier: `plane_${planeKey}` };
        }
      }
      break;
    }
  }

  // P3: keyword scoring
  const bioText = flattenFieldText(interpParsed.biometric_measurements || "").toLowerCase();
  const hasNeg = NEGATIVE_MEASUREMENT_PATTERNS.some(p => bioText.includes(p));
  let classScores;
  if (hasNeg) {
    const dampened = { ...FIELD_WEIGHTS, biometric_measurements: 0.3 };
    classScores = scoreClassesPerField(interpParsed, dampened);
  } else {
    classScores = scoreClassesPerField(interpParsed);
  }
  const groups = aggregateGroupScores(classScores, interpParsed);

  // Disambiguation: body_full vs doppler
  const ORGAN_KW = ["stomach","liver","hepat","artery","arter","vein","venous","gastric","ductus"];
  if (groups.length >= 2) {
    const gnames = {};
    groups.forEach((g, i) => gnames[g.name] = i);
    if (gnames.body_full !== undefined && gnames.doppler !== undefined && groups[0].name === "body_full") {
      let ft = "";
      for (const [k, v] of Object.entries(interpParsed)) {
        if (!k.startsWith("_") && v) ft += " " + flattenFieldText(v);
      }
      ft = ft.toLowerCase();
      if (ORGAN_KW.some(kw => ft.includes(kw))) {
        const di = gnames.doppler;
        [groups[0], groups[di]] = [groups[di], groups[0]];
      }
    }
  }

  if (groups.length > 0 && groups[0].gscore > 0) {
    const top = groups[0];
    let detClasses = [...top.detCls];
    let segClasses = top.segCls ? [...top.segCls] : null;
    let tier = `interp_${top.name}`;

    if (groups.length > 1) {
      const sec = groups[1];
      if (sec.gscore > 0 && sec.gscore >= top.gscore * 0.3 && top.compat.includes(sec.name)) {
        detClasses.push(...sec.detCls);
        if (sec.segCls) {
          if (!segClasses) segClasses = [];
          segClasses.push(...sec.segCls);
        }
        tier = `interp_${top.name}+${sec.name}`;
      }
    }
    detClasses = detClasses.slice(0, 5);
    if (segClasses) segClasses = segClasses.slice(0, 5);
    return { det: detClasses.join(", "), seg: segClasses ? segClasses.join(", ") : null, tier };
  }

  // P4: generic cls_label fallback
  if (clsLabel && CLASSIFY_TO_DETECT[clsLabel]) {
    return { det: CLASSIFY_TO_DETECT[clsLabel], seg: CLASSIFY_TO_SEGMENT[clsLabel] || null, tier: `cls_generic_${clsLabel}` };
  }
  if (clsLabel) {
    const hit = keywordLookup(clsLabel);
    if (hit) return { det: hit.det, seg: hit.seg, tier: `cls_keyword_${hit.kw}` };
  }

  // P5: default fallback
  return { det: DEFAULT_DETECT_CLASSES, seg: null, tier: "fallback_no_match" };
}