diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..062d3f8 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,47 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.11"] + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v3 + with: + enable-cache: true + cache-dependency-glob: "**/pyproject.toml" + - name: Set up Python ${{ matrix.python-version }} + run: uv python install ${{ matrix.python-version }} + - name: Install dependencies + run: uv sync --extra dev + - name: Lint (advisory) + run: uv run ruff check . + continue-on-error: true + - name: Format check + run: uv run ruff format --check . + - name: Run tests + run: uv run pytest -m "not integration" -q + + typecheck: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v3 + with: + enable-cache: true + cache-dependency-glob: "**/pyproject.toml" + - name: Set up Python + run: uv python install 3.11 + - name: Install dependencies + run: uv sync --extra dev + - name: Type check (advisory) + run: uv run mypy deepseek_ocr/ --ignore-missing-imports + continue-on-error: true diff --git a/README.md b/README.md index 925e5cc..77d7a8c 100644 --- a/README.md +++ b/README.md @@ -10,12 +10,12 @@ Command-line tool for OCR using DeepSeek vision models. Supports Ollama (local) - **Multi-backend**: Ollama (local, free) and vLLM (OpenAI-compatible API) - Supports PDFs and images (JPG, PNG, WEBP, GIF, BMP, TIFF) -- Per-document output folders with figures -- Batch processing with incremental resume (skips already-processed files) +- Canonical output via the shared [`ocr-output-contract`](https://github.com/r-uben/ocr-output-contract): one `///.md` per document under `## Page N` headers, dual `metadata.json` (per-doc sidecar + root index), input-relative keying (no basename collisions) +- Batch processing of directory trees with incremental resume (skips already-completed documents; re-runs when the input, model, backend, task, or prompt changes) +- Truncation detection: a length-truncated page is recorded `status=partial/failed`, never a silent `completed` - Retry with exponential backoff for transient failures -- Parallel page processing for faster PDF OCR -- `--dry-run` to preview files before processing -- Clean markdown output with HTML tables converted to markdown +- `--dry-run` to preview the exact documents that will be processed +- Clean markdown output with HTML tables converted to markdown (`--raw` keeps the model's verbatim text) ## Choosing an OCR tool @@ -67,10 +67,10 @@ deepseek-ocr document.jpg # Process a PDF deepseek-ocr paper.pdf -# Process all files in a directory -deepseek-ocr ./documents/ --recursive +# Process a directory tree (always walked recursively) +deepseek-ocr ./documents/ -# Preview files without processing +# Preview the documents that would be processed deepseek-ocr ./documents/ --dry-run # Custom output directory @@ -79,8 +79,11 @@ deepseek-ocr doc.pdf -o ./results/ # Use vLLM backend deepseek-ocr paper.pdf --backend vllm --vllm-url http://gpu-server:8000/v1 -# Parallel processing for faster PDF OCR -deepseek-ocr large-document.pdf -w 2 +# Raise the per-page token budget if dense pages truncate +deepseek-ocr large-document.pdf --max-tokens 16384 + +# Keep the model's verbatim output (skip the cleaner) +deepseek-ocr paper.pdf --raw # Extract and analyze embedded figures deepseek-ocr paper.pdf --analyze-figures @@ -95,23 +98,24 @@ deepseek-ocr paper.pdf -q deepseek-ocr [OPTIONS] INPUT_PATH Options: - -o, --output-dir PATH Output directory for results - -r, --recursive Recursively process directories + -o, --output-dir PATH Output root (default: /ocr/) + -r, --recursive Accepted for compatibility; batch trees are + ALWAYS walked recursively --model TEXT Model name (default: deepseek-ocr) - --prompt TEXT Custom prompt for OCR + --prompt TEXT Custom prompt for OCR (overrides --task) --task [convert|ocr|layout|extract|parse] OCR task type - --extract-images Extract and save page images from PDFs - --no-metadata Exclude metadata from output --dpi INTEGER PDF rendering DPI (default: 200) - -w, --workers INTEGER Parallel workers for PDF pages (default: 1) --analyze-figures Extract and analyze embedded figures with AI + --raw Keep verbatim model output (skip the cleaner) + --max-tokens INTEGER Max tokens per page (default: 8192). Raise if + dense pages truncate --max-dim INTEGER Max image dimension (default: 1920, 0 to disable) --backend [ollama|vllm] Backend to use (default: ollama) --vllm-url TEXT vLLM API URL (default: http://localhost:8000/v1) - --reprocess Force reprocessing of already-done files - --dry-run Preview files without processing - -q, --quiet Suppress output, print paths only + --reprocess Force reprocessing of already-done documents + --dry-run Preview documents without processing + -q, --quiet Suppress output, print one .md path per line --verbose Enable verbose output --help Show this message and exit. ``` @@ -138,36 +142,44 @@ deepseek-ocr info ## Output Format -Each document gets its own folder: +Output follows the shared `ocr-output-contract`. The default output root is +`/ocr/` for a single file and `/ocr/` for a directory +(override with `-o`). Each document gets its own folder, mirroring the input +subtree so same-named files in different directories never collide: ``` -output/ +ocr/ +├── metadata.json # root index, keyed by input-relative path └── document/ - ├── document.md # OCR markdown - └── figures/ # Extracted figures (if --analyze-figures) - └── page1_fig1.png + ├── document.md # OCR markdown + ├── metadata.json # per-document sidecar (provenance) + └── figures/ # extracted figures (if --analyze-figures) + └── figure_1_page1.png ``` -The markdown includes metadata: +The markdown body carries **no YAML frontmatter** — all provenance lives in the +JSON sidecars. Pages are separated by `## Page N` headers: ```markdown ---- -source: /path/to/document.pdf -processed: 2025-12-01T15:30:00 -pages: 3 -processing_time: 18.45s -model: deepseek-ocr -backend: ollama ---- - ## Page 1 +[Extracted content...] + +## Page 2 + [Extracted content...] ``` +The per-document `metadata.json` records the ratified schema (`status`, +`checksum`, `model`, `backend`, `processing_time`, `timestamp` (UTC), +`output_path`, `pages`, plus a run `fingerprint`). + ### Batch Resume -Batch processing saves `metadata.json` in the output directory. On re-run, already-processed files are skipped automatically. Use `--reprocess` to force reprocessing. +The root `metadata.json` records every processed document. On re-run, a document +is skipped only when the input is unchanged, its `.md` still exists on disk, and +the run configuration (model, backend, task, prompt) is unchanged. Use +`--reprocess` to force reprocessing. ## Configuration @@ -176,33 +188,30 @@ Create a `.env` file or set environment variables with `DEEPSEEK_OCR_` prefix: ```bash DEEPSEEK_OCR_BACKEND=ollama DEEPSEEK_OCR_MODEL_NAME=deepseek-ocr -DEEPSEEK_OCR_OUTPUT_DIR=output DEEPSEEK_OCR_OLLAMA_URL=http://localhost:11434 DEEPSEEK_OCR_VLLM_BASE_URL=http://localhost:8000/v1 DEEPSEEK_OCR_MAX_DIMENSION=1920 +DEEPSEEK_OCR_MAX_TOKENS=8192 DEEPSEEK_OCR_MAX_RETRIES=3 DEEPSEEK_OCR_RETRY_DELAY=1.0 +DEEPSEEK_OCR_LOG_LEVEL=INFO ``` ## Programmatic Usage ```python from pathlib import Path -from deepseek_ocr import create_backend, OCRProcessor +from deepseek_ocr import create_backend, process backend = create_backend(backend_type="ollama", model_name="deepseek-ocr") backend.load_model() -processor = OCRProcessor( - backend=backend, - output_dir=Path("./results"), - workers=2, -) - -result = processor.process_file(Path("document.pdf")) -print(result.output_text) +# process() routes all output through the contract and returns a RunOutcome. +outcome = process(Path("document.pdf"), backend, output_dir=Path("./results")) +for md_path in outcome.outputs: + print("wrote", md_path) +print("exit code:", outcome.exit_code) # nonzero if any document/page failed -processor.save_result(result) backend.unload_model() ``` diff --git a/deepseek_ocr/__init__.py b/deepseek_ocr/__init__.py index 9861567..1d6ebe2 100644 --- a/deepseek_ocr/__init__.py +++ b/deepseek_ocr/__init__.py @@ -5,10 +5,11 @@ __license__ = "MIT" from deepseek_ocr.backends import Backend, OllamaBackend, VLLMBackend, create_backend -from deepseek_ocr.processor import OCRProcessor +from deepseek_ocr.processor import discover_documents, process __all__ = [ - "OCRProcessor", + "process", + "discover_documents", "Backend", "OllamaBackend", "VLLMBackend", diff --git a/deepseek_ocr/__main__.py b/deepseek_ocr/__main__.py index 8b0a0da..e9abb65 100644 --- a/deepseek_ocr/__main__.py +++ b/deepseek_ocr/__main__.py @@ -2,7 +2,5 @@ from .cli import main - if __name__ == "__main__": main() - diff --git a/deepseek_ocr/backends/__init__.py b/deepseek_ocr/backends/__init__.py index 04fba21..e003205 100644 --- a/deepseek_ocr/backends/__init__.py +++ b/deepseek_ocr/backends/__init__.py @@ -13,6 +13,7 @@ def create_backend( max_dimension: int | None = None, max_retries: int | None = None, retry_delay: float | None = None, + max_tokens: int | None = None, **kwargs, ) -> Backend: """Factory function to create the appropriate backend. @@ -23,6 +24,7 @@ def create_backend( max_dimension: Maximum image dimension for resizing max_retries: Maximum number of retries for transient errors retry_delay: Base delay in seconds between retries + max_tokens: Maximum tokens per page response (truncation control) **kwargs: Additional backend-specific arguments Returns: @@ -36,6 +38,7 @@ def create_backend( max_dimension=max_dimension, max_retries=max_retries, retry_delay=retry_delay, + max_tokens=max_tokens, ollama_url=kwargs.get("ollama_url"), ) elif backend_type == "vllm": @@ -44,6 +47,7 @@ def create_backend( max_dimension=max_dimension, max_retries=max_retries, retry_delay=retry_delay, + max_tokens=max_tokens, base_url=kwargs.get("vllm_base_url"), ) else: diff --git a/deepseek_ocr/backends/base.py b/deepseek_ocr/backends/base.py index e2f41f6..b52f68f 100644 --- a/deepseek_ocr/backends/base.py +++ b/deepseek_ocr/backends/base.py @@ -5,7 +5,6 @@ import time from abc import ABC, abstractmethod from pathlib import Path -from typing import Union from PIL import Image @@ -68,7 +67,7 @@ def unload_model(self) -> None: @abstractmethod def process_image( self, - image: Union[Image.Image, Path, str], + image: Image.Image | Path | str, prompt: str | None = None, task: str = "convert", return_raw: bool = False, @@ -86,6 +85,28 @@ def process_image( """ ... + def process_image_with_meta( + self, + image: Image.Image | Path | str, + prompt: str | None = None, + task: str = "convert", + return_raw: bool = False, + ) -> tuple[str, object | None]: + """Process an image and return ``(text, finish_reason)``. + + The ``finish_reason`` is the model's completion stop reason (e.g. + ``"length"`` / ``"stop"`` / ``None``). The processor feeds it to the + contract's :func:`ocr_output_contract.is_truncated` so a length-truncated + non-empty page is recorded as a per-page failure (status=partial/failed) + rather than a silent ``completed`` (content loss). + + The default delegates to :meth:`process_image` and reports no finish + reason, so a backend that does not surface one (or a mock) keeps working + with truncation detection simply disabled. Backends that expose a stop + reason override this to return it. + """ + return self.process_image(image, prompt=prompt, task=task, return_raw=return_raw), None + def _retry(self, func, *args, **kwargs): """Execute func with exponential backoff retry on TransientError. @@ -105,7 +126,7 @@ def _retry(self, func, *args, **kwargs): except TransientError as e: last_error = e if attempt < self.max_retries: - delay = self.retry_delay * (2 ** attempt) + random.uniform(0, 0.5) + delay = self.retry_delay * (2**attempt) + random.uniform(0, 0.5) logger.warning( f"[{self.backend_name}] Transient error (attempt {attempt + 1}/" f"{self.max_retries + 1}): {e}. Retrying in {delay:.1f}s..." @@ -115,9 +136,7 @@ def _retry(self, func, *args, **kwargs): logger.error( f"[{self.backend_name}] Max retries ({self.max_retries}) exhausted: {e}" ) - raise RuntimeError( - f"Max retries ({self.max_retries}) exhausted: {last_error}" - ) + raise RuntimeError(f"Max retries ({self.max_retries}) exhausted: {last_error}") def process_images_batch( self, @@ -135,7 +154,7 @@ def process_images_batch( results.append(result) return results - def describe_figure(self, image: Union[Image.Image, Path, str]) -> str: + def describe_figure(self, image: Image.Image | Path | str) -> str: """Generate a description of a figure/chart/diagram.""" return self.process_image( image, diff --git a/deepseek_ocr/backends/ollama.py b/deepseek_ocr/backends/ollama.py index cbd73da..178e78b 100644 --- a/deepseek_ocr/backends/ollama.py +++ b/deepseek_ocr/backends/ollama.py @@ -4,7 +4,6 @@ import io import logging from pathlib import Path -from typing import Union import requests from PIL import Image @@ -31,6 +30,7 @@ def __init__( max_dimension: int | None = None, max_retries: int | None = None, retry_delay: float | None = None, + max_tokens: int | None = None, ): super().__init__( model_name=model_name, @@ -38,6 +38,7 @@ def __init__( max_retries=max_retries if max_retries is not None else settings.max_retries, retry_delay=retry_delay if retry_delay is not None else settings.retry_delay, ) + self.max_tokens = max_tokens if max_tokens is not None else settings.max_tokens self.ollama_url = ollama_url or settings.ollama_url or OLLAMA_API_URL logger.info(f"Initialized OllamaBackend with model: {self.model_name}") @@ -94,8 +95,14 @@ def _image_to_base64(self, image: Image.Image) -> str: image.save(buffer, format="JPEG", quality=95) return base64.b64encode(buffer.getvalue()).decode("utf-8") - def _call_ollama_api(self, image_b64: str, prompt: str) -> str: - """Make the Ollama API call, raising TransientError for retryable failures.""" + def _call_ollama_api(self, image_b64: str, prompt: str) -> tuple[str, str | None]: + """Make the Ollama API call, returning ``(text, finish_reason)``. + + Raises TransientError for retryable failures. Ollama reports the stop + reason as ``done_reason`` (e.g. ``"length"`` when the response hit + ``num_predict``); it is threaded out so the processor can flag truncated + pages instead of recording an incomplete page as a silent success. + """ try: response = requests.post( f"{self.ollama_url}/api/generate", @@ -106,6 +113,7 @@ def _call_ollama_api(self, image_b64: str, prompt: str) -> str: "stream": False, "options": { "num_ctx": 8192, + "num_predict": self.max_tokens, "temperature": 0.1, }, }, @@ -114,9 +122,7 @@ def _call_ollama_api(self, image_b64: str, prompt: str) -> str: except requests.exceptions.Timeout as e: raise TransientError("Ollama request timed out", original=e) except requests.exceptions.ConnectionError as e: - raise TransientError( - f"Lost connection to Ollama at {self.ollama_url}", original=e - ) + raise TransientError(f"Lost connection to Ollama at {self.ollama_url}", original=e) if response.status_code in _TRANSIENT_STATUS_CODES: raise TransientError( @@ -126,16 +132,30 @@ def _call_ollama_api(self, image_b64: str, prompt: str) -> str: if response.status_code != 200: raise RuntimeError(f"Ollama API error: {response.text}") - return response.json().get("response", "") + payload = response.json() + return payload.get("response", ""), payload.get("done_reason") def process_image( self, - image: Union[Image.Image, Path, str], + image: Image.Image | Path | str, prompt: str | None = None, task: str = "convert", return_raw: bool = False, ) -> str: """Process image and return OCR text.""" + text, _finish_reason = self.process_image_with_meta( + image, prompt=prompt, task=task, return_raw=return_raw + ) + return text + + def process_image_with_meta( + self, + image: Image.Image | Path | str, + prompt: str | None = None, + task: str = "convert", + return_raw: bool = False, + ) -> tuple[str, str | None]: + """Process image and return ``(text, finish_reason)``.""" if not self.model: raise RuntimeError("Model not loaded. Call load_model() first") @@ -155,7 +175,7 @@ def process_image( image = resize_image_if_needed(image, self.max_dimension) image_b64 = self._image_to_base64(image) - raw_text = self._retry(self._call_ollama_api, image_b64, prompt) + raw_text, finish_reason = self._retry(self._call_ollama_api, image_b64, prompt) if return_raw: - return raw_text - return clean_ocr_output(raw_text) + return raw_text, finish_reason + return clean_ocr_output(raw_text), finish_reason diff --git a/deepseek_ocr/backends/vllm.py b/deepseek_ocr/backends/vllm.py index 4beaaf9..506edcc 100644 --- a/deepseek_ocr/backends/vllm.py +++ b/deepseek_ocr/backends/vllm.py @@ -4,7 +4,6 @@ import io import logging from pathlib import Path -from typing import Union from PIL import Image @@ -27,6 +26,7 @@ def __init__( max_dimension: int | None = None, max_retries: int | None = None, retry_delay: float | None = None, + max_tokens: int | None = None, ): super().__init__( model_name=model_name, @@ -34,6 +34,7 @@ def __init__( max_retries=max_retries if max_retries is not None else settings.max_retries, retry_delay=retry_delay if retry_delay is not None else settings.retry_delay, ) + self.max_tokens = max_tokens if max_tokens is not None else settings.max_tokens self.base_url = base_url or getattr(settings, "vllm_base_url", None) or VLLM_DEFAULT_URL self._client = None logger.info(f"Initialized VLLMBackend with model: {self.model_name} at {self.base_url}") @@ -52,8 +53,7 @@ def load_model(self) -> None: import openai except ImportError: raise RuntimeError( - "openai package is required for vLLM backend. " - "Install with: pip install openai" + "openai package is required for vLLM backend. Install with: pip install openai" ) logger.info(f"Connecting to vLLM at {self.base_url}") @@ -76,8 +76,7 @@ def load_model(self) -> None: ) except Exception as e: raise RuntimeError( - f"Failed to connect to vLLM at {self.base_url}: {e}. " - "Is the vLLM server running?" + f"Failed to connect to vLLM at {self.base_url}: {e}. Is the vLLM server running?" ) self.model = True @@ -95,8 +94,13 @@ def _image_to_base64_url(self, image: Image.Image) -> str: b64 = base64.b64encode(buffer.getvalue()).decode("utf-8") return f"data:image/jpeg;base64,{b64}" - def _call_vllm_api(self, image_url: str, prompt: str) -> str: - """Make the vLLM API call, raising TransientError for retryable failures.""" + def _call_vllm_api(self, image_url: str, prompt: str) -> tuple[str, str | None]: + """Make the vLLM API call, returning ``(text, finish_reason)``. + + Raises TransientError for retryable failures. The ``finish_reason`` is + threaded out so the processor can detect length-truncated pages instead + of recording an incomplete page as a silent success. + """ import openai try: @@ -117,10 +121,11 @@ def _call_vllm_api(self, image_url: str, prompt: str) -> str: ], } ], - max_tokens=2048, + max_tokens=self.max_tokens, temperature=0.1, ) - return response.choices[0].message.content or "" + choice = response.choices[0] + return (choice.message.content or "", getattr(choice, "finish_reason", None)) except openai.APITimeoutError as e: raise TransientError(f"vLLM request timed out: {e}", original=e) @@ -133,20 +138,31 @@ def _call_vllm_api(self, image_url: str, prompt: str) -> str: except openai.APIStatusError as e: # 502, 503, 504 come through as APIStatusError if e.status_code in {502, 503, 504}: - raise TransientError( - f"vLLM HTTP {e.status_code}: {e}", original=e - ) + raise TransientError(f"vLLM HTTP {e.status_code}: {e}", original=e) # Non-transient (400, 404, etc.) — don't retry raise RuntimeError(f"vLLM API error (HTTP {e.status_code}): {e}") def process_image( self, - image: Union[Image.Image, Path, str], + image: Image.Image | Path | str, prompt: str | None = None, task: str = "convert", return_raw: bool = False, ) -> str: """Process image and return OCR text using vLLM.""" + text, _finish_reason = self.process_image_with_meta( + image, prompt=prompt, task=task, return_raw=return_raw + ) + return text + + def process_image_with_meta( + self, + image: Image.Image | Path | str, + prompt: str | None = None, + task: str = "convert", + return_raw: bool = False, + ) -> tuple[str, str | None]: + """Process image and return ``(text, finish_reason)`` using vLLM.""" if not self.model or self._client is None: raise RuntimeError("Model not loaded. Call load_model() first") @@ -166,7 +182,7 @@ def process_image( image = resize_image_if_needed(image, self.max_dimension) image_url = self._image_to_base64_url(image) - raw_text = self._retry(self._call_vllm_api, image_url, prompt) + raw_text, finish_reason = self._retry(self._call_vllm_api, image_url, prompt) if return_raw: - return raw_text - return clean_ocr_output(raw_text) + return raw_text, finish_reason + return clean_ocr_output(raw_text), finish_reason diff --git a/deepseek_ocr/cli.py b/deepseek_ocr/cli.py index 7c39de0..ec92f3b 100644 --- a/deepseek_ocr/cli.py +++ b/deepseek_ocr/cli.py @@ -1,8 +1,14 @@ -"""Command-line interface for DeepSeek OCR.""" +"""Command-line interface for DeepSeek OCR. + +Output is written through the shared ``ocr-output-contract`` package so deepseek's +output is byte-structure-identical to the rest of the engine family: default root +``/ocr/`` (``-o`` overrides, never required), one +``///.md`` per document under ``## Page N`` headers, dual +metadata sidecars, and a nonzero exit on any failure. +""" import sys from pathlib import Path -from typing import Optional import click from rich.console import Console @@ -11,8 +17,9 @@ from deepseek_ocr import __version__ from deepseek_ocr.backends import create_backend from deepseek_ocr.config import settings -from deepseek_ocr.processor import OCRProcessor -from deepseek_ocr.utils import collect_files, is_pdf_file, setup_logging +from deepseek_ocr.processor import _IMAGE_SUFFIXES, discover_documents +from deepseek_ocr.processor import process as run_process +from deepseek_ocr.utils import is_pdf_file, setup_logging console = Console() err_console = Console(stderr=True) @@ -23,7 +30,7 @@ def print_banner(quiet: bool = False) -> None: console.print(f"[dim]deepseek-ocr v{__version__}[/dim]") -def _format_size(size_bytes: int) -> str: +def _format_size(size_bytes: float) -> str: """Format byte count as human-readable string.""" for unit in ("B", "KB", "MB", "GB"): if size_bytes < 1024: @@ -42,9 +49,16 @@ def _get_pdf_page_count(path: Path) -> int: return count -def _run_dry_run(input_path: Path, recursive: bool, quiet: bool) -> None: - """List files that would be processed without actually processing them.""" - files = collect_files(input_path, recursive=recursive) +def _run_dry_run(input_path: Path, output_dir: Path | None, quiet: bool) -> None: + """List documents that would be processed, using the REAL discovery. + + Uses the same :func:`discover_documents` (and the same resolved output root, + so prior outputs are excluded) as the real run, so the preview can never + diverge from what is actually processed. + """ + files = discover_documents(input_path, output_dir) + if not files: + raise ValueError(f"no documents found at {input_path}") if quiet: for f in files: @@ -62,25 +76,34 @@ def _run_dry_run(input_path: Path, recursive: bool, quiet: bool) -> None: total_pages = 0 for idx, f in enumerate(files, 1): - size = f.stat().st_size - total_size += size - - if is_pdf_file(f): - try: - pages = _get_pdf_page_count(f) - except Exception: - pages = 0 - file_type = "PDF" + if f.is_dir(): + # An image directory is OCR'd as ONE document, one page per image. + # Count ONLY the image files the real run (_gather_images) processes, so + # the preview never over-reports pages for a dir holding stray non-image + # files (the dry-run/real-run divergence the review flagged). + images = [p for p in f.iterdir() if p.is_file() and p.suffix.lower() in _IMAGE_SUFFIXES] + size = sum(p.stat().st_size for p in images) + pages = len(images) + file_type = "IMG DIR" else: - pages = 1 - file_type = f.suffix.upper().lstrip(".") + size = f.stat().st_size + if is_pdf_file(f): + try: + pages = _get_pdf_page_count(f) + except Exception: + pages = 0 + file_type = "PDF" + else: + pages = 1 + file_type = f.suffix.upper().lstrip(".") + total_size += size total_pages += pages table.add_row(str(idx), f.name, file_type, _format_size(size), str(pages)) console.print(table) console.print( - f"\n[bold]{len(files)}[/bold] files, " + f"\n[bold]{len(files)}[/bold] documents, " f"[bold]{_format_size(total_size)}[/bold] total, " f"[bold]{total_pages}[/bold] pages" ) @@ -99,7 +122,8 @@ def cli(ctx: click.Context) -> None: deepseek-ocr ./papers/ --recursive deepseek-ocr document.pdf --dry-run - Supports Ollama (local, default) and vLLM (OpenAI-compatible API) backends. + Output goes to /ocr/ by default (-o overrides; never required). + Supports Ollama (local, default) and vLLM (OpenAI-compatible) backends. """ ctx.ensure_object(dict) @@ -110,118 +134,115 @@ def cli(ctx: click.Context) -> None: "-o", "--output-dir", type=click.Path(path_type=Path), - help="Output directory for results", + default=None, + help="Output root (default: /ocr/). Writes /.md per document.", ) @click.option( "-r", "--recursive", is_flag=True, - help="Recursively process directories", + help="Accepted for backward compatibility; batch directory trees are ALWAYS walked recursively.", ) @click.option( "--model", "model_name", type=str, - default="deepseek-ocr", - help="Ollama model name (default: deepseek-ocr)", + default=None, + help="Model name. CLI wins; else DEEPSEEK_OCR_MODEL_NAME / settings (default: deepseek-ocr).", ) @click.option( "--prompt", type=str, - help="Custom prompt for OCR", + default=None, + help="Custom prompt for OCR (overrides --task).", ) @click.option( "--task", type=click.Choice(["convert", "ocr", "layout", "extract", "parse"]), default="convert", - help="OCR task type", -) -@click.option( - "--extract-images", - is_flag=True, - help="Extract and save page images from PDFs", -) -@click.option( - "--no-metadata", - is_flag=True, - help="Exclude metadata from output", + show_default=True, + help="OCR task type; selects the backend prompt unless --prompt is given.", ) @click.option( "--dpi", type=int, default=200, - help="PDF rendering DPI (default: 200, higher=slower but better quality)", + show_default=True, + help="PDF rendering DPI (higher=slower but better quality).", ) @click.option( - "-w", - "--workers", - type=int, - default=1, - help="Parallel workers for PDF pages (default: 1). Note: Ollama processes sequentially, so >1 workers mainly overlap I/O, not GPU inference.", + "--analyze-figures", + is_flag=True, + help="Extract and describe embedded figures/images from PDFs.", ) @click.option( - "--analyze-figures", + "--raw", is_flag=True, - help="Extract and analyze embedded figures/images from PDFs with AI descriptions.", + help="Keep the model's verbatim output (skip clean_ocr_output, which strips [[d,d,d,d]] boxes and <...> spans).", +) +@click.option( + "--max-tokens", + type=int, + default=None, + help="Max tokens per page response (default: 8192 or DEEPSEEK_OCR_MAX_TOKENS). Raise if dense pages truncate.", ) @click.option( "--max-dim", "max_dimension", type=int, default=None, - help="Maximum image dimension (width or height). Larger images are resized to prevent timeouts. Default: 1920. Set to 0 to disable.", + help="Maximum image dimension. Larger images are resized to prevent timeouts. 0 disables.", ) @click.option( "--backend", type=click.Choice(["ollama", "vllm"]), default=None, - help="Backend to use: 'ollama' (local, default) or 'vllm' (OpenAI-compatible). Can also set DEEPSEEK_OCR_BACKEND env var.", + help="Backend: 'ollama' (local, default) or 'vllm'. Or set DEEPSEEK_OCR_BACKEND.", ) @click.option( "--vllm-url", "vllm_base_url", type=str, default=None, - help="vLLM API URL (default: http://localhost:8000/v1). Can also set DEEPSEEK_OCR_VLLM_BASE_URL env var.", + help="vLLM API URL (default: http://localhost:8000/v1). Or set DEEPSEEK_OCR_VLLM_BASE_URL.", ) @click.option( "--reprocess", is_flag=True, - help="Force reprocessing of already-processed files in batch mode.", + help="Re-OCR documents already recorded completed.", ) @click.option( "--dry-run", is_flag=True, - help="List files that would be processed without running OCR. Shows file count, sizes, and page counts.", + help="List files that would be processed without running OCR.", ) @click.option( "-q", "--quiet", is_flag=True, - help="Suppress non-error output. Only print file paths on success (useful for scripting).", + help="Suppress non-error output. Print one output .md path per line (for scripting).", ) @click.option( "--verbose", is_flag=True, - help="Enable verbose output", + help="Enable verbose output.", ) @click.pass_context def process( ctx: click.Context, input_path: Path, - output_dir: Optional[Path], + output_dir: Path | None, recursive: bool, - model_name: str, - prompt: Optional[str], + model_name: str | None, + prompt: str | None, task: str, - extract_images: bool, - no_metadata: bool, dpi: int, - workers: int, analyze_figures: bool, - max_dimension: Optional[int], - backend: Optional[str], - vllm_base_url: Optional[str], + raw: bool, + max_tokens: int | None, + max_dimension: int | None, + backend: str | None, + vllm_base_url: str | None, reprocess: bool, dry_run: bool, quiet: bool, @@ -229,9 +250,8 @@ def process( ) -> None: """Process documents and images with OCR. - INPUT_PATH can be a single file or a directory containing multiple files. - - Supported formats: PDF, JPG, PNG, WEBP, GIF, BMP, TIFF + INPUT_PATH can be a single file, a directory of page images (one document), or + a tree of documents (batch). Supported: PDF, JPG, PNG, WEBP, GIF, BMP, TIFF. \b Examples: @@ -246,7 +266,7 @@ def process( if dry_run: print_banner(quiet=quiet) try: - _run_dry_run(input_path, recursive=recursive, quiet=quiet) + _run_dry_run(input_path, output_dir=output_dir, quiet=quiet) except Exception as e: err_console.print(f"[red]error:[/red] {e}") sys.exit(1) @@ -254,78 +274,65 @@ def process( print_banner(quiet=quiet) - # Resolve backend from CLI flag or config + # Resolve backend from CLI flag or config. backend_type = backend or settings.backend - # For vLLM, default model is deepseek-vl2 unless explicitly specified - if backend_type == "vllm" and model_name == "deepseek-ocr": - model_name = "deepseek-vl2" + # Resolve the model: an explicit --model on the command line wins; otherwise + # fall back to settings.model_name (which honors DEEPSEEK_OCR_MODEL_NAME / .env), + # NOT a click default that would silently shadow the env var. We treat the model + # as "unset" when --model was not passed AND the env did not override the default. + model_from_cli = model_name is not None + resolved_model: str = model_name if model_name is not None else settings.model_name + model_is_default = not model_from_cli and resolved_model == "deepseek-ocr" + + # For vLLM, default model is deepseek-vl2 unless a model was explicitly chosen. + if backend_type == "vllm" and model_is_default: + resolved_model = "deepseek-vl2" try: backend_instance = create_backend( backend_type=backend_type, - model_name=model_name, + model_name=resolved_model, max_dimension=max_dimension, + max_tokens=max_tokens, ollama_url=settings.ollama_url, vllm_base_url=vllm_base_url or settings.vllm_base_url, ) - backend_instance.load_model() - - processor_kwargs = { - "backend": backend_instance, - "extract_images": extract_images, - "include_metadata": not no_metadata, - "dpi": dpi, - "workers": workers, - "analyze_figures": analyze_figures, - } - if output_dir: - processor_kwargs["output_dir"] = output_dir - - processor = OCRProcessor(**processor_kwargs) - - if input_path.is_file(): - result = processor.process_file(input_path, prompt=prompt, show_progress=not verbose and not quiet) - output_path = processor.save_result(result) - if quiet: - console.print(str(output_path)) - else: - console.print(f"[dim]->[/dim] {output_path}") - else: - results = processor.process_batch( - input_path, - recursive=recursive, - prompt=prompt, - show_progress=not verbose and not quiet, - reprocess=reprocess, - ) - if quiet: - for result in results: - base_name = result.input_path.stem - output_path = processor.output_dir / base_name / f"{base_name}.md" - console.print(str(output_path)) - else: - table = Table(show_header=True, header_style="bold") - table.add_column("File", style="cyan") - table.add_column("Pages", justify="right") - table.add_column("Time (s)", justify="right") - - for result in results: - table.add_row( - result.input_path.name, - str(result.page_count), - f"{result.processing_time:.2f}", - ) - - console.print(table) + outcome = run_process( + input_path, + backend_instance, + dpi=dpi, + task=task, + prompt=prompt, + output_dir=output_dir, + reprocess=reprocess, + analyze_figures=analyze_figures, + raw=raw, + ) backend_instance.unload_model() - except Exception as e: err_console.print(f"[red]error:[/red] {e}") sys.exit(1) + total = outcome.completed + outcome.failed + outcome.partial + if quiet: + for path in outcome.outputs: + console.print(path) + else: + console.print(f"[dim]deepseek-ocr: backend={backend_type} dpi={dpi}[/dim]") + console.print(f" {outcome.completed}/{total} document(s) completed") + if outcome.has_failures: + err_console.print( + f" {outcome.failed} failed, {outcome.partial} partial: " + + ", ".join(outcome.failures) + ) + + # Uniform exit policy (canon SYS-02): nonzero if any document/page failed. + if outcome.exit_code != 0: + sys.exit(outcome.exit_code) + @cli.command() def info() -> None: @@ -339,7 +346,6 @@ def info() -> None: sys_table.add_row("Python", f"{sys.version_info.major}.{sys.version_info.minor}") sys_table.add_row("Backend", settings.backend) - # Check Ollama status from deepseek_ocr.backends.ollama import OllamaBackend ollama_backend = OllamaBackend() @@ -358,8 +364,7 @@ def info() -> None: settings_table.add_column("Value", style="yellow") settings_table.add_row("Model", settings.model_name) - settings_table.add_row("Output Directory", str(settings.output_dir)) - settings_table.add_row("Extract Images", str(settings.extract_images)) + settings_table.add_row("Output root default", "/ocr/") settings_table.add_row("Max Image Dimension", str(settings.max_dimension)) console.print(settings_table) @@ -396,11 +401,7 @@ def main() -> None: candidate = argv[first_non_option_index] if candidate not in known_subcommands and Path(candidate).exists(): - argv = ( - argv[:first_non_option_index] - + ["process"] - + argv[first_non_option_index:] - ) + argv = argv[:first_non_option_index] + ["process"] + argv[first_non_option_index:] sys.argv = [sys.argv[0], *argv] cli(obj={}) diff --git a/deepseek_ocr/config.py b/deepseek_ocr/config.py index 4b403b9..dc9ce87 100644 --- a/deepseek_ocr/config.py +++ b/deepseek_ocr/config.py @@ -1,6 +1,5 @@ """Configuration management for DeepSeek OCR CLI.""" -from pathlib import Path from typing import Literal from pydantic import Field @@ -43,18 +42,10 @@ class Settings(BaseSettings): description="Maximum image dimension (width or height). Larger images are resized to prevent Ollama timeouts. Set to 0 to disable.", ) - # Output configuration - output_dir: Path = Field( - default=Path("output"), - description="Default output directory", - ) - extract_images: bool = Field( - default=False, - description="Extract and save images from documents", - ) - include_metadata: bool = Field( - default=True, - description="Include metadata in output markdown", + # Generation + max_tokens: int = Field( + default=8192, + description="Maximum tokens per page response. Dense full pages can exceed the old 2048 default and truncate (silent content loss); raise if pages truncate.", ) # Retry configuration diff --git a/deepseek_ocr/metadata.py b/deepseek_ocr/metadata.py deleted file mode 100644 index 480be7f..0000000 --- a/deepseek_ocr/metadata.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Metadata tracking for incremental batch processing.""" - -import hashlib -import json -import logging -import os -from datetime import datetime -from pathlib import Path -from typing import Any, Dict, Optional - -logger = logging.getLogger(__name__) - -METADATA_VERSION = "1" -METADATA_FILENAME = "metadata.json" - - -def _file_checksum(path: Path) -> str: - """Compute SHA256 checksum of a file.""" - h = hashlib.sha256() - with open(path, "rb") as f: - for chunk in iter(lambda: f.read(8192), b""): - h.update(chunk) - return f"sha256:{h.hexdigest()}" - - -class MetadataManager: - """Manages metadata.json for incremental batch processing.""" - - def __init__(self, output_dir: Path) -> None: - self.output_dir = Path(output_dir) - self.metadata_path = self.output_dir / METADATA_FILENAME - self._data: Dict[str, Any] = {"version": METADATA_VERSION, "files": {}} - self._load() - - def _load(self) -> None: - """Load existing metadata from disk.""" - if self.metadata_path.exists(): - try: - self._data = json.loads(self.metadata_path.read_text(encoding="utf-8")) - logger.debug(f"Loaded metadata with {len(self._data.get('files', {}))} entries") - except (json.JSONDecodeError, OSError) as e: - logger.warning(f"Failed to load metadata, starting fresh: {e}") - self._data = {"version": METADATA_VERSION, "files": {}} - - def save(self) -> None: - """Atomic write: write to tmp file then rename.""" - self.output_dir.mkdir(parents=True, exist_ok=True) - tmp_path = self.metadata_path.with_suffix(".tmp") - tmp_path.write_text( - json.dumps(self._data, indent=2, ensure_ascii=False), - encoding="utf-8", - ) - os.replace(str(tmp_path), str(self.metadata_path)) - - def is_processed(self, file_path: Path) -> bool: - """Check if a file has already been processed (matching name and checksum).""" - key = file_path.name - entry = self._data.get("files", {}).get(key) - if entry is None: - return False - if entry.get("status") != "completed": - return False - # Verify checksum matches (file hasn't changed) - current_checksum = _file_checksum(file_path) - return entry.get("checksum") == current_checksum - - def record( - self, - file_path: Path, - *, - pages: int, - processing_time: float, - model: str, - backend: str, - output_path: str, - ) -> None: - """Record metadata for a processed file and save incrementally.""" - key = file_path.name - self._data["files"][key] = { - "status": "completed", - "pages": pages, - "processing_time": round(processing_time, 2), - "model": model, - "backend": backend, - "timestamp": datetime.now().isoformat(), - "output_path": output_path, - "checksum": _file_checksum(file_path), - } - self.save() - - @property - def files(self) -> Dict[str, Any]: - """Access the files dict.""" - return self._data.get("files", {}) diff --git a/deepseek_ocr/processor.py b/deepseek_ocr/processor.py index e366636..4896761 100644 --- a/deepseek_ocr/processor.py +++ b/deepseek_ocr/processor.py @@ -1,25 +1,59 @@ -"""Document processing and batch handling for DeepSeek OCR.""" +"""Render inputs to page images, OCR each via a backend, write canonical Markdown. +deepseek is a *local, page-based* engine: it renders a PDF (or reads a directory of +page images) to per-page images and OCRs each page through one of its backends +(ollama / vllm). This module owns *how OCR happens* and deepseek's figure extraction; +the shared ``ocr-output-contract`` package owns *where the bytes go* and what shape +the metadata takes. + +The per-page text list this module produces is fed straight into the contract's +:func:`assemble_pages`, so every page of a document lands in ONE +``///.md`` under ``## Page N`` headers. An empty/whitespace +page response is a per-page FAILURE recorded in the sidecar, never a silent 0-byte +success reported as ok. The Markdown body carries NO YAML frontmatter — all provenance +lives in the per-document ``metadata.json`` sidecar (canon: single source of truth). +""" + +from __future__ import annotations + +import hashlib import io import logging -from concurrent.futures import ThreadPoolExecutor, as_completed +import time from dataclasses import dataclass, field -from datetime import datetime from pathlib import Path -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING import fitz # PyMuPDF +from ocr_output_contract import ( + UNREADABLE_CHECKSUM, + DocMetadata, + RootIndex, + RunOutcome, + Status, + assemble_pages, + doc_dir_for, + figure_filename, + figure_markdown_link, + figures_dir_for, + is_truncated, + is_within_output_root, + iter_input_files, + markdown_path_for, + relative_key, + resolve_output_root, + run_fingerprint, + safe_checksum, + sha256_checksum, + utc_timestamp, + write_doc_metadata, +) from PIL import Image -from tqdm import tqdm -from deepseek_ocr.config import settings -from deepseek_ocr.metadata import MetadataManager from deepseek_ocr.utils import ( - collect_files, + IMAGE_EXTENSIONS, ensure_dir, - is_pdf_file, load_image, - sanitize_filename, ) if TYPE_CHECKING: @@ -27,10 +61,14 @@ logger = logging.getLogger(__name__) +_IMAGE_SUFFIXES = {ext.lower() for ext in IMAGE_EXTENSIONS} +#: Inputs deepseek treats as a single source document to render page-by-page. +_DOC_SUFFIXES = {".pdf"} + @dataclass class FigureInfo: - """Container for extracted figure information.""" + """Container for an extracted embedded figure (opt-in via --analyze-figures).""" page_num: int figure_num: int @@ -40,463 +78,531 @@ class FigureInfo: format: str context: str = "" description: str = "" - saved_path: Optional[Path] = None - - -class OCRResult: - """Container for OCR processing results.""" - - def __init__( - self, - input_path: Path, - output_text: str, - page_count: int = 1, - processing_time: float = 0.0, - metadata: Optional[Dict] = None, - ): - self.input_path = input_path - self.output_text = output_text - self.page_count = page_count - self.processing_time = processing_time - self.metadata = metadata or {} - self.timestamp = datetime.now() - - def to_markdown(self, include_metadata: bool = True) -> str: - lines = [] - - if include_metadata: - lines.append("---") - lines.append(f"source: {self.input_path}") - lines.append(f"processed: {self.timestamp.isoformat()}") - lines.append(f"pages: {self.page_count}") - lines.append(f"processing_time: {self.processing_time:.2f}s") - for key, value in self.metadata.items(): - lines.append(f"{key}: {value}") - lines.append("---") - lines.append("") - - lines.append(self.output_text) - - return "\n".join(lines) - - -class OCRProcessor: - """Handles OCR processing for documents and images.""" - - def __init__( - self, - backend: Optional["Backend"] = None, - output_dir: Optional[Path] = None, - extract_images: bool = False, - include_metadata: bool = True, - dpi: int = 200, - workers: int = 1, - analyze_figures: bool = False, - ): - if backend is not None: - self._backend = backend - else: - from deepseek_ocr.backends import create_backend - - self._backend = create_backend() - - self.output_dir = output_dir or settings.output_dir - self.extract_images = extract_images or settings.extract_images - self.include_metadata = include_metadata and settings.include_metadata - self.dpi = dpi - self.workers = max(1, workers) - self.analyze_figures = analyze_figures - - ensure_dir(self.output_dir) - logger.info(f"OCRProcessor initialized with output_dir: {self.output_dir}, workers: {self.workers}") - - def _pdf_to_images(self, pdf_path: Path) -> List[Image.Image]: - logger.debug(f"Converting PDF to images: {pdf_path} at {self.dpi} DPI") - - try: - images = [] - pdf_document = fitz.open(pdf_path) - - for page_num in range(len(pdf_document)): - page = pdf_document[page_num] - - zoom = self.dpi / 72 - mat = fitz.Matrix(zoom, zoom) - pix = page.get_pixmap(matrix=mat) - - img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) - images.append(img) - - logger.debug(f"Converted page {page_num + 1}/{len(pdf_document)}") - - pdf_document.close() - logger.info(f"Converted {len(images)} pages from {pdf_path.name}") - - return images - - except Exception as e: - raise RuntimeError(f"Failed to convert PDF {pdf_path}: {e}") + saved_path: Path | None = None - def _save_images(self, images: List[Image.Image], base_name: str) -> Path: - images_dir = self.output_dir / base_name / "images" - ensure_dir(images_dir) - - for idx, image in enumerate(images, 1): - image_path = images_dir / f"page_{idx:04d}.png" - image.save(image_path, "PNG") - logger.debug(f"Saved image: {image_path}") - - logger.info(f"Saved {len(images)} images to {images_dir}") - return images_dir - - def _extract_figures_from_pdf(self, pdf_path: Path) -> List[FigureInfo]: - """Extract embedded figures/images from a PDF.""" - figures: List[FigureInfo] = [] +@dataclass +class DocResult: + """Result of OCR'ing one source document (a PDF or an image directory). + + ``pages`` holds the per-page markdown in order; ``page_errors`` maps a + 1-indexed page number to its error string for any page that failed (or whose + model response was empty). ``status`` maps to the contract enum. + """ + + source: Path + pages: list[str] + processing_time: float = 0.0 + page_errors: dict[int, str] = field(default_factory=dict) + error: str | None = None + figures_markdown: str = "" + + @property + def page_count(self) -> int: + return len(self.pages) + + @property + def status(self) -> Status: + """``completed`` = every page ok; ``partial`` = some ok; ``failed`` = none.""" + if self.pages and not self.page_errors and not self.error: + return Status.COMPLETED + succeeded = self.page_count - len(self.page_errors) + if succeeded > 0: + return Status.PARTIAL + return Status.FAILED + + +# --------------------------------------------------------------------------- +# Page OCR (image -> page-text; deepseek-owned). Output shape is the contract's job. +# --------------------------------------------------------------------------- + + +def _render_pdf(pdf_path: Path, dpi: int) -> list[Image.Image]: + """Render each PDF page to an in-memory RGB PIL image.""" + images: list[Image.Image] = [] + zoom = dpi / 72 + mat = fitz.Matrix(zoom, zoom) + with fitz.open(pdf_path) as doc: + for page_num in range(len(doc)): + pix = doc[page_num].get_pixmap(matrix=mat) + images.append(Image.frombytes("RGB", [pix.width, pix.height], pix.samples)) + return images + + +def _gather_images(directory: Path) -> list[Path]: + images = sorted(p for p in directory.iterdir() if p.suffix.lower() in _IMAGE_SUFFIXES) + if not images: + raise ValueError(f"no page images found in {directory}") + return images + + +def _ocr_pages( + source: Path, + images: list[Image.Image], + backend: Backend, + task: str, + prompt: str | None, + start: float, + raw: bool = False, +) -> DocResult: + """OCR an ordered list of page images into a DocResult. + + Each page is an independent backend call so page boundaries and per-page + failures are real. A page that errors OR returns empty/whitespace text is + recorded in ``page_errors`` (empty != success) and gets an explicit failure + marker in its slot, keeping the page count and ``## Page N`` numbering aligned. + + Truncation is a per-page failure too: when the backend reports it stopped on + a length/token limit (``finish_reason`` in the contract's + :data:`TRUNCATION_FINISH_REASONS`), the page text is non-empty but + incomplete, so recording it as ``completed`` would be silent content loss. + Such a page is recorded in ``page_errors`` (driving status=partial/failed) + while keeping the recovered-so-far text in its slot rather than discarding it. + + ``raw`` opts out of the backend's ``clean_ocr_output`` post-processing so the + model's verbatim output (including ``[[d,d,d,d]]`` boxes and ``<...>`` spans) + is preserved for faithful extraction. + """ + pages: list[str] = [] + page_errors: dict[int, str] = {} + for idx, image in enumerate(images, start=1): try: - doc = fitz.open(pdf_path) - + text, finish_reason = _ocr_one_image(backend, image, prompt, task, raw) + if not text.strip(): + raise ValueError("empty OCR response (no text returned)") + # Per-page truncation: 1 image -> 1 page, so the page-shortfall signal + # does not apply; the finish_reason length-limit signal does. Keep the + # recovered-so-far text but flag the page so it is not a silent success. + if is_truncated(finish_reason, parsed_pages=1, actual_pages=1): + page_errors[idx] = ( + f"truncated response (finish_reason={finish_reason!r}); " + "page content is incomplete" + ) + pages.append(text) + except Exception as exc: + logger.error("OCR failed for page %d of %s: %s", idx, source.name, exc) + page_errors[idx] = str(exc) + pages.append(f"*[OCR failed for page {idx}]*") + return DocResult( + source=source, + pages=pages, + processing_time=time.time() - start, + page_errors=page_errors, + ) + + +def _ocr_one_image( + backend: Backend, + image: Image.Image, + prompt: str | None, + task: str, + raw: bool = False, +) -> tuple[str, object | None]: + """Call the backend, returning ``(text, finish_reason)``. + + Backends that surface a completion ``finish_reason`` override + :meth:`Backend.process_image_with_meta`; the default falls back to + :meth:`Backend.process_image` with ``finish_reason=None`` (no truncation + signal), keeping the contract intact for engines/mocks that do not report it. + """ + return backend.process_image_with_meta(image, prompt=prompt, task=task, return_raw=raw) + + +def _ocr_one_document( + doc: Path, + backend: Backend, + task: str, + prompt: str | None, + dpi: int, + raw: bool = False, +) -> tuple[DocResult, list[Image.Image] | None]: + """OCR one document. Returns the result plus rendered images (for figure reuse).""" + start = time.time() + try: + if doc.is_dir(): + image_paths = _gather_images(doc) + images = [load_image(p) for p in image_paths] + return _ocr_pages(doc, images, backend, task, prompt, start, raw), None + if doc.suffix.lower() in _DOC_SUFFIXES: + images = _render_pdf(doc, dpi) + return _ocr_pages(doc, images, backend, task, prompt, start, raw), images + if doc.suffix.lower() in _IMAGE_SUFFIXES: + images = [load_image(doc)] + return _ocr_pages(doc, images, backend, task, prompt, start, raw), None + raise ValueError(f"unsupported input: {doc} (expected a .pdf, image, or directory)") + except Exception as exc: + logger.error("could not process %s: %s", doc, exc) + return ( + DocResult(source=doc, pages=[], processing_time=time.time() - start, error=str(exc)), + None, + ) + + +# --------------------------------------------------------------------------- +# Figure extraction (deepseek-owned, opt-in). Written under the contract doc dir. +# --------------------------------------------------------------------------- + + +def _extract_figures_from_pdf(pdf_path: Path) -> list[FigureInfo]: + """Extract embedded figures/images from a PDF.""" + figures: list[FigureInfo] = [] + try: + with fitz.open(pdf_path) as doc: for page_num in range(len(doc)): page = doc[page_num] page_text = page.get_text() - images = page.get_images(full=True) - - for img_idx, img_info in enumerate(images): + for img_idx, img_info in enumerate(page.get_images(full=True)): xref = img_info[0] - try: base_image = doc.extract_image(xref) - image_bytes = base_image["image"] - - pil_image = Image.open(io.BytesIO(image_bytes)) + pil_image = Image.open(io.BytesIO(base_image["image"])) if pil_image.mode != "RGB": pil_image = pil_image.convert("RGB") - - context = "" - img_rects = page.get_image_rects(xref) - if img_rects: - rect = img_rects[0] - expanded_rect = fitz.Rect( - max(0, rect.x0 - 50), - max(0, rect.y0 - 150), - rect.x1 + 50, - rect.y1 + 150 + context = page_text[:500].strip() if page_text else "" + figures.append( + FigureInfo( + page_num=page_num + 1, + figure_num=img_idx + 1, + image=pil_image, + width=base_image["width"], + height=base_image["height"], + format=base_image["ext"], + context=context, ) - context = page.get_text("text", clip=expanded_rect).strip() - - if not context: - context = page_text[:500].strip() if page_text else "" - - figure = FigureInfo( - page_num=page_num + 1, - figure_num=img_idx + 1, - image=pil_image, - width=base_image["width"], - height=base_image["height"], - format=base_image["ext"], - context=context, ) - figures.append(figure) - - except Exception as e: - logger.warning(f"Failed to extract image {img_idx + 1} from page {page_num + 1}: {e}") - continue - - doc.close() - logger.info(f"Extracted {len(figures)} figures from {pdf_path.name}") - - except Exception as e: - logger.error(f"Failed to extract figures from {pdf_path}: {e}") - - return figures - - def _save_figures(self, figures: List[FigureInfo], base_name: str) -> Path: - """Save extracted figures to disk.""" - figures_dir = self.output_dir / base_name / "figures" - ensure_dir(figures_dir) - - for fig in figures: - filename = f"page{fig.page_num}_fig{fig.figure_num}.{fig.format}" - fig_path = figures_dir / filename - fig.image.save(fig_path) - fig.saved_path = fig_path - logger.debug(f"Saved figure: {fig_path}") - - logger.info(f"Saved {len(figures)} figures to {figures_dir}") - return figures_dir - - def _analyze_single_figure( - self, figure: FigureInfo - ) -> Tuple[int, int, str, Optional[str]]: - """Analyze a single figure. Returns (page_num, fig_num, description, error).""" - try: - if figure.context: - prompt = ( - f"This figure appears in a document with the following context:\n" - f"---\n{figure.context[:500]}\n---\n\n" - f"Describe what this figure shows. Include details about any charts, " - f"graphs, diagrams, or visual elements. Explain what it represents " - f"in the context of the document." - ) - else: - prompt = ( - "Describe this figure in detail. Include information about any charts, " - "graphs, diagrams, tables, or visual elements. Explain what it represents." - ) - - description = self._backend.process_image(figure.image, prompt=prompt) - return (figure.page_num, figure.figure_num, description, None) - - except Exception as e: - logger.error(f"Failed to analyze figure {figure.figure_num} on page {figure.page_num}: {e}") - return (figure.page_num, figure.figure_num, "", str(e)) - - def _analyze_figures( - self, figures: List[FigureInfo], show_progress: bool = True - ) -> List[FigureInfo]: - """Analyze all figures and populate their descriptions.""" - if not figures: - return figures - - if self.workers == 1: - fig_iterator = ( - tqdm(figures, desc="Analyzing figures", unit="fig") - if show_progress and len(figures) > 1 - else figures - ) - for fig in fig_iterator: - _, _, description, error = self._analyze_single_figure(fig) - fig.description = description if not error else f"[Analysis Error: {error}]" - else: - with ThreadPoolExecutor(max_workers=self.workers) as executor: - futures = { - executor.submit(self._analyze_single_figure, fig): fig - for fig in figures - } - - if show_progress and len(figures) > 1: - pbar = tqdm(total=len(figures), desc=f"Analyzing figures ({self.workers}w)", unit="fig") - else: - pbar = None - - for future in as_completed(futures): - fig = futures[future] - _, _, description, error = future.result() - fig.description = description if not error else f"[Analysis Error: {error}]" - if pbar: - pbar.update(1) - - if pbar: - pbar.close() - - return figures - - def _figures_to_markdown(self, figures: List[FigureInfo]) -> str: - """Convert figure analyses to markdown.""" - if not figures: - return "" - - lines = ["\n\n---\n\n# Figures\n"] - - for fig in figures: - lines.append(f"\n## Figure {fig.figure_num} (Page {fig.page_num})\n") - if fig.saved_path: - lines.append(f"![Figure {fig.figure_num}](./figures/{fig.saved_path.name})\n") - lines.append(f"*Size: {fig.width}x{fig.height} ({fig.format})*\n") - lines.append(f"\n{fig.description}\n") - - return "\n".join(lines) - - def _process_single_page( - self, page_data: Tuple[int, Image.Image], prompt: Optional[str] = None - ) -> Tuple[int, str, Optional[str]]: - """Process a single page image. Returns (page_num, text, error).""" - page_num, image = page_data - try: - output = self._backend.process_image(image, prompt=prompt) - return (page_num, output, None) - except Exception as e: - logger.error(f"Failed to process page {page_num}: {e}") - return (page_num, "", str(e)) - - def process_file( - self, file_path: Path, prompt: Optional[str] = None, show_progress: bool = True - ) -> OCRResult: - """Process a single file (image or PDF).""" - if not file_path.exists(): - raise FileNotFoundError(f"File not found: {file_path}") - - logger.info(f"Processing file: {file_path}") - start_time = datetime.now() - - if self._backend.model is None: - self._backend.load_model() - + except Exception as exc: + logger.warning( + "Failed to extract image %d on page %d: %s", + img_idx + 1, + page_num + 1, + exc, + ) + except Exception as exc: + logger.error("Failed to extract figures from %s: %s", pdf_path, exc) + return figures + + +def _process_figures(doc: Path, doc_dir: Path, backend: Backend) -> str: + """Extract, save, and describe embedded figures; return their markdown block. + + Figures land under ``/figures/figure__page

.png`` (canon naming) + with resolvable links in the returned markdown. + """ + figures = _extract_figures_from_pdf(doc) + if not figures: + return "" + figures_dir = ensure_dir(figures_dir_for(doc_dir)) + lines = ["\n\n---\n\n# Figures\n"] + for fig in figures: + filename = figure_filename(fig.figure_num, fig.page_num) + fig_path = figures_dir / filename + fig.image.save(fig_path, "PNG") + fig.saved_path = fig_path try: - if is_pdf_file(file_path): - images = self._pdf_to_images(file_path) - page_count = len(images) - - if self.extract_images: - base_name = sanitize_filename(file_path.stem) - self._save_images(images, base_name) - - page_data = [(idx, img) for idx, img in enumerate(images, 1)] - - if self.workers == 1: - outputs = [] - page_iterator = ( - tqdm(page_data, total=page_count, desc="OCR pages", unit="page") - if show_progress and page_count > 1 - else page_data - ) - for idx, image in page_iterator: - logger.debug(f"Processing page {idx}/{page_count}") - output = self._backend.process_image(image, prompt=prompt) - outputs.append(f"## Page {idx}\n\n{output}") - else: - results_dict: Dict[int, str] = {} - errors: List[str] = [] - - with ThreadPoolExecutor(max_workers=self.workers) as executor: - futures = { - executor.submit(self._process_single_page, pd, prompt): pd[0] - for pd in page_data - } - - if show_progress and page_count > 1: - pbar = tqdm(total=page_count, desc=f"OCR pages ({self.workers}w)", unit="page") - else: - pbar = None - - for future in as_completed(futures): - page_num, text, error = future.result() - if error: - errors.append(f"Page {page_num}: {error}") - results_dict[page_num] = f"[OCR Error: {error}]" - else: - results_dict[page_num] = text - if pbar: - pbar.update(1) - - if pbar: - pbar.close() - - if errors: - logger.warning(f"Errors on {len(errors)} pages: {errors}") - - outputs = [ - f"## Page {i}\n\n{results_dict[i]}" - for i in sorted(results_dict.keys()) - ] - - output_text = "\n\n".join(outputs) - - if self.analyze_figures: - figures = self._extract_figures_from_pdf(file_path) - if figures: - base_name = sanitize_filename(file_path.stem) - self._save_figures(figures, base_name) - self._analyze_figures(figures, show_progress=show_progress) - output_text += self._figures_to_markdown(figures) - - else: - image = load_image(file_path) - output_text = self._backend.process_image(image, prompt=prompt) - page_count = 1 - - processing_time = (datetime.now() - start_time).total_seconds() - - metadata = { - "model": self._backend.model_name, - "backend": getattr(self._backend, "backend_name", "ollama"), - } - - result = OCRResult( - input_path=file_path, - output_text=output_text, - page_count=page_count, - processing_time=processing_time, - metadata=metadata, + fig.description = backend.describe_figure(fig.image) + except Exception as exc: + logger.error("figure description failed (page %d): %s", fig.page_num, exc) + fig.description = f"[Analysis Error: {exc}]" + lines.append(f"\n## Figure {fig.figure_num} (Page {fig.page_num})\n") + lines.append(figure_markdown_link(fig.figure_num, fig.page_num) + "\n") + lines.append(f"*Size: {fig.width}x{fig.height} ({fig.format})*\n") + lines.append(f"\n{fig.description}\n") + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Output writing (all routed through the ocr-output-contract package) +# --------------------------------------------------------------------------- + + +def _backend_model(backend: Backend) -> str: + return getattr(backend, "model_name", "") or "" + + +def _backend_name(backend: Backend) -> str: + return getattr(backend, "backend_name", "") or "" + + +def _doc_checksum(source: Path) -> str: + """Content checksum for idempotency. Hash a PDF/image's bytes, or an image dir's.""" + if source.is_file(): + return sha256_checksum(source) + h = hashlib.sha256() + for img in _gather_images(source): + h.update(img.name.encode("utf-8")) + with open(img, "rb") as f: + for chunk in iter(lambda: f.read(8192), b""): + h.update(chunk) + return f"sha256:{h.hexdigest()}" + + +def _safe_doc_checksum(source: Path) -> str | None: + """Like :func:`_doc_checksum`, but ``None`` instead of raising on an OSError. + + Mirrors the contract's :func:`safe_checksum` for deepseek's two input shapes: + a single file (delegating to ``safe_checksum``) and an image directory (which + may still raise if an image is deleted/unreadable between discovery and + processing). A ``None`` result means "input unreadable now" — the idempotency + pre-check must NOT skip (it processes, and the per-doc catch-all in + :func:`_ocr_one_document` records that one doc ``status=failed`` and the batch + CONTINUES, rather than an ``OSError`` aborting the whole run — the SYS-02 + "one bad file aborts the batch" failure mode). + """ + if source.is_file(): + # safe_checksum already returns None on OSError for the single-file case. + return safe_checksum(source) + try: + return _doc_checksum(source) + except OSError: + return None + + +def _build_doc_metadata( + result: DocResult, + markdown_path: Path, + output_root: Path, + backend: Backend, + fingerprint: str | None, +) -> DocMetadata: + """Assemble the per-document metadata record from a DocResult.""" + status = result.status + error = None + if status is not Status.COMPLETED: + if result.page_errors: + error = "; ".join(f"page {n}: {msg}" for n, msg in sorted(result.page_errors.items())) + elif result.error: + error = result.error + # Tolerant checksum: if the source became unreadable mid-run we still persist a + # status=failed record rather than letting the failure-metadata write itself throw + # and escape the per-doc error boundary. Fall back to the canonical ``sha256:`` + # UNREADABLE_CHECKSUM sentinel (v0.1.3) instead of "" — the conformance harness + # requires a ``sha256:`` checksum even on a failure record. The sentinel can never + # equal a real digest, so a later readable run gets a different checksum and + # reprocesses. (failure_checksum's generalization to deepseek's two input shapes: + # _safe_doc_checksum handles both the single-file and image-directory cases.) + checksum = _safe_doc_checksum(result.source) or UNREADABLE_CHECKSUM + return DocMetadata( + status=status, + checksum=checksum, + model=_backend_model(backend), + backend=_backend_name(backend), + processing_time=result.processing_time, + timestamp=utc_timestamp(), + output_path=str(markdown_path.relative_to(output_root)), + pages=result.page_count, + error=error, + fingerprint=fingerprint, + ) + + +def _write_document( + result: DocResult, + output_root: Path, + rel_key: str, + backend: Backend, + index: RootIndex, + fingerprint: str | None, +) -> tuple[DocMetadata, Path]: + """Write the aggregated markdown + BOTH metadata levels for one document. + + Output is always written (even on failure) so failures are recorded with + ``status=failed`` per the canon. The single ``/.md`` aggregates + every page under ``## Page N`` headers. NO YAML frontmatter — provenance is in + the sidecar only. + """ + doc_dir = doc_dir_for(output_root, rel_key) + doc_dir.mkdir(parents=True, exist_ok=True) + markdown_path = markdown_path_for(doc_dir, rel_key) + + body = assemble_pages(result.pages) if result.pages else "*[OCR Failed]*\n" + body += result.figures_markdown + markdown_path.write_text(body, encoding="utf-8") + + meta = _build_doc_metadata(result, markdown_path, output_root, backend, fingerprint) + write_doc_metadata(doc_dir, rel_key, meta) + index.record(rel_key, meta) + return meta, markdown_path + + +# --------------------------------------------------------------------------- +# Orchestration +# --------------------------------------------------------------------------- + + +def process( + source: Path, + backend: Backend, + dpi: int = 200, + task: str = "convert", + prompt: str | None = None, + output_dir: Path | None = None, + reprocess: bool = False, + analyze_figures: bool = False, + raw: bool = False, +) -> RunOutcome: + """Process an input (file or directory of documents) through the contract. + + ``source`` is either a single document (a ``.pdf``, an image, or a directory of + page images treated as ONE document) or a batch directory tree of such documents. + Output goes to ``resolve_output_root(source, output_dir)`` — default + ``/ocr/``; ``-o`` overrides; never required. + + ``raw`` opts out of ``clean_ocr_output`` so the model's verbatim text is kept. + + Returns a :class:`RunOutcome` whose ``exit_code`` is nonzero if any + document/page failed (uniform across single-file and batch). + """ + # Resolve the output root FIRST so discovery can exclude it. The canonical + # default for a directory input is ``/ocr/`` — INSIDE the scanned tree — + # so a naive recursive walk would re-ingest the engine's own .md/figure outputs + # as inputs on the next run. The contract's iter_input_files prunes that subtree. + output_root = resolve_output_root(source, output_dir) + + documents, scan_root = _discover_documents(source, output_root) + if not documents: + raise ValueError(f"no documents found at {source}") + + if backend.model is None or backend.model is False: + backend.load_model() + + output_root.mkdir(parents=True, exist_ok=True) + index = RootIndex(output_root) + # Run fingerprint keys the idempotency cache on everything that changes a + # document's OUTPUT for a given input. The dedicated task/prompt params handle + # the prompt selector; ``extra`` carries the remaining RESOLVED output-affecting + # flags (raw, dpi, max_tokens, analyze_figures) so a re-run with any of them + # changed reprocesses instead of silently reusing a stale cached result. + # + # When a custom --prompt is set the backends IGNORE --task (they only call + # get_prompt(task) when prompt is None), so task is dropped from the fingerprint + # to avoid needlessly reprocessing two same-prompt runs that differ only in task. + fingerprint = run_fingerprint( + model=_backend_model(backend), + backend=_backend_name(backend), + task=None if prompt is not None else task, + prompt=prompt, + extra={ + "raw": raw, + "dpi": dpi, + "max_tokens": getattr(backend, "max_tokens", None), + "analyze_figures": analyze_figures, + }, + ) + + outcome = RunOutcome() + for doc in documents: + rel_key = relative_key(doc, scan_root) + # Idempotency pre-check uses the SAFE checksum: an input that became + # unreadable between discovery and processing yields None, which can never + # match a recorded checksum, so the doc is NOT skipped — it falls through to + # _ocr_one_document, whose catch-all records it status=failed and the batch + # CONTINUES (no whole-run abort: the SYS-02 class the contract guards). + pre_checksum = _safe_doc_checksum(doc) + if ( + not reprocess + and pre_checksum is not None + and index.is_completed(rel_key, pre_checksum, fingerprint=fingerprint) + ): + logger.info("skip %s (already completed; use --reprocess)", rel_key) + # Quiet mode emits one .md path per processed doc; a cached/resumed doc + # is still "present", so emit its path too (compute + verify on disk). + skip_md = markdown_path_for(doc_dir_for(output_root, rel_key), rel_key) + outcome.add( + Status.COMPLETED, + output_path=str(skip_md) if skip_md.exists() else None, ) - - logger.info( - f"Processed {file_path.name} - {page_count} pages in {processing_time:.2f}s" - ) - - return result - - except Exception as e: - raise RuntimeError(f"Failed to process {file_path}: {e}") - - def process_batch( - self, - input_path: Path, - recursive: bool = False, - prompt: Optional[str] = None, - show_progress: bool = True, - reprocess: bool = False, - ) -> List[OCRResult]: - """Process multiple files from a directory. - - Args: - reprocess: If True, reprocess files even if already in metadata. - """ - files = collect_files(input_path, recursive=recursive) - logger.info(f"Found {len(files)} files to process") - - metadata = MetadataManager(self.output_dir) - - # Filter already-processed files unless reprocess is requested - if not reprocess: - to_process = [] - skipped = 0 - for f in files: - if metadata.is_processed(f): - skipped += 1 - else: - to_process.append(f) - if skipped: - logger.info(f"Skipping {skipped} already-processed files") - files = to_process - - if self._backend.model is None: - self._backend.load_model() - - results = [] - iterator = tqdm(files, desc="Processing files") if show_progress else files - - for file_path in iterator: - try: - result = self.process_file(file_path, prompt=prompt) - results.append(result) - output_path = self.save_result(result) - - metadata.record( - file_path, - pages=result.page_count, - processing_time=result.processing_time, - model=result.metadata.get("model", ""), - backend=result.metadata.get("backend", ""), - output_path=str(output_path.relative_to(self.output_dir)), - ) - - except Exception as e: - logger.error(f"Failed to process {file_path}: {e}") - continue - - logger.info(f"Successfully processed {len(results)}/{len(files)} files") - return results - - def save_result(self, result: OCRResult, output_path: Optional[Path] = None) -> Path: - if output_path is None: - base_name = sanitize_filename(result.input_path.stem) - output_path = self.output_dir / base_name / f"{base_name}.md" - - ensure_dir(output_path.parent) - - markdown_content = result.to_markdown(include_metadata=self.include_metadata) - output_path.write_text(markdown_content, encoding="utf-8") - - logger.info(f"Saved result to: {output_path}") - return output_path + continue + + result, _images = _ocr_one_document(doc, backend, task, prompt, dpi, raw) + + if ( + analyze_figures + and result.status is not Status.FAILED + and doc.suffix.lower() in _DOC_SUFFIXES + ): + doc_dir = doc_dir_for(output_root, rel_key) + doc_dir.mkdir(parents=True, exist_ok=True) + result.figures_markdown = _process_figures(doc, doc_dir, backend) + + meta, markdown_path = _write_document( + result, output_root, rel_key, backend, index, fingerprint + ) + outcome.add( + meta.status, + detail=None if meta.status is Status.COMPLETED else rel_key, + output_path=str(markdown_path), + ) + + return outcome + + +def discover_documents(source: Path, output_dir: Path | None = None) -> list[Path]: + """Public preview of what :func:`process` would OCR, in the SAME order. + + Resolves the output root exactly as the real run does and returns the list of + source documents discovered under ``source`` (with the output-root subtree + excluded). The CLI ``--dry-run`` uses this so the preview never diverges from + the real run (the dry-run/real-run divergence the review flagged). + """ + output_root = resolve_output_root(source, output_dir) + documents, _scan_root = _discover_documents(source, output_root) + return documents + + +def _is_image_dir_document(source: Path, output_root: Path) -> bool: + """True when ``source`` is unambiguously ONE document made of page images. + + socr renders a PDF to ``page_0001.png ...`` and passes the directory, expecting + deepseek to OCR it as a single document. That is only unambiguous when the + directory contains ONLY images: no PDFs and no subdirectories. The moment a PDF + or a subdirectory is present, ``source`` is a batch tree (the papers-library use + case) and must be walked recursively — treating it as one image-document would + silently drop every PDF and subdir (the HIGH bug this guard fixes). + + Classification is OUTPUT-ROOT-AWARE so it is STABLE across runs. With the + default output root ``/ocr/`` nested inside the scanned directory, the + first run creates that ``ocr/`` subtree; without this exclusion, a re-run would + see the new subdir, return False, and silently reclassify the SAME folder as a + batch tree (the run-1 aggregated ``scan/scan.md`` orphaned, per-image + ``page_0001_png/`` trees emitted on run 2, broken resume). The resolved output + root — and its own ``.md``/figure/metadata outputs — are excluded from the + scan so a first run and a re-run classify the same input identically. + """ + if not source.is_dir(): + return False + has_image = False + for p in source.iterdir(): + # Skip the engine's own output subtree (default root nests in the input): + # it must not shift classification between runs. + if is_within_output_root(p, output_root): + continue + if p.is_dir(): + return False + suffix = p.suffix.lower() + if suffix in _DOC_SUFFIXES: + return False + if suffix in _IMAGE_SUFFIXES: + has_image = True + return has_image + + +def _discover_documents(source: Path, output_root: Path) -> tuple[list[Path], Path]: + """Return ``(documents, scan_root)`` for ``source``, excluding the output root. + + * A bare ``.pdf`` or image file is one document (scan root = its parent). + * A directory that contains ONLY page images (no PDFs, no subdirs) is a SINGLE + image-dir document keyed by its folder name (scan root = its parent). + * Otherwise the directory is a batch tree: every ``.pdf``/image under it is one + document, discovered via the contract's :func:`iter_input_files`, which prunes + the resolved ``output_root`` subtree so prior outputs are never re-ingested. + """ + if source.is_file(): + return [source], source.parent + if not source.is_dir(): + return [], source + if _is_image_dir_document(source, output_root): + # The image-dir document keys on its own folder name (no file stem). + return [source], source.parent + suffixes = _DOC_SUFFIXES | _IMAGE_SUFFIXES + documents = list(iter_input_files(source, output_root, suffixes=suffixes)) + return documents, source diff --git a/deepseek_ocr/utils.py b/deepseek_ocr/utils.py index f5358be..81f5d7c 100644 --- a/deepseek_ocr/utils.py +++ b/deepseek_ocr/utils.py @@ -3,7 +3,6 @@ import logging import re from pathlib import Path -from typing import List from PIL import Image @@ -14,7 +13,16 @@ def setup_logging(level: str = "WARNING", verbose: bool = False) -> logging.Logger: - log_level = logging.DEBUG if verbose else logging.WARNING + """Configure logging. ``--verbose`` forces DEBUG; otherwise honor ``level``. + + ``level`` accepts a level name (e.g. ``"INFO"``, ``"WARNING"``) or number, so + ``DEEPSEEK_OCR_LOG_LEVEL`` is no longer a dead setting. + """ + if verbose: + log_level: int = logging.DEBUG + else: + resolved = logging.getLevelName(str(level).upper()) + log_level = resolved if isinstance(resolved, int) else logging.WARNING logging.basicConfig( level=log_level, @@ -37,11 +45,11 @@ def is_pdf_file(file_path: Path) -> bool: return file_path.suffix.lower() == PDF_EXTENSION -def collect_files(input_path: Path, recursive: bool = False) -> List[Path]: +def collect_files(input_path: Path, recursive: bool = False) -> list[Path]: if not input_path.exists(): raise FileNotFoundError(f"Path not found: {input_path}") - files: List[Path] = [] + files: list[Path] = [] if input_path.is_file(): if is_supported_file(input_path): @@ -127,14 +135,14 @@ def ensure_dir(directory: Path) -> Path: def _html_table_to_markdown(html_table: str) -> str: rows = [] - row_matches = re.findall(r']*>(.*?)', html_table, re.DOTALL | re.IGNORECASE) + row_matches = re.findall(r"]*>(.*?)", html_table, re.DOTALL | re.IGNORECASE) for row_html in row_matches: - cells = re.findall(r']*>(.*?)', row_html, re.DOTALL | re.IGNORECASE) + cells = re.findall(r"]*>(.*?)", row_html, re.DOTALL | re.IGNORECASE) cleaned_cells = [] for cell in cells: - cell = re.sub(r'<[^>]+>', '', cell) - cell = ' '.join(cell.split()) + cell = re.sub(r"<[^>]+>", "", cell) + cell = " ".join(cell.split()) cleaned_cells.append(cell) if cleaned_cells: rows.append(cleaned_cells) @@ -153,38 +161,38 @@ def _html_table_to_markdown(html_table: str) -> str: def clean_ocr_output(text: str) -> str: """Remove grounding annotations, convert HTML tables to markdown, decode entities.""" - text = re.sub(r'<\|ref\|>.*?<\|/ref\|>', '', text) - text = re.sub(r'<\|det\|>\[\[.*?\]\]<\|/det\|>', '', text) - text = re.sub(r'<\|[^|]+\|>', '', text) + text = re.sub(r"<\|ref\|>.*?<\|/ref\|>", "", text) + text = re.sub(r"<\|det\|>\[\[.*?\]\]<\|/det\|>", "", text) + text = re.sub(r"<\|[^|]+\|>", "", text) # Strip bare bounding box coordinates that leak through without tags # e.g., "text[[114, 531, 883, 619]]" → "text" - text = re.sub(r'\[\[\d+,\s*\d+,\s*\d+,\s*\d+\]\]', '', text) + text = re.sub(r"\[\[\d+,\s*\d+,\s*\d+,\s*\d+\]\]", "", text) def replace_table(match: re.Match) -> str: return _html_table_to_markdown(match.group(0)) - text = re.sub(r']*>.*?', replace_table, text, flags=re.DOTALL | re.IGNORECASE) + text = re.sub(r"]*>.*?", replace_table, text, flags=re.DOTALL | re.IGNORECASE) - text = re.sub(r'<(sup|sub)>([^<]*)', r'^\2', text, flags=re.IGNORECASE) - text = re.sub(r'

([^<]*)
', r'\1', text, flags=re.IGNORECASE) - text = re.sub(r'', '\n', text, flags=re.IGNORECASE) - text = re.sub(r'<[^>]+>', '', text) + text = re.sub(r"<(sup|sub)>([^<]*)", r"^\2", text, flags=re.IGNORECASE) + text = re.sub(r"
([^<]*)
", r"\1", text, flags=re.IGNORECASE) + text = re.sub(r"", "\n", text, flags=re.IGNORECASE) + text = re.sub(r"<[^>]+>", "", text) html_entities = { - '&': '&', - '<': '<', - '>': '>', - '"': '"', - ''': "'", - ' ': ' ', - ''': "'", - ''': "'", + "&": "&", + "<": "<", + ">": ">", + """: '"', + "'": "'", + " ": " ", + "'": "'", + "'": "'", } for entity, char in html_entities.items(): text = text.replace(entity, char) - text = re.sub(r'\n{3,}', '\n\n', text) - lines = [line.strip() for line in text.split('\n')] - text = '\n'.join(lines) + text = re.sub(r"\n{3,}", "\n\n", text) + lines = [line.strip() for line in text.split("\n")] + text = "\n".join(lines) text = text.strip() return text diff --git a/examples/basic_usage.py b/examples/basic_usage.py index e92b5b9..f378980 100644 --- a/examples/basic_usage.py +++ b/examples/basic_usage.py @@ -1,54 +1,50 @@ #!/usr/bin/env python -"""Basic usage example for DeepSeek OCR CLI library.""" +"""Basic usage example for DeepSeek OCR CLI library. + +Output is routed through the shared ``ocr-output-contract`` package: each document +lands in ``//.md`` under ``## Page N`` headers, with a +per-document ``metadata.json`` sidecar and a rolled-up root index. ``process()`` +returns a :class:`RunOutcome` whose ``outputs`` lists the written ``.md`` paths. +""" from pathlib import Path -from deepseek_ocr import create_backend, OCRProcessor + +from deepseek_ocr import create_backend, process def main() -> None: """Demonstrate basic usage of the library.""" print("DeepSeek OCR - Basic Usage Example\n") - # Initialize backend (Ollama by default) + # Initialize backend (Ollama by default). process() also loads it if needed. print("1. Connecting to Ollama...") backend = create_backend(backend_type="ollama", model_name="deepseek-ocr") backend.load_model() print(" Connected\n") - # Create processor - print("2. Creating OCR processor...") - processor = OCRProcessor( - backend=backend, - output_dir=Path("./output"), - include_metadata=True, - workers=2, - ) + output_dir = Path("./output") - # Example 1: Process a single image (if available) + # Example 1: Process a single image (if available). example_image = Path("test_image.jpg") if example_image.exists(): - print(f"3. Processing image: {example_image}") - result = processor.process_file(example_image) - print(f" Extracted {len(result.output_text)} characters") - print(f" Processing time: {result.processing_time:.2f}s") - - output_path = processor.save_result(result) - print(f" Saved to: {output_path}\n") + print(f"2. Processing image: {example_image}") + outcome = process(example_image, backend, output_dir=output_dir) + for md_path in outcome.outputs: + print(f" Saved to: {md_path}") + print(f" completed={outcome.completed} failed={outcome.failed}\n") else: - print(f"3. Skipping image processing (no {example_image} found)\n") + print(f"2. Skipping image processing (no {example_image} found)\n") - # Example 2: Process a PDF (if available) + # Example 2: Process a PDF (if available). example_pdf = Path("test_document.pdf") if example_pdf.exists(): - print(f"4. Processing PDF: {example_pdf}") - result = processor.process_file(example_pdf) - print(f" Processed {result.page_count} pages") - print(f" Processing time: {result.processing_time:.2f}s") - - output_path = processor.save_result(result) - print(f" Saved to: {output_path}\n") + print(f"3. Processing PDF: {example_pdf}") + outcome = process(example_pdf, backend, output_dir=output_dir) + for md_path in outcome.outputs: + print(f" Saved to: {md_path}") + print(f" exit_code={outcome.exit_code}\n") else: - print(f"4. Skipping PDF processing (no {example_pdf} found)\n") + print(f"3. Skipping PDF processing (no {example_pdf} found)\n") # Cleanup backend.unload_model() diff --git a/examples/batch_processing.py b/examples/batch_processing.py index dcf9ea7..9fba0c1 100644 --- a/examples/batch_processing.py +++ b/examples/batch_processing.py @@ -1,12 +1,19 @@ #!/usr/bin/env python -"""Batch processing example for DeepSeek OCR CLI library.""" +"""Batch processing example for DeepSeek OCR CLI library. + +A directory tree of documents is walked recursively (the resolved output root is +excluded from discovery, so prior outputs are never re-ingested). Each document is +keyed by its input-relative path, so same-named files in different subdirectories +never collide. Re-running skips already-completed documents (incremental resume). +""" from pathlib import Path -from deepseek_ocr import create_backend, OCRProcessor + +from deepseek_ocr import create_backend, process def main() -> None: - """Demonstrate batch processing of multiple files.""" + """Demonstrate batch processing of a directory tree.""" print("DeepSeek OCR - Batch Processing Example\n") input_dir = Path("./documents") @@ -17,40 +24,24 @@ def main() -> None: print("Please create ./documents/ and add some PDF/image files") return - # Initialize backend print("1. Connecting to Ollama...") backend = create_backend(backend_type="ollama", model_name="deepseek-ocr") backend.load_model() print(" Connected\n") - # Create processor - print("2. Setting up processor...") - processor = OCRProcessor( - backend=backend, - output_dir=output_dir, - extract_images=True, - include_metadata=True, - ) - - # Process all files in directory - print(f"3. Processing files in: {input_dir}") - results = processor.process_batch( - input_path=input_dir, - recursive=True, - show_progress=True, - ) - - # Print summary - print(f"\n4. Processing Summary:") - print(f" Files processed: {len(results)}") - if results: - total_pages = sum(r.page_count for r in results) - total_time = sum(r.processing_time for r in results) - print(f" Total pages: {total_pages}") - print(f" Total time: {total_time:.2f}s") - if total_pages > 0: - print(f" Average time per page: {total_time/total_pages:.2f}s") - print(f" Output directory: {output_dir}\n") + print(f"2. Processing documents under: {input_dir}") + outcome = process(input_dir, backend, output_dir=output_dir) + + total = outcome.completed + outcome.failed + outcome.partial + print("\n3. Processing Summary:") + print(f" Documents processed: {total}") + print(f" Completed: {outcome.completed}") + print(f" Partial: {outcome.partial}") + print(f" Failed: {outcome.failed}") + if outcome.failures: + print(f" Failures: {', '.join(outcome.failures)}") + print(f" Output directory: {output_dir}") + print(f" Exit code: {outcome.exit_code}\n") backend.unload_model() print("Done!") diff --git a/examples/figure_analysis.py b/examples/figure_analysis.py index ed706ed..b9acb6f 100644 --- a/examples/figure_analysis.py +++ b/examples/figure_analysis.py @@ -1,8 +1,14 @@ #!/usr/bin/env python -"""Example: Extract and analyze embedded figures from PDFs.""" +"""Example: Extract and analyze embedded figures from PDFs. + +With ``analyze_figures=True``, embedded figures are extracted, saved under +``//figures/figure__page

.png`` (canon naming), and an +AI description is appended to the document markdown. +""" from pathlib import Path -from deepseek_ocr import create_backend, OCRProcessor + +from deepseek_ocr import create_backend, process def main() -> None: @@ -12,30 +18,22 @@ def main() -> None: backend = create_backend(backend_type="ollama", model_name="deepseek-ocr") backend.load_model() - # Create processor with figure analysis enabled - processor = OCRProcessor( - backend=backend, - output_dir=Path("./output"), - analyze_figures=True, - workers=2, - ) - + output_dir = Path("./output") pdf_path = Path("test_document.pdf") if not pdf_path.exists(): print(f"Error: {pdf_path} not found") + backend.unload_model() return print(f"Processing: {pdf_path}") - result = processor.process_file(pdf_path) - - print(f"Pages: {result.page_count}") - print(f"Time: {result.processing_time:.2f}s") + outcome = process(pdf_path, backend, output_dir=output_dir, analyze_figures=True) - output_path = processor.save_result(result) - print(f"Output: {output_path}") + for md_path in outcome.outputs: + print(f"Output: {md_path}") + print(f"completed={outcome.completed} failed={outcome.failed}") - # Figures are saved in output/doc_name/figures/ - figures_dir = Path("./output") / pdf_path.stem / "figures" + # Figures are saved in //figures/. + figures_dir = output_dir / pdf_path.stem / "figures" if figures_dir.exists(): figures = list(figures_dir.glob("*")) print(f"Figures: {len(figures)} saved to {figures_dir}/") diff --git a/pyproject.toml b/pyproject.toml index 73c00de..6e8d567 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,14 +1,12 @@ -[tool.poetry] +[project] name = "deepseek-ocr-cli" version = "0.4.3" -description = "CLI tool for OCR using DeepSeek-OCR model via Ollama" -authors = ["Ruben Fernandez Fuertes "] +description = "CLI tool for OCR using DeepSeek-OCR model via Ollama or vLLM" readme = "README.md" -license = "MIT" -repository = "https://github.com/r-uben/deepseek-ocr-cli" -homepage = "https://github.com/r-uben/deepseek-ocr-cli" +requires-python = ">=3.11" +license = { text = "MIT" } +authors = [{ name = "Ruben Fernandez Fuertes", email = "fernandezfuertesruben@gmail.com" }] keywords = ["ocr", "deepseek", "cli", "pdf", "document-processing", "ollama"] -packages = [{include = "deepseek_ocr"}] classifiers = [ "Development Status :: 4 - Beta", "Intended Audience :: Developers", @@ -16,55 +14,63 @@ classifiers = [ "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Text Processing :: General", ] +dependencies = [ + "ocr-output-contract @ git+https://github.com/r-uben/ocr-output-contract@v0.1.3", + "pillow>=11.0.0", + "click>=8.1.7", + "python-dotenv>=1.0.0", + "tqdm>=4.67.0", + "pydantic>=2.10.0", + "pydantic-settings>=2.7.0", + "pymupdf>=1.25.1", + "rich>=13.9.4", + "requests>=2.31.0", +] -[tool.poetry.dependencies] -python = "^3.10" -pillow = "^11.0.0" -click = "^8.1.7" -python-dotenv = "^1.0.0" -tqdm = "^4.67.0" -pydantic = "^2.10.0" -pydantic-settings = "^2.7.0" -pymupdf = "^1.25.1" -rich = "^13.9.4" -requests = "^2.31.0" -openai = {version = "^1.0.0", optional = true} - -[tool.poetry.extras] -vllm = ["openai"] +[project.urls] +repository = "https://github.com/r-uben/deepseek-ocr-cli" +homepage = "https://github.com/r-uben/deepseek-ocr-cli" -[tool.poetry.group.dev.dependencies] -pytest = "^8.3.4" -pytest-cov = "^6.0.0" -black = "^24.10.0" -ruff = "^0.8.4" -mypy = "^1.13.0" -types-pillow = "^10.2.0" -types-requests = "^2.31.0" +[project.optional-dependencies] +vllm = ["openai>=1.0.0"] +dev = [ + "pytest>=8.3.4", + "pytest-cov>=6.0.0", + "ruff>=0.8.4", + "mypy>=1.13.0", + "types-pillow>=10.2.0", + "types-requests>=2.31.0", +] -[tool.poetry.scripts] +[project.scripts] deepseek-ocr = "deepseek_ocr.cli:main" -[tool.black] -line-length = 100 -target-version = ['py310'] +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["deepseek_ocr"] + +# Required so the git+https direct reference to ocr-output-contract is allowed. +[tool.hatch.metadata] +allow-direct-references = true [tool.ruff] line-length = 100 -target-version = "py310" +target-version = "py311" [tool.ruff.lint] select = ["E", "F", "I", "N", "W", "UP"] ignore = ["E501"] [tool.mypy] -python_version = "3.10" +python_version = "3.11" warn_return_any = true warn_unused_configs = true disallow_untyped_defs = true @@ -75,8 +81,5 @@ testpaths = ["tests"] python_files = ["test_*.py"] python_classes = ["Test*"] python_functions = ["test_*"] -addopts = "--cov=deepseek_ocr --cov-report=term-missing --cov-report=html" - -[build-system] -requires = ["poetry-core"] -build-backend = "poetry.core.masonry.api" +markers = ["integration: hits live models/servers; excluded from default run"] +addopts = "-m 'not integration'" diff --git a/tests/test_cli.py b/tests/test_cli.py index 8fa0f62..edb7d95 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,13 +1,23 @@ """Tests for CLI interface.""" -import os import tempfile from pathlib import Path +import fitz from click.testing import CliRunner from PIL import Image -from deepseek_ocr.cli import cli, _format_size, _run_dry_run +from deepseek_ocr.cli import _format_size, cli + + +def _make_pdf(path: Path, pages: int = 1) -> None: + """Write a tiny multi-page PDF for discovery/dry-run tests.""" + doc = fitz.open() + for i in range(pages): + page = doc.new_page(width=200, height=300) + page.insert_text((40, 60), f"page {i + 1}") + doc.save(str(path)) + doc.close() class TestFormatSize: @@ -43,22 +53,63 @@ def test_dry_run_single_image(self) -> None: assert "dry run" in result.output.lower() def test_dry_run_directory(self) -> None: - """Dry run on a directory lists all supported files.""" + """Dry run on a batch tree lists each document, excluding unsupported files. + + A directory that mixes PDFs with subdirectories is a BATCH tree, so each + document is listed individually (the dry-run preview uses the same real + discovery as the run, so it never diverges). + """ runner = CliRunner() with tempfile.TemporaryDirectory() as tmpdir: - for name in ("a.png", "b.jpg", "c.txt"): - p = Path(tmpdir) / name - if name.endswith(".txt"): - p.write_text("not an image") - else: - Image.new("RGB", (50, 50)).save(p) + root = Path(tmpdir) + # A subdir forces batch-tree semantics (not a single image-dir doc). + (root / "sub").mkdir() + _make_pdf(root / "a.pdf") + _make_pdf(root / "sub" / "b.pdf") + (root / "c.txt").write_text("not a document") result = runner.invoke(cli, ["process", tmpdir, "--dry-run"]) assert result.exit_code == 0 - assert "a.png" in result.output - assert "b.jpg" in result.output + assert "a.pdf" in result.output + assert "b.pdf" in result.output assert "c.txt" not in result.output + def test_dry_run_image_dir_is_one_document(self) -> None: + """A flat directory of only images is ONE image-dir document (socr case).""" + runner = CliRunner() + with tempfile.TemporaryDirectory() as tmpdir: + root = Path(tmpdir) / "scan" + root.mkdir() + for n in (1, 2): + Image.new("RGB", (50, 50)).save(root / f"page_{n:04d}.png") + + result = runner.invoke(cli, ["process", str(root), "--dry-run"]) + assert result.exit_code == 0 + # The directory itself is the single document; pages = image count. + assert "scan" in result.output + assert "IMG DIR" in result.output + + def test_dry_run_image_dir_page_count_ignores_stray_files(self) -> None: + """LOW fix: image-dir dry-run counts only image files, matching the real run. + + A directory of 2 images plus a stray non-image file is one image-dir + document of 2 pages (the real run filters by image suffix). The dry-run + preview must report 2 pages, not 3, so it never diverges from what is OCR'd. + """ + runner = CliRunner() + with tempfile.TemporaryDirectory() as tmpdir: + root = Path(tmpdir) / "scan" + root.mkdir() + for n in (1, 2): + Image.new("RGB", (50, 50)).save(root / f"page_{n:04d}.png") + (root / "notes.txt").write_text("not an image") + + result = runner.invoke(cli, ["process", str(root), "--dry-run"]) + assert result.exit_code == 0 + assert "IMG DIR" in result.output + # The summary line reports total pages = image count (2), not 3. + assert "2 pages" in result.output + def test_dry_run_quiet(self) -> None: """Dry run with --quiet outputs only file paths.""" runner = CliRunner() @@ -89,6 +140,63 @@ def test_quiet_suppresses_banner(self) -> None: assert "deepseek-ocr v" not in result.output +class TestModelPrecedence: + """LOW fix: DEEPSEEK_OCR_MODEL_NAME / settings.model_name wins when --model omitted.""" + + def _run_capturing_model(self, monkeypatch, argv: list[str]) -> str: + """Invoke `process` with backend creation + OCR stubbed; return the model_name.""" + import deepseek_ocr.cli as cli_mod + + captured: dict[str, str] = {} + + def fake_create_backend(*, backend_type, model_name, **kwargs): + captured["model_name"] = model_name + return _StubBackend() # OCR is stubbed; only unload_model() is called + + monkeypatch.setattr(cli_mod, "create_backend", fake_create_backend) + monkeypatch.setattr(cli_mod, "run_process", lambda *a, **k: _StubOutcome()) + runner = CliRunner() + with tempfile.TemporaryDirectory() as tmpdir: + pdf = Path(tmpdir) / "doc.pdf" + _make_pdf(pdf) + result = runner.invoke(cli, ["process", str(pdf), *argv]) + assert result.exit_code == 0, result.output + return captured["model_name"] + + def test_env_model_used_when_flag_omitted(self, monkeypatch) -> None: + """With no --model, settings.model_name (DEEPSEEK_OCR_MODEL_NAME) wins.""" + from deepseek_ocr.config import settings + + monkeypatch.setattr(settings, "model_name", "env-model") + assert self._run_capturing_model(monkeypatch, []) == "env-model" + + def test_cli_model_overrides_env(self, monkeypatch) -> None: + """An explicit --model wins over settings.model_name.""" + from deepseek_ocr.config import settings + + monkeypatch.setattr(settings, "model_name", "env-model") + assert self._run_capturing_model(monkeypatch, ["--model", "cli-model"]) == "cli-model" + + +class _StubOutcome: + """Minimal RunOutcome stand-in so the CLI can finish without real OCR.""" + + completed = 0 + failed = 0 + partial = 0 + has_failures = False + exit_code = 0 + outputs: list[str] = [] + failures: list[str] = [] + + +class _StubBackend: + """Backend stand-in: the CLI only calls unload_model() after a stubbed run.""" + + def unload_model(self) -> None: + pass + + class TestAutoInsertProcess: """Tests that the main() entry point auto-inserts 'process' subcommand.""" diff --git a/tests/test_config.py b/tests/test_config.py index f7b0709..3af62fe 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,7 +1,5 @@ """Tests for configuration management.""" -from pathlib import Path - from deepseek_ocr.config import Settings @@ -16,9 +14,7 @@ def test_default_settings(self) -> None: assert settings.ollama_url == "http://localhost:11434" assert settings.vllm_base_url == "http://localhost:8000/v1" assert settings.max_dimension == 1920 - assert settings.output_dir == Path("output") - assert settings.extract_images is False - assert settings.include_metadata is True + assert settings.max_tokens == 8192 assert settings.max_retries == 3 assert settings.retry_delay == 1.0 diff --git a/tests/test_metadata.py b/tests/test_metadata.py deleted file mode 100644 index db5eed3..0000000 --- a/tests/test_metadata.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Tests for metadata tracking.""" - -import json -from pathlib import Path - -from deepseek_ocr.metadata import MetadataManager, _file_checksum - - -class TestFileChecksum: - """Tests for file checksum computation.""" - - def test_checksum_deterministic(self, tmp_path: Path) -> None: - f = tmp_path / "test.txt" - f.write_text("hello world") - assert _file_checksum(f) == _file_checksum(f) - - def test_checksum_changes_with_content(self, tmp_path: Path) -> None: - f = tmp_path / "test.txt" - f.write_text("hello") - c1 = _file_checksum(f) - f.write_text("world") - c2 = _file_checksum(f) - assert c1 != c2 - - def test_checksum_prefix(self, tmp_path: Path) -> None: - f = tmp_path / "test.txt" - f.write_text("data") - assert _file_checksum(f).startswith("sha256:") - - -class TestMetadataManager: - """Tests for MetadataManager.""" - - def test_save_and_load(self, tmp_path: Path) -> None: - mgr = MetadataManager(tmp_path) - f = tmp_path / "doc.pdf" - f.write_bytes(b"%PDF-fake") - - mgr.record( - f, - pages=3, - processing_time=5.67, - model="deepseek-ocr", - backend="ollama", - output_path="doc.md", - ) - - # Load fresh from disk - mgr2 = MetadataManager(tmp_path) - assert "doc.pdf" in mgr2.files - entry = mgr2.files["doc.pdf"] - assert entry["status"] == "completed" - assert entry["pages"] == 3 - assert entry["processing_time"] == 5.67 - assert entry["model"] == "deepseek-ocr" - assert entry["output_path"] == "doc.md" - - def test_is_processed_true(self, tmp_path: Path) -> None: - mgr = MetadataManager(tmp_path) - f = tmp_path / "doc.pdf" - f.write_bytes(b"%PDF-fake") - - mgr.record( - f, - pages=1, - processing_time=1.0, - model="m", - backend="b", - output_path="doc.md", - ) - - assert mgr.is_processed(f) is True - - def test_is_processed_false_unknown_file(self, tmp_path: Path) -> None: - mgr = MetadataManager(tmp_path) - f = tmp_path / "new.pdf" - f.write_bytes(b"new content") - assert mgr.is_processed(f) is False - - def test_is_processed_false_after_content_change(self, tmp_path: Path) -> None: - mgr = MetadataManager(tmp_path) - f = tmp_path / "doc.pdf" - f.write_bytes(b"version1") - - mgr.record( - f, - pages=1, - processing_time=1.0, - model="m", - backend="b", - output_path="doc.md", - ) - - # Modify the file - f.write_bytes(b"version2") - assert mgr.is_processed(f) is False - - def test_atomic_write(self, tmp_path: Path) -> None: - """Verify no .tmp file remains after save.""" - mgr = MetadataManager(tmp_path) - mgr.save() - - assert (tmp_path / "metadata.json").exists() - assert not (tmp_path / "metadata.json.tmp").exists() - - def test_load_corrupt_json(self, tmp_path: Path) -> None: - """Corrupt metadata.json should not crash, just start fresh.""" - meta_file = tmp_path / "metadata.json" - meta_file.write_text("{bad json", encoding="utf-8") - - mgr = MetadataManager(tmp_path) - assert mgr.files == {} - - def test_empty_dir_no_metadata(self, tmp_path: Path) -> None: - mgr = MetadataManager(tmp_path) - assert mgr.files == {} - - def test_multiple_records(self, tmp_path: Path) -> None: - mgr = MetadataManager(tmp_path) - - for name in ["a.pdf", "b.pdf", "c.pdf"]: - f = tmp_path / name - f.write_bytes(name.encode()) - mgr.record( - f, - pages=1, - processing_time=0.5, - model="m", - backend="b", - output_path=f"{name}.md", - ) - - mgr2 = MetadataManager(tmp_path) - assert len(mgr2.files) == 3 - - def test_metadata_json_structure(self, tmp_path: Path) -> None: - mgr = MetadataManager(tmp_path) - f = tmp_path / "test.pdf" - f.write_bytes(b"test") - mgr.record( - f, - pages=2, - processing_time=3.14, - model="deepseek-ocr", - backend="ollama", - output_path="test.md", - ) - - raw = json.loads((tmp_path / "metadata.json").read_text()) - assert raw["version"] == "1" - assert "test.pdf" in raw["files"] - assert "checksum" in raw["files"]["test.pdf"] - assert "timestamp" in raw["files"]["test.pdf"] diff --git a/tests/test_output_contract.py b/tests/test_output_contract.py new file mode 100644 index 0000000..5bf1206 --- /dev/null +++ b/tests/test_output_contract.py @@ -0,0 +1,430 @@ +"""Engine-level conformance: deepseek's REAL output vs the shared contract. + +The contract *primitives* (path/key computation, page assembly, metadata writers, +the exit-code policy) are unit-tested inside the ``ocr-output-contract`` package +itself, so they are NOT re-tested here. + +What stays here is the engine-side proof: run deepseek's actual processor (with a +mocked backend — no Ollama, no GPU, no network) over real inputs and assert the +produced output tree conforms to the family-wide contract via the package's reusable +:func:`ocr_output_contract.conformance.assert_conforms` harness, including: + +* a multi-page PDF aggregating into ONE ``/.md`` with both pages; +* an empty/whitespace model response recorded as ``status=failed`` driving a nonzero exit; +* the HIGH fix: the ``.md`` body carries NO YAML frontmatter (provenance -> sidecar only); +* the HIGH fix: ``--task`` actually reaches the backend prompt. +""" + +from __future__ import annotations + +import json + +import fitz +from ocr_output_contract import UNREADABLE_CHECKSUM +from ocr_output_contract.conformance import ExpectedDoc, assert_conforms +from PIL import Image + +from deepseek_ocr.backends.base import Backend +from deepseek_ocr.processor import discover_documents, process + + +class FakeBackend(Backend): + """A mocked backend: deterministic per-page text, no network/model.""" + + def __init__(self, text="OCR page text", model="deepseek-ocr"): + super().__init__(model_name=model) + self.text = text + self.calls = 0 + self.seen_prompts: list[str] = [] + + @property + def backend_name(self) -> str: + return "ollama" + + def load_model(self) -> None: + self.model = True + + def unload_model(self) -> None: + self.model = False + + def process_image(self, image, prompt=None, task="convert", return_raw=False): + self.calls += 1 + resolved = prompt if prompt is not None else self.get_prompt(task) + self.seen_prompts.append(resolved) + return self.text + + +def _make_pdf(path, pages=2): + doc = fitz.open() + for i in range(pages): + page = doc.new_page(width=300, height=400) + page.insert_text((40, 60), f"Source page {i + 1}") + doc.save(path) + doc.close() + + +def test_multipage_pdf_conforms_no_data_loss(tmp_path): + """A 2-page PDF -> one /.md with both pages, no per-page folders.""" + pdf = tmp_path / "sample.pdf" + _make_pdf(pdf, pages=2) + out = tmp_path / "out" + + counter = {"n": 0} + + class PerPage(FakeBackend): + def process_image(self, image, prompt=None, task="convert", return_raw=False): + counter["n"] += 1 + return f"PAGE-{counter['n']}-CONTENT" + + outcome = process(pdf, PerPage(), dpi=120, output_dir=out) + assert outcome.exit_code == 0 + + assert_conforms( + out, + [ExpectedDoc(rel_key="sample.pdf", pages=2, status="completed")], + require_failures_nonzero_exit=outcome.exit_code != 0, + ) + + md = out / "sample" / "sample.md" + body = md.read_text() + assert "## Page 1" in body and "## Page 2" in body + assert "PAGE-1-CONTENT" in body and "PAGE-2-CONTENT" in body + assert not (out / "sample_p0001").exists() + assert not (out / "sample_p0002").exists() + + +def test_no_yaml_frontmatter_in_markdown(tmp_path): + """HIGH fix: the .md body must NOT carry YAML frontmatter; provenance is in the sidecar.""" + pdf = tmp_path / "doc.pdf" + _make_pdf(pdf, pages=1) + out = tmp_path / "out" + + outcome = process(pdf, FakeBackend(text="clean body text"), dpi=120, output_dir=out) + assert outcome.exit_code == 0 + + body = (out / "doc" / "doc.md").read_text() + # No leading YAML frontmatter delimiter, and none of the old provenance keys leaked. + assert not body.lstrip().startswith("---") + for key in ("source:", "processed:", "processing_time:", "model:", "backend:"): + assert key not in body + + # Provenance lives ONLY in the sidecar. + meta = json.loads((out / "doc" / "metadata.json").read_text()) + assert meta["backend"] == "ollama" + assert meta["model"] == "deepseek-ocr" + assert meta["status"] == "completed" + + +def test_task_reaches_backend_prompt(tmp_path): + """HIGH fix: --task is honoured -> the task's prompt reaches the backend.""" + pdf = tmp_path / "t.pdf" + _make_pdf(pdf, pages=1) + out = tmp_path / "out" + + be = FakeBackend(text="text") + process(pdf, be, dpi=120, task="ocr", output_dir=out) + + # "ocr" maps to a distinct prompt from the default "convert". + assert be.seen_prompts == [Backend.PROMPTS["ocr"]] + assert be.seen_prompts[0] != Backend.PROMPTS["convert"] + + +def test_empty_response_fails_and_exits_nonzero(tmp_path): + """An empty model response -> status=failed, nonzero exit (not a 0-byte success).""" + pdf = tmp_path / "blank.pdf" + _make_pdf(pdf, pages=1) + out = tmp_path / "out" + + outcome = process(pdf, FakeBackend(text=" "), dpi=120, output_dir=out) + assert outcome.exit_code != 0 + + assert_conforms( + out, + [ExpectedDoc(rel_key="blank.pdf", status="failed")], + require_failures_nonzero_exit=True, + ) + doc_meta = json.loads((out / "blank" / "metadata.json").read_text()) + assert doc_meta["status"] == "failed" + + +def test_nested_batch_conforms_no_basename_collision(tmp_path): + """Two same-basename PDFs in different subdirs both survive (input-relative key).""" + root = tmp_path / "in" + (root / "a").mkdir(parents=True) + (root / "b").mkdir(parents=True) + _make_pdf(root / "a" / "intro.pdf", pages=1) + _make_pdf(root / "b" / "intro.pdf", pages=1) + out = tmp_path / "out" + + outcome = process(root, FakeBackend(), dpi=120, output_dir=out) + assert outcome.exit_code == 0 + + assert_conforms( + out, + [ + ExpectedDoc(rel_key="a/intro.pdf", pages=1, status="completed"), + ExpectedDoc(rel_key="b/intro.pdf", pages=1, status="completed"), + ], + ) + + +def test_image_dir_document_conforms(tmp_path): + """A directory of page images is one conforming document keyed by folder name.""" + src = tmp_path / "scan" + src.mkdir() + for n in (1, 2): + Image.new("RGB", (60, 60), "white").save(src / f"page_{n:04d}.png") + out = tmp_path / "out" + + outcome = process(src, FakeBackend(), dpi=120, output_dir=out) + assert outcome.exit_code == 0 + + assert_conforms( + out, + [ExpectedDoc(rel_key="scan", pages=2, status="completed")], + ) + + +def test_mixed_dir_with_stray_image_processes_all_pdfs(tmp_path): + """HIGH fix: a batch dir holding PDFs + a stray image processes ALL PDFs. + + Previously any direct image misclassified the whole directory as ONE + image-document, silently dropping every PDF and subdirectory (the + papers-library batch use case). Now a directory that contains a PDF (or a + subdir) is a batch tree: every PDF AND the stray image are OCR'd as their own + documents, and subdirectories are recursed. + """ + root = tmp_path / "papers" + (root / "sub").mkdir(parents=True) + _make_pdf(root / "paper1.pdf", pages=1) + _make_pdf(root / "sub" / "paper2.pdf", pages=1) + Image.new("RGB", (60, 60), "white").save(root / "cover.png") + out = tmp_path / "out" + + outcome = process(root, FakeBackend(), dpi=120, output_dir=out) + assert outcome.exit_code == 0 + # 2 PDFs + 1 stray image = 3 documents, none dropped. + assert outcome.completed == 3 + + assert_conforms( + out, + [ + ExpectedDoc(rel_key="paper1.pdf", pages=1, status="completed"), + ExpectedDoc(rel_key="sub/paper2.pdf", pages=1, status="completed"), + ExpectedDoc(rel_key="cover.png", pages=1, status="completed"), + ], + ) + + +def test_rerun_does_not_reingest_default_output_root(tmp_path): + """HIGH fix: the default /ocr/ output is excluded from re-run discovery. + + With the default (nested) output root, a second run must not re-discover the + first run's own .md/figure outputs as fresh inputs. Discovery via + iter_input_files prunes the resolved output-root subtree. + """ + root = tmp_path / "papers" + root.mkdir() + _make_pdf(root / "doc.pdf", pages=1) + + # First run with the DEFAULT output root (/ocr/, inside the tree). + outcome1 = process(root, FakeBackend(), dpi=120) + assert outcome1.completed == 1 + assert (root / "ocr").is_dir() # default root sits inside the scanned tree + + # Second run: the only new document is still the one PDF; the ocr/ outputs + # (markdown + metadata) must NOT be re-discovered as inputs. + docs = discover_documents(root) + assert docs == [root / "doc.pdf"] + + +def test_truncated_page_recorded_partial_not_completed(tmp_path): + """MEDIUM fix: a length-truncated non-empty page is partial/failed, not completed. + + The backend reports a length finish_reason; the page text is non-empty but + incomplete, so recording it as completed would be silent content loss. + """ + pdf = tmp_path / "dense.pdf" + _make_pdf(pdf, pages=1) + out = tmp_path / "out" + + class TruncatingBackend(FakeBackend): + def process_image_with_meta(self, image, prompt=None, task="convert", return_raw=False): + # Non-empty text, but the model stopped on the token limit. + return "partial dense content cut off", "length" + + outcome = process(pdf, TruncatingBackend(), dpi=120, output_dir=out) + assert outcome.exit_code != 0 + + doc_meta = json.loads((out / "dense" / "metadata.json").read_text()) + assert doc_meta["status"] == "failed" # single page truncated -> no good pages + assert "truncat" in doc_meta["error"].lower() + + +def test_quiet_skip_emits_md_path_on_resume(tmp_path): + """MEDIUM fix: a cached/resumed doc still contributes its .md path to outputs. + + On the first run the doc is processed and its path emitted; on the second run + it is skipped (already completed) but must STILL appear in outcome.outputs so + `deepseek-ocr dir -q` does not silently drop cached docs. + """ + pdf = tmp_path / "doc.pdf" + _make_pdf(pdf, pages=1) + out = tmp_path / "out" + + outcome1 = process(pdf, FakeBackend(), dpi=120, output_dir=out) + md = out / "doc" / "doc.md" + assert outcome1.outputs == [str(md)] + + # Second run: skip branch, but the .md path is still emitted. + outcome2 = process(pdf, FakeBackend(), dpi=120, output_dir=out) + assert outcome2.completed == 1 + assert outcome2.outputs == [str(md)] + + +def test_rerun_under_different_task_reprocesses(tmp_path): + """The run fingerprint invalidates the cache when --task changes. + + A re-run under a different task must NOT silently reuse the prior output: the + fingerprint (model/backend/task/prompt) differs, so is_completed returns False. + """ + pdf = tmp_path / "doc.pdf" + _make_pdf(pdf, pages=1) + out = tmp_path / "out" + + be1 = FakeBackend(text="convert output") + process(pdf, be1, dpi=120, task="convert", output_dir=out) + assert be1.calls == 1 + + # Same input, different task -> fingerprint differs -> reprocessed (not skipped). + be2 = FakeBackend(text="ocr output") + process(pdf, be2, dpi=120, task="ocr", output_dir=out) + assert be2.calls == 1 + + +def test_image_dir_classification_stable_across_reruns_default_root(tmp_path): + """MEDIUM fix: an image-dir document classifies identically on run 1 and re-run. + + With the DEFAULT output root (``/ocr/``, nested INSIDE the scanned dir), + the first run of an image directory creates that ``ocr/`` subtree. Without + output-root-aware classification, the re-run would see the new subdir, decide + the folder is no longer a pure image dir, and silently reclassify it as a BATCH + tree (one aggregated ``scan/scan.md`` on run 1 -> per-image ``page_0001_png/`` + trees on run 2: non-idempotent, orphaned output, spurious index entries, broken + resume). Classification must be STABLE: a first run and a re-run see the same + single image-dir document. + """ + src = tmp_path / "scan" + src.mkdir() + for n in (1, 2): + Image.new("RGB", (60, 60), "white").save(src / f"page_{n:04d}.png") + + # Run 1 (default output root): classified as ONE image-dir document. + docs_run1 = discover_documents(src) + assert docs_run1 == [src] + + # Actually run it so the default ocr/ subtree is created inside the scanned dir. + outcome1 = process(src, FakeBackend(), dpi=120) + assert outcome1.completed == 1 + assert (src / "ocr").is_dir() # default root nests inside the scanned tree + + # Re-run classification: the engine's own ocr/ subtree must NOT flip the + # decision. Still exactly ONE image-dir document, identical to run 1. + docs_run2 = discover_documents(src) + assert docs_run2 == docs_run1 == [src] + + # And a real re-run stays idempotent: still one completed document, and NO + # per-image batch trees were emitted under the output root. + outcome2 = process(src, FakeBackend(), dpi=120) + assert outcome2.completed == 1 + out_root = src / "ocr" + assert (out_root / "scan" / "scan.md").exists() + assert not (out_root / "page_0001_png").exists() + assert not (out_root / "page_0002_png").exists() + + +def test_unreadable_file_recorded_failed_batch_continues(tmp_path): + """SYS-02 fix: an unreadable input fails ITSELF; the batch keeps going. + + A file that becomes unreadable between discovery and processing must be + recorded status=failed (via the safe checksum + per-doc catch-all), not raise + an OSError that aborts the whole batch and drops the good files. + """ + import os + import stat + + root = tmp_path / "papers" + root.mkdir() + _make_pdf(root / "a_good.pdf", pages=1) + bad = root / "b_bad.pdf" + _make_pdf(bad, pages=1) + _make_pdf(root / "c_good.pdf", pages=1) + out = tmp_path / "out" + + os.chmod(bad, 0) + try: + # Capture whether chmod(0) actually denied reads in THIS run (it does not + # when running as root, e.g. some CI), before permissions are restored. + bad_was_unreadable = not os.access(bad, os.R_OK) + outcome = process(root, FakeBackend(), dpi=120, output_dir=out) + finally: + os.chmod(bad, stat.S_IRUSR | stat.S_IWUSR) + + # The whole batch did NOT abort: two good files completed, the bad one failed. + assert outcome.completed == 2 + assert outcome.failed == 1 + assert outcome.exit_code != 0 + + # The failed doc has a durable status=failed record carrying a VALID ``sha256:`` + # checksum (v0.1.3): the conformance harness rejects a non-``sha256:`` checksum + # even on a failure record, so the unreadable-input fallback must be the + # ``sha256:`` UNREADABLE_CHECKSUM sentinel, never None/"". + bad_meta = json.loads((out / "b_bad" / "metadata.json").read_text()) + assert bad_meta["status"] == "failed" + assert bad_meta["checksum"].startswith("sha256:") + # When chmod(0) actually denied reads (i.e. not running as root, as in normal + # local runs) the digest is unobtainable, so the record carries the sentinel. + if bad_was_unreadable: + assert bad_meta["checksum"] == UNREADABLE_CHECKSUM + + +def test_rerun_under_different_raw_flag_reprocesses(tmp_path): + """The run fingerprint extra invalidates the cache when --raw changes. + + --raw genuinely changes page text (it skips clean_ocr_output), so a re-run with + a different --raw must reprocess, not silently reuse the prior cleaned output. + """ + pdf = tmp_path / "doc.pdf" + _make_pdf(pdf, pages=1) + out = tmp_path / "out" + + be1 = FakeBackend(text="output") + process(pdf, be1, dpi=120, raw=False, output_dir=out) + assert be1.calls == 1 + + # Same input/task/prompt, different --raw -> fingerprint extra differs -> reprocess. + be2 = FakeBackend(text="output") + process(pdf, be2, dpi=120, raw=True, output_dir=out) + assert be2.calls == 1 + + +def test_rerun_same_prompt_different_task_not_reprocessed(tmp_path): + """LOW fix: --task does NOT over-invalidate the cache when --prompt is set. + + The backends ignore --task when a custom --prompt is given (they only call + get_prompt(task) when prompt is None), so two runs with the SAME prompt but a + DIFFERENT task produce identical output and must be cache-skipped, not + needlessly reprocessed. task is dropped from the fingerprint when prompt is set. + """ + pdf = tmp_path / "doc.pdf" + _make_pdf(pdf, pages=1) + out = tmp_path / "out" + + be1 = FakeBackend(text="output") + process(pdf, be1, dpi=120, prompt="my custom prompt", task="convert", output_dir=out) + assert be1.calls == 1 + + # Same prompt, different task -> same fingerprint -> skipped (no reprocess). + be2 = FakeBackend(text="output") + process(pdf, be2, dpi=120, prompt="my custom prompt", task="ocr", output_dir=out) + assert be2.calls == 0 diff --git a/tests/test_retry.py b/tests/test_retry.py index 34ed57d..180b891 100644 --- a/tests/test_retry.py +++ b/tests/test_retry.py @@ -5,7 +5,7 @@ import pytest import requests -from deepseek_ocr.backends.base import Backend, TransientError +from deepseek_ocr.backends.base import TransientError from deepseek_ocr.backends.ollama import OllamaBackend @@ -33,6 +33,7 @@ def _make_backend(self, max_retries: int = 3, retry_delay: float = 0.0) -> Ollam backend.max_dimension = None backend.max_retries = max_retries backend.retry_delay = retry_delay + backend.max_tokens = 8192 backend.model = True backend.ollama_url = "http://localhost:11434" return backend @@ -118,6 +119,7 @@ def _make_backend(self) -> OllamaBackend: backend.max_dimension = None backend.max_retries = 2 backend.retry_delay = 0.0 + backend.max_tokens = 8192 backend.model = True backend.ollama_url = "http://localhost:11434" return backend diff --git a/tests/test_utils.py b/tests/test_utils.py index a2bdca4..37ac2dd 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,14 +1,13 @@ """Tests for utility functions.""" -import pytest from pathlib import Path + from deepseek_ocr.utils import ( - is_supported_file, + IMAGE_EXTENSIONS, is_image_file, is_pdf_file, + is_supported_file, sanitize_filename, - IMAGE_EXTENSIONS, - PDF_EXTENSION, ) @@ -49,21 +48,21 @@ class TestFilenameSanitization: def test_sanitize_removes_invalid_chars(self) -> None: """Test that invalid characters are replaced.""" - assert sanitize_filename('test.pdf') == 'test_file_.pdf' - assert sanitize_filename('test/file\\name.pdf') == 'test_file_name.pdf' - assert sanitize_filename('test:file|name.pdf') == 'test_file_name.pdf' + assert sanitize_filename("test.pdf") == "test_file_.pdf" + assert sanitize_filename("test/file\\name.pdf") == "test_file_name.pdf" + assert sanitize_filename("test:file|name.pdf") == "test_file_name.pdf" def test_sanitize_handles_empty(self) -> None: """Test that empty filenames get a default.""" - assert sanitize_filename('') == 'untitled' - assert sanitize_filename(' ') == 'untitled' + assert sanitize_filename("") == "untitled" + assert sanitize_filename(" ") == "untitled" def test_sanitize_strips_spaces_and_dots(self) -> None: """Test that leading/trailing spaces and dots are removed.""" - assert sanitize_filename(' test.pdf ') == 'test.pdf' - assert sanitize_filename('.test.pdf.') == 'test.pdf' + assert sanitize_filename(" test.pdf ") == "test.pdf" + assert sanitize_filename(".test.pdf.") == "test.pdf" def test_sanitize_preserves_valid_names(self) -> None: """Test that valid filenames are preserved.""" - assert sanitize_filename('valid_filename.pdf') == 'valid_filename.pdf' - assert sanitize_filename('test-file-123.png') == 'test-file-123.png' + assert sanitize_filename("valid_filename.pdf") == "valid_filename.pdf" + assert sanitize_filename("test-file-123.png") == "test-file-123.png"