Ask Your Question

Revision history [back]

click to hide/show revision 1
initial version

How can I use a custom tensorflow model with the CV2 DNN module?

I retrained an object detection model based on Google's Tensorflow object detection API. I exported it as a frozen inference graph. I would like to use it with CV2's DNN module:

cap = cv2.VideoCapture(URL)

cvNet = cv2.dnn.readNetFromTensorflow('graph.pb', 'graph.pbtxt')

while True:
    ret, img = cap.read()
    rows = img.shape[0]
    cols = img.shape[1]
    cvNet.setInput(cv2.dnn.blobFromImage(img, 1.0/127.5, (300, 300), (127.5, 127.5, 127.5), swapRB=True, crop=False))
    cvOut = cv2Net.forward()

    for detection in cv2Out[0,0,:,:]:
        score = float(detection[2])
        if score > 0.3:
            left = detection[3] * cols
            top = detection[4] * rows
            right = detection[5] * cols
            bottom = detection[6] * rows
            cv2.rectangle(img, (int(left), int(top)), (int(right), int(bottom)), (23, 230, 210), thickness=2)
    cv2.imshow('img', img)
    if cv2.waitKey(1) ==27:
        exit(0)

I get this error: Const input blob for weights not found in function getConstBlob

From my research, I believe I have to optimize the inference graph. I can't find any documentation as how to do this.

If anyone could point me in the right direction, it would be very much appreciated.