Jump to content

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.

Link to comment
Share on other sites

Link to post
Share on other sites

Create an account or sign in to comment

You need to be a member in order to leave a comment

Create an account

Sign up for a new account in our community. It's easy!

Register a new account

Sign in

Already have an account? Sign in here.

Sign In Now

×