Aman
a32b7508c7
PaddleOCR with preprocessing, scispaCy NER, figure/table detection, citation extraction, and chunked Markdown output with frontmatter. Includes watch mode and notebook reprocessing.
148 lines
5.3 KiB
Python
148 lines
5.3 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from ocr_pipeline.config import settings
|
|
from ocr_pipeline.utils.logging import ProgressTracker, get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ProcessedImage:
|
|
path: Path
|
|
image_hash: str
|
|
ocr_result: Any
|
|
preprocessing_time: float
|
|
ocr_time: float
|
|
success: bool
|
|
error: str | None = None
|
|
|
|
|
|
def _process_single_image(args: tuple[Path, dict]) -> ProcessedImage:
|
|
"""Worker function for multiprocessing - must be at module level."""
|
|
path, config_dict = args
|
|
from ocr_pipeline.ocr.engine import ImagePreprocessor, OCREngine, compute_image_hash
|
|
|
|
preprocessor = ImagePreprocessor()
|
|
ocr_engine = OCREngine()
|
|
|
|
try:
|
|
image_hash = compute_image_hash(path)
|
|
prep_start = time.time()
|
|
image = preprocessor.preprocess(path)
|
|
preprocessing_time = time.time() - prep_start
|
|
|
|
ocr_start = time.time()
|
|
result = ocr_engine.process(image)
|
|
ocr_time = time.time() - ocr_start
|
|
|
|
return ProcessedImage(
|
|
path=path,
|
|
image_hash=image_hash,
|
|
ocr_result=result,
|
|
preprocessing_time=preprocessing_time,
|
|
ocr_time=ocr_time,
|
|
success=True,
|
|
)
|
|
except Exception as e:
|
|
return ProcessedImage(
|
|
path=path,
|
|
image_hash="",
|
|
ocr_result=None,
|
|
preprocessing_time=0,
|
|
ocr_time=0,
|
|
success=False,
|
|
error=str(e),
|
|
)
|
|
|
|
|
|
class ParallelProcessor:
|
|
def __init__(self, workers: int | None = None):
|
|
self.workers = workers or settings.processing.workers
|
|
self.batch_size = settings.processing.batch_size
|
|
self.retry_attempts = settings.processing.retry_attempts
|
|
self.timeout = settings.processing.timeout_per_image
|
|
|
|
def find_images(self) -> list[Path]:
|
|
all_images = []
|
|
for base_path_str in settings.input.paths:
|
|
base_path = Path(base_path_str).expanduser()
|
|
if not base_path.exists():
|
|
logger.warning("input_path_not_found", path=str(base_path))
|
|
continue
|
|
|
|
for pattern in settings.input.patterns:
|
|
if settings.input.recursive:
|
|
images = list(base_path.rglob(pattern))
|
|
else:
|
|
images = list(base_path.glob(pattern))
|
|
all_images.extend(images)
|
|
|
|
unique_images = []
|
|
seen = set()
|
|
for img in all_images:
|
|
if img.is_file() and img.suffix.lower() in (".png", ".jpg", ".jpeg", ".tiff", ".bmp"):
|
|
if img not in seen:
|
|
seen.add(img)
|
|
unique_images.append(img)
|
|
|
|
logger.info("found_images", count=len(unique_images))
|
|
return sorted(unique_images)
|
|
|
|
def process(self, image_paths: list[Path] | None = None) -> list[ProcessedImage]:
|
|
if image_paths is None:
|
|
image_paths = self.find_images()
|
|
|
|
if not image_paths:
|
|
logger.warning("no_images_to_process")
|
|
return []
|
|
|
|
results = []
|
|
failed = []
|
|
|
|
with ProgressTracker(len(image_paths), "OCR Processing") as tracker:
|
|
with ProcessPoolExecutor(max_workers=self.workers) as executor:
|
|
future_to_path = {
|
|
executor.submit(_process_single_image, (path, {})): path for path in image_paths
|
|
}
|
|
|
|
for future in as_completed(future_to_path):
|
|
path = future_to_path[future]
|
|
try:
|
|
result = future.result(timeout=self.timeout)
|
|
if not result.success and self.retry_attempts > 0:
|
|
for attempt in range(self.retry_attempts):
|
|
logger.warning("retrying", path=str(path), attempt=attempt + 1)
|
|
retry_result = _process_single_image((path, {}))
|
|
if retry_result.success:
|
|
result = retry_result
|
|
break
|
|
results.append(result)
|
|
tracker.update(success=result.success)
|
|
tracker.log_result(path, result.success, result.error)
|
|
except Exception as e:
|
|
logger.error("processing_failed", path=str(path), error=str(e))
|
|
results.append(
|
|
ProcessedImage(
|
|
path=path,
|
|
image_hash="",
|
|
ocr_result=None,
|
|
preprocessing_time=0,
|
|
ocr_time=0,
|
|
success=False,
|
|
error=str(e),
|
|
)
|
|
)
|
|
tracker.update(success=False)
|
|
tracker.log_result(path, False, str(e))
|
|
|
|
successful = [r for r in results if r.success]
|
|
failed = [r for r in results if not r.success]
|
|
logger.info("processing_complete", successful=len(successful), failed=len(failed))
|
|
return results
|