fast_neural_style.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. from __future__ import print_function
  2. import cv2 as cv
  3. import numpy as np
  4. import argparse
  5. parser = argparse.ArgumentParser(
  6. description='This script is used to run style transfer models from '
  7. 'https://github.com/jcjohnson/fast-neural-style using OpenCV')
  8. parser.add_argument('--input', help='Path to image or video. Skip to capture frames from camera')
  9. parser.add_argument('--model', help='Path to .t7 model')
  10. parser.add_argument('--width', default=-1, type=int, help='Resize input to specific width.')
  11. parser.add_argument('--height', default=-1, type=int, help='Resize input to specific height.')
  12. parser.add_argument('--median_filter', default=0, type=int, help='Kernel size of postprocessing blurring.')
  13. args = parser.parse_args()
  14. net = cv.dnn.readNetFromTorch(cv.samples.findFile(args.model))
  15. net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
  16. if args.input:
  17. cap = cv.VideoCapture(args.input)
  18. else:
  19. cap = cv.VideoCapture(0)
  20. cv.namedWindow('Styled image', cv.WINDOW_NORMAL)
  21. while cv.waitKey(1) < 0:
  22. hasFrame, frame = cap.read()
  23. if not hasFrame:
  24. cv.waitKey()
  25. break
  26. inWidth = args.width if args.width != -1 else frame.shape[1]
  27. inHeight = args.height if args.height != -1 else frame.shape[0]
  28. inp = cv.dnn.blobFromImage(frame, 1.0, (inWidth, inHeight),
  29. (103.939, 116.779, 123.68), swapRB=False, crop=False)
  30. net.setInput(inp)
  31. out = net.forward()
  32. out = out.reshape(3, out.shape[2], out.shape[3])
  33. out[0] += 103.939
  34. out[1] += 116.779
  35. out[2] += 123.68
  36. out /= 255
  37. out = out.transpose(1, 2, 0)
  38. t, _ = net.getPerfProfile()
  39. freq = cv.getTickFrequency() / 1000
  40. print(t / freq, 'ms')
  41. if args.median_filter:
  42. out = cv.medianBlur(out, args.median_filter)
  43. cv.imshow('Styled image', out)