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)