磐创AI分享

来源 | OpenCV学堂

作者 | gloomyfish

Hello大家好，这篇文章给大家详细介绍一下pytorch中最重要的组件torchvision，它包含了常见的数据集、模型架构与预训练模型权重文件、常见图像变换、计算机视觉任务训练。可以是说是pytorch中非常有用的模型迁移学习神器。本文将会介绍如何使用torchvison的预训练模型ResNet50实现图像分类。

模型

Torchvision.models包里面包含了常见的各种基础模型架构，主要包括：

AlexNet

VGG

ResNet

SqueezeNet

DenseNet

Inception v3

GoogLeNet

ShuffleNet v2

MobileNet v2

ResNeXt

Wide ResNet

MNASNet

这里我选择了ResNet50，基于ImageNet训练的基础网络来实现图像分类， 网络模型下载与加载如下：

model = torchvision.models.resnet50(pretrained=True).eval().cuda()

tf = 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]

        )])

使用模型实现图像分类

这里首先需要加载ImageNet的分类标签，目的是最后显示分类的文本标签时候使用。然后对输入图像完成预处理，使用ResNet50模型实现分类预测，对预测结果解析之后，显示标签文本，完整的代码演示如下：

1with open(imagenet_classes.txt) as

 f:

2    labels = [line.strip() for line in

 f.readlines()]

3 4src = cv.imread("D:/images/space_shuttle.jpg") # aeroplane.jpg 5image = cv.resize(src, (224, 224

))

6image = np.float32(image) / 255.0 7image[:,:,] -= (np.float32(0.485), np.float32(0.456), np.float32(0.406

))

8image[:,:,] /= (np.float32(0.229), np.float32(0.224), np.float32(0.225

))

9image = image.transpose((2, 0, 1

))

10input_x = torch.from_numpy(image).unsqueeze(0

)

11

print(input_x.size())

12

pred = model(input_x.cuda())

13pred_index = torch.argmax(pred, 1

).cpu().detach().numpy()

14

print(pred_index)

15print("current predict class name : %s"%labels[pred_index[0

]])

16cv.putText(src, labels[pred_index[0]], (50, 50), cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2

)

17cv.imshow("input"

, src)

18cv.waitKey(0

)

19cv.destroyAllWindows()

运行结果如下：

转ONNX支持

在torchvision中的模型基本上都可以转换为ONNX格式，而且被OpenCV DNN模块所支持，所以，很方便的可以对torchvision自带的模型转为ONNX，实现OpenCV DNN的调用，首先转为ONNX模型，直接使用torch.onnx.export即可转换(还不知道怎么转，快点看前面的例子)。转换之后使用OpenCV DNN调用的代码如下：

1with open(imagenet_classes.txt) as

 f:

2    labels = [line.strip() for line in

 f.readlines()]

3net = cv.dnn.readNetFromONNX("resnet.onnx"

)

4src = cv.imread("D:/images/messi.jpg")  # aeroplane.jpg 5image = cv.resize(src, (224, 224

))

6image = np.float32(image) / 255.0 7image[:, :, ] -= (np.float32(0.485), np.float32(0.456), np.float32(0.406

))

8image[:, :, ] /= (np.float32(0.229), np.float32(0.224), np.float32(0.225

))

9blob = cv.dnn.blobFromImage(image, 1.0, (224, 224), (0, 0, 0), False

)

10

net.setInput(blob)

11

probs = net.forward()

12

index = np.argmax(probs)

13cv.putText(src, labels[index], (50, 50), cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2

)

14cv.imshow("input"

, src)

15cv.waitKey(0

)

16cv.destroyAllWindows()

 运行结果见上图，这里就不再贴了。