轻松学Pytorch-使用ResNet50实现图像分类

在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)
10net.setInput(blob)
11probs = net.forward()
12index = 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()

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

✄————————————————

看到这里，说明你喜欢这篇文章，请点击「
在看
」或顺手「
转发
」「
点赞
」。

欢迎微信搜索「
panchuangxx
」，添加小编
磐小小仙
微信，每日朋友圈更新一篇高质量推文（无广告），为您提供更多精彩内容。

▼  
▼
  
 扫描二维码添加小编
  ▼

  
▼