使用pytorch预训练分割模型做预测

        专栏的第一篇博客介绍了语义分割，这篇博客主要是使用开源的模型做一些效果测试，先跑个demo。pytorch官方开源了一些预训练模型（开源预训练模型网址），参数信息如下：

        pytorch提供了DeepLabV3、FCN、LRASPP三个种类的预训练模型，我们根据官方提供的demo对DeepLabV3模型做效果测试。

测试模型效果

• 准备图片数据，在使用模型做测试之前我们先要下载一些图片数据，我是直接百度搜索小狗的图片，然后保存在本地。
• 本地推理，推理代码如下：

from torchvision.io.image import read_image
from torchvision.models.segmentation import fcn_resnet50, FCN_ResNet50_Weights
from torchvision.transforms.functional import to_pil_image

img = read_image("小狗图片位置")
# 设置模型类型
weights = FCN_ResNet50_Weights.DEFAULT
model = fcn_resnet50(weights=weights)
model.eval()

# 初始化图片预处理函数
preprocess = weights.transforms()

# 图像预处理
batch = preprocess(img).unsqueeze(0)

# 使用模型做预测
prediction = model(batch)["out"]
normalized_masks = prediction.softmax(dim=1)
class_to_idx = {cls: idx for (idx, cls) in enumerate(weights.meta["categories"])}
mask = normalized_masks[0, class_to_idx["dog"]]
to_pil_image(img).show()
to_pil_image(mask).show()

• 效果展示

结尾

        以上是使用pytorch预训练模型做的测试，后续出标注数据的内容。感谢您的观看，觉得写的还可以请帮忙点个赞和收藏！