专栏的第一篇博客介绍了语义分割,这篇博客主要是使用开源的模型做一些效果测试,先跑个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预训练模型做的测试,后续出标注数据的内容。感谢您的观看,觉得写的还可以请帮忙点个赞和收藏!