baseline for page evaluation and cleaning of metrics section

Author: christianabbet

main.py

@@ -302,7 +302,7 @@ def main(
 
     # Start MLFlow tracking
     if mlflow_tracking:
-        setup_mlflow(input_path, ground_truth_path, model_path, matching_params, classifier_name)
+        setup_mlflow(input_path, matching_params, ground_truth_path, model_path, classifier_name)
 
     # Process pages
     pdf_files = get_pdf_files(input_path)

src/evaluation.py

@@ -6,11 +6,14 @@
 
 import pandas as pd
 from dotenv import load_dotenv
+from extraction.evaluation.benchmark.metrics import OverallMetrics
+from extraction.evaluation.evaluation_dataclasses import Metrics
+from extraction.evaluation.utility import evaluate_single
 from pydantic import TypeAdapter
 
 from src.page_classes import PageClasses
 from src.page_structure import ProcessorDocumentEntities
-from src.schemas import DocumentGroundTruth
+from src.schemas import DocumentGroundTruth, DocumentPage
 
 load_dotenv()
 mlflow_tracking = os.getenv("MLFLOW_TRACKING").lower() == "true"
@@ -35,35 +38,37 @@ def load_ground_truth(ground_truth_path: Path) -> list[DocumentGroundTruth] | No
         return None
 
 
-def compute_confusion_stats(
-    predictions: list[DocumentGroundTruth], ground_truth: list[DocumentGroundTruth]
-) -> tuple[dict, int, int]:
-    """Computes confusion matrix entries, total pages and files processed for evaluating classification results."""
-    stats = {label: {"true_positives": 0, "false_negatives": 0, "false_positives": 0} for label in LABELS}
+def groundtruth_doc_to_pages(documents: list[DocumentGroundTruth]) -> dict[str, DocumentPage]:
+    """Convert list of documents to list of keyed pages.
 
-    pred_keys = set([pred.filename for pred in predictions])
-    gt_keys = set([pred.filename for pred in ground_truth])
+    Args:
+        documents (list[DocumentGroundTruth]): Documents with pages to flatten
 
-    # Evaluate on the intersection so we don't crash when pages are missing
-    common_keys = pred_keys & gt_keys
+    Returns:
+        dict[str, DocumentPage]: Keyed pages.
+    """
+    return {f"{doc.filename}-{page.page}": page for doc in documents for page in doc.pages}
 
-    missing_in_pred = gt_keys - pred_keys
-    missing_in_gt = pred_keys - gt_keys
-    if missing_in_pred:
-        logger.info(f"{len(missing_in_pred)} GT pages have no prediction (e.g., {next(iter(missing_in_pred))}).")
-    if missing_in_gt:
-        logger.info(f"{len(missing_in_gt)} predicted pages missing in GT (e.g., {next(iter(missing_in_gt))}).")
 
-    # TODO from here - finish evaluation
-    total_pages = len(common_keys)
-    total_files = len({fname for (fname, _page) in common_keys})
+# def compute_title_metric(title_gt: str, title_pred: str) -> bool | None:
+#     if title_gt is None:
+#         return None
+#     else:
+#         return title_gt.lower().strip() == title_pred.lower().strip()
+
+
+def compute_classification_stats(
+    predictions: dict[str, DocumentGroundTruth], ground_truth: dict[str, DocumentGroundTruth]
+) -> dict:
+    stats = {label: {"true_positives": 0, "false_negatives": 0, "false_positives": 0} for label in LABELS}
+    common_keys = predictions.keys() & ground_truth.keys()
 
     for key in common_keys:
         pred_page = predictions.get(key, {})
         gt_page = ground_truth.get(key, {})
         for label in LABELS:
-            pred = int(pred_page.get(label, 0))
-            gt = int(gt_page.get(label, 0))
+            pred = int(pred_page.classification.get(label, 0))
+            gt = int(gt_page.classification.get(label, 0))
 
             if gt == 1 and pred == 1:
                 stats[label]["true_positives"] += 1
@@ -72,24 +77,64 @@ def compute_confusion_stats(
             elif gt == 0 and pred == 1:
                 stats[label]["false_positives"] += 1
 
-    return stats, total_files, total_pages
+    return stats
 
 
-def save_confusion_stats(stats: dict, output_dir: Path) -> Path:
-    """Saves confusion matrix to output directory."""
-    csv_path = output_dir / "evaluation_metrics.csv"
+def compute_title_stats(
+    predictions: dict[str, DocumentGroundTruth], ground_truth: dict[str, DocumentGroundTruth]
+) -> dict:
+    stats = {"true_positives": 0, "false_negatives": 0, "false_positives": 0}
+    common_keys = predictions.keys() & ground_truth.keys()
 
+    for key in common_keys:
+        pred_title = predictions[key].title
+        gt_title = ground_truth[key].title
+        logger.info(f"{key}: {gt_title} == {pred_title}")
+        if pred_title and gt_title and pred_title == gt_title:
+            stats["true_positives"] += 1
+        else:
+            stats["false_negatives"] += 1
+            stats["false_positives"] += 1
+
+    return {"title": stats}
+
+
+def compute_stats(
+    predictions: list[DocumentGroundTruth], ground_truths: list[DocumentGroundTruth]
+) -> tuple[dict, dict]:
+    """Computes confusion matrix entries, total pages and files processed for evaluating classification results."""
+    pred_keyed = groundtruth_doc_to_pages(predictions)
+    gt_keyed = groundtruth_doc_to_pages(ground_truths)
+
+    # Evaluate on the intersection so we don't crash when pages are missing
+    pred_keys, gt_keys = set(pred_keyed.keys()), set(gt_keyed.keys())
+
+    missing_in_pred = gt_keys - pred_keys
+    missing_in_gt = pred_keys - gt_keys
+    if missing_in_pred:
+        logger.info(f"{len(missing_in_pred)} GT pages have no prediction (e.g., {next(iter(missing_in_pred))}).")
+    if missing_in_gt:
+        logger.info(f"{len(missing_in_gt)} predicted pages missing in GT (e.g., {next(iter(missing_in_gt))}).")
+
+    classification_stats = compute_classification_stats(pred_keyed, gt_keyed)
+    title_stats = compute_title_stats(pred_keyed, gt_keyed)
+
+    return classification_stats, title_stats
+
+
+def save_stats(stats_classification: list, csv_path: Path) -> Path:
+    """Saves confusion matrix to output directory."""
     with open(csv_path, "w", newline="") as f:
         writer = csv.writer(f)
         writer.writerow(
             [
-                "Class",
+                "Label",
                 "True_Positives",
                 "False_Negatives",
                 "False_Positives",
             ]
         )
-        for label, s in stats.items():
+        for label, s in stats_classification.items():
             writer.writerow(
                 [
                     label,
@@ -101,15 +146,27 @@ def save_confusion_stats(stats: dict, output_dir: Path) -> Path:
     return csv_path
 
 
-def log_metrics_to_mlflow(stats: dict, total_files: int, total_pages: int) -> None:
+def log_metrics_to_mlflow(stats_classification: dict, stats_title: dict) -> None:
     """Calculates and logs F1, precision and recall to MLflow."""
     if not mlflow_tracking:
         return None
 
+    # Log metrics for title extraction
+    tp, fp, fn = [stats_title["title"][label] for label in ["true_positives", "false_negatives", "false_positives"]]
+    precision = tp / (tp + fp) if (tp + fp) else 0.0
+    recall = tp / (tp + fn) if (tp + fn) else 0.0
+    f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) else 0.0
+    mlflow.log_metric("title/f1", f1)
+    mlflow.log_metric("title/precision", precision)
+    mlflow.log_metric("title/recall", recall)
+
+    logger.info(f"Title: F1={f1:.2%}, Precision={precision:.2%}, Recall={recall:.2%}")
+
+    # Log metrics for classification output
     precisions = []
     recalls = []
     f1_scores = []
-    for label, s in stats.items():
+    for label, s in stats_classification.items():
         tp, fn, fp = s["true_positives"], s["false_negatives"], s["false_positives"]
         precision = tp / (tp + fp) if (tp + fp) else 0.0
         recall = tp / (tp + fn) if (tp + fn) else 0.0
@@ -119,121 +176,39 @@ def log_metrics_to_mlflow(stats: dict, total_files: int, total_pages: int) -> No
         recalls.append(recall)
         f1_scores.append(f1)
 
-        mlflow.log_metric(f"F1 {label}", f1)
-        mlflow.log_metric(f"{label}_precision", precision)
-        mlflow.log_metric(f"{label}_recall", recall)
+        mlflow.log_metric(f"classification/{label}_f1", f1)
+        mlflow.log_metric(f"classification/{label}_precision", precision)
+        mlflow.log_metric(f"classification/{label}_recall", recall)
 
         logger.info(f"{label}: F1={f1:.2%}, Precision={precision:.2%}, Recall={recall:.2%}")
 
     macro_precision = sum(precisions) / len(precisions) if precisions else 0.0
     macro_recall = sum(recalls) / len(recalls) if recalls else 0.0
     macro_f1 = sum(f1_scores) / len(f1_scores) if f1_scores else 0.0
 
-    mlflow.log_metric("Macro Avg Precision", macro_precision)
-    mlflow.log_metric("Macro Avg Recall", macro_recall)
-    mlflow.log_metric("Macro Avg F1", macro_f1)
-
-    logger.info(f"Macro Avg: F1={macro_f1:.2%}, Precision={macro_precision:.2%}, Recall={macro_recall:.2%}")
-
-    mlflow.log_metric("total_pages", total_pages)
-    mlflow.log_metric("total_files", total_files)
-
-
-def create_page_comparison(pred_dict: dict, gt_dict: dict, output_dir: Path) -> pd.DataFrame:
-    """Create a per-page comparison CSV/DF for pages present in both predictions and ground truth (intersection)."""
-    output_dir.mkdir(parents=True, exist_ok=True)
-    report_path = output_dir / "per_page_comparison.csv"
-
-    columns = (
-        ["Filename", "Page"]
-        + [f"{label}_pred" for label in LABELS]
-        + [f"{label}_gt" for label in LABELS]
-        + [f"{label}_match" for label in LABELS]
-        + ["All_labels_match", "Status"]
-    )
-
-    pred_keys = set(pred_dict.keys())
-    gt_keys = set(gt_dict.keys())
-
-    # Only evaluate files/pages that are in predictions.
-    keys = pred_keys & gt_keys
-
-    rows = []
-    for filename, page_num in sorted(keys, key=lambda k: (k[0], k[1])):
-        pred_page = pred_dict[(filename, page_num)]
-        gt_page = gt_dict[(filename, page_num)]
-
-        preds = [int(pred_page.get(label, 0)) for label in LABELS]
-        gts = [int(gt_page.get(label, 0)) for label in LABELS]
-        matches = [int(p == g) for p, g in zip(preds, gts, strict=True)]
-        all_match = int(all(matches))
-
-        # Only keep misclassifications
-        if not all_match:
-            status = "mismatch"
-            row = [filename, page_num] + preds + gts + matches + [all_match, status]
-            rows.append(row)
-
-    df = pd.DataFrame(rows, columns=columns)
-
-    # Write to csv file
-    with open(report_path, "w", newline="", encoding="utf-8") as f:
-        writer = csv.writer(f)
-        writer.writerow(columns)
-        writer.writerows(rows)
-
-    if mlflow_tracking:
-        mlflow.log_artifact(str(report_path))
-    logger.info(f"Logged misclassifications to {report_path}")
-
-    return df
-
-
-def save_misclassifications(df: pd.DataFrame, output_dir: Path) -> None:
-    """Save misclassified pages and per-class CSVs."""
-
-    def get_active_labels(row, suffix):
-        return [label for label in LABELS if row[f"{label}_{suffix}"] == 1]
+    mlflow.log_metric("classification/macro_precision", macro_precision)
+    mlflow.log_metric("classification/macro_recall", macro_recall)
+    mlflow.log_metric("classification/marco_f1", macro_f1)
 
-    df["Predicted_labels"] = df.apply(lambda row: get_active_labels(row, suffix="pred"), axis=1)
-    df["Ground_truth_labels"] = df.apply(lambda row: get_active_labels(row, suffix="gt"), axis=1)
-
-    misclassified = df[df["All_labels_match"] == 0][["Filename", "Page", "Ground_truth_labels", "Predicted_labels"]]
-
-    mis_path = output_dir / "misclassifications.csv"
-    misclassified.to_csv(mis_path, index=False)
-    if mlflow_tracking:
-        mlflow.log_artifact(str(mis_path))
-
-    for true_class in LABELS:
-        class_mis = misclassified[
-            misclassified["Ground_truth_labels"].apply(lambda labels, cls=true_class: cls in labels)
-        ]
-        if not class_mis.empty:
-            path = output_dir / f"misclassified_{true_class}.csv"
-            class_mis.to_csv(path, index=False)
-            if mlflow_tracking:
-                mlflow.log_artifact(str(path))
+    logger.info(f"Classification Macro: F1={macro_f1:.2%}, Precision={macro_precision:.2%}, Recall={macro_recall:.2%}")
 
 
 def evaluate_results(
     predictions: list[ProcessorDocumentEntities], ground_truth_path: Path, output_dir: Path = Path("evaluation")
-) -> dict | None:
+) -> tuple[Path, Path]:
     """Evaluate classification predictions against ground truth."""
     output_dir.mkdir(parents=True, exist_ok=True)
 
     gt_list = load_ground_truth(ground_truth_path)
     pred_list = [pred.to_ground_truth() for pred in predictions]
 
-    stats, total_files, total_pages = compute_confusion_stats(pred_list, gt_list)
-    stats_path = save_confusion_stats(stats, output_dir)
+    stats_classification, stats_title = compute_stats(pred_list, gt_list)
+    stats_classification_path = save_stats(stats_classification, output_dir / "evaluation_metrics_classification.csv")
+    stats_title_path = save_stats(stats_title, output_dir / "evaluation_metrics_title.csv")
 
     if mlflow_tracking:
-        log_metrics_to_mlflow(stats, total_files, total_pages)
-        mlflow.log_artifact(str(stats_path))
-
-    # TODO
-    # comparison_data = create_page_comparison(pred_dict, gt_dict, output_dir)
-    # save_misclassifications(comparison_data, output_dir)
+        log_metrics_to_mlflow(stats_classification, stats_title)
+        mlflow.log_artifact(str(stats_classification_path))
+        mlflow.log_artifact(str(stats_title_path))
 
-    return stats
+    return stats_classification, stats_title_path