import glob import numpy as np import torch import os import cv2 from model.unet_model import UNet if __name__ == "__main__": # 选择设备，有cuda用cuda，没有就用cpu device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 加载网络，图片单通道，分类为1。 net = UNet(n_channels=1, n_classes=1) # 将网络拷贝到deivce中 net.to(device=device) # 加载模型参数 net.load_state_dict(torch.load('best_model.pth', map_location=device)) # 测试模式 net.eval() # 读取所有图片路径 tests_path = glob.glob('../data/data/test/test_image/*.png') # 遍历所有图片 for i in tests_path: # 保存结果地址 save_res_path = '../data/test/test_mask/*res.png' # 读取图片 img = cv2.imread(i) # 转为灰度图 img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) # 转为batch为1，通道为1，大小为512*512的数组 img = img.reshape(1, 1, img.shape[0], img.shape[1]) # 转为tensor img_tensor = torch.from_numpy(img) # 将tensor拷贝到device中，只用cpu就是拷贝到cpu中，用cuda就是拷贝到cuda中。 img_tensor = img_tensor.to(device=device, dtype=torch.float32) # 预测 pred = net(img_tensor) # 提取结果 pred = np.array(pred.data.cpu()[0])[0] # 处理结果 pred[pred >= 0.5] = 255 pred[pred < 0.5] = 0 # 保存图片 cv2.imwrite(save_res_path, pred) # print(pred) print("successfully save")

这段代码是一个用于图像分割的UNet模型的理代码。首先，它会检查是否有可用的CUDA设备，然后加载UNet模型并将其移动到相应的设备上。接下来，它会遍历所有测试图片的路径，读取并预处理每张图片。然后，使用加载的模型对图片进行预测，并将预测结果进行后处理，得到二值化的分割结果。最后，将结果保存到指定路径上。

值得注意的是，上述代码中的`save_res_path`变量应该是一个具体的文件路径，而不是一个通配符表达式。你需要将`save_res_path`修改为一个具体的文件路径，以保存每张图片的分割结果。

如果你有更多的问题，请继续提问。