Guest User

Untitled

a guest
Feb 20th, 2018
81
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 0.88 KB | None | 0 0
  1. import argparse
  2.  
  3. from PIL import Image
  4. from coremltools.models import MLModel
  5. import cv2
  6. import numpy as np
  7.  
  8. def main():
  9. parser = argparse.ArgumentParser(
  10. description='Stylize image using CoreML'
  11. )
  12.  
  13. parser.add_argument('-model', required=True, help='CoreML model path')
  14. args = parser.parse_args()
  15.  
  16. capture = cv2.VideoCapture(0)
  17. capture.set(3, 720)
  18. capture.set(4, 720)
  19.  
  20. net = MLModel(args.model)
  21.  
  22. while True:
  23. ret, frame = capture.read()
  24. if not ret:
  25. continue
  26.  
  27. image = Image.fromarray(frame)
  28. image = image.resize((720, 720))
  29. stylized_image = net.predict({'inputImage': image})['outputImage']
  30. stylized_image = np.array(stylized_image.convert('RGB'))
  31. cv2.imshow('', stylized_image)
  32.  
  33. k = cv2.waitKey(10)
  34. if k == 27:
  35. break
  36.  
  37.  
  38. if __name__ == "__main__":
  39. main()
Add Comment
Please, Sign In to add comment