forked from BVLC/caffe
-
Notifications
You must be signed in to change notification settings - Fork 1.7k
/
create_list.py
121 lines (113 loc) · 5.13 KB
/
create_list.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
import argparse
import os
from random import shuffle
import shutil
import subprocess
import sys
HOMEDIR = os.path.expanduser("~")
CURDIR = os.path.dirname(os.path.realpath(__file__))
# If true, re-create all list files.
redo = True
# The root directory which holds all information of the dataset.
data_dir = "{}/data/coco".format(HOMEDIR)
# The directory name which holds the image sets.
imgset_dir = "ImageSets"
# The direcotry which contains the images.
img_dir = "images"
img_ext = "jpg"
# The directory which contains the annotations.
anno_dir = "Annotations"
anno_ext = "json"
train_list_file = "{}/train.txt".format(CURDIR)
minival_list_file = "{}/minival.txt".format(CURDIR)
testdev_list_file = "{}/testdev.txt".format(CURDIR)
test_list_file = "{}/test.txt".format(CURDIR)
# Create training set.
# We follow Ross Girschick's split.
if redo or not os.path.exists(train_list_file):
datasets = ["train2014", "valminusminival2014"]
img_files = []
anno_files = []
for dataset in datasets:
imgset_file = "{}/{}/{}.txt".format(data_dir, imgset_dir, dataset)
with open(imgset_file, "r") as f:
for line in f.readlines():
name = line.strip("\n")
subset = name.split("_")[1]
img_file = "{}/{}/{}.{}".format(img_dir, subset, name, img_ext)
assert os.path.exists("{}/{}".format(data_dir, img_file)), \
"{}/{} does not exist".format(data_dir, img_file)
anno_file = "{}/{}/{}.{}".format(anno_dir, subset, name, anno_ext)
assert os.path.exists("{}/{}".format(data_dir, anno_file)), \
"{}/{} does not exist".format(data_dir, anno_file)
img_files.append(img_file)
anno_files.append(anno_file)
# Shuffle the images.
idx = [i for i in xrange(len(img_files))]
shuffle(idx)
with open(train_list_file, "w") as f:
for i in idx:
f.write("{} {}\n".format(img_files[i], anno_files[i]))
if redo or not os.path.exists(minival_list_file):
datasets = ["minival2014"]
subset = "val2014"
img_files = []
anno_files = []
for dataset in datasets:
imgset_file = "{}/{}/{}.txt".format(data_dir, imgset_dir, dataset)
with open(imgset_file, "r") as f:
for line in f.readlines():
name = line.strip("\n")
img_file = "{}/{}/{}.{}".format(img_dir, subset, name, img_ext)
assert os.path.exists("{}/{}".format(data_dir, img_file)), \
"{}/{} does not exist".format(data_dir, img_file)
anno_file = "{}/{}/{}.{}".format(anno_dir, subset, name, anno_ext)
assert os.path.exists("{}/{}".format(data_dir, anno_file)), \
"{}/{} does not exist".format(data_dir, anno_file)
img_files.append(img_file)
anno_files.append(anno_file)
with open(minival_list_file, "w") as f:
for i in xrange(len(img_files)):
f.write("{} {}\n".format(img_files[i], anno_files[i]))
if redo or not os.path.exists(testdev_list_file):
datasets = ["test-dev2015"]
subset = "test2015"
img_files = []
anno_files = []
for dataset in datasets:
imgset_file = "{}/{}/{}.txt".format(data_dir, imgset_dir, dataset)
with open(imgset_file, "r") as f:
for line in f.readlines():
name = line.strip("\n")
img_file = "{}/{}/{}.{}".format(img_dir, subset, name, img_ext)
assert os.path.exists("{}/{}".format(data_dir, img_file)), \
"{}/{} does not exist".format(data_dir, img_file)
anno_file = "{}/{}/{}.{}".format(anno_dir, subset, name, anno_ext)
assert os.path.exists("{}/{}".format(data_dir, anno_file)), \
"{}/{} does not exist".format(data_dir, anno_file)
img_files.append(img_file)
anno_files.append(anno_file)
with open(testdev_list_file, "w") as f:
for i in xrange(len(img_files)):
f.write("{} {}\n".format(img_files[i], anno_files[i]))
if redo or not os.path.exists(test_list_file):
datasets = ["test2015"]
subset = "test2015"
img_files = []
anno_files = []
for dataset in datasets:
imgset_file = "{}/{}/{}.txt".format(data_dir, imgset_dir, dataset)
with open(imgset_file, "r") as f:
for line in f.readlines():
name = line.strip("\n")
img_file = "{}/{}/{}.{}".format(img_dir, subset, name, img_ext)
assert os.path.exists("{}/{}".format(data_dir, img_file)), \
"{}/{} does not exist".format(data_dir, img_file)
anno_file = "{}/{}/{}.{}".format(anno_dir, subset, name, anno_ext)
assert os.path.exists("{}/{}".format(data_dir, anno_file)), \
"{}/{} does not exist".format(data_dir, anno_file)
img_files.append(img_file)
anno_files.append(anno_file)
with open(test_list_file, "w") as f:
for i in xrange(len(img_files)):
f.write("{} {}\n".format(img_files[i], anno_files[i]))