scripts/eval_metrics.py

#!/usr/bin/env python3
"""
eval_metrics.py — Compute PSNR, SSIM, LPIPS, FID for Generated CXR Images
==============================================================================

This script computes four standard metrics for evaluating generated medical images:
  • SSIM  — Structural Similarity (scikit-image)
  • PSNR  — Peak Signal-to-Noise Ratio (scikit-image)
  • LPIPS — Learned Perceptual Image Patch Similarity (lpips package)
  • FID   — Fréchet Inception Distance (torchmetrics)

All metrics are computed on paired (generated, ground truth) images from corresponding
folder structures.

Usage:
  python eval_metrics.py \\
    --generated_dir /path/to/generated/images \\
    --ground_truth_dir /path/to/ground_truth/images \\
    --output_dir /path/to/results

  # Example:
  python eval_metrics.py \\
    --generated_dir ./outputs/eval_eval_2026.04.29_03.59.44 \\
    --ground_truth_dir ./dataset/cxr_radiomics/test_images \\
    --output_dir ./results/metrics
"""

import os
import sys
import gc
import json
import logging
import argparse
from typing import List, Dict, Optional, Tuple
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from tqdm import tqdm


# ============================================================
# Setup Logging
# ============================================================

def setup_logging(output_dir: str):
    """Configure logging to both console and file."""
    os.makedirs(output_dir, exist_ok=True)
    log_file = os.path.join(output_dir, "metrics_eval.log")
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        handlers=[
            logging.FileHandler(log_file, mode="a"),
            logging.StreamHandler(sys.stdout),
        ],
        force=True,
    )
    return logging.getLogger(__name__)


# ============================================================
# Image Discovery & Pairing
# ============================================================

def find_image_pairs(
    generated_dir: str,
    ground_truth_dir: str,
) -> List[Tuple[str, str]]:
    """
    Scan both directories and match images by relative path (patient_id/image_name.png).
    
    Returns:
        List of (generated_path, ground_truth_path) tuples for paired images.
    """
    valid_pairs: List[Tuple[str, str]] = []
    
    # Find all generated images
    gen_images = {}
    for root, dirs, files in os.walk(generated_dir):
        for fname in files:
            if fname.lower().endswith(('.png', '.jpg', '.jpeg')):
                full_path = os.path.join(root, fname)
                relative_path = os.path.relpath(full_path, generated_dir)
                gen_images[relative_path] = full_path
    
    # Find corresponding GT images
    for relative_path, gen_path in gen_images.items():
        gt_path = os.path.join(ground_truth_dir, relative_path)
        if os.path.exists(gt_path):
            valid_pairs.append((gen_path, gt_path))
    
    return valid_pairs


# ============================================================
# Metrics Computation
# ============================================================

def compute_metrics(
    generated_dir: str,
    ground_truth_dir: str,
    output_dir: str,
    max_samples: Optional[int] = None,
    device: torch.device = None,
):
    """
    Compute SSIM, PSNR, LPIPS, FID for all paired images.
    
    Args:
        generated_dir: Path to generated images (organized as patient_id/image_name.png)
        ground_truth_dir: Path to ground truth images (same structure)
        output_dir: Where to save metrics report and logs
        max_samples: Limit evaluation to first N image pairs (for testing)
        device: PyTorch device (cuda:0 or cpu)
    """
    if device is None:
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    logger = setup_logging(output_dir)
    
    # ---- Import metric libraries ----
    try:
        from skimage.metrics import structural_similarity as skimage_ssim
        from skimage.metrics import peak_signal_noise_ratio as skimage_psnr
    except ImportError:
        logger.error("ERROR: scikit-image is required. pip install scikit-image")
        return None
    
    try:
        import lpips as lpips_pkg
    except ImportError:
        logger.error("ERROR: lpips is required. pip install lpips")
        return None
    
    try:
        from torchmetrics.image.fid import FrechetInceptionDistance
    except ImportError:
        logger.error("ERROR: torchmetrics[image] is required. pip install 'torchmetrics[image]'")
        return None
    
    logger.info(f"\n{'='*70}")
    logger.info(f"  Metrics Evaluation — device={device}")
    logger.info(f"{'='*70}")
    logger.info(f"  Generated images : {generated_dir}")
    logger.info(f"  Ground truth     : {ground_truth_dir}")
    logger.info(f"  Results output   : {output_dir}")
    
    # ---- Find image pairs ----
    valid_pairs = find_image_pairs(generated_dir, ground_truth_dir)
    logger.info(f"\n  Found {len(valid_pairs)} valid image pairs")
    
    if len(valid_pairs) == 0:
        logger.error("  No valid pairs found. Check directory structure.")
        return None
    
    if max_samples is not None and len(valid_pairs) > max_samples:
        valid_pairs = valid_pairs[:max_samples]
        logger.info(f"  Limited to first {max_samples} pairs (--max_samples)")
    
    # ---- Initialize metrics ----
    logger.info("\n  Initializing metric functions...")
    lpips_fn = lpips_pkg.LPIPS(net="alex").to(device).eval()
    fid_metric = FrechetInceptionDistance(feature=2048).to(device)
    
    ssim_scores: List[float] = []
    psnr_scores: List[float] = []
    lpips_scores: List[float] = []
    
    # Per-patient tracking
    patient_metrics: Dict[str, Dict[str, List[float]]] = {}
    
    # Batched evaluation buffers
    EVAL_BATCH = 32
    gt_fid_buf: List[torch.Tensor] = []
    gen_fid_buf: List[torch.Tensor] = []
    gt_lpips_buf: List[torch.Tensor] = []
    gen_lpips_buf: List[torch.Tensor] = []
    
    def _flush_fid():
        """Push accumulated buffers into torchmetrics FID state."""
        nonlocal gt_fid_buf, gen_fid_buf
        if gt_fid_buf:
            fid_metric.update(torch.stack(gt_fid_buf).to(device), real=True)
            fid_metric.update(torch.stack(gen_fid_buf).to(device), real=False)
            gt_fid_buf, gen_fid_buf = [], []
    
    def _flush_lpips():
        """Compute LPIPS for accumulated buffer."""
        nonlocal gt_lpips_buf, gen_lpips_buf
        if gt_lpips_buf:
            with torch.no_grad():
                gt_batch = torch.stack(gt_lpips_buf).to(device)
                gen_batch = torch.stack(gen_lpips_buf).to(device)
                d = lpips_fn(gen_batch, gt_batch)  # [B, 1, 1, 1]
                lpips_scores.extend(d.view(-1).cpu().tolist())
            gt_lpips_buf, gen_lpips_buf = [], []
    
    # ---- Main evaluation loop ----
    logger.info(f"\n  Starting evaluation on {len(valid_pairs)} image pairs...\n")
    error_count = 0
    
    for idx, (gen_path, gt_path) in enumerate(tqdm(valid_pairs, desc="Evaluating")):
        try:
            # Load images
            gen_pil = Image.open(gen_path).convert("RGB")
            gt_pil = Image.open(gt_path).convert("RGB")  # ensure 3-channel
            
            # ---- SSIM & PSNR (grayscale, uint8, data_range=255) ----
            gen_gray = np.array(gen_pil.convert("L"))      # [H, W] uint8
            gt_gray = np.array(gt_pil.convert("L"))
            
            ssim_val = skimage_ssim(gt_gray, gen_gray, data_range=255)
            psnr_val = skimage_psnr(gt_gray, gen_gray, data_range=255)
            ssim_scores.append(ssim_val)
            psnr_scores.append(psnr_val)
            
            # ---- Patient-level tracking ----
            # Extract patient ID from path (parent directory name)
            patient_id = os.path.basename(os.path.dirname(gen_path))
            if patient_id not in patient_metrics:
                patient_metrics[patient_id] = {"ssim": [], "psnr": [], "lpips": []}
            patient_metrics[patient_id]["ssim"].append(ssim_val)
            patient_metrics[patient_id]["psnr"].append(psnr_val)
            
            # ---- Prepare tensors for batched LPIPS/FID ----
            gen_np = np.array(gen_pil)  # [H, W, 3] uint8
            gt_np = np.array(gt_pil)
            
            # LPIPS: float32, [-1, 1], [3, H, W]
            gen_lpips_buf.append(
                torch.from_numpy(gen_np).permute(2, 0, 1).float() / 127.5 - 1.0
            )
            gt_lpips_buf.append(
                torch.from_numpy(gt_np).permute(2, 0, 1).float() / 127.5 - 1.0
            )
            
            # FID: uint8, [0, 255], [3, H, W]
            gen_fid_buf.append(torch.from_numpy(gen_np).permute(2, 0, 1))
            gt_fid_buf.append(torch.from_numpy(gt_np).permute(2, 0, 1))
            
            # Flush buffers when full
            if len(gt_fid_buf) >= EVAL_BATCH:
                _flush_lpips()
                _flush_fid()
        
        except Exception as e:
            error_count += 1
            logger.warning(f"  Error processing {gen_path}: {e}")
            continue
    
    # ---- Flush remaining buffers ----
    _flush_lpips()
    _flush_fid()
    
    # ---- Compute final aggregate metrics ----
    fid_score = float(fid_metric.compute().item()) if len(valid_pairs) > 0 else float("nan")
    
    results = {
        "SSIM": float(np.mean(ssim_scores)) if ssim_scores else 0.0,
        "PSNR": float(np.mean(psnr_scores)) if psnr_scores else 0.0,
        "LPIPS": float(np.mean(lpips_scores)) if lpips_scores else 0.0,
        "FID": fid_score,
        "num_evaluated": len(ssim_scores),
        "num_total": len(valid_pairs),
        "num_errors": error_count,
    }
    
    # ---- Per-patient summary ----
    per_patient = {}
    for pid, m in sorted(patient_metrics.items()):
        per_patient[pid] = {
            "SSIM": float(np.mean(m["ssim"])) if m["ssim"] else 0.0,
            "PSNR": float(np.mean(m["psnr"])) if m["psnr"] else 0.0,
            "LPIPS": float(np.mean(m["lpips"])) if m["lpips"] else 0.0,
            "count": len(m["ssim"]),
        }
    results["per_patient"] = per_patient
    
    # ---- Print and save report ----
    logger.info(f"\n{'='*70}")
    logger.info(f"  EVALUATION RESULTS")
    logger.info(f"{'='*70}")
    logger.info(f"  SSIM  ↑ : {results['SSIM']:.4f}")
    logger.info(f"  PSNR  ↑ : {results['PSNR']:.2f} dB")
    logger.info(f"  LPIPS ↓ : {results['LPIPS']:.4f}")
    logger.info(f"  FID   ↓ : {results['FID']:.2f}")
    logger.info(f"  Evaluated : {results['num_evaluated']} / {results['num_total']}")
    if error_count > 0:
        logger.info(f"  Errors : {error_count}")
    logger.info(f"{'='*70}")
    
    logger.info("\n  Per-patient breakdown (top 10):")
    for i, (pid, pm) in enumerate(sorted(per_patient.items())[:10]):
        logger.info(
            f"    {pid:20s}  SSIM={pm['SSIM']:.4f}  PSNR={pm['PSNR']:.2f}  "
            f"LPIPS={pm['LPIPS']:.4f}  (n={pm['count']})"
        )
    if len(per_patient) > 10:
        logger.info(f"    ... and {len(per_patient) - 10} more patients")
    
    # Save to disk
    report_path = os.path.join(output_dir, "metrics_report.json")
    with open(report_path, "w") as f:
        json.dump(results, f, indent=2)
    logger.info(f"\n  Report saved to {report_path}")
    
    return results


# ============================================================
# Main Entry Point
# ============================================================

def parse_args():
    parser = argparse.ArgumentParser(
        description="Compute PSNR, SSIM, LPIPS, FID for generated images",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--generated_dir", type=str, required=True,
        help="Path to folder with generated images (organized as patient_id/image.png)",
    )
    parser.add_argument(
        "--ground_truth_dir", type=str, required=True,
        help="Path to folder with ground truth images (same folder structure as generated_dir)",
    )
    parser.add_argument(
        "--output_dir", type=str, required=True,
        help="Directory to save metrics report and logs",
    )
    parser.add_argument(
        "--max_samples", type=int, default=None,
        help="Limit evaluation to first N image pairs (for quick testing)",
    )
    parser.add_argument(
        "--device", type=str, default="cuda:0",
        help="PyTorch device to use (cuda:0, cuda:1, cpu, etc.)",
    )
    
    return parser.parse_args()


def main():
    args = parse_args()
    
    # Validate paths
    if not os.path.isdir(args.generated_dir):
        print(f"ERROR: Generated directory not found: {args.generated_dir}")
        sys.exit(1)
    if not os.path.isdir(args.ground_truth_dir):
        print(f"ERROR: Ground truth directory not found: {args.ground_truth_dir}")
        sys.exit(1)
    
    # Setup device
    device = torch.device(args.device)
    
    # Compute metrics
    results = compute_metrics(
        generated_dir=args.generated_dir,
        ground_truth_dir=args.ground_truth_dir,
        output_dir=args.output_dir,
        max_samples=args.max_samples,
        device=device,
    )
    
    if results is not None:
        print("\nMetrics computation completed successfully!")
    else:
        print("\nMetrics computation failed!")
        sys.exit(1)


if __name__ == "__main__":
    main()