Walk-Assistant/make_data.py at master · YoongiKim/Walk-Assistant · 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
111
112
113
114
115
116
117
118
119
120
121
"""
Walk-Assistant : Recognizing sidewalk for the visually impaired
Copyright (C) 2018 Yoongi Kim (devlifecode@outlook.com)

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""


from opt_flow import OptFlow
import glob
import cv2
import os
import numpy as np
import threading
from tqdm import tqdm, trange

QUEUE_SIZE = 50
DATA_PATH = "data/videos/*.mp4"

# 영상이 중간에 끊기면 코덱 문제이므로 ffmpeg로 mute 해야 합니다.
# Linux: ffmpeg -i data/videos/test.mp4 -c copy -an data/videos/test_mute.mp4
# Windows: ffmpeg.exe -i data/videos/test.mp4 -c copy -an data/videos/test_mute.mp4

START_SKIP = 500  # 10초 생략
END_SKIP = 500  # 10초 생략

OUTPUT_PATH = "data/frames"
LABEL_PATH = "data/opt_flow.txt"
os.makedirs(OUTPUT_PATH, exist_ok=True)

files = glob.glob(DATA_PATH)

flow = OptFlow(height_start=0.5, height_end=1.0)

label_lines = []

for file in files:
    print(file)
    file_name = str(file).replace('\\', '/').split('/')[-1].replace('.mp4', '')

    cap = cv2.VideoCapture(file)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    for i in range(START_SKIP):
        cap.read()

    img_queue = []
    flow_queue = []  # (x, y) moved position of previous frame

    _, img1 = cap.read()
    img_queue.append(img1)
    succeed, img2 = cap.read()
    img_queue.append(img2)

    with trange(START_SKIP, total-END_SKIP-1) as t:
        for i in t:
            x, y = flow.get_direction(img1, img2, show=False)
            flow_queue.append((x, y))

            if len(img_queue) >= QUEUE_SIZE:
                pop_img = cv2.resize(img_queue.pop(0), (1280, 720))
                pop_flow = flow_queue.pop(0)

                # way_visual = flow.draw_way(pop_img, flow_queue, size=40)
                way = flow.draw_way(pop_img, flow_queue, size=200)
                way = cv2.resize(way, (16, 9))

                way_flatten = np.reshape(np.array(np.array(way)/255).astype(np.int), 144)

                if np.sum(way_flatten) > 2:  # 횡단보도 대기 무시
                    way_encode = ''
                    for b in way_flatten:
                        way_encode += str(b)
                    label_lines.append('%s_%d.jpg,%s\n' % (file_name, i, way_encode))
                    threading.Thread(target=cv2.imwrite, args=('%s/%s_%d.jpg' % (OUTPUT_PATH, file_name, i), pop_img)).start()
                    # t.write('%s_%d.jpg,%s\n' % (file_name, i, way_encode))

                    visual = cv2.resize(way, (1280, 720))
                    visual = cv2.cvtColor(np.array(visual).astype(np.uint8), cv2.COLOR_GRAY2BGR)
                    visual = cv2.add(pop_img, visual)

                    # way_visual = cv2.cvtColor(np.array(way_visual).astype(np.uint8), cv2.COLOR_GRAY2BGR)
                    # way_visual[:,:,0]=0
                    # visual = cv2.add(visual, way_visual)

                    cv2.imshow('visual', visual)
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        print('User Interrupted')
                        break

            img1 = img2
            succeed, img2 = cap.read()
            img_queue.append(img2)

print('Processed %d frames' % len(label_lines))

last_labels = []
if os.path.exists(LABEL_PATH):
    print('Found existing label file. Appending labels...')
    with open(LABEL_PATH, 'r') as p:
        last_labels = p.readlines()

total_labels = last_labels + label_lines
set_labels = set(total_labels)
print('Found %d duplicated labels' % (len(total_labels) - len(set_labels)))

with open(LABEL_PATH, 'w') as f:
    f.writelines(set_labels)

print('Saved %d labels' % len(set_labels))