SHARE
TWEET

Untitled

a guest Oct 21st, 2019 78 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import torch
  2. import cv2
  3. import numpy as np
  4. import json
  5. from torchvision import transforms
  6. from PIL import Image
  7.  
  8. model = torch.hub.load('pytorch/vision', 'mobilenet_v2', pretrained=True).cuda()
  9. model.eval()
  10.  
  11. with open("labels.json") as labels:
  12.     labels = json.load(labels)
  13.  
  14.  
  15. preprocess = transforms.Compose([
  16.     transforms.Resize(256),
  17.     transforms.CenterCrop(224),
  18.     transforms.ToTensor(),
  19.     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
  20. ])
  21.  
  22. cap = cv2.VideoCapture(2)
  23.  
  24. while(True):
  25.     # Capture frame-by-frame
  26.     ret, frame = cap.read()
  27.     # Our operations on the 1frame come here
  28.     frame2 = Image.fromarray(frame)
  29.  
  30.     frame2 = preprocess(frame2).cuda()
  31.     output = model(frame2.unsqueeze(0))
  32.     print(output.argmax())
  33.     output = int(output.argmax().cpu().numpy())
  34.     print(labels[str(output)])
  35.     cv2.imshow("frame",frame)
  36.     if cv2.waitKey(1) & 0xFF == ord('q'):
  37.         break
  38.  
  39.  
  40. cap.release()
  41. cv2.destroyAllWindows()
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top