quickdraw_app/static/app.js

const classSelect = document.querySelector("#class-select");
const form = document.querySelector("#generate-form");
const guidance = document.querySelector("#guidance");
const guidanceValue = document.querySelector("#guidance-value");
const statusEl = document.querySelector("#status");
const imageStage = document.querySelector("#image-stage");
const resultTitle = document.querySelector("#result-title");
const message = document.querySelector("#message");
const button = document.querySelector("#generate-button");
const downloadLink = document.querySelector("#download-link");
const recognizeButton = document.querySelector("#recognize-button");
const recognitionList = document.querySelector("#recognition-list");
const cnnStatus = document.querySelector("#cnn-status");

let currentImage = null;

function setMessage(text, isError = false) {
  message.textContent = text;
  message.classList.toggle("error", isError);
}

function setLoading(isLoading) {
  button.disabled = isLoading;
  button.querySelector("span:last-child").textContent = isLoading ? "生成中" : "生成";
}

async function loadStatus() {
  const response = await fetch("/api/status");
  const status = await response.json();
  statusEl.textContent = status.checkpoint_exists
    ? `${status.device} · ${status.image_size}x${status.image_size} · step ${status.step ?? "not loaded"}`
    : "checkpoint missing";
}

async function loadCnnStatus() {
  try {
    const response = await fetch("/api/cnn/classes");
    const data = await response.json();
    if (!response.ok) {
      throw new Error(data.detail || "CNN offline");
    }
    const classes = Array.isArray(data) ? data : data.classes;
    cnnStatus.textContent = `CNN: ${classes.length} classes`;
    cnnStatus.title = classes.join(", ");
  } catch (error) {
    cnnStatus.textContent = "CNN: offline";
    cnnStatus.title = error.message;
  }
}

async function loadClasses() {
  const response = await fetch("/api/classes");
  const data = await response.json();
  classSelect.innerHTML = "";
  for (const name of data.classes) {
    const option = document.createElement("option");
    option.value = name;
    option.textContent = name;
    classSelect.append(option);
  }
  classSelect.value = data.classes.includes("cat") ? "cat" : data.classes[0];
}

guidance.addEventListener("input", () => {
  guidanceValue.value = Number(guidance.value).toFixed(2);
});

form.addEventListener("submit", async (event) => {
  event.preventDefault();
  const formData = new FormData(form);
  const seedValue = formData.get("seed");
  const payload = {
    class_name: formData.get("class_name"),
    count: Number(formData.get("count")),
    guidance_scale: Number(formData.get("guidance_scale")),
    seed: seedValue ? Number(seedValue) : null,
  };

  setLoading(true);
  setMessage("模型采样大约需要几十秒，MPS/CPU 会更慢。");

  try {
    const response = await fetch("/api/generate", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(payload),
    });
    const data = await response.json();
    if (!response.ok) {
      throw new Error(data.detail || "生成失败");
    }

    imageStage.innerHTML = "";
    const image = document.createElement("img");
    image.src = data.image;
    image.alt = `${data.class_name} generated sketch samples`;
    imageStage.append(image);
    currentImage = data.image;
    resultTitle.textContent = data.class_name;
    downloadLink.href = data.image;
    downloadLink.setAttribute("aria-disabled", "false");
    recognizeButton.disabled = false;
    recognitionList.innerHTML = "";
    setMessage(`完成 · ${data.device} · step ${data.step} · CFG ${data.guidance_scale}`);
    await loadStatus();
  } catch (error) {
    setMessage(error.message, true);
  } finally {
    setLoading(false);
  }
});

recognizeButton.addEventListener("click", async () => {
  if (!currentImage) {
    return;
  }

  recognizeButton.disabled = true;
  recognizeButton.textContent = "识别中";
  setMessage("正在请求 CNN 图片识别服务。");

  try {
    const response = await fetch("/api/recognize", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ image: currentImage, top_k: 5, predictor: "quickdraw100" }),
    });
    const data = await response.json();
    if (!response.ok) {
      throw new Error(data.detail || "识别失败");
    }

    recognitionList.innerHTML = "";
    for (const item of data.top || []) {
      const row = document.createElement("div");
      row.className = "recognition-item";

      const label = document.createElement("strong");
      label.textContent = item.label;

      const track = document.createElement("div");
      track.className = "confidence-track";
      const fill = document.createElement("div");
      fill.className = "confidence-fill";
      fill.style.width = `${Math.max(0, Math.min(1, item.confidence)) * 100}%`;
      track.append(fill);

      const value = document.createElement("span");
      value.className = "confidence-value";
      value.textContent = `${Math.round(item.confidence * 100)}%`;

      row.append(label, track, value);
      recognitionList.append(row);
    }
    const top = data.prediction;
    setMessage(top ? `CNN top-1: ${top.label} · ${(top.confidence * 100).toFixed(1)}%` : "CNN 返回为空");
  } catch (error) {
    setMessage(error.message, true);
  } finally {
    recognizeButton.disabled = false;
    recognizeButton.textContent = "识别当前图片";
  }
});

Promise.all([loadStatus(), loadClasses(), loadCnnStatus()]).catch((error) => {
  setMessage(error.message, true);
});