Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import time | |
| from abc import ABC, abstractmethod | |
| from typing import List, Dict, Any, Tuple | |
| from ..data.schema import OCRWord | |
| class BaseOCREngine(ABC): | |
| def __init__(self, engine_name: str, cache_dir: str = None): | |
| self.engine_name = engine_name | |
| self.cache_dir = os.path.join(cache_dir, engine_name) if cache_dir else None | |
| if self.cache_dir: | |
| os.makedirs(self.cache_dir, exist_ok=True) | |
| def get_cached_result(self, doc_id: str) -> List[OCRWord]: | |
| if not self.cache_dir or not doc_id: | |
| return None | |
| cache_path = os.path.join(self.cache_dir, f"{doc_id}.json") | |
| if os.path.exists(cache_path): | |
| with open(cache_path, "r", encoding="utf-8") as f: | |
| data = json.load(f) | |
| return [OCRWord(**w) for w in data.get("words", [])] | |
| return None | |
| def save_cached_result(self, doc_id: str, words: List[OCRWord], latency_sec: float = 0.0): | |
| if not self.cache_dir or not doc_id: | |
| return | |
| cache_path = os.path.join(self.cache_dir, f"{doc_id}.json") | |
| with open(cache_path, "w", encoding="utf-8") as f: | |
| json.dump({ | |
| "doc_id": doc_id, | |
| "engine": self.engine_name, | |
| "latency_sec": latency_sec, | |
| "words": [w.model_dump() for w in words] | |
| }, f, indent=2) | |
| def extract_with_cache(self, image_path: str, doc_id: str = None) -> Tuple[List[OCRWord], float]: | |
| """Checks cache first; if missing, runs OCR, caches and records latency.""" | |
| cached = self.get_cached_result(doc_id) | |
| if cached is not None: | |
| return cached, 0.0 | |
| t0 = time.time() | |
| words = self.extract_words(image_path) | |
| latency = time.time() - t0 | |
| self.save_cached_result(doc_id, words, latency) | |
| return words, latency | |
| def extract_words(self, image_path: str) -> List[OCRWord]: | |
| pass | |