js/visualizer.js

/**
 * Canvas-based annotation renderer.
 * Ported from hf_space/visualizer.py (PIL -> Canvas 2D).
 */
import { CLASS_COLORS, _RAW_PALETTE } from "./constants.js";

function hexToRgb(hex) {
  hex = hex.replace("#", "");
  return [parseInt(hex.slice(0,2),16), parseInt(hex.slice(2,4),16), parseInt(hex.slice(4,6),16)];
}

function getColor(label) {
  let hex = CLASS_COLORS[label];
  if (!hex) hex = _RAW_PALETTE[Math.abs(hashCode(label)) % _RAW_PALETTE.length];
  return hexToRgb(hex);
}

function hashCode(s) {
  let h = 0;
  for (let i = 0; i < s.length; i++) h = ((h << 5) - h + s.charCodeAt(i)) | 0;
  return h;
}

/**
 * Draw annotations on a canvas.
 * @param {HTMLCanvasElement} canvas - Target canvas (must be sized to image dimensions).
 * @param {HTMLImageElement|ImageBitmap} image - Source image.
 * @param {object} parsed - Rescaled ParsedOutput with pixel coordinates.
 * @param {object} opts - { lineWidth, fontSize, maskAlpha }
 */
export function drawAnnotations(canvas, image, parsed, opts = {}) {
  const lw = opts.lineWidth || 3;
  const fs = opts.fontSize || 14;
  const maskAlpha = opts.maskAlpha || 0.25;

  canvas.width = image.width;
  canvas.height = image.height;
  const ctx = canvas.getContext("2d");
  ctx.drawImage(image, 0, 0);

  // Detections: solid bounding boxes
  for (const det of parsed.detections) {
    const [r,g,b] = getColor(det.label);
    const [x1,y1,x2,y2] = det.bbox;
    ctx.strokeStyle = `rgb(${r},${g},${b})`;
    ctx.lineWidth = lw;
    ctx.setLineDash([]);
    ctx.strokeRect(x1, y1, x2-x1, y2-y1);
    drawLabel(ctx, det.label, x1, y1, [r,g,b], fs);
  }

  // Segmentations: filled polygon + dashed bbox
  for (const seg of parsed.segmentations) {
    const [r,g,b] = getColor(seg.label);
    const [x1,y1,x2,y2] = seg.bbox;
    // Filled polygon
    if (seg.polygon && seg.polygon.length >= 3) {
      ctx.fillStyle = `rgba(${r},${g},${b},${maskAlpha})`;
      ctx.strokeStyle = `rgb(${r},${g},${b})`;
      ctx.lineWidth = 1;
      ctx.setLineDash([]);
      ctx.beginPath();
      ctx.moveTo(seg.polygon[0][0], seg.polygon[0][1]);
      for (let i = 1; i < seg.polygon.length; i++) ctx.lineTo(seg.polygon[i][0], seg.polygon[i][1]);
      ctx.closePath();
      ctx.fill();
      ctx.stroke();
    }
    // Dashed bbox
    ctx.strokeStyle = `rgb(${r},${g},${b})`;
    ctx.lineWidth = lw;
    ctx.setLineDash([10, 6]);
    ctx.strokeRect(x1, y1, x2-x1, y2-y1);
    ctx.setLineDash([]);
    drawLabel(ctx, seg.label, x1, y1, [r,g,b], fs);
  }

  // Keypoints
  for (const kp of parsed.keypoints) {
    const [r,g,b] = getColor(kp.label);
    const [x1,y1,x2,y2] = kp.bbox;
    // Dotted bbox
    ctx.strokeStyle = `rgb(${r},${g},${b})`;
    ctx.lineWidth = lw;
    ctx.setLineDash([6, 4]);
    ctx.strokeRect(x1, y1, x2-x1, y2-y1);
    ctx.setLineDash([]);
    drawLabel(ctx, kp.label, x1, y1, [r,g,b], fs);

    const rad = Math.max(3, Math.min(canvas.width, canvas.height) / 150);
    const visPts = [];
    kp.keypoints.forEach(([kx, ky, vis], idx) => {
      if (vis === 0) return;
      visPts.push([kx, ky]);
      ctx.fillStyle = `rgb(${r},${g},${b})`;
      ctx.strokeStyle = "white";
      ctx.lineWidth = 1;
      ctx.beginPath();
      ctx.arc(kx, ky, rad, 0, Math.PI * 2);
      ctx.fill();
      ctx.stroke();
      ctx.fillStyle = `rgb(${r},${g},${b})`;
      ctx.font = `${fs}px sans-serif`;
      ctx.fillText(String(idx + 1), kx + rad + 2, ky - rad + fs);
    });
    // Connect consecutive visible keypoints
    if (visPts.length >= 2) {
      ctx.strokeStyle = `rgb(${r},${g},${b})`;
      ctx.lineWidth = Math.max(1, lw - 1);
      ctx.beginPath();
      ctx.moveTo(visPts[0][0], visPts[0][1]);
      for (let i = 1; i < visPts.length; i++) ctx.lineTo(visPts[i][0], visPts[i][1]);
      ctx.stroke();
    }
  }

  // Classification labels at top
  let yOff = 10;
  ctx.font = `bold ${fs}px sans-serif`;
  for (const cls of parsed.classifications) {
    const [r,g,b] = getColor(cls.label);
    const text = `View: ${cls.label}`;
    const tm = ctx.measureText(text);
    const pad = 4;
    ctx.fillStyle = "rgba(50,50,50,0.78)";
    ctx.fillRect(10 - pad, yOff - pad, tm.width + 2*pad, fs + 2*pad);
    ctx.fillStyle = "white";
    ctx.fillText(text, 10, yOff + fs);
    yOff += fs + 2*pad + 6;
  }
}

function drawLabel(ctx, label, x, y, [r,g,b], fs) {
  ctx.font = `bold ${fs}px sans-serif`;
  const tm = ctx.measureText(label);
  const pad = 3;
  const ly = Math.max(0, y - fs - 2*pad);
  ctx.fillStyle = `rgb(${r},${g},${b})`;
  ctx.fillRect(x - pad, ly, tm.width + 2*pad, fs + 2*pad);
  ctx.fillStyle = "white";
  ctx.fillText(label, x, ly + fs + pad);
}