diff --git a/examples/efficientnet.py b/examples/efficientnet.py index e742bb4548..5f006b88ba 100644 --- a/examples/efficientnet.py +++ b/examples/efficientnet.py @@ -6,6 +6,7 @@ import os GPU = os.getenv("GPU", None) is not None import sys import io +import time import numpy as np np.set_printoptions(suppress=True) @@ -88,7 +89,7 @@ class EfficientNet: x = x.pad2d(padding=(0,1,0,1)) x = swish(self._bn0(x.conv2d(self._conv_stem, stride=2))) for block in self._blocks: - print(x.shape) + #print(x.shape) x = block(x) x = swish(self._bn1(x.conv2d(self._conv_head))) x = x.avg_pool2d(kernel_size=x.shape[2:4]) @@ -119,45 +120,30 @@ class EfficientNet: if GPU: mv.cuda_() -if __name__ == "__main__": - # instantiate my net - model = EfficientNet() - model.load_weights_from_torch() - - # load image and preprocess - from PIL import Image - if len(sys.argv) > 1: - url = sys.argv[1] - else: - url = "https://raw.githubusercontent.com/karpathy/micrograd/master/puppy.jpg" - img = Image.open(io.BytesIO(fetch(url))) +def infer(model, img): + # preprocess image aspect_ratio = img.size[0] / img.size[1] img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0)))) img = np.array(img) y0,x0=(np.asarray(img.shape)[:2]-224)//2 - img = img[y0:y0+224, x0:x0+224] + retimg = img = img[y0:y0+224, x0:x0+224] + + # if you want to look at the image + """ + import matplotlib.pyplot as plt + plt.imshow(img) + plt.show() + """ + + # low level preprocess img = np.moveaxis(img, [2,0,1], [0,1,2]) img = img.astype(np.float32).reshape(1,3,224,224) img /= 255.0 img -= np.array([0.485, 0.456, 0.406]).reshape((1,-1,1,1)) img /= np.array([0.229, 0.224, 0.225]).reshape((1,-1,1,1)) - # if you want to look at the micrograd puppy - """ - import matplotlib.pyplot as plt - plt.imshow(img[0].mean(axis=0)) - plt.show() - """ - - # category labels - import ast - lbls = fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt") - lbls = ast.literal_eval(lbls.decode('utf-8')) - # run the net - import time - st = time.time() if GPU: out = model.forward(Tensor(img).cuda()).cpu() else: @@ -169,8 +155,45 @@ if __name__ == "__main__": plt.plot(out.data[0]) plt.show() """ + return out, retimg - print("did inference in %.2f s" % (time.time()-st)) - print(np.argmax(out.data), np.max(out.data), lbls[np.argmax(out.data)]) +if __name__ == "__main__": + # instantiate my net + model = EfficientNet() + model.load_weights_from_torch() + + # category labels + import ast + lbls = fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt") + lbls = ast.literal_eval(lbls.decode('utf-8')) + + # load image and preprocess + from PIL import Image + url = sys.argv[1] + if url == 'webcam': + import pygame + pygame.init() + SCALE = 3 + screen = pygame.display.set_mode((224*SCALE, 224*SCALE)) + pygame.display.set_caption("capture") + + import cv2 + cap = cv2.VideoCapture(0) + while 1: + ret, frame = cap.read() + frame = Image.fromarray(frame[:, :, [2,1,0]]) + out, retimg = infer(model, frame) + simg = cv2.resize(retimg, (224*SCALE, 224*SCALE)) + pygame.surfarray.blit_array(screen, np.array(simg).swapaxes(0,1)) + pygame.display.update() + for e in pygame.event.get(): + pass + print(np.argmax(out.data), np.max(out.data), lbls[np.argmax(out.data)]) + else: + img = Image.open(io.BytesIO(fetch(url))) + out, _ = infer(model, img) + st = time.time() + print(np.argmax(out.data), np.max(out.data), lbls[np.argmax(out.data)]) + print("did inference in %.2f s" % (time.time()-st)) #print("NOT", np.argmin(out.data), np.min(out.data), lbls[np.argmin(out.data)])