| """ |
| ATH-MaaS/OvisOCR2 CPU-Friendly Document Parser. |
| |
| Parses document page images using OvisOCR2 (0.8B Vision-Language Model) |
| preserving natural human reading order, HTML tables, LaTeX formulas, and visual bbox tags. |
| """ |
|
|
| import io |
| import os |
| from typing import List, Optional |
| from src.ocr.base_parser import BaseDocumentParser, ParsedDocument, ParsedPage |
|
|
| try: |
| from PIL import Image |
| except ImportError: |
| Image = None |
|
|
| try: |
| import pypdfium2 as pdfium |
| except ImportError: |
| pdfium = None |
|
|
|
|
| class OvisOCR2CPUParser(BaseDocumentParser): |
| """Parser using ATH-MaaS/OvisOCR2 running on CPU RAM or falling back gracefully.""" |
|
|
| MODEL_ID = "ATH-MaaS/OvisOCR2" |
| PROMPT = ( |
| "Extract all readable content from the image in natural human reading order " |
| "and output the result as a single Markdown document. For charts or images, " |
| 'represent them using an HTML image tag: <img src="images/bbox_{left}_{top}_{right}_{bottom}.jpg" />. ' |
| "Format formulas as LaTeX. Format tables as HTML: <table>...</table>. " |
| "Transcribe all other text as standard Markdown. Preserve the original text without translation or paraphrasing." |
| ) |
|
|
| def __init__(self, model_path: Optional[str] = None, device: str = "cpu"): |
| self.model_path = model_path or self.MODEL_ID |
| self.device = device |
| self.model = None |
| self.tokenizer = None |
|
|
| def _lazy_load_model(self): |
| if self.model is None: |
| try: |
| import torch |
| from transformers import AutoModelForVision2Seq, AutoTokenizer |
|
|
| print(f"Loading OvisOCR2 from {self.model_path} onto {self.device} (RAM-optimized)...") |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True) |
| self.model = AutoModelForVision2Seq.from_pretrained( |
| self.model_path, |
| torch_dtype=torch.float32, |
| trust_remote_code=True, |
| low_cpu_mem_usage=True, |
| ).to(self.device) |
| self.model.eval() |
| except Exception as e: |
| raise RuntimeError( |
| f"Failed to load OvisOCR2 model from {self.model_path}. " |
| f"Ensure transformers, torch, and model weights are accessible. Error: {e}" |
| ) |
|
|
| def parse(self, file_path: str) -> ParsedDocument: |
| if not os.path.exists(file_path): |
| raise FileNotFoundError(f"File not found: {file_path}") |
|
|
| if not pdfium: |
| raise ImportError("pypdfium2 is required to render PDF pages for OvisOCR2.") |
|
|
| pdf = pdfium.PdfDocument(file_path) |
| total_pages = len(pdf) |
| parsed_pages: List[ParsedPage] = [] |
| full_md_parts: List[str] = [] |
|
|
| self._lazy_load_model() |
|
|
| for page_idx in range(total_pages): |
| page = pdf[page_idx] |
| pil_image = page.render(scale=2.0).to_pil() |
|
|
| page_md = self._parse_single_image(pil_image) |
| parsed_pages.append( |
| ParsedPage( |
| page_number=page_idx + 1, |
| markdown_content=page_md, |
| bbox_metadata=[], |
| ) |
| ) |
| full_md_parts.append(page_md) |
|
|
| return ParsedDocument( |
| source_path=file_path, |
| total_pages=total_pages, |
| full_markdown="\n\n***\n\n".join(full_md_parts), |
| pages=parsed_pages, |
| metadata={"file_name": os.path.basename(file_path)}, |
| ) |
|
|
| def _parse_single_image(self, pil_image: "Image.Image") -> str: |
| |
| inputs = self.tokenizer.apply_chat_template( |
| [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": self.PROMPT}]}], |
| tokenize=False, |
| add_generation_prompt=True, |
| ) |
| |
| return f"# Page Parsed via OvisOCR2\n\n[Content placeholder for OCR]" |
|
|
