Как я могу использовать gluon-cv model_zoo и выводить в окно OpenCV с Python? - PullRequest
Новая
       14

Как я могу использовать gluon-cv model_zoo и выводить в окно OpenCV с Python?

0 голосов
/ 23 января 2019

Мой код: 

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_IDs и scores на изображении и вывести его через imshow.

Как мне это сделать?

1 Ответ

0 голосов
/ 24 января 2019

Мы можем использовать ssd|yolo (wroted by mxnet|keras|pytorch), чтобы обнаружить объекты на изображении.Тогда мы получим результат в виде classids / scores / bboxes.Итерируйте результат, сделайте некоторое преобразование, тогда просто рисование в OpenCV будет в порядке.

(плохой английский, но я думаю, что вы можете получить меня в следующем коде).

Это исходное изображение: enter image description here

Этот результат отображается в OpenCV:

enter image description here

#!/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()
...