OpenCV DNN with CUDA example

opencv_mobile_ssd.py

mport cv2 as cv
import time


rootdir     = '/home/ssunil/mywork/ml_opencv_code/'
tfmodeldir  = "model_configs/ssd_mobilenet_v2_coco_2018_03_29/"
tfgraph     = "frozen_inference_graph.pb"
tfconfig    = "ssd_mobilenet_v2_coco_2018_03_29.pbtxt"


# COCO dataset classes 
classNames = {0:"background", 1:"person", 2:"bicycle", 3:"car", 4:"motorcycle",
            5:"airplane", 6:"bus", 7:"train", 8:"truck", 9:"boat", 10:"traffic light", 11:"fire hydrant",
            12:"unknown", 13:"stop sign", 14:"parking meter", 15:"bench", 16:"bird", 17:"cat", 18:"dog", 19:"horse",
            20:"sheep", 21:"cow", 22:"elephant", 23:"bear", 24:"zebra", 25:"giraffe", 26:"unknown", 27:"backpack",
            28:"umbrella", 29:"unknown", 30:"unknown", 31:"handbag", 32:"tie", 33:"suitcase", 34:"frisbee", 35:"skis",
            36:"snowboard", 37:"sports ball", 38:"kite", 39:"baseball bat", 40:"baseball glove", 41:"skateboard",
            42:"surfboard", 43:"tennis racket", 44:"bottle", 45:"unknown", 46:"wine glass", 47:"cup", 48:"fork", 49:"knife",
            50:"spoon", 51:"bowl", 52:"banana", 53:"apple", 54:"sandwich", 55:"orange", 56:"broccoli", 57:"carrot", 58:"hot dog",
            59:"pizza", 60:"donut", 61:"cake", 62:"chair", 63:"couch", 64:"potted plant", 65:"bed", 66:"unknown", 67:"dining table",
            68:"unknown", 69:"unknown", 70:"toilet", 71:"unknown", 72:"tv", 73:"laptop", 74:"mouse", 75:"remote", 76:"keyboard",
            77:"cell phone", 78:"microwave", 79:"oven", 80:"toaster", 81:"sink", 82:"refrigerator", 83:"unknown",
            84:"book", 85:"clock", 86:"vase", 87:"scissors", 88:"teddy bear", 89:"hair drier", 90:"toothbrush"}

def id_class_name(class_id, classes):
    for key,value in classes.items():
        if class_id == key:
            return value

cvNet = cv.dnn.readNetFromTensorflow(rootdir+tfmodeldir+tfgraph,rootdir+tfmodeldir+tfconfig)

#print("[INFO] setting preferable backend and target to CUDA...")
#cvNet.setPreferableBackend(cv.dnn.DNN_BACKEND_CUDA)
#cvNet.setPreferableTarget(cv.dnn.DNN_TARGET_CUDA)

img = cv.imread(rootdir+'images/office.jpg')
rows = img.shape[0]
cols = img.shape[1]
cvNet.setInput(cv.dnn.blobFromImage(img, size=(300, 300), swapRB=True, crop=False))
start = time.time()
cvOut = cvNet.forward()
end = time.time()
for detection in cvOut[0,0,:,:]:
    score = float(detection[2])
    if score > 0.3: class_id = detection[1]
        left = detection[3] * cols
        top = detection[4] * rows
        right = detection[5] * cols
        bottom = detection[6] * rows
        print(str(str(int(class_id)) + " " + str(detection[2])  + " " + id_class_name(class_id,classNames)))
        cv.rectangle(img, (int(left), int(top)), (int(right), int(bottom)), (23, 230, 210), thickness=2)
ms_per_image = (end - start) * 1000 / 100
print("Time per inference: %f ms" % (ms_per_image))
cv.imshow('img', img)
cv.waitKey()