fix(ocr): Arabic+Latin font fix and added tests

src/js/utils/ocr.ts

@@ -1,12 +1,14 @@
 import Tesseract from 'tesseract.js';
-import { PDFDocument, StandardFonts, rgb, PDFFont } from 'pdf-lib';
+import { PDFDocument, StandardFonts, rgb, PDFFont, degrees } from 'pdf-lib';
 import fontkit from '@pdf-lib/fontkit';
 import { getFontForLanguage } from './font-loader.js';
+import { languageToFontFamily } from '../config/font-mappings.js';
 import { OcrPage, OcrLine } from '@/types';
 import {
   parseHocrDocument,
   calculateWordTransform,
   calculateSpaceTransform,
+  scaleOcrPageToPdfPoints,
 } from './hocr-transform.js';
 import { getPDFDocument } from './helpers.js';
 import { loadPdfDocument } from './load-pdf-document.js';
@@ -18,13 +20,22 @@ export interface OcrOptions {
   binarize: boolean;
   whitelist: string;
   embedFullFonts?: boolean;
+  psm?: Tesseract.PSM;
+  perPageTimeoutMs?: number;
   onProgress?: (status: string, progress: number) => void;
 }
 
+export interface OcrWarning {
+  page: number;
+  kind: 'recognize-timeout' | 'recognize-error' | 'draw-error' | 'page-error';
+  message: string;
+}
+
 export interface OcrResult {
   pdfBytes: Uint8Array;
   pdfDoc: PDFDocument;
   fullText: string;
+  warnings: OcrWarning[];
 }
 
 function binarizeCanvas(ctx: CanvasRenderingContext2D) {
@@ -45,24 +56,27 @@ function drawOcrTextLayer(
   pageHeight: number,
   primaryFont: PDFFont,
   latinFont: PDFFont
-): void {
+): { droppedWords: number; droppedSpaces: number } {
+  let droppedWords = 0;
+  let droppedSpaces = 0;
   ocrPage.lines.forEach(function (line: OcrLine) {
-    const words = line.words;
+    const words =
+      line.direction === 'rtl' ? [...line.words].reverse() : line.words;
+    const lineRotation = degrees(
+      -line.textangle + Math.atan(line.baseline.slope) * (180 / Math.PI)
+    );
 
     for (let i = 0; i < words.length; i++) {
       const word = words[i];
-      const text = word.text.replace(
-        /[\u0000-\u001F\u007F-\u009F\u200E\u200F\u202A-\u202E\uFEFF]/g,
-        ''
-      );
+      const text = word.text.replace(/[\p{Cc}\p{Cf}]/gu, '');
 
       if (!text.trim()) continue;
 
-      const hasNonLatin = /[^\u0000-\u007F]/.test(text);
+      const hasNonLatin = /[^\p{ASCII}]/u.test(text);
       const font = hasNonLatin ? primaryFont : latinFont;
 
       if (!font) {
-        console.warn('Font not available for text: "' + text + '"');
+        droppedWords++;
         continue;
       }
 
@@ -87,11 +101,12 @@ function drawOcrTextLayer(
           y: transform.y,
           font,
           size: transform.fontSize,
-          color: rgb(0, 0, 0),
-          opacity: 0,
+          color: rgb(1, 0, 0),
+          opacity: 1,
+          rotate: lineRotation,
         });
-      } catch (error) {
-        console.warn(`Could not draw text "${text}":`, error);
+      } catch {
+        droppedWords++;
       }
 
       if (line.injectWordBreaks && i < words.length - 1) {
@@ -117,16 +132,43 @@ function drawOcrTextLayer(
               y: spaceTransform.y,
               font,
               size: spaceTransform.fontSize,
-              color: rgb(0, 0, 0),
-              opacity: 0,
+              color: rgb(1, 0, 0),
+              opacity: 1,
+              rotate: lineRotation,
             });
           } catch {
-            console.warn(`Could not draw space between words`);
+            droppedSpaces++;
           }
         }
       }
     }
   });
+  return { droppedWords, droppedSpaces };
+}
+
+const DEFAULT_PER_PAGE_TIMEOUT_MS = 5 * 60 * 1000;
+const PDF_CLEANUP_INTERVAL_PAGES = 25;
+
+function recognizeWithTimeout(
+  worker: Tesseract.Worker,
+  canvas: HTMLCanvasElement,
+  timeoutMs: number
+): Promise<Tesseract.RecognizeResult> {
+  return new Promise((resolve, reject) => {
+    const timer = setTimeout(() => {
+      reject(new Error(`OCR recognize timed out after ${timeoutMs}ms`));
+    }, timeoutMs);
+    worker
+      .recognize(canvas, {}, { text: true, hocr: true })
+      .then((result) => {
+        clearTimeout(timer);
+        resolve(result);
+      })
+      .catch((err) => {
+        clearTimeout(timer);
+        reject(err);
+      });
+  });
 }
 
 export async function performOcr(
@@ -139,9 +181,13 @@ export async function performOcr(
     binarize,
     whitelist,
     embedFullFonts,
+    psm,
+    perPageTimeoutMs,
     onProgress,
   } = options;
   const progress = onProgress || (() => {});
+  const pageTimeoutMs = perPageTimeoutMs ?? DEFAULT_PER_PAGE_TIMEOUT_MS;
+  const warnings: OcrWarning[] = [];
 
   const worker = await createConfiguredTesseractWorker(
     language,
@@ -151,140 +197,187 @@ export async function performOcr(
     }
   );
 
-  await worker.setParameters({
-    tessjs_create_hocr: '1',
-    tessedit_pageseg_mode: Tesseract.PSM.AUTO,
-  });
-
-  if (whitelist) {
-    await worker.setParameters({
-      tessedit_char_whitelist: whitelist,
-    });
-  }
-
-  const sourcePdfDoc = await loadPdfDocument(pdfBytes);
-  const pdf = await getPDFDocument({ data: pdfBytes }).promise;
-  const newPdfDoc = await PDFDocument.create();
-
-  newPdfDoc.registerFontkit(fontkit);
-
-  progress('Loading fonts...', 0);
-
-  const selectedLangs = language.split('+');
-  const cjkLangs = ['jpn', 'chi_sim', 'chi_tra', 'kor'];
-  const indicLangs = [
-    'hin',
-    'ben',
-    'guj',
-    'kan',
-    'mal',
-    'ori',
-    'pan',
-    'tam',
-    'tel',
-    'sin',
-  ];
-  const priorityLangs = [...cjkLangs, ...indicLangs, 'ara', 'rus', 'ukr'];
-
-  const primaryLang =
-    selectedLangs.find((l) => priorityLangs.includes(l)) ||
-    selectedLangs[0] ||
-    'eng';
-
-  const hasCJK = selectedLangs.some((l) => cjkLangs.includes(l));
-  const hasIndic = selectedLangs.some((l) => indicLangs.includes(l));
-  const hasLatin =
-    selectedLangs.some((l) => !priorityLangs.includes(l)) ||
-    selectedLangs.includes('eng');
-  const isIndicPlusLatin = hasIndic && hasLatin && !hasCJK;
-
-  let primaryFont: PDFFont;
-  let latinFont: PDFFont;
+  let fullText = '';
+  let newPdfDoc: PDFDocument | undefined;
 
   try {
-    if (isIndicPlusLatin) {
-      const [scriptFontBytes, latinFontBytes] = await Promise.all([
-        getFontForLanguage(primaryLang),
-        getFontForLanguage('eng'),
-      ]);
-      primaryFont = await newPdfDoc.embedFont(scriptFontBytes, {
-        subset: !embedFullFonts,
-      });
-      latinFont = await newPdfDoc.embedFont(latinFontBytes, {
-        subset: !embedFullFonts,
-      });
-    } else {
-      const fontBytes = await getFontForLanguage(primaryLang);
-      primaryFont = await newPdfDoc.embedFont(fontBytes, {
-        subset: !embedFullFonts,
+    await worker.setParameters({
+      tessedit_pageseg_mode: psm ?? Tesseract.PSM.AUTO,
+    });
+
+    if (whitelist) {
+      await worker.setParameters({
+        tessedit_char_whitelist: whitelist,
       });
+    }
+
+    const sourcePdfDoc = await loadPdfDocument(pdfBytes);
+    const pdf = await getPDFDocument({ data: pdfBytes }).promise;
+    newPdfDoc = await PDFDocument.create();
+
+    newPdfDoc.registerFontkit(fontkit);
+
+    progress('Loading fonts...', 0);
+
+    const selectedLangs = language.split('+');
+    const cjkLangs = ['jpn', 'chi_sim', 'chi_tra', 'kor'];
+    const indicLangs = [
+      'hin',
+      'ben',
+      'guj',
+      'kan',
+      'mal',
+      'ori',
+      'pan',
+      'tam',
+      'tel',
+      'sin',
+    ];
+    const priorityLangs = [...cjkLangs, ...indicLangs, 'ara', 'rus', 'ukr'];
+
+    const primaryLang =
+      selectedLangs.find((l) => priorityLangs.includes(l)) ||
+      selectedLangs[0] ||
+      'eng';
+
+    const hasCJK = selectedLangs.some((l) => cjkLangs.includes(l));
+    const hasLatin =
+      selectedLangs.some((l) => !priorityLangs.includes(l)) ||
+      selectedLangs.includes('eng');
+    const primaryFontFamily = languageToFontFamily[primaryLang] ?? 'Noto Sans';
+    const latinFontFamily = languageToFontFamily['eng'] ?? 'Noto Sans';
+    const sharesLatinFamily = primaryFontFamily === latinFontFamily;
+    const needsSeparateLatinFont =
+      priorityLangs.includes(primaryLang) &&
+      hasLatin &&
+      !hasCJK &&
+      !sharesLatinFamily;
+
+    let primaryFont: PDFFont;
+    let latinFont: PDFFont;
+
+    try {
+      if (needsSeparateLatinFont) {
+        const [scriptFontBytes, latinFontBytes] = await Promise.all([
+          getFontForLanguage(primaryLang),
+          getFontForLanguage('eng'),
+        ]);
+        primaryFont = await newPdfDoc.embedFont(scriptFontBytes, {
+          subset: !embedFullFonts,
+        });
+        latinFont = await newPdfDoc.embedFont(latinFontBytes, {
+          subset: !embedFullFonts,
+        });
+      } else {
+        const fontBytes = await getFontForLanguage(primaryLang);
+        primaryFont = await newPdfDoc.embedFont(fontBytes, {
+          subset: !embedFullFonts,
+        });
+        latinFont = primaryFont;
+      }
+    } catch (e) {
+      console.error('Font loading failed, falling back to Helvetica', e);
+      primaryFont = await newPdfDoc.embedFont(StandardFonts.Helvetica);
       latinFont = primaryFont;
     }
-  } catch (e) {
-    console.error('Font loading failed, falling back to Helvetica', e);
-    primaryFont = await newPdfDoc.embedFont(StandardFonts.Helvetica);
-    latinFont = primaryFont;
-  }
 
-  let fullText = '';
-
-  try {
     for (let i = 1; i <= pdf.numPages; i++) {
       progress(
         `Processing page ${i} of ${pdf.numPages}`,
         (i - 1) / pdf.numPages
       );
 
-      const page = await pdf.getPage(i);
-      const viewport = page.getViewport({ scale: resolution });
+      const pdfJsPage = await pdf.getPage(i);
+      try {
+        const viewport = pdfJsPage.getViewport({ scale: resolution });
 
-      const canvas = document.createElement('canvas');
-      canvas.width = viewport.width;
-      canvas.height = viewport.height;
-      const context = canvas.getContext('2d');
-      if (!context) throw new Error('Failed to create canvas context');
+        const canvas = document.createElement('canvas');
+        canvas.width = viewport.width;
+        canvas.height = viewport.height;
+        const context = canvas.getContext('2d');
+        if (!context) throw new Error('Failed to create canvas context');
 
-      await page.render({ canvasContext: context, viewport, canvas }).promise;
+        await pdfJsPage.render({ canvasContext: context, viewport, canvas })
+          .promise;
 
-      if (binarize) {
-        binarizeCanvas(context);
+        if (binarize) {
+          binarizeCanvas(context);
+        }
+
+        let pageText = '';
+        let pageHocr = '';
+        try {
+          const result = await recognizeWithTimeout(
+            worker,
+            canvas,
+            pageTimeoutMs
+          );
+          pageText = result.data.text ?? '';
+          pageHocr = result.data.hocr ?? '';
+        } catch (err) {
+          const message = err instanceof Error ? err.message : String(err);
+          warnings.push({
+            page: i,
+            kind: /timed out/i.test(message)
+              ? 'recognize-timeout'
+              : 'recognize-error',
+            message,
+          });
+        }
+
+        canvas.width = 0;
+        canvas.height = 0;
+
+        const [copiedPage] = await newPdfDoc.copyPages(sourcePdfDoc, [i - 1]);
+        newPdfDoc.addPage(copiedPage);
+
+        if (pageHocr) {
+          const ocrPage = parseHocrDocument(pageHocr);
+          const pdfPageWidth = copiedPage.getWidth();
+          const pdfPageHeight = copiedPage.getHeight();
+          scaleOcrPageToPdfPoints(ocrPage, pdfPageWidth, pdfPageHeight);
+          const drawStats = drawOcrTextLayer(
+            copiedPage,
+            ocrPage,
+            pdfPageHeight,
+            primaryFont,
+            latinFont
+          );
+          if (drawStats.droppedWords > 0 || drawStats.droppedSpaces > 0) {
+            warnings.push({
+              page: i,
+              kind: 'draw-error',
+              message: `${drawStats.droppedWords} word(s) and ${drawStats.droppedSpaces} space(s) failed to render in the text layer`,
+            });
+          }
+        }
+
+        fullText += pageText + '\n\n';
+      } catch (err) {
+        const message = err instanceof Error ? err.message : String(err);
+        warnings.push({ page: i, kind: 'page-error', message });
+      } finally {
+        pdfJsPage.cleanup();
       }
 
-      const result = await worker.recognize(
-        canvas,
-        {},
-        { text: true, hocr: true }
-      );
-      const data = result.data;
-
-      canvas.width = 0;
-      canvas.height = 0;
-
-      const [copiedPage] = await newPdfDoc.copyPages(sourcePdfDoc, [i - 1]);
-      newPdfDoc.addPage(copiedPage);
-
-      if (data.hocr) {
-        const ocrPage = parseHocrDocument(data.hocr);
-        drawOcrTextLayer(
-          copiedPage,
-          ocrPage,
-          viewport.height,
-          primaryFont,
-          latinFont
-        );
+      if (i % PDF_CLEANUP_INTERVAL_PAGES === 0) {
+        pdf.cleanup();
       }
-
-      fullText += data.text + '\n\n';
     }
   } finally {
     await worker.terminate();
   }
 
+  if (!newPdfDoc) {
+    throw new Error('OCR aborted before any page was processed');
+  }
+
   const savedBytes = await newPdfDoc.save();
 
   return {
     pdfBytes: new Uint8Array(savedBytes),
     pdfDoc: newPdfDoc,
     fullText,
+    warnings,
   };
 }