diff --git a/samples/dnn/googlenet_python.py b/samples/dnn/googlenet_python.py index 53ccdc5518..1db2cbf612 100644 --- a/samples/dnn/googlenet_python.py +++ b/samples/dnn/googlenet_python.py @@ -4,31 +4,21 @@ import cv2 from cv2 import dnn import timeit -def prepare_image(img): - img = cv2.resize(img, (224, 224)) - #convert interleaved image (RGBRGB) to planar(RRGGBB) - blob = np.moveaxis(img, 2, 0) - blob = np.reshape(blob.astype(np.float32), (-1, 3, 224, 224)) - return blob - def timeit_forward(net): - print("OpenCL:", cv2.ocl.useOpenCL()) print("Runtime:", timeit.timeit(lambda: net.forward(), number=10)) def get_class_list(): with open('synset_words.txt', 'rt') as f: - return [ x[x.find(" ") + 1 :] for x in f ] + return [x[x.find(" ") + 1:] for x in f] -blob = prepare_image(cv2.imread('space_shuttle.jpg')) +blob = dnn.blobFromImage(cv2.imread('space_shuttle.jpg'), 1, (224, 224), (104, 117, 123)) print("Input:", blob.shape, blob.dtype) -cv2.ocl.setUseOpenCL(True) #Disable OCL if you want net = dnn.readNetFromCaffe('bvlc_googlenet.prototxt', 'bvlc_googlenet.caffemodel') -net.setBlob(".data", blob) -net.forward() +net.setInput(blob) +prob = net.forward() #timeit_forward(net) #Uncomment to check performance -prob = net.getBlob("prob") print("Output:", prob.shape, prob.dtype) classes = get_class_list() print("Best match", classes[prob.argmax()]) \ No newline at end of file