import argparse import os from datetime import datetime from pathlib import Path import cv2 import numpy as np import torch from monai.metrics import DiceMetric, MeanIoU, HausdorffDistanceMetric from monai.transforms import ( Compose, LoadImaged, ScaleIntensityd, EnsureChannelFirstd, ToTensord, Resized, Lambdad ) from torch.utils.data import DataLoader from datasets.PolypDetectionDataset.PolypDetectionDataset import PolypDetectionDataset from lib.model.model import Wavelet_FFT_SwinUNETR def parse_args(): parser = argparse.ArgumentParser(description="Polyp Segmentation Model Evaluation Script") # ==================== Dataset-related Parameters ==================== parser.add_argument( "--dataset_name", type=str, required=True, help="Dataset name" ) parser.add_argument( "--data_root", type=str, default=r"./data/Polyp-Detection-Dataset", help="Root directory path of the dataset" ) parser.add_argument( "--num_workers", type=int, default=0, help="Number of worker processes for data loader" ) parser.add_argument( "--pin_memory", type=bool, default=True, help="Whether to enable pinned memory" ) parser.add_argument( "--target_spatial_size", type=tuple, default=(512, 512), help="Target spatial size, format like: '512,512' or '(512,512)'" ) parser.add_argument( "--batch_size", type=int, default=1, help="Batch size (smaller batch size recommended for evaluation)" ) # ==================== Model-related Parameters ==================== parser.add_argument( "--in_channels", type=int, default=3, help="Number of input image channels" ) parser.add_argument( "--out_channels", type=int, default=1, help="Number of output foreground channels" ) parser.add_argument( "--feature_size", type=int, default=48, help="Network feature dimension" ) parser.add_argument( "--spatial_dims", type=int, default=2, choices=[2, 3], help="Spatial dimension (2D or 3D)" ) parser.add_argument( "--use_wavelet", type=bool, default=True, help="Whether to enable wavelet enhancement module" ) parser.add_argument( "--wavelet_J", type=int, default=2, help="Wavelet decomposition levels" ) parser.add_argument( "--wavelet_wave", type=str, default="db4", help="Wavelet basis type" ) parser.add_argument( "--wavelet_reduction", type=int, default=16, help="Wavelet attention compression ratio" ) parser.add_argument( "--use_fft", type=bool, default=True, help="Whether to enable FFT enhancement module" ) parser.add_argument( "--use_v2", type=bool, default=True, help="Whether to enable Swin-UNETR v2 module" ) # ==================== Model Loading Parameters ==================== parser.add_argument( "--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Training device (cuda or cpu)" ) # ==================== Other Parameters ==================== parser.add_argument( "--save_results", type=bool, default=False, help="是否保存预测结果" ) parser.add_argument( "--dir_flag", type=str, default="_v4_minute", help="Prediction result save file suffix" ) parser.add_argument( "--results_dir", type=str, default="./evaluation_results", help="Directory for saving prediction results" ) parser.add_argument( "--outputs_dir", type=str, default="./outputs", help="是否保存预测结果" ) parser.add_argument( "--save_visualization", type=bool, default=True, help="Whether to save visualization results" ) parser.add_argument( "--vis_num_samples", type=int, default=1000, help="Number of samples to save for visualization" ) parser.add_argument( "--best_metric", type=str, default=False, help="Load best overall model, False means load best Dice model by default" ) return parser.parse_args() def create_val_transform(target_spatial_size=(512, 512)): """ Create validation set transformations Args: target_spatial_size: Target spatial size Returns: Compose: Validation transformation composition """ def convert_label_to_single_channel(label_tensor): """Convert RGB labels to single-channel binary mask""" single_channel = label_tensor[0:1, :, :] binary_label = (single_channel > 127).float() return binary_label val_transforms = Compose([ LoadImaged(keys=["image", "label"]), EnsureChannelFirstd(keys=["image", "label"]), Lambdad(keys=["label"], func=convert_label_to_single_channel), Resized(keys=["image", "label"], spatial_size=target_spatial_size, mode=("bilinear", "nearest")), ScaleIntensityd(keys=["image"]), ToTensord(keys=["image", "label"]), ]) return val_transforms def create_dataloader(args, dataset): """ Create data loader Args: args: Command line arguments dataset: Dataset Returns: DataLoader: Data loader """ loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=args.pin_memory, drop_last=False ) return loader def load_model(args, checkpoint_path): """ Load model Args: args: Command line arguments checkpoint_path: Checkpoint path Returns: model: Model with loaded weights """ print(f"\nLoading model: {checkpoint_path}") # Create model model = Wavelet_FFT_SwinUNETR( in_channels=args.in_channels, out_channels=args.out_channels, feature_size=args.feature_size, spatial_dims=args.spatial_dims, wavelet_enhancement=args.use_wavelet, wavelet_J=args.wavelet_J, wavelet_wave=args.wavelet_wave, wavelet_mode='symmetric', wavelet_reduction=args.wavelet_reduction, fft_enhancement=args.use_fft, use_v2=args.use_v2 ) # Load weights if not os.path.exists(checkpoint_path): raise FileNotFoundError(f"Model file does not exist: {checkpoint_path}") checkpoint = torch.load(checkpoint_path, map_location=args.device) # Check checkpoint format if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: model.load_state_dict(checkpoint["model_state_dict"]) print("✓ Model weights loaded from training checkpoint") else: model.load_state_dict(checkpoint) print("✓ Model weights loaded directly") model = model.to(args.device) model.eval() print(f"✓ Model loaded and set to evaluation mode") print(f"✓ 使用设备：{args.device}") return model def evaluate_model(model, dataloader, args): """ Evaluate model performance Args: model: Model dataloader: Data loader args: Command line arguments Returns: dict: Dictionary containing various metrics """ print("\n" + "=" * 60) print("Starting evaluation...") print("=" * 60) # Initialize metrics dice_metric = DiceMetric(reduction="mean") iou_metric = MeanIoU(reduction="mean") hd_metric = HausdorffDistanceMetric(reduction="mean") hd95_metric = HausdorffDistanceMetric(reduction="mean", percentile=95) total_samples = 0 flag = 0 saved_vis_count = 0 vis_dir = None # Create visualization save directory if args.save_visualization: vis_dir = os.path.join(args.results_dir, f"visualization_{args.dataset_name}") os.makedirs(vis_dir, exist_ok=True) print(f"✓ Visualization results will be saved to: {vis_dir}") with torch.no_grad(): dice_metric.reset() iou_metric.reset() hd_metric.reset() hd95_metric.reset() for batch_idx, batch_data in enumerate(dataloader): images = batch_data["image"].to(args.device) labels = batch_data["label"].to(args.device) # Forward propagation outputs = model(images) # [B, 1, H, W] # Post-processing outputs = torch.sigmoid(outputs) outputs = (outputs > 0.5).int() if flag == 0: flag = 1 print(f"\n{'=' * 60}") print(f"[Detailed Debug Info - Batch 0]") print(f"{'=' * 60}") print(f"Input image size: {images.shape}") print(f"Output image size after post-processing: {outputs.shape}") print(f"Unique values: {torch.unique(outputs)}") print(f"\n--- Labels ---") print(f"Labels shape: {labels.shape}") print(f"Labels unique values: {torch.unique(labels)}") print(f"{'=' * 60}\n") # Calculate metrics for current batch - use directly without additional processing dice_metric(y_pred=outputs, y=labels) iou_metric(y_pred=outputs, y=labels) hd_metric(y_pred=outputs, y=labels) hd95_metric(y_pred=outputs, y=labels) # Save visualization results if args.save_visualization and saved_vis_count < args.vis_num_samples: save_visualization( images=images, labels=labels, predictions=outputs, save_dir=vis_dir, batch_idx=batch_idx, max_samples=args.vis_num_samples - saved_vis_count ) saved_vis_count += images.shape[0] # Print progress if (batch_idx + 1) % 10 == 0 or (batch_idx + 1) == len(dataloader): print(f"进度：{batch_idx + 1}/{len(dataloader)} batches") total_samples += images.shape[0] # Aggregate all metrics mean_dice = dice_metric.aggregate().item() mean_iou = iou_metric.aggregate().item() mean_hd = hd_metric.aggregate().item() mean_hd95 = hd95_metric.aggregate().item() results = { "mDice": mean_dice, "mIoU": mean_iou, "mHD": mean_hd, "mHD95": mean_hd95, "total_samples": total_samples, } return results def save_visualization(images, labels, predictions, save_dir, batch_idx, max_samples): """ Save visualization results: original image, ground truth label, prediction label combined Args: images: Input image batch [B, C, H, W] labels: Ground truth labels [B, 1, H, W] predictions: Prediction labels [B, 1, H, W] save_dir: Save directory batch_idx: Batch index max_samples: Maximum samples to save """ for i in range(min(images.shape[0], max_samples)): try: # Extract single sample image = images[i].cpu() label = labels[i].cpu() prediction = predictions[i].cpu() # Image processing: de-normalize and convert to RGB if image.shape[0] == 1: # 灰度图 image_np = image[0].numpy() * 255 image_rgb = np.stack([image_np] * 3, axis=-1).astype(np.uint8) else: # RGB image image_np = image.numpy().transpose(1, 2, 0) # De-normalize (assuming z-score normalization, approximate handling) image_np = np.clip((image_np - image_np.min()) / (image_np.max() - image_np.min() + 1e-8) * 255, 0, 255) image_rgb = image_np.astype(np.uint8) # Label processing: convert to binary mask label_np = label[0].numpy() if label.shape[0] == 1 else label.numpy() label_binary = (label_np > 0.5).astype(np.float32) # Prediction processing: convert to binary mask pred_np = prediction[0].numpy() if prediction.shape[0] == 1 else prediction.numpy() pred_binary = (pred_np > 0.5).astype(np.float32) # Create pure black and white label image label_bw = (label_binary * 255).astype(np.uint8) label_bw_3ch = np.stack([label_bw] * 3, axis=-1) # Convert to 3 channels for concatenation # Create pure black and white prediction image pred_bw = (pred_binary * 255).astype(np.uint8) pred_bw_3ch = np.stack([pred_bw] * 3, axis=-1) # Convert to 3 channels for concatenation # Horizontally concatenate three images: original, ground truth B&W, prediction B&W combined = np.hstack([image_rgb, label_bw_3ch, pred_bw_3ch]) # Add text annotations (at the top of the image) h, w = image_rgb.shape[:2] # Calculate text width for centering orig_text_size = cv2.getTextSize('Original', cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0] gt_text_size = cv2.getTextSize('Ground Truth', cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0] pred_text_size = cv2.getTextSize('Prediction', cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0] combined = cv2.putText(combined, 'Original', ((w - orig_text_size[0]) // 2, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) combined = cv2.putText(combined, 'Ground Truth', (w + (w - gt_text_size[0]) // 2, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) combined = cv2.putText(combined, 'Prediction', (2 * w + (w - pred_text_size[0]) // 2, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) # Save sample_idx = batch_idx * images.shape[0] + i save_path = os.path.join(save_dir, f'sample_{sample_idx:04d}.png') cv2.imwrite(save_path, cv2.cvtColor(combined, cv2.COLOR_RGB2BGR)) except Exception as e: print(f"⚠️ Error saving sample {i}: {e}") continue def print_results(results, dataset_name, checkpoint_path, model): """ Print evaluation results Args: results: Evaluation results dictionary dataset_name: Dataset name checkpoint_path: Model checkpoint path model: Model """ print("\n" + "=" * 60) print("Evaluation Results") print("=" * 60) print(f"Model parameters: {sum(p.numel() for p in model.parameters())}") print(f"Dataset: {dataset_name}") print(f"Model: {checkpoint_path}") print(f"Number of samples: {results['total_samples']}") print("-" * 60) print(f"mDice (Mean Dice Coefficient): {results['mDice']:.3f}") print(f"mIoU (Mean Intersection over Union): {results['mIoU']:.3f}") print(f"mHD (Mean Hausdorff Distance): {results['mHD']:.3f}") print(f"mHD95 (95% Hausdorff Distance): {results['mHD95']:.3f}") print("=" * 60) def save_results(results, dataset_name, checkpoint_path, results_dir, model): """ Save evaluation results to file Args: results: Evaluation results dictionary dataset_name: Dataset name checkpoint_path: Model checkpoint path results_dir: Results save directory model: Model """ os.makedirs(results_dir, exist_ok=True) # Generate filename result_file = os.path.join(results_dir, f"eval_{dataset_name}.txt") with open(result_file, 'w', encoding='utf-8') as f: f.write("=" * 60 + "\n") f.write("Polyp Segmentation Model Evaluation Report\n") f.write("=" * 60 + "\n\n") f.write(f"Model parameters: {sum(p.numel() for p in model.parameters())}\n") f.write(f"Dataset: {dataset_name}\n") f.write(f"Model checkpoint: {checkpoint_path}\n") f.write(f"Evaluation time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") f.write(f"Number of samples: {results['total_samples']}\n\n") f.write("-" * 60 + "\n") f.write("Evaluation Metrics:\n") f.write("-" * 60 + "\n") f.write(f"mDice (Mean Dice Coefficient): {results['mDice']:.3f}\n") f.write(f"mIoU (Mean Intersection over Union): {results['mIoU']:.3f}\n") f.write(f"mHD (Mean Hausdorff Distance): {results['mHD']:.3f}\n") f.write(f"mHD95 (95% Hausdorff Distance): {results['mHD95']:.3f}\n") f.write("-" * 60 + "\n") print(f"\n✓ Evaluation results saved to: {result_file}") def main(): """ Main evaluation function """ # ==================== Step 1: Parse arguments ==================== args = parse_args() checkpoint_path = Path(args.outputs_dir + args.dir_flag) / f"best_dice_model_{args.dataset_name}.pt" if args.best_metric: checkpoint_path = Path(args.outputs_dir + args.dir_flag) / f"best_metric_model_{args.dataset_name}.pt" print("\n" + "=" * 60) print("Polyp Segmentation Model Evaluation") print("=" * 60) print(f"Dataset: {args.dataset_name}") print(f"Model checkpoint: {checkpoint_path}") # ==================== Step 2: Create validation set and data loader ==================== print("\nLoading validation set...") val_transform = create_val_transform(args.target_spatial_size) val_dataset = PolypDetectionDataset( root_dir=Path(args.data_root) / args.dataset_name, flag='val', transform=val_transform ) val_loader = create_dataloader(args, val_dataset) print(f"✓ Validation set size: {len(val_dataset)} samples") print(f"✓ Data loader: {len(val_loader)} batches") # ==================== Step 3: Load model ==================== model = load_model(args, checkpoint_path) # ==================== Step 4: Evaluate model ==================== results = evaluate_model(model, val_loader, args) # ==================== Step 5: Print and save results ==================== print_results(results, args.dataset_name, checkpoint_path, model) if args.save_results: save_results(results, args.dataset_name, checkpoint_path, args.results_dir + args.dir_flag, model) print("\n" + "=" * 60) print("Evaluation completed!") print("=" * 60) if __name__ == "__main__": main()