def test_pascal_voc_cache_writer(self):
images_dir, annotations_dir, label_map = test_util.create_pascal_voc(
self.get_temp_dir())
cache_writer = dataloader_util.PascalVocCacheFilesWriter(label_map,
images_dir,
num_shards=1)
cache_files = dataloader_util.get_cache_files(
cache_dir=self.get_temp_dir(), cache_prefix_filename='pascal')
cache_writer.write_files(cache_files, annotations_dir)
# Checks the TFRecord file.
tfrecord_files = cache_files.tfrecord_files
self.assertTrue(os.path.isfile(tfrecord_files[0]))
self.assertGreater(os.path.getsize(tfrecord_files[0]), 0)
# Checks the annotation json file.
annotations_json_file = cache_files.annotations_json_file
self.assertTrue(os.path.isfile(annotations_json_file))
self.assertGreater(os.path.getsize(annotations_json_file), 0)
expected_json_file = test_util.get_test_data_path('annotations.json')
self.assertTrue(filecmp.cmp(annotations_json_file, expected_json_file))
# Checks the meta_data file.
meta_data_file = cache_files.meta_data_file
self.assertTrue(os.path.isfile(meta_data_file))
self.assertGreater(os.path.getsize(meta_data_file), 0)
with tf.io.gfile.GFile(meta_data_file, 'r') as f:
meta_data_dict = yaml.load(f, Loader=yaml.FullLoader)
self.assertEqual(meta_data_dict['size'], 1)
self.assertEqual(meta_data_dict['label_map'], label_map)
# Checks xml_dict from `_get_xml_dict` function.
xml_dict = next(cache_writer._get_xml_dict(annotations_dir))
expected_xml_dict = {
'filename':
'2012_12.jpg',
'folder':
'',
'object': [{
'difficult': '1',
'bndbox': {
'xmin': '64',
'ymin': '64',
'xmax': '192',
'ymax': '192',
},
'name': 'person',
'truncated': '0',
'pose': '',
}],
'size': {
'width': '256',
'height': '256',
}
}
self.assertEqual(xml_dict, expected_xml_dict)
def test_get_cache_files(self):
cache_files = dataloader_util.get_cache_files(
cache_dir='/tmp/', cache_prefix_filename='train', num_shards=1)
self.assertEqual(cache_files.cache_prefix, '/tmp/train')
self.assertLen(cache_files.tfrecord_files, 1)
self.assertEqual(cache_files.tfrecord_files[0],
'/tmp/train-00000-of-00001.tfrecord')
self.assertEqual(cache_files.meta_data_file,
'/tmp/train_meta_data.yaml')
self.assertEqual(cache_files.annotations_json_file,
'/tmp/train_annotations.json')
def from_pascal_voc(
cls,
images_dir: str,
annotations_dir: str,
label_map: Union[List[str], Dict[int, str], str],
annotation_filenames: Optional[Collection[str]] = None,
ignore_difficult_instances: bool = False,
num_shards: int = 100,
max_num_images: Optional[int] = None,
cache_dir: Optional[str] = None,
cache_prefix_filename: Optional[str] = None) -> DetectorDataLoader:
"""Loads from dataset with PASCAL VOC format.
Refer to
https://towardsdatascience.com/coco-data-format-for-object-detection-a4c5eaf518c5#:~:text=Pascal%20VOC%20is%20an%20XML,for%20training%2C%20testing%20and%20validation
for the description of PASCAL VOC data format.
LabelImg Tool (https://github.com/tzutalin/labelImg) can annotate the image
and save annotations as XML files in PASCAL VOC data format.
Annotations are in the folder: ${annotations_dir}.
Raw images are in the foloder: ${images_dir}.
Args:
images_dir: Path to directory that store raw images.
annotations_dir: Path to the annotations directory.
label_map: Variable shows mapping label integers ids to string label
names. 0 is the reserved key for `background`. Label names can't be
duplicated. Supported format: 1. Dict, map label integers ids to string
label names, e.g.
{1: 'person', 2: 'notperson'}. 2. List, a list of label names. e.g.
['person', 'notperson'] which is
the same as setting label_map={1: 'person', 2: 'notperson'}.
3. String, name for certain dataset. Accepted values are: 'coco', 'voc'
and 'waymo'. 4. String, yaml filename that stores label_map.
annotation_filenames: Collection of annotation filenames (strings) to be
loaded. For instance, if there're 3 annotation files [0.xml, 1.xml,
2.xml] in `annotations_dir`, setting annotation_filenames=['0', '1']
makes this method only load [0.xml, 1.xml].
ignore_difficult_instances: Whether to ignore difficult instances.
`difficult` can be set inside `object` item in the annotation xml file.
num_shards: Number of shards for output file.
max_num_images: Max number of imags to process.
cache_dir: The cache directory to save TFRecord, metadata and json file.
When cache_dir is not set, a temporary folder will be created and will
not be removed automatically after training which makes it can be used
later.
cache_prefix_filename: The cache prefix filename. If not set, will
automatically generate it based on `image_dir`, `annotations_dir` and
`annotation_filenames`.
Returns:
ObjectDetectorDataLoader object.
"""
label_map = _get_label_map(label_map)
# If `cache_prefix_filename` is None, automatically generates a hash value.
if cache_prefix_filename is None:
cache_prefix_filename = util.get_cache_prefix_filename_from_pascal(
images_dir=images_dir,
annotations_dir=annotations_dir,
annotation_filenames=annotation_filenames,
num_shards=num_shards)
cache_files = util.get_cache_files(
cache_dir=cache_dir,
cache_prefix_filename=cache_prefix_filename,
num_shards=num_shards)
# If not cached, writes data into tfrecord_file_paths and
# annotations_json_file_path.
# If `num_shards` differs, it's still not cached.
if not util.is_cached(cache_files):
cache_writer = util.PascalVocCacheFilesWriter(
label_map=label_map,
images_dir=images_dir,
num_shards=num_shards,
max_num_images=max_num_images,
ignore_difficult_instances=ignore_difficult_instances)
cache_writer.write_files(
cache_files=cache_files,
annotations_dir=annotations_dir,
annotation_filenames=annotation_filenames)
return cls.from_cache(cache_files.cache_prefix)
def test_csv_cache_writer(self):
label_map = {1: 'Baked Goods', 2: 'Cheese', 3: 'Salad'}
images_dir = test_util.get_test_data_path('salad_images')
cache_writer = dataloader_util.CsvCacheFilesWriter(label_map,
images_dir,
num_shards=1)
csv_file = test_util.get_test_data_path('salads_ml_use.csv')
for set_name, size in [('TRAIN', 1), ('TEST', 2)]:
with tf.io.gfile.GFile(csv_file, 'r') as f:
lines = [
line for line in csv.reader(f)
if line[0].startswith(set_name)
]
cache_files = dataloader_util.get_cache_files(
cache_dir=self.get_temp_dir(), cache_prefix_filename='csv')
cache_writer.write_files(cache_files, lines)
# Checks the TFRecord file.
tfrecord_files = cache_files.tfrecord_files
self.assertTrue(os.path.isfile(tfrecord_files[0]))
self.assertGreater(os.path.getsize(tfrecord_files[0]), 0)
# Checks the annotation json file.
annotations_json_file = cache_files.annotations_json_file
self.assertTrue(os.path.isfile(annotations_json_file))
self.assertGreater(os.path.getsize(annotations_json_file), 0)
expected_json_file = test_util.get_test_data_path(
set_name.lower() + '_annotations.json')
self.assertTrue(
filecmp.cmp(annotations_json_file, expected_json_file))
# Checks the meta_data file.
meta_data_file = cache_files.meta_data_file
self.assertTrue(os.path.isfile(meta_data_file))
self.assertGreater(os.path.getsize(meta_data_file), 0)
with tf.io.gfile.GFile(meta_data_file, 'r') as f:
meta_data_dict = yaml.load(f, Loader=yaml.FullLoader)
self.assertEqual(meta_data_dict['size'], size)
self.assertEqual(meta_data_dict['label_map'], label_map)
# Checks xml_dict from `_get_xml_dict` function.
xml_dict = next(cache_writer._get_xml_dict(lines))
expected_xml_dict = {
'filename':
'279324025_3e74a32a84_o.jpg',
'object': [{
'name': 'Baked Goods',
'bndbox': {
'xmin': 9.1888,
'ymin': 101.982,
'xmax': 908.0176,
'ymax': 882.8832,
'name': 'Baked Goods',
},
'difficult': 0,
'truncated': 0,
'pose': 'Unspecified',
}],
'size': {
'width': 1600,
'height': 1200,
}
}
self.assertEqual(xml_dict, expected_xml_dict)