268 lines
10 KiB
TypeScript
268 lines
10 KiB
TypeScript
import { NextRequest, NextResponse } from "next/server";
|
|
import fs from "fs";
|
|
import path from "path";
|
|
import { Client } from "@gradio/client";
|
|
import { query } from "../../../db";
|
|
|
|
export const maxDuration = 120; // Allow up to 120 seconds for slow model inference
|
|
|
|
export async function GET(req: NextRequest) {
|
|
try {
|
|
const { searchParams } = new URL(req.url);
|
|
const action = searchParams.get("action") || "list";
|
|
const runId = searchParams.get("runId");
|
|
const imageType = searchParams.get("imageType"); // 'do', 'product' or null for all
|
|
|
|
if (runId) {
|
|
const runRes = await query(`
|
|
SELECT id, image_path, engine, status, ocr_result, time_elapsed_ms, image_type, created_at
|
|
FROM arena_runs
|
|
WHERE id = $1
|
|
`, [parseInt(runId)]);
|
|
|
|
if (runRes.rowCount === 0) {
|
|
return NextResponse.json({ success: false, error: "Run not found" }, { status: 404 });
|
|
}
|
|
return NextResponse.json({ success: true, run: runRes.rows[0] });
|
|
}
|
|
|
|
if (action === "stats") {
|
|
let queryText = `
|
|
SELECT
|
|
engine,
|
|
COUNT(*)::integer as total_runs,
|
|
COUNT(CASE WHEN status = 'done' THEN 1 END)::integer as success_runs,
|
|
COUNT(CASE WHEN status = 'failed' THEN 1 END)::integer as failed_runs,
|
|
ROUND(AVG(CASE WHEN status = 'done' THEN time_elapsed_ms END))::integer as avg_time_ms,
|
|
MIN(CASE WHEN status = 'done' THEN time_elapsed_ms END)::integer as min_time_ms,
|
|
MAX(CASE WHEN status = 'done' THEN time_elapsed_ms END)::integer as max_time_ms
|
|
FROM arena_runs
|
|
`;
|
|
const params: any[] = [];
|
|
if (imageType === "do" || imageType === "product") {
|
|
queryText += ` WHERE image_type = $1`;
|
|
params.push(imageType);
|
|
}
|
|
queryText += ` GROUP BY engine`;
|
|
|
|
const statsRes = await query(queryText, params);
|
|
return NextResponse.json({ success: true, stats: statsRes.rows });
|
|
}
|
|
|
|
const limit = parseInt(searchParams.get("limit") || "50");
|
|
let queryText = `
|
|
SELECT id, image_path, engine, status, time_elapsed_ms, image_type, created_at
|
|
FROM arena_runs
|
|
`;
|
|
const params: any[] = [];
|
|
if (imageType === "do" || imageType === "product") {
|
|
queryText += ` WHERE image_type = $1`;
|
|
params.push(imageType);
|
|
}
|
|
queryText += ` ORDER BY created_at DESC LIMIT $${params.length + 1}`;
|
|
params.push(limit);
|
|
|
|
const runsRes = await query(queryText, params);
|
|
return NextResponse.json({ success: true, runs: runsRes.rows });
|
|
} catch (error: any) {
|
|
console.error("Failed to fetch arena runs/stats:", error);
|
|
return NextResponse.json({ success: false, error: error.message }, { status: 500 });
|
|
}
|
|
}
|
|
|
|
export async function POST(req: NextRequest) {
|
|
const startTime = Date.now();
|
|
let engine: string | undefined;
|
|
let image: string | undefined;
|
|
let imageType = "do";
|
|
try {
|
|
const body = await req.json().catch(() => ({}));
|
|
engine = body.engine;
|
|
image = body.image;
|
|
|
|
if (!engine || !image) {
|
|
return NextResponse.json({ error: "Missing engine or image" }, { status: 400 });
|
|
}
|
|
|
|
imageType = body.imageType || "do";
|
|
if (typeof image === "string") {
|
|
if (image.startsWith("/produk-pfm/") || image.includes("produk-pfm") || image.includes("Product")) {
|
|
imageType = "product";
|
|
} else if (image.startsWith("/do-pfm/") || image.includes("do-pfm")) {
|
|
imageType = "do";
|
|
}
|
|
}
|
|
|
|
let imageBuffer: Buffer;
|
|
let base64Image = "";
|
|
|
|
// 1. Resolve image (local file or base64)
|
|
if (typeof image === "string" && (image.startsWith("/do-pfm/") || image.startsWith("/produk-pfm/"))) {
|
|
// Resolve path in public folder
|
|
const cleanPath = image.startsWith("/") ? image.slice(1) : image;
|
|
const filePath = path.join(process.cwd(), "public", cleanPath);
|
|
|
|
if (!fs.existsSync(filePath)) {
|
|
return NextResponse.json({ error: `File not found on server: ${image}` }, { status: 404 });
|
|
}
|
|
imageBuffer = fs.readFileSync(filePath);
|
|
base64Image = `data:image/jpeg;base64,${imageBuffer.toString("base64")}`;
|
|
} else if (typeof image === "string" && image.startsWith("data:")) {
|
|
// Base64 data URI
|
|
base64Image = image;
|
|
const base64Data = image.split(",")[1];
|
|
imageBuffer = Buffer.from(base64Data, "base64");
|
|
} else if (typeof image === "string") {
|
|
// Raw base64 string
|
|
base64Image = `data:image/jpeg;base64,${image}`;
|
|
imageBuffer = Buffer.from(image, "base64");
|
|
} else {
|
|
return NextResponse.json({ error: "Invalid image format" }, { status: 400 });
|
|
}
|
|
|
|
let outputText = "";
|
|
|
|
// 2. Route to the requested OCR engine
|
|
if (engine === "deepseek") {
|
|
const blob = new Blob([new Uint8Array(imageBuffer)], { type: "image/jpeg" });
|
|
const gradioUrl = process.env.DEEPSEEK_GRADIO_URL || "http://host.docker.internal:7873/v2/";
|
|
const client = await Client.connect(gradioUrl);
|
|
const result = await client.predict(2, [blob, "Default", "Markdown", ""]);
|
|
const data = result.data as any[];
|
|
outputText = data[1] || data[0] || "";
|
|
|
|
} else if (engine === "lightonocr") {
|
|
const url = process.env.LIGHTONOCR_API_URL || "http://host.docker.internal:7678/layout-parsing";
|
|
const res = await fetch(url, {
|
|
method: "POST",
|
|
headers: { "Content-Type": "application/json" },
|
|
body: JSON.stringify({
|
|
file: base64Image,
|
|
useLayoutDetection: false
|
|
})
|
|
});
|
|
if (!res.ok) {
|
|
throw new Error(`LightOnOCR backend error: ${res.status} ${await res.text()}`);
|
|
}
|
|
const data = await res.json();
|
|
outputText = data.result?.layoutParsingResults?.[0]?.markdown?.text || "";
|
|
|
|
} else if (engine === "nemotron") {
|
|
const url = process.env.NEMOTRON_API_URL || "http://host.docker.internal:8009/layout-parsing";
|
|
const res = await fetch(url, {
|
|
method: "POST",
|
|
headers: { "Content-Type": "application/json" },
|
|
body: JSON.stringify({
|
|
file: base64Image,
|
|
model: "Multilingual (en, zh, ja, ko, ru, …)",
|
|
merge_level: "layout"
|
|
})
|
|
});
|
|
if (!res.ok) {
|
|
throw new Error(`Nemotron backend error: ${res.status} ${await res.text()}`);
|
|
}
|
|
const data = await res.json();
|
|
outputText = data.result?.layoutParsingResults?.[0]?.markdown?.text || "";
|
|
|
|
} else if (engine === "paddle") {
|
|
const url = process.env.PIPELINE_URL || "http://paddleocr-pipeline-api:8090/layout-parsing";
|
|
const rawB64 = base64Image.includes(",") ? base64Image.split(",")[1] : base64Image;
|
|
const res = await fetch(url, {
|
|
method: "POST",
|
|
headers: { "Content-Type": "application/json" },
|
|
body: JSON.stringify({
|
|
file: rawB64,
|
|
matchHistoryJob: false,
|
|
useLayoutDetection: true,
|
|
fileType: 1,
|
|
useDocUnwarping: false,
|
|
useDocOrientationClassify: false
|
|
})
|
|
});
|
|
if (!res.ok) {
|
|
throw new Error(`PaddleOCR backend error: ${res.status} ${await res.text()}`);
|
|
}
|
|
const data = await res.json();
|
|
const pipelineResult = data.result || data;
|
|
outputText = pipelineResult?.layoutParsingResults?.[0]?.markdown?.text || "";
|
|
|
|
} else if (engine === "dots") {
|
|
// Calling python API directly
|
|
const url = process.env.DOTS_API_URL || "http://host.docker.internal:7872/layout-parsing";
|
|
const res = await fetch(url, {
|
|
method: "POST",
|
|
headers: { "Content-Type": "application/json" },
|
|
body: JSON.stringify({
|
|
file: base64Image,
|
|
promptLabel: "ocr",
|
|
useLayoutDetection: true
|
|
})
|
|
});
|
|
if (!res.ok) {
|
|
throw new Error(`Dots OCR backend error: ${res.status} ${await res.text()}`);
|
|
}
|
|
const data = await res.json();
|
|
outputText = data.result?.layoutParsingResults?.[0]?.markdown?.text || "";
|
|
|
|
} else if (engine === "glm") {
|
|
const gradioUrl = process.env.GLM_GRADIO_URL || "http://host.docker.internal:7875/";
|
|
const client = await Client.connect(gradioUrl);
|
|
const result = await client.predict(2, ["Text", base64Image, 1024, 60]);
|
|
const data = result.data as any[];
|
|
outputText = data[0] || "";
|
|
|
|
} else {
|
|
return NextResponse.json({ error: `Unknown engine: ${engine}` }, { status: 400 });
|
|
}
|
|
|
|
const elapsedMs = Date.now() - startTime;
|
|
|
|
// Record successful run
|
|
try {
|
|
const loggedImagePath = (typeof image === "string" && image.startsWith("data:"))
|
|
? `[Base64 Upload: ${image.length} chars]`
|
|
: (typeof image === "string" && image.length > 500)
|
|
? `[Raw Base64: ${image.length} chars]`
|
|
: image;
|
|
await query(
|
|
`INSERT INTO arena_runs (image_path, engine, status, ocr_result, time_elapsed_ms, image_type)
|
|
VALUES ($1, $2, $3, $4, $5, $6)`,
|
|
[loggedImagePath, engine, "done", outputText, elapsedMs, imageType]
|
|
);
|
|
} catch (dbErr) {
|
|
console.error("Failed to log success to arena_runs:", dbErr);
|
|
}
|
|
|
|
return NextResponse.json({
|
|
success: true,
|
|
text: outputText,
|
|
elapsedMs
|
|
});
|
|
|
|
} catch (error: any) {
|
|
console.error("OCR Arena proxy error:", error);
|
|
const elapsedMs = Date.now() - startTime;
|
|
|
|
// Record failed run
|
|
try {
|
|
const loggedImagePath = (typeof image === "string" && image.startsWith("data:"))
|
|
? `[Base64 Upload: ${image.length} chars]`
|
|
: (typeof image === "string" && image.length > 500)
|
|
? `[Raw Base64: ${image.length} chars]`
|
|
: image;
|
|
await query(
|
|
`INSERT INTO arena_runs (image_path, engine, status, ocr_result, time_elapsed_ms, image_type)
|
|
VALUES ($1, $2, $3, $4, $5, $6)`,
|
|
[loggedImagePath || "unknown", engine || "unknown", "failed", error.message || "Unknown error", elapsedMs, imageType]
|
|
);
|
|
} catch (dbErr) {
|
|
console.error("Failed to log failure to arena_runs:", dbErr);
|
|
}
|
|
|
|
return NextResponse.json({
|
|
success: false,
|
|
error: error.message || "Failed to process OCR request"
|
|
}, { status: 500 });
|
|
}
|
|
}
|