-
Notifications
You must be signed in to change notification settings - Fork 0
/
categorize.py
86 lines (79 loc) · 3.27 KB
/
categorize.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
from argparse import ArgumentParser
import os
import csv
import shutil
def main():
parser = ArgumentParser(description="Move images to folders named for corresponding action")
parser.add_argument("-v", "--verbose", action="store_true", help="print extra messages (default: False")
args = parser.parse_args()
verbose = args.verbose
os.chdir("data/")
while True:
try:
os.chdir("categorized")
break
except FileNotFoundError:
os.mkdir("categorized")
# make sure a directory exists for each label
existing_subds = os.listdir()
possible_labels = ["forward", "forward left", "forward right", "backward",
"backward left", "backward right", "left", "right", "stop", "none"]
for possible_label in possible_labels:
if possible_label not in existing_subds:
os.mkdir(possible_label)
# get list of raw data directories
dirs = os.listdir("../all data")
# for each directory, open data.csv
for d in dirs:
print("processing {} images".format(d))
with open("../all data/{}/data.csv".format(d)) as data:
data_reader = csv.reader(data)
# get header
header = next(data_reader)
up_idx = header.index("wheels: up key")
down_idx = header.index("wheels: down key")
left_idx = header.index("wheels: left key")
right_idx = header.index("wheels: right key")
stop_idx = header.index("stop")
row_number = 1
for row in data_reader:
forward = int(row[up_idx])
backward = int(row[down_idx])
left = int(row[left_idx])
right = int(row[right_idx])
stop = int(row[stop_idx])
labels = []
if forward:
labels.append("forward")
if backward:
labels.append("backward")
if left:
labels.append("left")
if right:
labels.append("right")
# stop overrides all other labels
if stop:
labels = ["stop"]
else:
# remove contradictory labels
if "forward" in labels and "backward" in labels:
labels.remove("forward")
labels.remove("backward")
if "left" in labels and "right" in labels:
labels.remove("left")
labels.remove("right")
# if there are no labels, assign label "none"
if len(labels) == 0:
labels = ["none"]
label = ' '.join(labels)
if label == "forward left":
print(label)
# rename corresponding image & copy to proper folder
# example new file name: NE4 1.jpg
if verbose:
print("moving {}.jpg to {}/{} {}.jpg".format(row_number, label, d, row_number))
shutil.copy("../all data/{}/img/{}.jpg".format(d, row_number),
"./{}/{} {}.jpg".format(label, d, row_number))
row_number += 1
if __name__ == "__main__":
main()