Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import argparse
- from PIL import Image
- from coremltools.models import MLModel
- import cv2
- import numpy as np
- def main():
- parser = argparse.ArgumentParser(
- description='Stylize image using CoreML'
- )
- parser.add_argument('-model', required=True, help='CoreML model path')
- args = parser.parse_args()
- capture = cv2.VideoCapture(0)
- capture.set(3, 720)
- capture.set(4, 720)
- net = MLModel(args.model)
- while True:
- ret, frame = capture.read()
- if not ret:
- continue
- image = Image.fromarray(frame)
- image = image.resize((720, 720))
- stylized_image = net.predict({'inputImage': image})['outputImage']
- stylized_image = np.array(stylized_image.convert('RGB'))
- cv2.imshow('', stylized_image)
- k = cv2.waitKey(10)
- if k == 27:
- break
- if __name__ == "__main__":
- main()
Add Comment
Please, Sign In to add comment