feat: integrate OCR visualization (#121)

Signed-off-by: samiullahchattha <[email protected]>
Co-authored-by: samiullahchattha <[email protected]>

Author: samiuc

docling_eval/cli/main.py

@@ -65,6 +65,7 @@
 from docling_eval.evaluators.ocr_evaluator import (
     OcrDatasetEvaluationResult,
     OCREvaluator,
+    OCRVisualizer,
 )
 from docling_eval.evaluators.readingorder_evaluator import (
     DatasetReadingOrderEvaluation,
@@ -820,8 +821,18 @@ def visualize(
                 fd.write(f"F1 Score: {ocr_evaluation.f1_score:.2f}\n")
                 fd.write(f"Recall: {ocr_evaluation.recall:.2f}\n")
                 fd.write(f"Precision: {ocr_evaluation.precision:.2f}\n")
+
+            _log.info(f"OCR evaluation stats saved to {log_filename}")
+
+            ocr_visualizer = OCRVisualizer()
+            ocr_visualizer(
+                dataset_path=idir,
+                ocr_evaluation_report_path=metrics_filename,
+                output_directory=odir,
+                data_split_name=split,
+            )
         except Exception as e:
-            _log.error(f"Error processing markdown text evaluation: {str(e)}")
+            _log.error(f"Error processing OCR evaluation: {str(e)}")
 
     else:
         _log.error(f"Unsupported modality for visualization: {modality}")

docling_eval/evaluators/ocr_evaluator.py

@@ -248,21 +248,6 @@ def __call__(
         )
         visualizations_output_path.mkdir(parents=True, exist_ok=True)
 
-        document_evaluations_map: Dict[str, DocumentEvaluationEntry] = {}
-        if ocr_evaluation_report_path and ocr_evaluation_report_path.exists():
-            with open(ocr_evaluation_report_path, "r") as report_file:
-                report_content: Dict[str, Any] = json.load(report_file)
-                for eval_item_data in report_content.get("evaluations", []):
-                    try:
-                        doc_entry = DocumentEvaluationEntry.model_validate(
-                            eval_item_data
-                        )
-                        document_evaluations_map[doc_entry.doc_id] = doc_entry
-                    except Exception as e_parse:
-                        _log.warning(
-                            f"Failed to parse document evaluation item: {eval_item_data}. Error: {e_parse}"
-                        )
-
         path_to_parquet_files: str = str(dataset_path / data_split_name / "*.parquet")
         hf_dataset: Dataset = load_dataset(
             "parquet", data_files={data_split_name: path_to_parquet_files}
@@ -283,20 +268,6 @@ def __call__(
                     BenchMarkColumns.GROUNDTRUTH_PAGE_IMAGES
                 )
 
-                page_image_bytes_list: List[Dict[str, bytes]] = []
-                if isinstance(page_images_data, list) and page_images_data:
-                    if (
-                        isinstance(page_images_data[0], dict)
-                        and "bytes" in page_images_data[0]
-                    ):
-                        page_image_bytes_list = page_images_data
-
-                if (
-                    ocr_evaluation_report_path
-                    and doc_id_val not in document_evaluations_map
-                ):
-                    continue
-
                 ground_truth_segmented_pages: Dict[int, SegmentedPage] = {}
                 prediction_segmented_pages: Dict[int, SegmentedPage] = {}
 
@@ -316,16 +287,9 @@ def __call__(
                     if parsed_pred_pages:
                         prediction_segmented_pages = parsed_pred_pages
 
-                if not page_image_bytes_list:
-                    _log.warning(
-                        f"No page images found for document {doc_id_val}. Skipping visualization."
-                    )
-                    continue
-
-                image_raw_bytes: bytes = page_image_bytes_list[0]["bytes"]
-                base_image: Image.Image = Image.open(BytesIO(image_raw_bytes)).convert(
-                    "RGB"
-                )
+                base_image: Image.Image = page_images_data[0]
+                if base_image.mode != "RGB":
+                    base_image = base_image.convert("RGB")
 
                 comparison_image: Image.Image = self._render_ocr_comparison_on_image(
                     doc_id_val,