webcam support

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)])