Pytorch预训练模型、内置模型实现图像分类、检测和分割

原创：余晓龙

Pytorch中提供了很多已经在ImageNet数据集上训练好的模型了，可以直接被加载到模型中进行预测任务。预训练模型存放在Pytorch的torchvision中库，在torchvision库的models模块下可以查看内置的模型，models模块中的模型包含四大类，如图所示：

一、图像分类代码实现

# coding: utf-8

from PIL import Image
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 步骤一（替换sans-serif字体）
plt.rcParams['axes.unicode_minus'] = False   # 步骤二（解决坐标轴负数的负号显示问题）
import json
import numpy as np

import torch
import torch.nn.functional as F
from torchvision import models, transforms

# 1.下载并加载预训练模型
model = models.resnet18(pretrained=True)
model = model.eval()

# 2.加载标签并对输入数据进行处理
labels_path = './imagenet_class_index.json'
with open(labels_path) as json_data:
    idx2labels = json.load(json_data)
# print(idx2labels)


def getone(onestr):
    return onestr.replace(',', ' ')


# 加载中文标签
with open('./zh_label.csv', 'r+', encoding='gbk') as f:
    # print(f)
    # print(map(getone, list(f)))
    zh_labels = list(map(getone, list(f)))

    print(len(zh_labels), type(zh_labels), zh_labels[:5])


transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

# 3.使用模型进行预测
def preimg(img):
    if img.mode == 'RGBA':
        ch = 4
        a = np.asarray(img)[:, :, :3]
        img = Image.fromarray(a)
    return img


im = preimg(Image.open('panda.jpg'))
transformed_img = transform(im)

inputimg = transformed_img.unsqueeze(0)

output = model(inputimg)
output = F.softmax(output, dim=1)

prediction_score, pred_label_i