Upload inference.py with huggingface_hub

inference.py

@@ -0,0 +1,217 @@
+"""
+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()