Add field-level confusion matrix evaluation script

evaluate_model.py

@@ -0,0 +1,264 @@
+"""
+Evaluate the fine-tuned Donut model and generate a Field-Level Confusion Matrix.
+Run this on the Workbench where the model and datasets are located.
+
+Usage:
+    python scripts/evaluate_model.py \
+        --model_path outputs/receipt_donut_gcp_enterprise/best_model \
+        --config configs/gcp_l4_enterprise.yaml \
+        --output_dir evaluation_results
+
+Outputs:
+    - evaluation_results/field_confusion_matrix.png
+    - evaluation_results/field_accuracy.json
+    - evaluation_results/error_analysis.html
+"""
+
+import os
+import sys
+import json
+import argparse
+import Levenshtein
+from pathlib import Path
+from collections import defaultdict
+
+import numpy as np
+import torch
+from PIL import Image
+import matplotlib.pyplot as plt
+from transformers import DonutProcessor, VisionEncoderDecoderModel
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from core.unified_dataset import UnifiedReceiptDataset
+
+
+FIELDS = ["merchant", "date", "subtotal", "tax", "total", "address"]
+
+
+def normalize_text(text):
+    """Lowercase and strip whitespace for fair comparison."""
+    if text is None:
+        return ""
+    return str(text).lower().strip().replace("$", "").replace(",", "")
+
+
+def categorize_match(gt, pred):
+    """
+    Categorize a single field prediction into:
+      - correct: exact match after normalization
+      - minor_typo: < 20% Levenshtein distance
+      - incorrect: everything else
+    """
+    gt_norm = normalize_text(gt)
+    pred_norm = normalize_text(pred)
+
+    if not gt_norm and not pred_norm:
+        return "correct"  # Both missing = agreement
+    if not gt_norm or not pred_norm:
+        return "incorrect"  # One missing, one present
+
+    if gt_norm == pred_norm:
+        return "correct"
+
+    dist = Levenshtein.distance(gt_norm, pred_norm)
+    max_len = max(len(gt_norm), len(pred_norm))
+    ratio = dist / max_len if max_len > 0 else 0
+
+    if ratio < 0.20:
+        return "minor_typo"
+    return "incorrect"
+
+
+def run_inference(model, processor, image_path, device):
+    """Run model inference on a single image and return parsed JSON dict."""
+    img = Image.open(image_path).convert("RGB")
+    pixel_values = processor(img, return_tensors="pt").pixel_values.to(device)
+    decoder_input_ids = torch.tensor([[model.config.decoder_start_token_id]]).to(device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            pixel_values,
+            decoder_input_ids=decoder_input_ids,
+            max_length=512,
+            pad_token_id=processor.tokenizer.pad_token_id,
+            eos_token_id=processor.tokenizer.eos_token_id,
+            use_cache=True,
+            bad_words_ids=[[processor.tokenizer.unk_token_id]],
+        )
+
+    seq = processor.tokenizer.batch_decode(outputs.sequences)[0]
+    seq = seq.replace(processor.tokenizer.eos_token, "").replace(
+        processor.tokenizer.pad_token, ""
+    )
+    seq = seq.replace(
+        processor.tokenizer.decode([model.config.decoder_start_token_id]), ""
+    ).strip()
+
+    try:
+        return json.loads(seq)
+    except json.JSONDecodeError:
+        return {}
+
+
+def evaluate(model, processor, dataset, device, max_samples=None):
+    """
+    Evaluate the model on a dataset and return per-field statistics.
+    """
+    counts = {field: {"correct": 0, "minor_typo": 0, "incorrect": 0} for field in FIELDS}
+    errors = []
+
+    n = min(len(dataset), max_samples) if max_samples else len(dataset)
+    print(f"Evaluating on {n} samples...")
+
+    for i in range(n):
+        sample = dataset[i]
+        image_path = sample["image_path"]
+        gt = sample["ground_truth"]
+
+        pred = run_inference(model, processor, image_path, device)
+
+        sample_error = {"image": image_path, "gt": gt, "pred": pred, "fields": {}}
+        all_correct = True
+
+        for field in FIELDS:
+            gt_val = gt.get(field, "")
+            pred_val = pred.get(field, "")
+            cat = categorize_match(gt_val, pred_val)
+            counts[field][cat] += 1
+            sample_error["fields"][field] = cat
+            if cat != "correct":
+                all_correct = False
+
+        if not all_correct:
+            errors.append(sample_error)
+
+        if (i + 1) % 50 == 0:
+            print(f"  Processed {i + 1}/{n}")
+
+    return counts, errors
+
+
+def plot_confusion_matrix(counts, output_dir):
+    """Generate a stacked bar chart confusion matrix per field."""
+    categories = ["correct", "minor_typo", "incorrect"]
+    colors = ["#4CAF50", "#FFC107", "#F44336"]
+
+    fig, ax = plt.subplots(figsize=(10, 6))
+    x = np.arange(len(FIELDS))
+    width = 0.25
+
+    for i, cat in enumerate(categories):
+        values = [counts[f][cat] for f in FIELDS]
+        ax.bar(x + i * width, values, width, label=cat.replace("_", " ").title(), color=colors[i])
+
+    ax.set_xlabel("Field")
+    ax.set_ylabel("Count")
+    ax.set_title("Field-Level Confusion Matrix (Validation/Test Set)")
+    ax.set_xticks(x + width)
+    ax.set_xticklabels(FIELDS, rotation=15, ha="right")
+    ax.legend()
+    ax.grid(axis="y", linestyle="--", alpha=0.5)
+    plt.tight_layout()
+
+    save_path = os.path.join(output_dir, "field_confusion_matrix.png")
+    plt.savefig(save_path, dpi=150)
+    print(f"Saved confusion matrix to {save_path}")
+    plt.close()
+
+
+def save_accuracy_json(counts, output_dir):
+    """Save numerical accuracy breakdown per field."""
+    results = {}
+    for field in FIELDS:
+        total = sum(counts[field].values())
+        results[field] = {
+            "correct_pct": round(counts[field]["correct"] / total * 100, 1),
+            "minor_typo_pct": round(counts[field]["minor_typo"] / total * 100, 1),
+            "incorrect_pct": round(counts[field]["incorrect"] / total * 100, 1),
+            "counts": counts[field],
+        }
+
+    save_path = os.path.join(output_dir, "field_accuracy.json")
+    with open(save_path, "w") as f:
+        json.dump(results, f, indent=2)
+    print(f"Saved accuracy JSON to {save_path}")
+
+
+def save_error_html(errors, output_dir, max_display=50):
+    """Generate an HTML file showing side-by-side GT vs Pred errors."""
+    html = ["<html><head><style>",
+            "body{font-family:sans-serif;margin:20px}",
+            "table{border-collapse:collapse;width:100%}",
+            "th,td{border:1px solid #ccc;padding:8px;text-align:left}",
+            "th{background:#f0f0f0}",
+            ".correct{color:green}.minor{color:orange}.incorrect{color:red}",
+            "</style></head><body>",
+            f"<h1>Error Analysis ({min(len(errors), max_display)} of {len(errors)} failures)</h1>",
+            "<table><tr><th>Image</th><th>Field</th><th>Ground Truth</th><th>Predicted</th><th>Status</th></tr>"]
+
+    for err in errors[:max_display]:
+        img_name = os.path.basename(err["image"])
+        for field in FIELDS:
+            status = err["fields"][field]
+            if status == "correct":
+                continue
+            css_class = "correct" if status == "correct" else ("minor" if status == "minor_typo" else "incorrect")
+            html.append(f"<tr><td>{img_name}</td><td>{field}</td>"
+                        f"<td>{err['gt'].get(field, 'N/A')}</td>"
+                        f"<td>{err['pred'].get(field, 'N/A')}</td>"
+                        f"<td class='{css_class}'>{status}</td></tr>")
+
+    html.append("</table></body></html>")
+
+    save_path = os.path.join(output_dir, "error_analysis.html")
+    with open(save_path, "w") as f:
+        f.write("\n".join(html))
+    print(f"Saved error analysis HTML to {save_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate Donut receipt model")
+    parser.add_argument("--model_path", required=True, help="Path to fine-tuned model")
+    parser.add_argument("--config", default="configs/gcp_l4_enterprise.yaml", help="Training config YAML")
+    parser.add_argument("--output_dir", default="evaluation_results", help="Where to save results")
+    parser.add_argument("--max_samples", type=int, default=None, help="Limit evaluation samples")
+    parser.add_argument("--split", default="test", choices=["train", "val", "test"], help="Which split to evaluate")
+    args = parser.parse_args()
+
+    import yaml
+    with open(args.config, "r") as f:
+        config = yaml.safe_load(f)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    print(f"Loading model from {args.model_path}...")
+    processor = DonutProcessor.from_pretrained(args.model_path)
+    model = VisionEncoderDecoderModel.from_pretrained(args.model_path)
+    model.to(device).eval()
+
+    print(f"Loading dataset split: {args.split}")
+    dataset = UnifiedReceiptDataset(
+        root=config["data"]["dataset_root"],
+        split=args.split,
+        processor=None,
+        include_datasets=config["data"].get("include_datasets"),
+    )
+
+    counts, errors = evaluate(model, processor, dataset, device, args.max_samples)
+    plot_confusion_matrix(counts, args.output_dir)
+    save_accuracy_json(counts, args.output_dir)
+    save_error_html(errors, args.output_dir)
+
+    print("\n=== Evaluation Complete ===")
+    for field in FIELDS:
+        total = sum(counts[field].values())
+        c = counts[field]["correct"]
+        m = counts[field]["minor_typo"]
+        i = counts[field]["incorrect"]
+        print(f"  {field:12s}: Correct={c}/{total} ({c/total*100:.1f}%) | "
+              f"Minor={m} | Incorrect={i}")
+
+
+if __name__ == "__main__":
+    main()