def get_image_train_eval(data_path):
"""get image list when data struct is
{
|-- data_url
|-- train
|-- Images
|-- a.jpg
...
|-- Annotations
|-- a.txt (or a.xml)
|-- label_map_dict (optional)
|-- eval
|-- Images
|-- b.jpg
...
|-- Annotations
|-- b.txt (or b.xml)
...
|-- label_map_dict (optional)
|-- label_map_dict (optional)
}
:param data_path: data store url
Returns:
train_data_list,
eval_data_list,
"""
image_list_train = []
# get all labeled train data
image_list_set = file.list_directory(os.path.join(data_path, 'train', 'Images'))
assert not image_list_set == [], 'there is no file in data url'
for i in image_list_set:
if file.exists(os.path.join(data_path, 'train', 'Annotations', os.path.splitext(i)[0] + '.xml')):
image_list_train.append([os.path.join(data_path, 'train', 'Images', i),
os.path.join(data_path, 'train', 'Annotations', os.path.splitext(i)[0] + '.xml')])
elif file.exists(os.path.join(data_path, 'train', 'Annotations', os.path.splitext(i)[0] + '.txt')):
image_list_train.append([os.path.join(data_path, 'train', 'Images', i),
file.read(os.path.join(data_path, 'train',
'Annotations',
os.path.splitext(i)[0] + '.txt'))])
# get all labeled eval data
image_list_eval = []
image_list_set = []
image_list_set = file.list_directory(os.path.join(data_path, 'eval', 'Images'))
assert not image_list_set == [], 'there is no file in data url'
for i in image_list_set:
if file.exists(os.path.join(data_path, 'eval', 'Annotations', os.path.splitext(i)[0] + '.xml')):
image_list_eval.append([os.path.join(data_path, 'eval', 'Images', i),
os.path.join(data_path, 'eval', 'Annotations', os.path.splitext(i)[0] + '.xml')])
elif file.exists(os.path.join(data_path, 'eval', 'Annotations', os.path.splitext(i)[0] + '.txt')):
image_list_eval.append([os.path.join(data_path, 'eval', 'Images', i),
file.read(os.path.join(data_path, 'eval',
'Annotations',
os.path.splitext(i)[0] + '.txt'))])
return image_list_train, image_list_eval
def _get_label_names(self, index_file):
classes = []
if index_file is not None and file.exists(index_file):
label_file = h5py.File(index_file, 'r')
classes_name = label_file['labels_list'][:]
label_file.close()
classes = [name.decode('utf-8') for name in classes_name]
else:
for data_file in self.img_list:
annotation_file = data_file[1]
tree = ET.ElementTree()
parser = ET.XMLParser(target=ET.TreeBuilder())
parser.feed(file.read(annotation_file, binary=True))
tree._root = parser.close()
objs = tree.findall('object')
non_diff_objs = [
obj for obj in objs if int(obj.find('difficult').text) == 0
]
objs = non_diff_objs
for obj in objs:
class_name = obj.find('name').text.lower().strip()
if class_name not in classes:
classes.append(class_name)
if index_file is not None:
index_file_path = index_file
else:
index_file_path = INDEX_FILE_NAME
if file.exists(index_file_path):
file.remove(index_file_path)
label_file = h5py.File(index_file, 'w')
label_file.create_dataset('labels_list',
data=[item.encode() for item in classes])
label_file.close()
return classes
def get_image_images_annotation(data_path, split_spec):
"""get image list when data struct is
{
|-- data_url
|-- Images
|-- a.jpg
|-- b.jpg
...
|-- Annotations
|-- a.txt (or a.xml)
|-- b.txt (or b.xml)
...
|-- label_map_dict (optional)
}
:param data_path: data store url
:param split_spec: split train percent if data doesn't have evaluation data
Returns:
train_data_list,
eval_data_list,
"""
image_set = []
label_dict = {}
label_num = 0
class_name = []
# get all labeled data
image_list_set = file.list_directory(os.path.join(data_path, 'Images'))
assert not image_list_set == [], 'there is no file in data url'
for i in image_list_set:
if file.exists(os.path.join(data_path, 'Annotations', os.path.splitext(i)[0] + '.xml')):
image_set.append([os.path.join(data_path, 'Images', i),
os.path.join(data_path, 'Annotations', os.path.splitext(i)[0] + '.xml')])
elif file.exists(os.path.join(data_path, 'Annotations', os.path.splitext(i)[0] + '.txt')):
label_name = file.read(os.path.join(data_path, 'Annotations',
os.path.splitext(i)[0] + '.txt'))
if label_name not in label_dict.keys():
label_dict[label_name] = label_num
class_name.append(label_name)
label_num = label_num + 1
image_set.append([os.path.join(data_path, 'Images', i),
label_dict[label_name]])
# split data to train and eval
num_examples = len(image_set)
train_num = int(num_examples * split_spec)
shuffle_list = list(range(num_examples))
random.shuffle(shuffle_list)
image_list_train = []
image_list_eval = []
for idx, item in enumerate(shuffle_list):
if idx < train_num:
image_list_train.append(image_set[item])
else:
image_list_eval.append(image_set[item])
return image_list_train, image_list_eval, class_name
def _load_label(self, idx):
"""Parse xml file and return labels."""
anno_path = self.img_list[idx][1]
tree = ET.ElementTree()
parser = ET.XMLParser(target=ET.TreeBuilder())
parser.feed(file.read(anno_path, binary=True))
tree._root = parser.close()
root = tree.getroot()
size = root.find('size')
width = float(size.find('width').text)
height = float(size.find('height').text)
if idx not in self._im_shapes:
# store the shapes for later usage
self._im_shapes[idx] = (width, height)
label = []
for obj in root.iter('object'):
difficult = int(obj.find('difficult').text)
cls_name = obj.find('name').text.strip().lower()
if cls_name not in self.classes:
self.classes.append(cls_name)
cls_id = self.classes.index(cls_name)
xml_box = obj.find('bndbox')
xmin = (float(xml_box.find('xmin').text) - 1)
ymin = (float(xml_box.find('ymin').text) - 1)
xmax = (float(xml_box.find('xmax').text) - 1)
ymax = (float(xml_box.find('ymax').text) - 1)
xmin = 0 if xmin < 0 else xmin
ymin = 0 if ymin < 0 else ymin
xmax = 0 if xmax < 0 else xmax
ymax = 0 if ymax < 0 else ymax
try:
self._validate_label(xmin, ymin, xmax, ymax, width, height)
except AssertionError as e:
raise RuntimeError("Invalid label at {}, {}".format(
anno_path, e))
label.append([xmin, ymin, xmax, ymax, cls_id, difficult])
return np.array(label)
def get_data_iter(data_path, train_file=None, val_file=None, split_spec=1,
hyper_train={}, hyper_val={}, **kwargs):
train_set = None
val_set = None
train_list = None
val_list = None
if train_file is not None:
assert file.exists(train_file), 'not found train file'
train_path = file.read(train_file).split("\n")[0:-1]
train_list = [path.replace('\r', '').split(' ') for path in train_path]
train_list = [[os.path.join(data_path, path[0]),
os.path.join(data_path, path[1])] for path in train_list]
if val_file is not None:
assert file.exists(val_file), 'not found val file'
val_path = file.read(val_file).split("\n")[0:-1]
val_list = [path.replace('\r', '').split(' ') for path in val_path]
val_list = [[os.path.join(data_path, path[0]),
os.path.join(data_path, path[1])] for path in val_list]
if train_file is None and val_file is None:
train_list, val_list, _ = get_image_list(data_path, split_spec)
if 'anchors' not in kwargs:
kwargs['anchors'] = [[116, 90, 156, 198, 373, 326],
[30, 61, 62, 45, 59, 119],
[10, 13, 16, 30, 33, 23]]
if 'offsets' not in kwargs:
kwargs['offsets'] = [(13, 13), (26, 26), (52, 52)]
if train_list is not None and len(train_list) > 0:
dataset = Detection_dataset(img_list=train_list,
index_file=hyper_train.get(
'index_file', None),
width=hyper_train.get('width', 416),
height=hyper_train.get('height', 416),
is_train=True,
** kwargs)
max_gt_box_number = max([len(item) for item in dataset.label_cache])
batch_size = hyper_train.get('batch_size', 32)
train_set = gluon.data.DataLoader(
dataset=dataset,
batch_size=batch_size,
shuffle=hyper_train.get('shuffle', True),
batchify_fn=_train_batchify_fn(max_gt_box_number),
last_batch='rollover',
num_workers=hyper_train.get('preprocess_threads', 4))
next_data_batch = next(iter(train_set))
setattr(train_set, 'reset', _reset)
setattr(train_set, 'provide_data', _get_provide_data(next_data_batch))
setattr(train_set, 'provide_label', _get_provide_label(
next_data_batch, (batch_size, max_gt_box_number, 4), is_train=True))
if val_list is not None and len(val_list) > 0:
assert 'index_file' in hyper_val and file.exists(
hyper_val['index_file']), 'not found label name file'
dataset = Detection_dataset(img_list=val_list,
index_file=hyper_val.get(
'index_file'),
width=hyper_val.get('width', 416),
height=hyper_val.get('height', 416),
is_train=False,
** kwargs)
max_gt_box_number = max([len(item) for item in dataset.label_cache])
batch_size = hyper_val.get('batch_size', 32)
val_set = gluon.data.DataLoader(
dataset=dataset,
batch_size=batch_size,
shuffle=hyper_val.get('shuffle', True),
batchify_fn=_val_batchify_fn(max_gt_box_number),
last_batch='keep',
num_workers=hyper_val.get('preprocess_threads', 4))
next_data_batch = next(iter(val_set))
setattr(val_set, 'reset', _reset)
setattr(val_set, 'provide_data', _get_provide_data(next_data_batch))
setattr(val_set, 'provide_label', _get_provide_label(
next_data_batch, is_train=False))
return train_set, val_set
def __getitem__(self, idx):
img_path = self.img_list[idx][0]
label = self.label_cache[idx]
img = file.read(img_path, binary=True)
img = mx_img.imdecode(img, 1)
return self.yolo_transform(img, label)