### 使用 PyTorch 实现视频 PSNR 和 SSIM 的计算 为了实现视频的 PSNR 和 SSIM 指标计算，可以利用 PyTorch 提供的功能来逐帧处理视频并调用相应的度量函数。以下是详细的说明和代码示例。 #### 背景介绍 PSNR（Peak Signal-to-Noise Ratio）用于衡量两幅图像之间的差异程度[^3]，而 SSIM（Structural Similarity Index Measure）则更注重于感知质量评估[^1]。对于视频而言，可以通过逐帧比较的方式计算每一对对应帧的 PSNR 和 SSIM 值，并最终取平均值作为整个视频的质量评价标准。 #### 依赖库安装 在实现之前，请确保已安装必要的 Python 库： ```bash pip install torch torchvision scikit-image ``` #### PyTorch 中的 PSNR 和 SSIM 计算方法 PyTorch 并未直接提供内置的 PSNR 函数，但可以通过简单的数学运算实现它。而对于 SSIM，则可以直接使用 `torchvision` 或第三方库如 `pytorch-msssim` 来完成[^4]。 #### 代码实现 以下是一个完整的代码示例： ```python import torch from skimage.metrics import structural_similarity as ssim_metric import numpy as np def calculate_psnr(img1, img2): mse = torch.mean((img1 - img2) ** 2).item() if mse == 0: return float('inf') max_pixel_value = 1.0 # Assuming images are normalized to [0, 1] psnr = 20 * torch.log10(max_pixel_value / torch.sqrt(mse)) return psnr.item() def calculate_ssim(img1, img2): img1_np = img1.permute(1, 2, 0).cpu().numpy() # Convert to HWC format and move to CPU img2_np = img2.permute(1, 2, 0).cpu().numpy() # Ensure the range is between 0 and 1 for both images img1_np = (img1_np - img1_np.min()) / (img1_np.max() - img1_np.min()) img2_np = (img2_np - img2_np.min()) / (img2_np.max() - img2_np.min()) ssim_val = ssim_metric( img1_np, img2_np, multichannel=True, data_range=img2_np.max() - img2_np.min(), gaussian_weights=True, sigma=1.5, use_sample_covariance=False ) return ssim_val def evaluate_video(video_gt_frames, video_pred_frames): """ Evaluate PSNR and SSIM metrics frame-by-frame. :param video_gt_frames: List of ground truth frames (each a tensor with shape CxHxW). :param video_pred_frames: List of predicted frames (same structure as above). :return: Average PSNR and SSIM values across all frames. """ total_psnr = 0.0 total_ssim = 0.0 num_frames = min(len(video_gt_frames), len(video_pred_frames)) # Ensure equal length for i in range(num_frames): gt_frame = video_gt_frames[i].unsqueeze(0) # Add batch dimension pred_frame = video_pred_frames[i].unsqueeze(0) psnr_val = calculate_psnr(gt_frame, pred_frame) ssim_val = calculate_ssim(gt_frame.squeeze(), pred_frame.squeeze()) total_psnr += psnr_val total_ssim += ssim_val avg_psnr = total_psnr / num_frames avg_ssim = total_ssim / num_frames return avg_psnr, avg_ssim # Example usage if __name__ == "__main__": device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Generate dummy video frames (normalized tensors in [0, 1]) video_gt_frames = [torch.rand(3, 256, 256).to(device) for _ in range(10)] video_pred_frames = [torch.rand(3, 256, 256).to(device) for _ in range(10)] avg_psnr, avg_ssim = evaluate_video(video_gt_frames, video_pred_frames) print(f"Average PSNR: {avg_psnr:.4f}, Average SSIM: {avg_ssim:.4f}") ``` 上述代码实现了以下几个功能： 1. 定义了一个 `calculate_psnr` 方法来计算两张图片间的 PSNR 值。 2. 利用了 `skimage.metrics.structural_similarity` 来计算 SSIM 值。 3. 将输入的视频分解成一系列帧，并通过循环逐一计算各帧的 PSNR 和 SSIM 值，最后返回均值。 #### 注意事项 - 输入的图像应被标准化到 `[0, 1]` 范围内以匹配大多数预定义的损失函数或度量函数的要求。 - 如果需要支持批量处理，可以在现有基础上进一步优化逻辑结构。 ---