screengrab.py 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. #!/usr/local/bin/python
  2. # -*- coding: utf-8 -*-
  3. import numpy as np
  4. import cv2
  5. from mss import mss
  6. from PIL import Image
  7. import threading
  8. import time
  9. from random import randint
  10. import tensorflow as tf
  11. from tensorflow.keras import layers, models, initializers
  12. import sys
  13. sys.path.append('../hostLib/')
  14. from hostLib.layers.conv2d import Conv2D as Conv2DFPGA
  15. bounding_box = {'top': 0, 'left': 0, 'width': 2560, 'height': 1440}
  16. width, height = 228, 228
  17. sct = mss()
  18. stop = 0
  19. a = layers.Input(dtype=tf.float32, shape=(width, height, 3))
  20. z = Conv2DFPGA(1, kernel_initializer=initializers.Constant(1/25))(a)
  21. model = models.Model(inputs=a, outputs=z)
  22. #model.compile(loss=tf.keras.losses.categorical_crossentropy,
  23. # optimizer=tf.keras.optimizers.Adadelta(),
  24. # metrics=['accuracy'])
  25. sct_img = sct.grab(bounding_box)
  26. np_img = np.array(sct_img)
  27. resized_image = cv2.resize(np_img, (width, height))
  28. resized_image = cv2.cvtColor(resized_image, cv2.COLOR_BGRA2BGR)
  29. while True:
  30. cv2.ellipse(
  31. resized_image,
  32. (randint(0,width), randint(0,height)),
  33. (randint(0,width/2), randint(0,height/2)),
  34. randint(0,360),
  35. 0,
  36. 360,
  37. [randint(0,256), randint(0,256), randint(0,256)],
  38. 10
  39. )
  40. img32 = tf.cast(resized_image, tf.float32)
  41. #img32 = np.expand_dims(img32, axis=2)
  42. cv2.imshow('screen', resized_image)
  43. x,y,w,h = cv2.getWindowImageRect('screen')
  44. batch = np.expand_dims(img32, axis=0)
  45. batch = tf.tile(batch, [1,1,1,1])
  46. print(batch.shape)
  47. predictions = model.predict(batch)
  48. pred8 = tf.cast(predictions, tf.uint8)
  49. for i in range(pred8.shape[0]):
  50. name = 'conv_{}'.format(i)
  51. cv2.imshow(name, pred8.numpy()[i])
  52. cv2.moveWindow(name, x+((i+1)*300)%1500, y+int((i+1)/5)*300)
  53. if (cv2.waitKey(1) & 0xFF) == ord('x') or stop:
  54. cv2.destroyAllWindows()
  55. stop = True
  56. break