Recorder-Tabber/src/lib/recognize.ts

import { spawn } from "node:child_process";
import { existsSync, readdirSync } from "node:fs";
import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import path from "node:path";
import { parseMusicXml } from "@/lib/musicxml";
import type { Score } from "@/lib/score";

export type RecognizerInfo = {
  available: boolean;
  name: string;
  detail: string;
};

export type RecognizeSuccess = {
  ok: true;
  score: Score;
  warnings: string[];
  recognizer: string;
};

export type RecognizeFailure = {
  ok: false;
  error: string;
  status: number;
  recognizer: string;
};

export function resolvePython(): string | null {
  const fromEnv = process.env.OEMER_PYTHON;
  if (fromEnv && existsSync(fromEnv)) return fromEnv;
  return null;
}

function modelsReady(python: string): boolean {
  const lib = path.resolve(path.dirname(python), "..", "lib");
  if (!existsSync(lib)) return false;
  for (const entry of readdirSync(lib)) {
    const model = path.join(
      lib,
      entry,
      "site-packages",
      "oemer",
      "checkpoints",
      "unet_big",
      "model.onnx",
    );
    if (existsSync(model)) return true;
  }
  return false;
}

export function recognizerInfo(): RecognizerInfo {
  const python = resolvePython();
  if (!python) {
    return {
      available: false,
      name: "oemer",
      detail:
        "The local reader (oemer) is not installed. Sample tunes and the note editor still work. No photo is turned into guessed notes.",
    };
  }
  if (!modelsReady(python)) {
    return {
      available: false,
      name: "oemer 0.1.8",
      detail:
        "oemer is installed, but its model files are missing, so a photo cannot be read. Sample tunes and the editor still work.",
    };
  }
  return {
    available: true,
    name: "oemer 0.1.8",
    detail:
      "Photos are read on this computer by oemer. It misses notes. Correct the list; nothing is sent away, and a failed read does not invent notes.",
  };
}

function scriptPath(): string {
  const candidates = [
    path.join(process.cwd(), "scripts", "run_oemer.py"),
    path.join(process.cwd(), "..", "scripts", "run_oemer.py"),
  ];
  return candidates.find((candidate) => existsSync(candidate)) ?? candidates[0];
}

function runProcess(
  command: string,
  args: string[],
  timeoutMs: number,
): Promise<{ code: number; stdout: string; stderr: string }> {
  return new Promise((resolve, reject) => {
    const child = spawn(command, args, {
      env: {
        ...process.env,
        MPLBACKEND: "Agg",
        PYTHONUNBUFFERED: "1",
      },
    });
    let stdout = "";
    let stderr = "";
    const timer = setTimeout(() => {
      child.kill("SIGKILL");
      reject(new Error("timeout"));
    }, timeoutMs);
    child.stdout.setEncoding("utf8");
    child.stderr.setEncoding("utf8");
    child.stdout.on("data", (chunk) => {
      stdout += chunk;
    });
    child.stderr.on("data", (chunk) => {
      stderr += chunk;
    });
    child.on("error", (error) => {
      clearTimeout(timer);
      reject(error);
    });
    child.on("close", (code) => {
      clearTimeout(timer);
      resolve({ code: code ?? 1, stdout, stderr });
    });
  });
}

function shortFailure(stderr: string, stdout: string): string {
  const lines = `${stderr}\n${stdout}`
    .split("\n")
    .map((line) => line.trim())
    .filter(Boolean)
    .filter((line) => !line.startsWith("MUSICXML="));
  const errorLine = [...lines].reverse().find((line) => /Error|Exception|assert|too small|could not/i.test(line));
  const text = errorLine ?? lines.at(-1) ?? "The recognizer stopped without a result.";
  return text.slice(0, 280);
}

export async function recognizeImage(file: File): Promise<RecognizeSuccess | RecognizeFailure> {
  const info = recognizerInfo();
  if (!info.available) {
    return { ok: false, error: info.detail, status: 503, recognizer: info.name };
  }
  const python = resolvePython();
  if (!python) {
    return { ok: false, error: info.detail, status: 503, recognizer: info.name };
  }

  const directory = await mkdtemp(path.join(tmpdir(), "recorder-tab-"));
  try {
    const inputPath = path.join(directory, "upload.bin");
    await writeFile(inputPath, Buffer.from(await file.arrayBuffer()));
    const script = scriptPath();
    let result: { code: number; stdout: string; stderr: string };
    try {
      result = await runProcess(python, [script, inputPath, directory], 240_000);
    } catch (error) {
      const timedOut = error instanceof Error && error.message === "timeout";
      return {
        ok: false,
        status: 422,
        recognizer: info.name,
        error: timedOut
          ? "Reading this page took too long and was stopped. No notes were invented. Try a clearer crop of one line, or enter the notes by hand."
          : "The local reader could not be started. No notes were invented.",
      };
    }

    const marker = result.stdout
      .split("\n")
      .map((line) => line.trim())
      .find((line) => line.startsWith("MUSICXML="));
    if (result.code !== 0 || !marker) {
      return {
        ok: false,
        status: 422,
        recognizer: info.name,
        error: `This image could not be read. ${shortFailure(result.stderr, result.stdout)} No notes were invented.`,
      };
    }

    const xmlPath = marker.slice("MUSICXML=".length);
    const xml = await readFile(xmlPath, "utf8");
    const parsed = parseMusicXml(xml);
    if (!parsed.ok) {
      return { ok: false, status: 422, recognizer: info.name, error: parsed.error };
    }

    const title = file.name.replace(/\.[^.]+$/, "").trim();
    return {
      ok: true,
      recognizer: info.name,
      warnings: parsed.warnings,
      score: {
        ...parsed.score,
        title: title || parsed.score.title,
        origin: "upload",
        attribution: `Read locally by ${info.name}. Correct any wrong pitch or rhythm before you play.`,
      },
    };
  } finally {
    await rm(directory, { recursive: true, force: true });
  }
}