spelunkai/labeling/frontend/app.js

"use strict";

// Point this at the backend if it's not running on the default local port.
const API_BASE = window.SPELUNKAI_API_BASE || "http://127.0.0.1:8000";

let currentSetId = null;
const subClassNamesById = new Map();
let activeSubClassId = null;

let frames = []; // [{frame, status}, ...] from GET /sets/{id}/frames
let currentIndex = 0;
let currentLabels = [];
let selectedLabelId = null;

// ---- API helpers ----------------------------------------------------------

async function apiRequest(method, path, body) {
  const resp = await fetch(`${API_BASE}${path}`, {
    method,
    headers: body !== undefined ? { "Content-Type": "application/json" } : undefined,
    body: body !== undefined ? JSON.stringify(body) : undefined,
  });
  if (!resp.ok) {
    const text = await resp.text().catch(() => "");
    throw new Error(`${method} ${path} -> ${resp.status} ${text}`);
  }
  if (resp.status === 204) return null;
  return resp.json();
}

const apiGet = (path) => apiRequest("GET", path);
const apiPost = (path, body) => apiRequest("POST", path, body);
const apiPatch = (path, body) => apiRequest("PATCH", path, body);
const apiPut = (path, body) => apiRequest("PUT", path, body);
const apiDelete = (path) => apiRequest("DELETE", path);

// ---- misc helpers -----------------------------------------------------

function getUser() {
  let user = localStorage.getItem("spelunkai_labeling_user");
  if (!user) {
    user = (window.prompt("Your username (for label attribution):") || "anonymous").trim() || "anonymous";
    localStorage.setItem("spelunkai_labeling_user", user);
  }
  return user;
}

// Deterministic per-class color so the same sub-class always renders the
// same way across boxes, the class picker, and different frames.
function colorForId(id) {
  const hue = (Number(id) * 47) % 360;
  return `hsl(${hue} 70% 50%)`;
}

function imagePathToUrl(imagePath) {
  return imagePath.split("/").map(encodeURIComponent).join("/");
}

// ---- sets & classes ---------------------------------------------------

async function loadSets() {
  const sets = await apiGet("/sets");
  const select = document.getElementById("set-select");
  select.innerHTML = "";
  if (sets.length === 0) {
    select.innerHTML = '<option value="">No sets yet - create one via the API</option>';
    return;
  }
  for (const s of sets) {
    const option = document.createElement("option");
    option.value = s.id;
    option.textContent = s.name;
    select.appendChild(option);
  }
  await onSetChange();
}

async function onSetChange() {
  currentSetId = document.getElementById("set-select").value;
  if (!currentSetId) return;
  const labelSet = await apiGet(`/sets/${currentSetId}`);
  renderClassList(labelSet);
  await loadFrames();
}

function renderClassList(labelSet) {
  subClassNamesById.clear();
  activeSubClassId = null;

  const container = document.getElementById("class-list");
  container.innerHTML = "";
  if (labelSet.main_classes.length === 0) {
    container.innerHTML = '<p class="hint">No classes yet for this set.</p>';
    return;
  }

  for (const mainClass of labelSet.main_classes) {
    const heading = document.createElement("div");
    heading.className = "main-class-label";
    heading.textContent = mainClass.name;
    container.appendChild(heading);

    for (const subClass of mainClass.sub_classes) {
      subClassNamesById.set(subClass.id, subClass.name);

      const chip = document.createElement("button");
      chip.type = "button";
      chip.className = "class-chip";
      chip.textContent = subClass.name;
      chip.style.setProperty("--chip-color", colorForId(subClass.id));
      chip.addEventListener("click", () => {
        activeSubClassId = subClass.id;
        document.querySelectorAll(".class-chip").forEach((c) => c.classList.remove("active"));
        chip.classList.add("active");
      });
      container.appendChild(chip);
    }
  }
}

// ---- frames -------------------------------------------------------------

async function loadFrames() {
  const sessionName = document.getElementById("session-filter").value.trim();
  const status = document.getElementById("status-filter").value;
  const params = new URLSearchParams();
  if (sessionName) params.set("session_name", sessionName);
  if (status) params.set("status", status);
  const query = params.toString() ? `?${params}` : "";

  frames = await apiGet(`/sets/${currentSetId}/frames${query}`);
  currentIndex = 0;

  const statusSelect = document.getElementById("frame-status");
  if (frames.length === 0) {
    document.getElementById("frame-image").removeAttribute("src");
    document.getElementById("bbox-layer").innerHTML = "";
    document.getElementById("frame-position").textContent = "No frames match the current filters.";
    statusSelect.disabled = true;
    currentLabels = [];
    return;
  }
  statusSelect.disabled = false;
  await showFrame(0);
}

function goToFrame(index) {
  if (frames.length === 0) return;
  showFrame(Math.max(0, Math.min(index, frames.length - 1)));
}

async function showFrame(index) {
  currentIndex = Math.max(0, Math.min(index, frames.length - 1));
  const entry = frames[currentIndex];
  selectedLabelId = null;

  const img = document.getElementById("frame-image");
  img.src = `${API_BASE}/images/${imagePathToUrl(entry.frame.image_path)}`;

  document.getElementById("frame-position").textContent =
    `${currentIndex + 1} / ${frames.length} — ${entry.frame.session_name}#${entry.frame.frame_index}`;
  document.getElementById("frame-status").value = entry.status;

  await loadLabelsForCurrentFrame();
}

async function loadLabelsForCurrentFrame() {
  const entry = frames[currentIndex];
  currentLabels = await apiGet(`/frames/${entry.frame.id}/sets/${currentSetId}/labels`);

  const img = document.getElementById("frame-image");
  if (img.complete) {
    renderBoxes();
  } else {
    img.onload = () => renderBoxes();
  }
}

async function onStatusChange(event) {
  const status = event.target.value;
  const entry = frames[currentIndex];
  try {
    await apiPut(`/frames/${entry.frame.id}/sets/${currentSetId}/status`, { status, updated_by: getUser() });
    entry.status = status;
  } catch (err) {
    alert(`Failed to update status: ${err.message}`);
  }
}

// ---- bounding boxes -------------------------------------------------------

function renderBoxes() {
  const layer = document.getElementById("bbox-layer");
  layer.innerHTML = "";

  const frame = frames[currentIndex].frame;
  const containerRect = document.getElementById("image-container").getBoundingClientRect();
  const toScreen = containerRect.width / frame.width;

  for (const label of currentLabels) {
    const div = document.createElement("div");
    div.className = "bbox" + (label.id === selectedLabelId ? " selected" : "");
    div.style.setProperty("--bbox-color", colorForId(label.sub_class_id));
    div.style.left = `${label.x * toScreen}px`;
    div.style.top = `${label.y * toScreen}px`;
    div.style.width = `${label.width * toScreen}px`;
    div.style.height = `${label.height * toScreen}px`;

    const tag = document.createElement("span");
    tag.className = "bbox-tag";
    tag.textContent = subClassNamesById.get(label.sub_class_id) || "?";
    div.appendChild(tag);

    const handle = document.createElement("div");
    handle.className = "bbox-handle";
    div.appendChild(handle);

    div.addEventListener("pointerdown", (e) => onBoxPointerDown(e, label, div, handle));
    layer.appendChild(div);
  }
}

function applyPendingBoxStyle(label, div) {
  const frame = frames[currentIndex].frame;
  const containerRect = document.getElementById("image-container").getBoundingClientRect();
  const toScreen = containerRect.width / frame.width;

  const x = label._pendingX ?? label.x;
  const y = label._pendingY ?? label.y;
  const w = label._pendingWidth ?? label.width;
  const h = label._pendingHeight ?? label.height;
  div.style.left = `${x * toScreen}px`;
  div.style.top = `${y * toScreen}px`;
  div.style.width = `${w * toScreen}px`;
  div.style.height = `${h * toScreen}px`;
}

function onBoxPointerDown(e, label, div, handle) {
  e.stopPropagation();
  selectedLabelId = label.id;
  document.querySelectorAll(".bbox").forEach((b) => b.classList.remove("selected"));
  div.classList.add("selected");

  const isResize = e.target === handle;
  const startScreenX = e.clientX;
  const startScreenY = e.clientY;
  const startX = label.x;
  const startY = label.y;
  const startW = label.width;
  const startH = label.height;
  const frame = frames[currentIndex].frame;
  const containerRect = document.getElementById("image-container").getBoundingClientRect();
  const toImage = frame.width / containerRect.width;

  function onMove(ev) {
    const dxImg = (ev.clientX - startScreenX) * toImage;
    const dyImg = (ev.clientY - startScreenY) * toImage;
    if (isResize) {
      label._pendingWidth = Math.max(4, startW + dxImg);
      label._pendingHeight = Math.max(4, startH + dyImg);
    } else {
      label._pendingX = Math.max(0, startX + dxImg);
      label._pendingY = Math.max(0, startY + dyImg);
    }
    applyPendingBoxStyle(label, div);
  }

  function onUp() {
    document.removeEventListener("pointermove", onMove);
    document.removeEventListener("pointerup", onUp);

    const patch = {};
    if (label._pendingX !== undefined) {
      patch.x = label._pendingX;
      patch.y = label._pendingY;
    }
    if (label._pendingWidth !== undefined) {
      patch.width = label._pendingWidth;
      patch.height = label._pendingHeight;
    }
    delete label._pendingX;
    delete label._pendingY;
    delete label._pendingWidth;
    delete label._pendingHeight;

    if (Object.keys(patch).length > 0) {
      Object.assign(label, patch);
      apiPatch(`/labels/${label.id}`, patch).catch((err) => {
        alert(`Failed to save box: ${err.message}`);
        loadLabelsForCurrentFrame();
      });
    }
  }

  document.addEventListener("pointermove", onMove);
  document.addEventListener("pointerup", onUp);
}

function onLayerPointerDown(e) {
  const layer = document.getElementById("bbox-layer");
  if (e.target !== layer) return; // a box handled its own pointerdown

  if (!activeSubClassId) {
    selectedLabelId = null;
    renderBoxes();
    return;
  }

  const containerRect = document.getElementById("image-container").getBoundingClientRect();
  const startX = e.clientX;
  const startY = e.clientY;

  const preview = document.createElement("div");
  preview.className = "bbox";
  preview.style.setProperty("--bbox-color", colorForId(activeSubClassId));
  layer.appendChild(preview);

  function onMove(ev) {
    const x0 = Math.min(startX, ev.clientX) - containerRect.left;
    const y0 = Math.min(startY, ev.clientY) - containerRect.top;
    const w = Math.abs(ev.clientX - startX);
    const h = Math.abs(ev.clientY - startY);
    preview.style.left = `${x0}px`;
    preview.style.top = `${y0}px`;
    preview.style.width = `${w}px`;
    preview.style.height = `${h}px`;
    preview.dataset.x0 = x0;
    preview.dataset.y0 = y0;
    preview.dataset.w = w;
    preview.dataset.h = h;
  }

  function onUp() {
    document.removeEventListener("pointermove", onMove);
    document.removeEventListener("pointerup", onUp);

    const w = Number(preview.dataset.w || 0);
    const h = Number(preview.dataset.h || 0);
    preview.remove();

    if (w < 4 || h < 4) {
      selectedLabelId = null;
      renderBoxes();
      return;
    }

    const frame = frames[currentIndex].frame;
    const toImage = frame.width / containerRect.width;
    createLabel({
      sub_class_id: activeSubClassId,
      x: Number(preview.dataset.x0) * toImage,
      y: Number(preview.dataset.y0) * toImage,
      width: w * toImage,
      height: h * toImage,
      created_by: getUser(),
    });
  }

  document.addEventListener("pointermove", onMove);
  document.addEventListener("pointerup", onUp);
}

async function createLabel(payload) {
  const setId = currentSetId;
  const frameId = frames[currentIndex].frame.id;
  try {
    const label = await apiPost(`/frames/${frameId}/sets/${setId}/labels`, payload);
    currentLabels.push(label);
    selectedLabelId = label.id;
    renderBoxes();
  } catch (err) {
    alert(`Failed to create label: ${err.message}`);
  }
}

async function deleteSelectedLabel() {
  const id = selectedLabelId;
  try {
    await apiDelete(`/labels/${id}`);
    currentLabels = currentLabels.filter((l) => l.id !== id);
    selectedLabelId = null;
    renderBoxes();
  } catch (err) {
    alert(`Failed to delete label: ${err.message}`);
  }
}

// ---- keyboard shortcuts -----------------------------------------------

function onKeyDown(e) {
  if (e.target.tagName === "INPUT" || e.target.tagName === "SELECT") return;

  if ((e.key === "Delete" || e.key === "Backspace") && selectedLabelId != null) {
    e.preventDefault();
    deleteSelectedLabel();
  } else if (e.key === "ArrowLeft") {
    goToFrame(currentIndex - 1);
  } else if (e.key === "ArrowRight") {
    goToFrame(currentIndex + 1);
  }
}

// ---- init -----------------------------------------------------------------

document.addEventListener("DOMContentLoaded", () => {
  document.getElementById("set-select").addEventListener("change", onSetChange);
  document.getElementById("apply-filters").addEventListener("click", loadFrames);
  document.getElementById("prev-frame").addEventListener("click", () => goToFrame(currentIndex - 1));
  document.getElementById("next-frame").addEventListener("click", () => goToFrame(currentIndex + 1));
  document.getElementById("frame-status").addEventListener("change", onStatusChange);
  document.getElementById("bbox-layer").addEventListener("pointerdown", onLayerPointerDown);
  document.addEventListener("keydown", onKeyDown);

  loadSets().catch((err) => alert(`Failed to load sets: ${err.message}`));
});