Spaces:
Running
Running
File size: 4,825 Bytes
b611f38 | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """
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
}
|