``` def main(): num = 800 # Total generated picture num args = create_argparser().parse_args() dist_util.setup_dist() logger.configure() logger.log("creating model...") model, diffusion = sr_create_model_and_diffusion( **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) ) model.load_state_dict( dist_util.load_state_dict(args.model_path, map_location="cuda") ) model.to(dist_util.dev()) model.eval() logger.log("loading data...") data = load_superres_data(args.batch_size) logger.log("creating samples...") MSE, SSIM, PSNR, LPIP = [], [], [], [] lpip_loss = lpips.LPIPS(net="alex").to(dist_util.dev()) l2loss = th.nn.MSELoss().to(dist_util.dev()) radon_288_736 = para_prepare_parallel(2.5) radon_72_736 = para_prepare_parallel(8.5) radon_36_736 = para_prepare_parallel(16.5) helper = {"fbp_para_288_736": radon_288_736, "fbp_para_36_736": radon_36_736, "fbp_para_72_736": radon_72_736} for i in range(0, num//args.batch_size):# model_kwargs = next(data) raw_img = model_kwargs.pop('raw_img').to("cuda") index = model_kwargs.pop('index') model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()} model_kwargs["fbp_para_36_736"] = radon_36_736 model_kwargs["fbp_para_288_736"] = radon_288_736 sample_fn = p_sample_loop_super_res sample, sample_72_288 = sample_fn( model, (args.batch_size, 1, 288, 736), #args.large_size, args.large_size # clip_denoised=args.clip_denoised, model_kwargs=model_kwargs, ) model_72_sino = F.interpolate(sample_72_288, [72, 736], mode="nearest") model_72_fbp = run_reco(model_72_sino + 1., helper["fbp_para_72_736"])[:,:,112:624,112:624] model_72_fbp_npy = model_72_fbp.cpu().detach().numpy() model_output_fbp = run_reco(sample + 1., helper["fbp_para_288_736"])[:,:,112:624,112:624] target_fbp = run_reco(raw_img + 1., helper["fbp_para_288_736"])[:,:,112:624,112:624] output_fbp_npy = model_output_fbp.cpu().detach().numpy() for j in range(0, args.batch_size): l2loss_value = l2loss(model_output_fbp[j], target_fbp[j]).item() print("index:", index[j], "MSELoss:", l2loss_value) MSE.append(l2loss_value) raw_npy = target_fbp.cpu().detach().numpy() ssim_value = ssim(np.squeeze(output_fbp_npy[j]),np.squeeze( raw_npy[j]), data_range = raw_npy[j].max() - raw_npy[j].min()) psnr_value = psnr(np.squeeze(output_fbp_npy[j]),np.squeeze( raw_npy[j]), data_range = raw_npy[j].max() - raw_npy[j].min()) print("index:", index[j], "SSIM:", ssim_value) SSIM.append(ssim_value) PSNR.append(psnr_value) lpip_value = lpip_loss(model_output_fbp[j], target_fbp[j]) print("lpips:", lpip_value.item()) LPIP.append(lpip_value.item()) print("mse mean:", np.mean(MSE)) print("ssim mean:", np.mean(SSIM)) print("psnr mean:", np.mean(PSNR)) print("lpip mean:", np.mean(LPIP))```解释代码
时间: 2025-04-05 08:22:59 浏览: 32
这段代码的主要目的是通过超分辨率扩散模型生成一系列高质量图像,并评估其相对于真实目标图像的质量指标(例如MSE、SSIM、PSNR和LPIPS)。接下来我们分模块逐一解读代码的核心功能及其工作流程。
---
### **函数概述**
`main()` 是整个脚本的主入口点,负责完成以下几个核心任务:
1. 初始化模型及扩散过程配置。
2. 加载数据集并逐批处理。
3. 调用采样函数生成超分辨率图像。
4. 计算多项质量评价指标以量化生成效果。
---
### **详细分解**
#### 1. 参数初始化与日志记录
```python
num = 800 # 需要生成的总图片数
args = create_argparser().parse_args()
dist_util.setup_dist() # 设置分布式训练环境
logger.configure() # 日志工具配置
logger.log("creating model...")
```
- **`create_argparser()`**: 解析命令行参数,提供灵活性以便用户调整实验条件。
- **`dist_util.setup_dist()`**: 如果使用的是多GPU分布式的场景下准备必要的通信设置。
- **`logger.configure()` & `logger.log(...)`**: 提供标准化的日志管理方便跟踪程序进度信息。
#### 2. 创建模型和加载权重
```python
model, diffusion = sr_create_model_and_diffusion(
**args_to_dict(args, sr_model_and_diffusion_defaults().keys())
)
model.load_state_dict(dist_util.load_state_dict(args.model_path, map_location="cuda"))
model.to(dist_util.dev())
model.eval()
```
- **创建模型**:
- 根据传入参数动态构建适合当前任务需求的超级分辨率模型以及对应的去噪扩散管道(diffusion process)。
- 此处还包含了将模型部署到CUDA设备上运行的操作。
- **加载已保存的模型权值**:
- 来源于指定路径下的文件(由 `args.model_path` 确定位置)。
#### 3. 数据加载
```python
data = load_superres_data(args.batch_size)
```
调用了专门设计的数据读取器获取测试所需批量化的低分辨率影像资料。
#### 4. 定义辅助变量用于后续计算
```python
radon_288_736 = para_prepare_parallel(2.5)
radon_72_736 = para_prepare_parallel(8.5)
radon_36_736 = para_prepare_parallel(16.5)
helper = {
"fbp_para_288_736": radon_288_736,
"fbp_para_36_736": radon_36_736,
"fbp_para_72_736": radon_72_736
}
```
此处定义了几个与拉东变换(Radon Transform)相关的参数集合存储起来备用。
#### 5. 图像重建循环部分
对于每一组输入数据进行如下操作:
##### (a)提取单次批次的信息
```python
for i in range(0, num // args.batch_size):
model_kwargs = next(data)
raw_img = model_kwargs.pop('raw_img').to("cuda")
index = model_kwargs.pop('index')
model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()}
model_kwargs["fbp_para_36_736"] = radon_36_736
model_kwargs["fbp_para_288_736"] = radon_288_736
sample_fn = p_sample_loop_super_res
sample, sample_72_288 = sample_fn(
model,
(args.batch_size, 1, 288, 736),
model_kwargs=model_kwargs,
)
```
这里进行了两件重要事情:
- 把从迭代器取得的一帧或多帧样本准备好送进网络做预测。
- 实际调用了前面提到过的 `p_sample_loop_super_res`, 输出包含两个维度不同的恢复结果。
##### (b)二次加工输出特征图谱
```python
model_72_sino = F.interpolate(sample_72_288, [72, 736], mode="nearest")
...
output_fbp_npy = model_output_fbp.cpu().detach().numpy()
```
通过对初步产生的高频细节映射施加额外插值及其他数学反投影手段得到更直观可见形态的内容表示形式。
#### 6. 统计各项误差度量值
遍历每一个单独个体像素块之间差异程度衡量标准得出总体表现情况概览。
---
### 结果汇总打印
最终会把收集的所有数值求均值得出整体效能报告结论展示出来便于比较分析之需所用。
---
阅读全文
相关推荐

















