File: googlenet_python.py

package info (click to toggle)
opencv 3.2.0%2Bdfsg-6
  • links: PTS, VCS
  • area: main
  • in suites: buster
  • size: 238,480 kB
  • sloc: xml: 901,650; cpp: 703,419; lisp: 20,142; java: 17,843; python: 17,641; ansic: 603; cs: 601; sh: 516; perl: 494; makefile: 117
file content (34 lines) | stat: -rw-r--r-- 1,070 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from __future__ import print_function
import numpy as np
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 ]

blob = prepare_image(cv2.imread('space_shuttle.jpg'))
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()
#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()])