base/
base.rs

1use candle_core::Tensor;
2use candle_nn::Module;
3use dehazing::model::DehazeNet;
4
5fn main() {
6    let device = candle_core::Device::cuda_if_available(0).unwrap();
7    let base_dir = env!("CARGO_MANIFEST_DIR");
8
9    let model = DehazeNet::with_device(&device).unwrap();
10
11    let img = image::open(format!("{base_dir}/testdata/test2.png")).unwrap();
12
13    let raw = img.to_rgb8().into_vec();
14    let data = Tensor::from_vec(
15        raw,
16        (img.height() as usize, img.width() as usize, 3),
17        &device,
18    )
19    .unwrap()
20    .to_dtype(candle_core::DType::F32)
21    .unwrap()
22    .broadcast_div(&Tensor::new(255f32, &device).unwrap())
23    .unwrap()
24    .permute((2, 0, 1))
25    .unwrap()
26    .unsqueeze(0)
27    .unwrap();
28
29    println!("{data:?}");
30
31    let out = model.forward(&data).unwrap();
32
33    // 处理输出张量
34    let out = out.squeeze(0).unwrap(); // 移除批次维度 [c, h, w]
35
36    let (_, height, width) = out.dims3().unwrap();
37
38    let image_data: Vec<u8> = out
39        .permute((1, 2, 0))
40        .unwrap() // [H, W, C] 符合图像布局
41        .flatten_all()
42        .unwrap()
43        .to_vec1::<f32>()
44        .unwrap()
45        .iter()
46        .map(|&v| (v.clamp(0.0, 1.0) * 255.0) as u8)
47        .collect();
48
49    // 保存图像
50    let img_out =
51        image::RgbImage::from_raw(width as u32, height as u32, image_data).expect("创建图像失败");
52
53    img_out.save("result/dehazed_output.jpg").expect("保存图像失败");
54    println!("去雾结果已保存为 result/dehazed_output.jpg");
55}
base/base.rs

base/
base.rs
