-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathcollect-data.py
127 lines (115 loc) · 5.08 KB
/
collect-data.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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import cv2
import numpy as np
import os
# Create the directory structure
if not os.path.exists("data"):
os.makedirs("data")
os.makedirs("data/train")
os.makedirs("data/test")
os.makedirs("data/train/0")
os.makedirs("data/train/1")
os.makedirs("data/train/2")
os.makedirs("data/train/3")
os.makedirs("data/train/4")
os.makedirs("data/train/5")
os.makedirs("data/train/6")
os.makedirs("data/train/7")
os.makedirs("data/train/8")
os.makedirs("data/train/9")
os.makedirs("data/test/0")
os.makedirs("data/test/1")
os.makedirs("data/test/2")
os.makedirs("data/test/3")
os.makedirs("data/test/4")
os.makedirs("data/test/5")
os.makedirs("data/test/6")
os.makedirs("data/test/7")
os.makedirs("data/test/8")
os.makedirs("data/test/9")
# Train or test
mode = 'train'
directory = 'data/'+mode+'/'
cap = cv2.VideoCapture(0)
while True:
_, frame = cap.read()
# Simulating mirror image
frame = cv2.flip(frame, 1)
# Getting count of existing images
count = {'zero': len(os.listdir(directory+"/0")),
'one': len(os.listdir(directory+"/1")),
'two': len(os.listdir(directory+"/2")),
'three': len(os.listdir(directory+"/3")),
'four': len(os.listdir(directory+"/4")),
'five': len(os.listdir(directory+"/5")),
'six': len(os.listdir(directory+"/6")),
'seven': len(os.listdir(directory+"/7")),
'eight': len(os.listdir(directory+"/8")),
'nine': len(os.listdir(directory+"/9"))}
# Printing the count in each set to the screen
cv2.putText(frame, "MODE : "+mode, (10, 50), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
cv2.putText(frame, "IMAGE COUNT", (10, 100), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
cv2.putText(frame, "ZERO : "+str(count['zero']), (10, 120), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
cv2.putText(frame, "ONE : "+str(count['one']), (10, 140), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
cv2.putText(frame, "TWO : "+str(count['two']), (10, 160), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
cv2.putText(frame, "THREE : "+str(count['three']), (10, 180), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
cv2.putText(frame, "FOUR : "+str(count['four']), (10, 200), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
cv2.putText(frame, "FIVE : "+str(count['five']), (10, 220), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
cv2.putText(frame, "SIX : "+str(count['six']), (10, 240), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
cv2.putText(frame, "SEVEN : "+str(count['seven']), (10, 260), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
cv2.putText(frame, "EIGHT : "+str(count['eight']), (10, 280), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
cv2.putText(frame, "NINE : "+str(count['nine']), (10, 300), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,255), 1)
# Coordinates of the ROI
x1 = int(0.5*frame.shape[1])
y1 = 10
x2 = frame.shape[1]-10
y2 = int(0.5*frame.shape[1])
# Drawing the ROI
# The increment/decrement by 1 is to compensate for the bounding box
cv2.rectangle(frame, (x1-1, y1-1), (x2+1, y2+1), (255,0,0) ,1)
# Extracting the ROI
roi = frame[y1:y2, x1:x2]
roi = cv2.resize(roi, (64, 64))
cv2.imshow("Frame", frame)
#_, mask = cv2.threshold(mask, 200, 255, cv2.THRESH_BINARY)
#kernel = np.ones((1, 1), np.uint8)
#img = cv2.dilate(mask, kernel, iterations=1)
#img = cv2.erode(mask, kernel, iterations=1)
# do the processing after capturing the image!
roi = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
_, roi = cv2.threshold(roi, 120, 255, cv2.THRESH_BINARY)
cv2.imshow("ROI", roi)
interrupt = cv2.waitKey(10)
if interrupt & 0xFF == 27: # esc key
break
if interrupt & 0xFF == ord('0'):
cv2.imwrite(directory+'0/'+str(count['zero'])+'.jpg', roi)
if interrupt & 0xFF == ord('1'):
cv2.imwrite(directory+'1/'+str(count['one'])+'.jpg', roi)
if interrupt & 0xFF == ord('2'):
cv2.imwrite(directory+'2/'+str(count['two'])+'.jpg', roi)
if interrupt & 0xFF == ord('3'):
cv2.imwrite(directory+'3/'+str(count['three'])+'.jpg', roi)
if interrupt & 0xFF == ord('4'):
cv2.imwrite(directory+'4/'+str(count['four'])+'.jpg', roi)
if interrupt & 0xFF == ord('5'):
cv2.imwrite(directory+'5/'+str(count['five'])+'.jpg', roi)
if interrupt & 0xFF == ord('6'):
cv2.imwrite(directory+'6/'+str(count['six'])+'.jpg', roi)
if interrupt & 0xFF == ord('7'):
cv2.imwrite(directory+'7/'+str(count['seven'])+'.jpg', roi)
if interrupt & 0xFF == ord('8'):
cv2.imwrite(directory+'8/'+str(count['eight'])+'.jpg', roi)
if interrupt & 0xFF == ord('9'):
cv2.imwrite(directory+'9/'+str(count['nine'])+'.jpg', roi)
cap.release()
cv2.destroyAllWindows()
"""
d = "old-data/test/0"
newd = "data/test/0"
for walk in os.walk(d):
for file in walk[2]:
roi = cv2.imread(d+"/"+file)
roi = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
_, mask = cv2.threshold(roi, 120, 255, cv2.THRESH_BINARY)
cv2.imwrite(newd+"/"+file, mask)
"""