111 lines
4.0 KiB
TypeScript
111 lines
4.0 KiB
TypeScript
import type { ToolInput } from "../../core/io/input-types";
|
|
import type { ToolResult } from "../../core/io/output-types";
|
|
import type { ToolContext } from "../../core/plugins/plugin-types";
|
|
import type { PdfToDocxOcrOptions } from "./index";
|
|
|
|
function abortIfNeeded(signal: AbortSignal) {
|
|
if (signal.aborted) throw new DOMException("Operation cancelled", "AbortError");
|
|
}
|
|
|
|
export async function runPdfToDocxOcr(
|
|
input: ToolInput,
|
|
options: PdfToDocxOcrOptions,
|
|
context: ToolContext
|
|
): Promise<ToolResult> {
|
|
const file = input.files?.[0];
|
|
if (!file) throw new Error("Add a PDF document.");
|
|
if (file.type && file.type !== "application/pdf") throw new Error("The selected file is not a PDF.");
|
|
|
|
abortIfNeeded(context.signal);
|
|
context.reportProgress({ percentage: 1, message: "Opening PDF locally..." });
|
|
|
|
const [pdfjs, { createWorker }] = await Promise.all([
|
|
import("pdfjs-dist"),
|
|
import("tesseract.js"),
|
|
]);
|
|
pdfjs.GlobalWorkerOptions.workerSrc = new URL(
|
|
"pdfjs-dist/build/pdf.worker.min.mjs",
|
|
import.meta.url
|
|
).toString();
|
|
|
|
const loadingTask = pdfjs.getDocument({ data: new Uint8Array(await file.arrayBuffer()) });
|
|
const pdf = await loadingTask.promise;
|
|
const pageTexts: string[] = [];
|
|
const scale = options.quality === "high" ? 2.5 : 1.75;
|
|
|
|
const worker = await createWorker(options.language, undefined, {
|
|
logger: (message) => {
|
|
if (message.status === "recognizing text" && typeof message.progress === "number") {
|
|
const completedPages = pageTexts.length;
|
|
context.reportProgress({
|
|
percentage: 8 + ((completedPages + message.progress) / pdf.numPages) * 88,
|
|
message: `Recognizing page ${completedPages + 1} of ${pdf.numPages}...`,
|
|
});
|
|
}
|
|
},
|
|
});
|
|
|
|
const cancelWorker = () => void worker.terminate();
|
|
context.signal.addEventListener("abort", cancelWorker, { once: true });
|
|
|
|
try {
|
|
for (let pageNumber = 1; pageNumber <= pdf.numPages; pageNumber += 1) {
|
|
abortIfNeeded(context.signal);
|
|
context.reportProgress({
|
|
percentage: 5 + ((pageNumber - 1) / pdf.numPages) * 88,
|
|
message: `Rendering page ${pageNumber} of ${pdf.numPages} locally...`,
|
|
});
|
|
|
|
const page = await pdf.getPage(pageNumber);
|
|
const viewport = page.getViewport({ scale });
|
|
const canvas = document.createElement("canvas");
|
|
canvas.width = Math.ceil(viewport.width);
|
|
canvas.height = Math.ceil(viewport.height);
|
|
const canvasContext = canvas.getContext("2d", { alpha: false });
|
|
if (!canvasContext) throw new Error("Your browser could not render this PDF page.");
|
|
await page.render({ canvas, canvasContext, viewport }).promise;
|
|
abortIfNeeded(context.signal);
|
|
|
|
const result = await worker.recognize(canvas);
|
|
pageTexts.push(result.data.text.trim());
|
|
page.cleanup();
|
|
canvas.width = 1;
|
|
canvas.height = 1;
|
|
}
|
|
} finally {
|
|
context.signal.removeEventListener("abort", cancelWorker);
|
|
await worker.terminate().catch(() => undefined);
|
|
await loadingTask.destroy();
|
|
}
|
|
|
|
abortIfNeeded(context.signal);
|
|
context.reportProgress({ percentage: 97, message: "Building editable DOCX..." });
|
|
const { Document, Packer, Paragraph } = await import("docx");
|
|
const children: Array<InstanceType<typeof Paragraph>> = [];
|
|
pageTexts.forEach((text, pageIndex) => {
|
|
const lines = text.split(/\r?\n/);
|
|
lines.forEach((line, lineIndex) => {
|
|
children.push(new Paragraph({
|
|
text: line,
|
|
pageBreakBefore: options.pageBreaks && pageIndex > 0 && lineIndex === 0,
|
|
spacing: { after: line.trim() ? 120 : 0 },
|
|
}));
|
|
});
|
|
});
|
|
|
|
const docxDocument = new Document({ sections: [{ children }] });
|
|
const blob = await Packer.toBlob(docxDocument);
|
|
context.reportProgress({ percentage: 100, message: "DOCX ready." });
|
|
|
|
return {
|
|
type: "files",
|
|
files: [{
|
|
name: `${file.name.replace(/\.pdf$/i, "")}_ocr.docx`,
|
|
mimeType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
blob,
|
|
sizeBefore: file.size,
|
|
sizeAfter: blob.size,
|
|
}],
|
|
};
|
|
}
|