Chiro/record.py at master · NWChen/Chiro · GitHub

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
import numpy as np
import cv2
import os
import time

# Takes in a camera object and any overlay information.
# Returns the raw RGB image from the camera.
def record(camera, overlay_text):
    # Continually display raw output of webcam
    ret, raw_frame = camera.read() # We could convert color spaces to HSV here
    frame = raw_frame
    #cv2.putText(frame, overlay_text, (0, 25), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
    cv2.imshow('frame', raw_frame)
    return raw_frame

def name_this_dir(cwd):
    max_index = 0
    for dir in os.listdir(cwd):
        index = int(dir)
        if index > max_index:
            max_index = index
    return str(max_index)

if __name__ == '__main__':
    three = '''
    .----------------.
    | .--------------. |
    | |    ______    | |
    | |   / ____ `.  | |
    | |   `'  __) |  | |
    | |   _  |__ '.  | |
    | |  | \____) |  | |
    | |   \______.'  | |
    | |              | |
    | '--------------' |
    '----------------'
    '''
    two = '''
     .----------------.
    | .--------------. |
    | |    _____     | |
    | |   / ___ `.   | |
    | |  |_/___) |   | |
    | |   .'____.'   | |
    | |  / /____     | |
    | |  |_______|   | |
    | |              | |
    | '--------------' |
    '----------------'
    '''
    one = '''
    .----------------.
    | .--------------. |
    | |     __       | |
    | |    /  |      | |
    | |    `| |      | |
    | |     | |      | |
    | |    _| |_     | |
    | |   |_____|    | |
    | |              | |
    | '--------------' |
    '----------------'
    '''
    goodjob = '''
     _____  ____   ____  _____         _  ____  ____
    / ____|/ __ \ / __ \|  __ \       | |/ __ \|  _ \
   | |  __| |  | | |  | | |  | |      | | |  | | |_) |
   | | |_ | |  | | |  | | |  | |  _   | | |  | |  _ <
   | |__| | |__| | |__| | |__| | | |__| | |__| | |_) |
    \_____|\____/ \____/|_____/   \____/ \____/|____(_)
    '''
    sleeps = [three, two, one]

    # Get user input for data type
    label = ''
    while label not in ('good', 'bad'):
        label = raw_input('LABEL FOR THIS DATASET (good / bad): ')

    # Generate training set directory
    temp_dir = '%s/data/%s' % (os.getcwd(), label)
    dir = '%s/%s/' % (temp_dir, name_this_dir(temp_dir))
    if not os.path.exists(dir):
        os.makedirs(dir)
    print dir
    index = 0

    # Count down to shot
    camera = cv2.VideoCapture(0)
    for i in sleeps:
        print i
        time.sleep(1)

    start_time = time.time()
    time.clock()

    elapsed = 0
    while elapsed < 10:
        elapsed = time.time() - start_time
        frame = record(camera, label)
        filename = dir + str(index) + '.jpg'
        height, width, channels = frame.shape # Ensure we actually took an image.
        if width and height:
            cv2.imwrite(filename, frame) # Save the image.
            index += 1
        if cv2.waitKey(1) & 0xFF == ord('q'): # Quit on 'q' and clean up.
            break

    # Destroy camera object and GUI
    camera.release()
    cv2.destroyAllWindows()