""" RadioUNet V3 推理脚本 使用训练好的模型对SoundMapDiff数据集进行推理 """ import os import sys import argparse import torch import numpy as np from PIL import Image import matplotlib.pyplot as plt from pathlib import Path from skimage.metrics import structural_similarity as ssim sys.path.append(os.path.join(os.path.dirname(__file__), 'lib')) from lib.modules import RadioWNet from lib.soundmap_loader import SoundMapDataset from torch.utils.data import DataLoader def calculate_metrics(pred, target): """计算评估指标""" pred_np = pred.cpu().numpy().squeeze() target_np = target.cpu().numpy().squeeze() # MSE mse = np.mean((pred_np - target_np) ** 2) # MAE mae = np.mean(np.abs(pred_np - target_np)) # RMSE rmse = np.sqrt(mse) # SSIM ssim_val = ssim(pred_np, target_np, data_range=1.0) # PSNR if mse > 0: psnr = 10 * np.log10(1.0 / mse) else: psnr = float('inf') return { 'mse': mse, 'mae': mae, 'rmse': rmse, 'ssim': ssim_val, 'psnr': psnr } def visualize_prediction(inputs, target, pred, metrics, save_path): """可视化预测结果""" fig, axes = plt.subplots(2, 2, figsize=(12, 12)) # 建筑物布局 axes[0, 0].imshow(inputs[0].cpu().numpy(), cmap='gray') axes[0, 0].set_title('Building Layout', fontsize=14) axes[0, 0].axis('off') # 声源位置 axes[0, 1].imshow(inputs[1].cpu().numpy(), cmap='hot') axes[0, 1].set_title('Sound Source', fontsize=14) axes[0, 1].axis('off') # 真实热力图 - 使用viridis颜色方案（紫→蓝→绿→黄） im1 = axes[1, 0].imshow(target.cpu().numpy().squeeze(), cmap='viridis', vmin=0, vmax=1) axes[1, 0].set_title('Ground Truth', fontsize=14) axes[1, 0].axis('off') plt.colorbar(im1, ax=axes[1, 0], fraction=0.046, pad=0.04) # 预测热力图 - 使用viridis颜色方案（紫→蓝→绿→黄） im2 = axes[1, 1].imshow(pred.cpu().numpy().squeeze(), cmap='viridis', vmin=0, vmax=1) axes[1, 1].set_title(f"Prediction (SSIM: {metrics['ssim']:.4f})", fontsize=14) axes[1, 1].axis('off') plt.colorbar(im2, ax=axes[1, 1], fraction=0.046, pad=0.04) # 添加指标信息 metrics_text = f"MSE: {metrics['mse']:.6f} | MAE: {metrics['mae']:.4f} | SSIM: {metrics['ssim']:.4f} | PSNR: {metrics['psnr']:.2f} dB" fig.suptitle(metrics_text, fontsize=12, y=0.02) plt.tight_layout() plt.savefig(save_path, dpi=150, bbox_inches='tight') plt.close() def main(): parser = argparse.ArgumentParser(description='RadioUNet V3 推理脚本') parser.add_argument('--checkpoint', type=str, default='outputs/radiounet_v3/checkpoints/best_model.pth', help='模型检查点路径') parser.add_argument('--dataset_dir', type=str, default='/home/djk/generate/dataset/SoundMapDiff', help='数据集目录') parser.add_argument('--output_dir', type=str, default='outputs/radiounet_v3/inference', help='输出目录') parser.add_argument('--num_samples', type=int, default=20, help='推理样本数量') parser.add_argument('--img_size', type=int, default=256, help='图像尺寸') args = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"使用设备: {device}") # 创建输出目录 output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) # 加载模型 print(f"加载模型: {args.checkpoint}") model = RadioWNet(inputs=2, phase="firstU").to(device) checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False) if 'model_state_dict' in checkpoint: model.load_state_dict(checkpoint['model_state_dict']) print(f"加载Epoch {checkpoint.get('epoch', 'unknown')}的模型") else: model.load_state_dict(checkpoint) model.eval() # 加载测试数据集 print(f"加载数据集: {args.dataset_dir}") test_dataset = SoundMapDataset( dataset_dir=args.dataset_dir, phase="test", img_size=args.img_size ) # 均匀采样 total_samples = len(test_dataset) indices = np.linspace(0, total_samples - 1, args.num_samples, dtype=int) print(f"测试样本数: {total_samples}, 采样数: {args.num_samples}") print(f"\n{'='*60}") print("开始推理...") print(f"{'='*60}\n") all_metrics = [] with torch.no_grad(): for i, idx in enumerate(indices): inputs, target = test_dataset[idx] inputs = inputs.unsqueeze(0).to(device) target = target.unsqueeze(0).to(device) # 推理 outputs = model(inputs) if isinstance(outputs, list): outputs = outputs[0] # 计算指标 metrics = calculate_metrics(outputs.squeeze(0), target.squeeze(0)) all_metrics.append(metrics) # 可视化 save_path = output_dir / f'prediction_{i+1}_idx{idx}.png' visualize_prediction(inputs.squeeze(0), target.squeeze(0), outputs.squeeze(0), metrics, save_path) print(f"样本 {i+1}/{args.num_samples} (idx={idx}): " f"SSIM={metrics['ssim']:.4f}, MSE={metrics['mse']:.6f}, PSNR={metrics['psnr']:.2f}dB") # 计算平均指标 avg_metrics = { 'mse': np.mean([m['mse'] for m in all_metrics]), 'mae': np.mean([m['mae'] for m in all_metrics]), 'rmse': np.mean([m['rmse'] for m in all_metrics]), 'ssim': np.mean([m['ssim'] for m in all_metrics]), 'psnr': np.mean([m['psnr'] for m in all_metrics]) } print(f"\n{'='*60}") print("平均评估指标") print(f"{'='*60}") print(f" 平均 MSE: {avg_metrics['mse']:.6f}") print(f" 平均 MAE: {avg_metrics['mae']:.4f}") print(f" 平均 RMSE: {avg_metrics['rmse']:.4f}") print(f" 平均 SSIM: {avg_metrics['ssim']:.4f}") print(f" 平均 PSNR: {avg_metrics['psnr']:.2f} dB") print(f"{'='*60}") # 保存报告 report_path = output_dir / 'evaluation_report.txt' with open(report_path, 'w', encoding='utf-8') as f: f.write("RadioUNet V3 评估报告\n") f.write("=" * 60 + "\n\n") f.write(f"模型: {args.checkpoint}\n") f.write(f"测试样本数: {args.num_samples}\n\n") for i, (idx, m) in enumerate(zip(indices, all_metrics)): f.write(f"样本 {i+1} (索引 {idx}):\n") f.write(f" MSE: {m['mse']:.6f}\n") f.write(f" MAE: {m['mae']:.4f}\n") f.write(f" SSIM: {m['ssim']:.4f}\n") f.write(f" PSNR: {m['psnr']:.2f} dB\n\n") f.write("=" * 60 + "\n") f.write("平均指标:\n") f.write("=" * 60 + "\n") f.write(f" 平均 MSE: {avg_metrics['mse']:.6f}\n") f.write(f" 平均 MAE: {avg_metrics['mae']:.4f}\n") f.write(f" 平均 RMSE: {avg_metrics['rmse']:.4f}\n") f.write(f" 平均 SSIM: {avg_metrics['ssim']:.4f}\n") f.write(f" 平均 PSNR: {avg_metrics['psnr']:.2f} dB\n") print(f"\n✅ 推理完成！结果保存在: {output_dir}") if __name__ == '__main__': main()