Spaces:
Sleeping
Sleeping
| import os | |
| import time | |
| from typing import List, Dict, Any | |
| import numpy as np | |
| import pandas as pd | |
| import mlflow | |
| from .ocr_space_engine import OCRSpaceEngine | |
| def compute_cer(reference: str, hypothesis: str) -> float: | |
| """Computes Character Error Rate (CER) via edit distance.""" | |
| r = reference.lower().strip() | |
| h = hypothesis.lower().strip() | |
| if not r: | |
| return 0.0 if not h else 1.0 | |
| n, m = len(r), len(h) | |
| dp = [[0] * (m + 1) for _ in range(n + 1)] | |
| for i in range(n + 1): | |
| dp[i][0] = i | |
| for j in range(m + 1): | |
| dp[0][j] = j | |
| for i in range(1, n + 1): | |
| for j in range(1, m + 1): | |
| if r[i - 1] == h[j - 1]: | |
| dp[i][j] = dp[i - 1][j - 1] | |
| else: | |
| dp[i][j] = 1 + min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1]) | |
| dist = dp[n][m] | |
| return float(dist / max(len(r), 1)) | |
| def run_ocr_comparison( | |
| sample_records: List[Any], | |
| cache_dir: str = "data/processed/ocr_cache", | |
| tracking_uri: str = "sqlite:///mlflow.db" | |
| ) -> pd.DataFrame: | |
| """ | |
| Runs comparison between OCR engines and logs results to MLflow under 'ocr_comparison'. | |
| """ | |
| mlflow.set_tracking_uri(tracking_uri) | |
| exp = mlflow.set_experiment("ocr_comparison") | |
| engines = { | |
| "ocr_space_engine2": OCRSpaceEngine(ocr_engine="2", cache_dir=cache_dir), | |
| "ocr_space_engine1": OCRSpaceEngine(ocr_engine="1", cache_dir=cache_dir) | |
| } | |
| comparison_rows = [] | |
| with mlflow.start_run(run_name="ocr_engine_benchmark", experiment_id=exp.experiment_id): | |
| for name, engine in engines.items(): | |
| runtimes = [] | |
| confidences = [] | |
| cers = [] | |
| for rec in sample_records: | |
| img_path = rec.image_path if hasattr(rec, "image_path") else rec.get("image_path") | |
| doc_id = rec.doc_id if hasattr(rec, "doc_id") else rec.get("doc_id") | |
| gt_words = rec.ocr_words if hasattr(rec, "ocr_words") else rec.get("ocr_words", []) | |
| gt_text = " ".join([w.text if hasattr(w, "text") else w.get("text", "") for w in gt_words]) | |
| t0 = time.time() | |
| words, _ = engine.extract_with_cache(img_path, doc_id=doc_id) | |
| dur = time.time() - t0 | |
| runtimes.append(dur) | |
| if words: | |
| confs = [w.conf for w in words] | |
| confidences.append(float(np.mean(confs))) | |
| pred_text = " ".join([w.text for w in words]) | |
| if gt_text: | |
| cers.append(compute_cer(gt_text, pred_text)) | |
| else: | |
| cers.append(0.0) | |
| else: | |
| confidences.append(0.0) | |
| cers.append(1.0) | |
| avg_runtime = float(np.mean(runtimes)) if runtimes else 0.0 | |
| avg_conf = float(np.mean(confidences)) if confidences else 0.0 | |
| avg_cer = float(np.mean(cers)) if cers else 0.0 | |
| mlflow.log_metric(f"{name}_avg_runtime_sec", avg_runtime) | |
| mlflow.log_metric(f"{name}_avg_confidence", avg_conf) | |
| mlflow.log_metric(f"{name}_cer", avg_cer) | |
| comparison_rows.append({ | |
| "engine": name, | |
| "avg_runtime_sec": round(avg_runtime, 4), | |
| "avg_confidence": round(avg_conf, 4), | |
| "cer": round(avg_cer, 4) | |
| }) | |
| df_res = pd.DataFrame(comparison_rows) | |
| csv_path = "reports/tables/ocr_engine_comparison.csv" | |
| os.makedirs(os.path.dirname(csv_path), exist_ok=True) | |
| df_res.to_csv(csv_path, index=False) | |
| mlflow.log_artifact(csv_path) | |
| return df_res | |
| if __name__ == "__main__": | |
| from ..data.schema import DocumentRecord, OCRWord | |
| dummy_records = [ | |
| DocumentRecord( | |
| doc_id="sample_0", | |
| source_dataset="test", | |
| doc_class="invoice", | |
| image_path="", | |
| ocr_words=[OCRWord(text="INVOICE TOTAL 100", bbox=[0, 0, 10, 10], conf=1.0)], | |
| fields={} | |
| ) | |
| ] | |
| df = run_ocr_comparison(dummy_records) | |
| print(df) | |