dsh-ab-ocr 0.1.0

package/python/source.py

@@ -0,0 +1,182 @@
+"""Document sources: one still image, or a PDF rendered one page at a time.
+
+Every source reports how many pages it carries and renders one page on demand,
+so the caller can OCR a page and drop its pixels before asking for the next. No
+source keeps more than the page currently being processed.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Any
+
+#: Bytes read at a time while digesting a source, so a large scan is not held twice.
+DIGEST_CHUNK = 1024 * 1024
+
+#: Extensions handled as a still image (a multi-frame file is read frame by frame).
+IMAGE_SUFFIXES = frozenset({
+    '.png', '.jpg', '.jpeg', '.jpe', '.bmp', '.dib', '.tif', '.tiff',
+    '.webp', '.gif', '.jp2', '.pbm', '.pgm', '.ppm', '.tga', '.ico',
+})
+
+#: Extensions handled as a PDF.
+PDF_SUFFIXES = frozenset({'.pdf'})
+
+#: Smallest render scale, so an absurd maxPixels still yields a readable page.
+MIN_SCALE = 0.2
+
+
+class UnsupportedDocument(Exception):
+    """Raised when a path is neither a readable image nor a readable PDF."""
+
+
+def source_digest(path: str) -> str:
+    """Digest a source file's bytes.
+
+    The caller names a recognition after this, so the name follows the document's
+    content and not merely its length: two files that happen to weigh the same are
+    still two recognitions, and a file that has not changed is always the same one.
+    @param path: the source file to read.
+    @returns: the lowercase hex digest of its bytes.
+    """
+    import hashlib
+
+    digest = hashlib.sha256()
+    with open(path, 'rb') as handle:
+        for chunk in iter(lambda: handle.read(DIGEST_CHUNK), b''):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+@dataclass
+class RenderedPage:
+    """One page's pixels, in the channel order the OCR engine expects."""
+
+    index: int
+    width: int
+    height: int
+    image: Any
+
+
+class Document:
+    """A page source the worker renders and releases one page at a time."""
+
+    def page_count(self) -> int:
+        """Return the number of pages this source carries."""
+        raise NotImplementedError
+
+    def render(self, index: int) -> RenderedPage:
+        """Render one zero-based page into a BGR array."""
+        raise NotImplementedError
+
+    def close(self) -> None:
+        """Release the underlying file handle or document."""
+
+
+def document_kind(path: str) -> str:
+    """Classify a path as a PDF, an image, or unsupported."""
+    suffix = os.path.splitext(path)[1].lower()
+    if suffix in PDF_SUFFIXES:
+        return 'pdf'
+    if suffix in IMAGE_SUFFIXES:
+        return 'image'
+    return 'unsupported'
+
+
+class ImageDocument(Document):
+    """A still image, or every frame of a multi-frame image file."""
+
+    def __init__(self, path: str, max_pixels: int) -> None:
+        from PIL import Image
+
+        self._image = Image.open(path)
+        self._frames = int(getattr(self._image, 'n_frames', 1))
+        self._index = 0
+        self._max_pixels = max_pixels
+        if self._frames > 1:
+            # Force a decode now so a corrupt later frame fails before any OCR.
+            self._image.seek(0)
+
+    def page_count(self) -> int:
+        return self._frames
+
+    def render(self, index: int) -> RenderedPage:
+        # Imported here so the worker's protocol helpers load on an interpreter
+        # that carries no OCR package, which is where the unit suite runs.
+        import numpy as np
+
+        if index != self._index:
+            self._image.seek(index)
+            self._index = index
+        frame = self._image.convert('RGB')
+        array = np.asarray(frame)
+        array = np.ascontiguousarray(array[:, :, ::-1])
+        pixels = self._max_pixels
+        if pixels > 0 and array.shape[0] * array.shape[1] > pixels:
+            step = max(1, int(np.ceil(np.sqrt(array.shape[0] * array.shape[1] / pixels))))
+            array = np.ascontiguousarray(array[::step, ::step])
+        return RenderedPage(index=index + 1, width=array.shape[1], height=array.shape[0], image=array)
+
+    def close(self) -> None:
+        self._image.close()
+
+
+class PdfDocument(Document):
+    """A PDF, rendered with PDFium at a scale derived from the requested DPI."""
+
+    def __init__(self, path: str, dpi: int, max_pixels: int) -> None:
+        import pypdfium2 as pdfium
+
+        self._pdfium = pdfium
+        self._document = pdfium.PdfDocument(path)
+        self._dpi = dpi
+        self._max_pixels = max_pixels
+
+    def page_count(self) -> int:
+        return len(self._document)
+
+    def render(self, index: int) -> RenderedPage:
+        import numpy as np
+
+        page = self._document[index]
+        try:
+            scale = self._dpi / 72.0
+            if self._max_pixels > 0:
+                width_points, height_points = page.get_size()
+                pixels = (width_points * scale) * (height_points * scale)
+                if pixels > self._max_pixels:
+                    scale = max(MIN_SCALE, scale * (self._max_pixels / pixels) ** 0.5)
+            bitmap = page.render(scale=scale)
+            try:
+                # PDFium's default bitmap is BGRA, which is the channel order the
+                # OCR engine's detector and recognizer were trained on.
+                array = np.ascontiguousarray(bitmap.to_numpy())
+            finally:
+                bitmap.close()
+        finally:
+            # Closing each page immediately is what keeps a long document's
+            # resident cost flat: PDFium holds per-page caches otherwise.
+            page.close()
+        return RenderedPage(index=index + 1, width=array.shape[1], height=array.shape[0], image=array)
+
+    def close(self) -> None:
+        self._document.close()
+
+
+def open_document(path: str, dpi: int, max_pixels: int) -> Document:
+    """Open a path as a page source.
+
+    @param path: path to a PDF or a still image.
+    @param dpi: render resolution for PDF pages.
+    @param max_pixels: ceiling on one rendered page's pixel count; 0 disables it.
+    @returns: a document that renders one page per call.
+    """
+    kind = document_kind(path)
+    if kind == 'pdf':
+        return PdfDocument(path, dpi, max_pixels)
+    if kind == 'image':
+        return ImageDocument(path, max_pixels)
+    raise UnsupportedDocument(
+        'unsupported document: ' + os.path.basename(path) + ' is neither a PDF nor a supported image'
+    )

package/scripts/setup.mjs

@@ -0,0 +1,251 @@
+/**
+ * Create the Python environment the OCR worker runs in.
+ *
+ * The OCR engine and its model weights ship inside the RapidOCR wheel, so
+ * "installing the model" and "installing the dependencies" are the same act:
+ * a virtual environment holding `python/requirements.txt`. This script performs
+ * that act, verifies it with the worker's own `--self-test`, and reports what it
+ * found — including the model files the wheel actually carries, so a failure to
+ * recognize can be told apart from a failure to install.
+ *
+ * It is idempotent: a run against a working environment verifies and exits
+ * without touching the network. `--force` rebuilds the environment from scratch.
+ * @module @deepseek-ai/dsh-ab-ocr/scripts/setup
+ */
+
+import { spawnSync } from 'node:child_process'
+import { existsSync, readFileSync, rmSync } from 'node:fs'
+import { dirname, join, resolve } from 'node:path'
+import { fileURLToPath } from 'node:url'
+
+const HERE = dirname(fileURLToPath(import.meta.url))
+const PACKAGE = resolve(HERE, '..')
+const PYTHON_DIR = join(PACKAGE, 'python')
+const REQUIREMENTS = join(PYTHON_DIR, 'requirements.txt')
+const WORKER = join(PYTHON_DIR, 'ocr_worker.py')
+const VENV = join(PYTHON_DIR, '.venv')
+
+/** Interpreter names to try, in order, when no path is given. */
+const CANDIDATES = ['python', 'python3', 'py']
+
+/**
+ * Run one command and capture its streams.
+ * @param command - the executable.
+ * @param args - its arguments.
+ * @returns the exit status and the two streams.
+ */
+function run(command, args) {
+  const result = spawnSync(command, args, { encoding: 'utf8', windowsHide: true })
+  return {
+    status: result.status,
+    stdout: result.stdout ?? '',
+    stderr: result.stderr ?? '',
+    error: result.error,
+  }
+}
+
+/**
+ * Find a base interpreter that runs.
+ * @param requested - an explicit path from the command line or the environment.
+ * @returns the working command, or undefined when none of the candidates run.
+ */
+function findInterpreter(requested) {
+  const tried = requested === undefined ? CANDIDATES : [requested]
+  for (const candidate of tried) {
+    const probe = run(candidate, ['--version'])
+    if (probe.error === undefined && probe.status === 0) return candidate
+  }
+  return undefined
+}
+
+/**
+ * The interpreter inside the virtual environment.
+ * @returns its path, which differs by platform.
+ */
+function venvInterpreter() {
+  return process.platform === 'win32'
+    ? join(VENV, 'Scripts', 'python.exe')
+    : join(VENV, 'bin', 'python')
+}
+
+/**
+ * Ask the worker to report its environment.
+ * @param python - the interpreter to ask.
+ * @returns the parsed report, or undefined when the worker did not answer.
+ */
+function selfTest(python) {
+  const result = run(python, [WORKER, '--self-test'])
+  const line = result.stdout.split('\n').find(entry => entry.trim() !== '')
+  if (line === undefined) return undefined
+  try {
+    return JSON.parse(line)
+  } catch {
+    // A worker that cannot start prints nothing parseable; the caller reports the
+    // absence rather than this parse.
+    return undefined
+  }
+}
+
+/**
+ * The model files the installed wheel carries.
+ * @param python - the interpreter holding RapidOCR.
+ * @returns file names mapped to their size in bytes, or an empty record.
+ */
+function modelFiles(python) {
+  const script = [
+    'import json, os, rapidocr',
+    "d = os.path.join(os.path.dirname(rapidocr.__file__), 'models')",
+    "names = sorted(n for n in os.listdir(d) if n.endswith('.onnx')) if os.path.isdir(d) else []",
+    'print(json.dumps({n: os.path.getsize(os.path.join(d, n)) for n in names}))',
+  ].join('\n')
+  const result = run(python, ['-c', script])
+  if (result.status !== 0) return {}
+  try {
+    return JSON.parse(result.stdout.trim())
+  } catch {
+    return {}
+  }
+}
+
+/**
+ * Parse the command line.
+ * @param argv - the arguments after the script name.
+ * @returns the options, or a message describing what is wrong.
+ */
+function parseArgs(argv) {
+  const options = { force: false, prefetch: false, json: false, python: process.env.DSH_OCR_PYTHON }
+  for (let index = 0; index < argv.length; index += 1) {
+    const flag = argv[index]
+    if (flag === '--force') options.force = true
+    else if (flag === '--prefetch') options.prefetch = true
+    else if (flag === '--json') options.json = true
+    else if (flag === '--python') {
+      const value = argv[index += 1]
+      if (value === undefined || value.startsWith('--')) return { error: '--python needs an interpreter path' }
+      options.python = value
+    } else return { error: `unknown option ${JSON.stringify(flag)}` }
+  }
+  return options
+}
+
+/**
+ * Lay out one version of a requirement line for the report.
+ * @param requirements - the requirement file's text.
+ * @returns the requirement lines, comments dropped.
+ */
+function requirements(requirements) {
+  return readFileSync(requirements, 'utf8')
+    .split('\n')
+    .map(line => line.trim())
+    .filter(line => line !== '' && !line.startsWith('#'))
+}
+
+/**
+ * Install one environment and verify it.
+ * @param options - the parsed command line.
+ * @returns the process exit status.
+ */
+function main(options) {
+  const say = (message) => {
+    if (!options.json) console.log(message)
+  }
+  const base = findInterpreter(options.python)
+  if (base === undefined) {
+    console.error('setup: no Python interpreter found; pass --python <path> or set DSH_OCR_PYTHON')
+    return 2
+  }
+
+  const python = venvInterpreter()
+  const existing = existsSync(python)
+  say(`setup: base interpreter ${base}`)
+  say(`setup: environment ${VENV}${existing ? '' : ' (creating)'}`)
+
+  if (options.force && existsSync(VENV)) {
+    rmSync(VENV, { recursive: true, force: true })
+    say('setup: removed the existing environment (--force)')
+  }
+
+  if (!existsSync(python)) {
+    const created = run(base, ['-m', 'venv', VENV])
+    if (created.status !== 0) {
+      console.error(`setup: could not create the environment\n${created.stderr.trim()}`)
+      return 3
+    }
+  }
+
+  let report = selfTest(python)
+  if (report?.ready !== true) {
+    say('setup: installing ' + requirements(REQUIREMENTS).join(', '))
+    // Prefer a local wheels cache shipped with the plugin (python/wheels) so
+    // setup works offline; pip still falls back to the configured index for any
+    // wheel not present locally.
+    const wheelsDir = join(PYTHON_DIR, 'wheels')
+    const findLinks = existsSync(wheelsDir)
+      ? ['--find-links', wheelsDir]
+      : []
+    if (findLinks.length) say(`setup: using local wheel cache ${wheelsDir}`)
+    // Bootstrap the build tools from the local cache so --no-build-isolation can
+    // compile any source distribution (e.g. antlr4-python3-runtime) without
+    // reaching the network. Standard venvs ship these, but an offline venv may
+    // not, and a build-isolation subprocess would otherwise hit the index.
+    if (findLinks.length) {
+      const bootstrap = run(python, ['-m', 'pip', 'install', '--disable-pip-version-check', ...findLinks, 'setuptools', 'wheel'])
+      if (bootstrap.status !== 0) {
+        console.error(`setup: could not bootstrap build tools\n${bootstrap.stderr.trim()}`)
+        return 3
+      }
+    }
+    const install = run(python, ['-m', 'pip', 'install', '--disable-pip-version-check', '--no-build-isolation', ...findLinks, '-r', REQUIREMENTS])
+    if (install.status !== 0) {
+      console.error(`setup: pip failed\n${install.stderr.trim()}`)
+      return 3
+    }
+    report = selfTest(python)
+  }
+  if (report?.ready !== true) {
+    console.error('setup: the environment still fails its own check')
+    for (const [name, version] of Object.entries(report?.modules ?? {})) {
+      if (version === null || String(name).endsWith('Error')) console.error(`  ${name}: ${version ?? 'missing'}`)
+    }
+    return 3
+  }
+
+  const models = modelFiles(python)
+  if (options.prefetch) {
+    say('setup: loading the engine and its models')
+    const warm = run(python, ['-c', 'import logging; logging.disable(logging.CRITICAL); from rapidocr import RapidOCR; RapidOCR()'])
+    if (warm.status !== 0) {
+      console.error(`setup: the engine did not load\n${warm.stderr.trim()}`)
+      return 3
+    }
+  }
+
+  const summary = {
+    python,
+    base,
+    version: report.python,
+    modules: report.modules,
+    models,
+    prefetched: options.prefetch,
+  }
+  if (options.json) {
+    console.log(JSON.stringify(summary, null, 2))
+    return 0
+  }
+  say(`setup: ready — Python ${report.python}`)
+  for (const [name, version] of Object.entries(report.modules)) say(`  ${name} ${version}`)
+  const names = Object.keys(models)
+  const bytes = Object.values(models).reduce((total, size) => total + size, 0)
+  say(names.length === 0
+    ? 'setup: no model files found beside the wheel'
+    : `setup: ${names.length} model file(s), ${(bytes / 1048576).toFixed(1)} MB — ${names.join(', ')}`)
+  return 0
+}
+
+const OPTIONS = parseArgs(process.argv.slice(2))
+if ('error' in OPTIONS) {
+  console.error(`setup: ${OPTIONS.error}`)
+  process.exitCode = 2
+} else {
+  process.exitCode = main(OPTIONS)
+}

package/tsconfig.json

@@ -0,0 +1,30 @@
+{
+  "compilerOptions": {
+    "target": "es2024",
+    "module": "esnext",
+    "moduleResolution": "bundler",
+    "declaration": true,
+    "declarationMap": true,
+    "sourceMap": true,
+    "skipLibCheck": true,
+    "esModuleInterop": true,
+    "allowImportingTsExtensions": true,
+    "rewriteRelativeImportExtensions": true,
+    "strict": true,
+    "noUncheckedIndexedAccess": true,
+    "exactOptionalPropertyTypes": true,
+    "noImplicitOverride": true,
+    "noFallthroughCasesInSwitch": true,
+    "noUnusedLocals": true,
+    "noUnusedParameters": true,
+    "jsx": "react-jsx",
+    "types": [
+      "node"
+    ],
+    "rootDir": "src",
+    "outDir": "lib"
+  },
+  "include": [
+    "src"
+  ]
+}

package/tsdown.config.ts

@@ -0,0 +1,18 @@
+/** Bundle config: one node entry (host tool). */
+import { defineConfig } from 'tsdown'
+
+export default defineConfig({
+  name: "dsh-ab-ocr",
+  entry: ['lib/index.js'],
+  outDir: 'lib',
+  format: ['esm'],
+  platform: 'node',
+  target: 'es2024',
+  fixedExtension: false,
+  dts: false,
+  clean: false,
+  deps: {
+    neverBundle: [/^@deepseek-ai\//],
+    alwaysBundle: (specifier) => !specifier.startsWith('@deepseek-ai/') && !specifier.startsWith('node:'),
+  },
+})