#!/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()