-
Notifications
You must be signed in to change notification settings - Fork 6
/
gettrainingdata.py
58 lines (48 loc) · 1.25 KB
/
gettrainingdata.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
import grabscreen
import csv
import os
import time
import cv2
from getkeys import key_check
import numpy as np
from PIL import Image
def imageprocessing(image):
processed_image = cv2.cvtColor(image,cv2.COLOR_BGR2GRAY)
#processed_image = cv2.Canny(processed_image,threshold1=200,threshold2=200)
return processed_image
def writecsv(o1):
with open('log.csv','a',newline='') as fp:
writer = csv.writer(fp,delimiter=',')
writer.writerow(o1)
def getkey():
key = key_check()
output = [0,0,0]
#AWD
if 'A' in key:
output[0] = 1
elif 'D' in key:
output[2] = 1
else:
output[1] = 1
return output
if __name__ == "__main__":
try:
os.mkdir("Images")
except:
pass
i = 1
#dataarray = []
for j in range(5):
print(j + 1)
time.sleep(1)
while True:
image1 = grabscreen.grab_screen([0, 350, 800, 650])
image1 = imageprocessing(image1)
image1 = cv2.resize(image1, (400, 150))
image_name = str(i)+ '.png'
full_name_image = os.path.join('Images', image_name)
keypressed = getkey()
writecsv([full_name_image,keypressed])
cv2.imwrite(full_name_image,image1)
print(i)
i += 1