def load_wpif(self, dataset_dir):
"""Load a subset of the COCO dataset.
dataset_dir: The root directory of the dataset.
"""
# get all labeled json file under dataset
examples_list = []
for example in os.listdir(dataset_dir):
if example[-5:] == '.json':
examples_list.append(os.path.join(dataset_dir, example))
# Add classes
self.add_class("WPIF", self.class_names.index('pad'), 'pad')
for example_name in examples_list:
label_file = LabelFile()
label_file.load(example_name, dataset_dir)
# Add images
image_path = label_file.imagePath
image_name = os.path.split(image_path)[-1]
image_id = int(image_name.split('.')[0][6:12])
width = label_file.width
height = label_file.height
annotations = []
for label, points, line_color, fill_color in label_file.shapes:
annotations.append([label, points])
self.add_image("WPIF", image_id=image_id,
path=image_path,
width=width, height=height,
annotations=annotations)
