语义分割模型预测结果保存代码
时间: 2023-08-06 10:12:48 浏览: 66
以下是一个使用 PyTorch 框架实现语义分割模型预测结果保存的代码示例:
```python
import torch
import torchvision.transforms as transforms
from PIL import Image
# 加载模型
model = torch.load('model.pt')
# 转换输入图像
transform = transforms.Compose([
transforms.Resize((256, 256)), # 调整图像大小
transforms.ToTensor(), # 转换为张量
])
img = Image.open('input.jpg')
img_tensor = transform(img)
# 进行预测
model.eval()
with torch.no_grad():
output = model(img_tensor.unsqueeze(0))
prediction = torch.argmax(output, dim=1).squeeze(0)
# 将预测结果保存为图像
prediction_img = transforms.ToPILImage()(prediction.byte())
prediction_img.save('output.png')
```
在代码中,首先加载已经训练好的模型,然后使用 `transforms` 转换输入图像,接着进行模型预测,最后将预测结果保存为图像。需要注意的是,预测结果为张量,需要使用 `ToPILImage()` 将其转换为图像后才能保存。
相关推荐
![7z](https://img-home.csdnimg.cn/images/20210720083312.png)
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)