How can I use the gluon-cv model_zoo and output to an OpenCV window with Python?

Question

My code is:

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)

I somehow need to draw the bounding_codes, class_IDs, and scores onto the image and output it via imshow.

How can I accomplish this?

Answer 1

We can use ssd|yolo (wroted by mxnet|keras|pytorch) to detect the objects in the image. Then we will get the result as a form of classids/scores/bboxes. Iterator the result, do some transform, then just drawing in OpenCV will be OK.

(Poor English, but I think you can get me in the following code).

---

This is the source image:

This the result displayed in 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()