首页 > 解决方案 > 如何使用 gluon-cv model_zoo 并通过 Python 输出到 OpenCV 窗口?

问题描述

我的代码是:

import gluoncv as gcv

net = gcv.model_zoo.get_model('ssd_512_mobilenet1.0_voc', pretrained=True)

windowName = "ssdObject"
cv2.namedWindow(windowName, cv2.WINDOW_NORMAL)
cv2.resizeWindow(windowName, 1280, 720)
cv2.moveWindow(windowName, 0, 0)
cv2.setWindowTitle(windowName, "SSD Object Detection")
while True:
    # Check to see if the user closed the window
    if cv2.getWindowProperty(windowName, 0) < 0:
        # This will fail if the user closed the window; Nasties get printed to the console
        break
    ret_val, frame = video_capture.read()

    frame = mx.nd.array(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).astype('uint8')
    rgb_nd, frame = gcv.data.transforms.presets.ssd.transform_test(frame, short=512, max_size=700)

    # # Run frame through network
    class_IDs, scores, bounding_boxes = net(rgb_nd)

    displayBuf = frame
    cv2.imshow(windowName, displayBuf)
    cv2.waitKey(0)

我不知何故需要在图像上绘制bounding_codes,class_IDsscores并通过imshow.

我怎样才能做到这一点?

标签: pythonopencvobject-detectionmxnet

解决方案


我们可以使用ssd|yolo(由 编写mxnet|keras|pytorch)来检测图像中的对象。然后我们将以classids/scores/bboxes的形式获得结果。迭代结果,做一些变换,然后在 OpenCV 中绘图就可以了。

(英语很差,但我想你可以在下面的代码中找到我)。


这是源图像: 在此处输入图像描述

这是 OpenCV 中显示的结果:

在此处输入图像描述


#!/usr/bin/python3
# 2019/01/24 09:05
# 2019/01/24 10:25

import gluoncv as gcv
import mxnet as mx
import cv2
import numpy as np
# https://github.com/pjreddie/darknet/blob/master/data/dog.jpg

## (1) Create network 
net = gcv.model_zoo.get_model('ssd_512_mobilenet1.0_voc', pretrained=True)

## (2) Read the image and preprocess 
img = cv2.imread("dog.jpg")
rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

xrgb = mx.nd.array(rgb).astype('uint8')
rgb_nd, xrgb = gcv.data.transforms.presets.ssd.transform_test(xrgb, short=512, max_size=700)

## (3) Interface 
class_IDs, scores, bounding_boxes = net(rgb_nd)

## (4) Display 
for i in range(len(scores[0])):
    #print(class_IDs.reshape(-1))
    #print(scores.reshape(-1))
    cid = int(class_IDs[0][i].asnumpy())
    cname = net.classes[cid]
    score = float(scores[0][i].asnumpy())
    if score < 0.5:
        break
    x,y,w,h = bbox =  bounding_boxes[0][i].astype(int).asnumpy()
    print(cid, score, bbox)
    tag = "{}; {:.4f}".format(cname, score)
    cv2.rectangle(img, (x,y), (w, h), (0, 255, 0), 2)
    cv2.putText(img, tag, (x, y-20),  cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,0,255), 1)

cv2.imshow("ssd", img);
cv2.waitKey()

推荐阅读