Spaces:
Sleeping
Sleeping
File size: 1,978 Bytes
c1b1f3a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | 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
@abstractmethod
def extract_words(self, image_path: str) -> List[OCRWord]:
pass
|