Spaces:
Running
Running
| """ | |
| Sequential Local Benchmark Runner. | |
| Executes all 6 OCR models sequentially, captures accurate inference times, | |
| generates bounding box visualizations, and computes diff & consensus metrics. | |
| """ | |
| import os | |
| import time | |
| import logging | |
| from typing import Dict, Any, List, Optional | |
| from PIL import Image | |
| from core.models import OCRModelOutput, Region, ALL_REGION_TYPES | |
| from core.region_classifier import compute_region_summary | |
| from core.visualizer import save_visualization, image_to_base64_jpeg | |
| from adapters import AVAILABLE_MODELS, MODEL_CATALOG, get_adapter_by_name | |
| from utils.memory_manager import model_manager | |
| from utils.diff_engine import generate_consensus_matrix, compute_diff, compute_similarity_ratio | |
| logger = logging.getLogger("BenchmarkRunner") | |
| def run_sequential_local_benchmark( | |
| image: Image.Image, | |
| output_dir: str = "outputs/visualizations", | |
| document_name: str = "document.png" | |
| ) -> Dict[str, Any]: | |
| """ | |
| Executes all 6 models sequentially on local hardware, | |
| generating bounding box annotated images, performance metrics, and diff matrices. | |
| """ | |
| os.makedirs(output_dir, exist_ok=True) | |
| benchmark_start = time.perf_counter() | |
| results: List[Dict[str, Any]] = [] | |
| valid_texts: Dict[str, str] = {} | |
| img_w, img_h = image.size | |
| logger.info(f"Starting sequential local benchmark on '{document_name}' ({img_w}x{img_h} px)") | |
| for name in AVAILABLE_MODELS: | |
| logger.info(f"==> Benchmarking Model: {name}") | |
| adapter = get_adapter_by_name(name) | |
| # Execute adapter inference with timing & error capture | |
| res = adapter.process(image) | |
| is_success = (res.get("status") == "SUCCESS") | |
| regions_data: List[Region] = res.get("regions", []) | |
| # Convert dictionary regions to Region objects if needed | |
| parsed_regions: List[Region] = [] | |
| for r in regions_data: | |
| if isinstance(r, Region): | |
| parsed_regions.append(r) | |
| elif isinstance(r, dict): | |
| parsed_regions.append(Region( | |
| box=r.get("box", [0, 0, 0, 0]), | |
| text=r.get("text", ""), | |
| region_type=r.get("region_type", "Text"), | |
| confidence=r.get("confidence"), | |
| details=r.get("details") | |
| )) | |
| # Compute word count & region distribution | |
| raw_text = res.get("text") or res.get("markdown") or "" | |
| word_count = len(raw_text.split()) if raw_text else 0 | |
| region_counts = compute_region_summary(parsed_regions) if is_success else {} | |
| # Render Annotated Visualization Image | |
| annotated_path = os.path.join(output_dir, f"{name.lower().replace('-', '_')}_annotated.png") | |
| annotated_base64 = None | |
| if is_success and parsed_regions: | |
| try: | |
| save_visualization(image, parsed_regions, annotated_path) | |
| annotated_img = Image.open(annotated_path) | |
| annotated_base64 = image_to_base64_jpeg(annotated_img, quality=85) | |
| logger.info(f"Rendered {len(parsed_regions)} bounding boxes for {name} -> {annotated_path}") | |
| except Exception as ve: | |
| logger.warning(f"Visualization rendering error for {name}: {ve}") | |
| output_item = OCRModelOutput( | |
| model_name=name, | |
| model_id=MODEL_CATALOG[name]["hf_model_id"], | |
| status=res.get("status", "ERROR"), | |
| inference_time_seconds=res.get("inference_time_seconds"), | |
| inference_time_str=res.get("inference_time_str", "N/A"), | |
| text=res.get("text"), | |
| markdown=res.get("markdown"), | |
| json=res.get("json"), | |
| output_type=res.get("output_type", MODEL_CATALOG[name]["preferred_output"]), | |
| word_count=word_count, | |
| regions=parsed_regions, | |
| region_counts=region_counts, | |
| annotated_image_path=annotated_path if (is_success and parsed_regions) else None, | |
| annotated_image_base64=annotated_base64, | |
| error=res.get("error") | |
| ) | |
| results.append(output_item.to_dict()) | |
| if is_success and raw_text.strip(): | |
| valid_texts[name] = raw_text.strip() | |
| total_time = round(time.perf_counter() - benchmark_start, 2) | |
| logger.info(f"Completed local benchmark across 6 models in {total_time}s") | |
| # Pairwise Consensus Matrix on successful outputs only | |
| consensus_data = generate_consensus_matrix(valid_texts) | |
| return { | |
| "success": True, | |
| "document_name": document_name, | |
| "image_dimensions": {"width": img_w, "height": img_h}, | |
| "total_benchmark_time_seconds": total_time, | |
| "available_region_types": ALL_REGION_TYPES, | |
| "results": results, | |
| "consensus_matrix": consensus_data | |
| } | |