-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcollect_imgs.py
More file actions
52 lines (41 loc) · 1.29 KB
/
collect_imgs.py
File metadata and controls
52 lines (41 loc) · 1.29 KB
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
import os
import cv2
DATA_DIR = './data'
if not os.path.exists(DATA_DIR):
os.makedirs(DATA_DIR)
number_of_classes = 26
dataset_size = 200
cap = None
for i in range(10):
cap = cv2.VideoCapture(i)
if cap.isOpened():
break
if not cap.isOpened():
print("Error: Could not open video stream.")
exit()
for j in range(number_of_classes):
if not os.path.exists(os.path.join(DATA_DIR, str(j))):
os.makedirs(os.path.join(DATA_DIR, str(j)))
print(f'Collecting data for class {chr(j + 65)}')
while True:
ret, frame = cap.read()
if not ret:
print("Error: Failed to capture frame.")
break
cv2.putText(frame, f'Ready to collect images for class {chr(j + 65)}? Press "Q" ! :)', (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2,
cv2.LINE_AA)
cv2.imshow('frame', frame)
if cv2.waitKey(25) == ord('q'):
break
counter = 0
while counter < dataset_size:
ret, frame = cap.read()
if not ret:
print("Error: Failed to capture frame.")
break
cv2.imshow('frame', frame)
cv2.waitKey(25)
cv2.imwrite(os.path.join(DATA_DIR, str(j), f'{counter}.jpg'), frame)
counter += 1
cap.release()
cv2.destroyAllWindows()