def test_load_imagenet_train(self):
imagenet_tf_record_data_path = '/localdata/datasets/imagenet-data'
# if the path doesn't exist the function will throw an error
ds, img_shape, num_examples, num_classes = DataGenerator.get_imagenet(
imagenet_tf_record_data_path, 'train')
self.assertTrue(isinstance(ds, tf.data.Dataset))
self.assertEqual(img_shape, (224, 224, 3))
self.assertEqual(num_examples, 1281167)
self.assertEqual(num_classes, 1000)
def test_build_imagenet_validation(self):
imagenet_original_data_path = '/localdata/datasets/imagenet-raw-data'
output_directory = '/tmp/temporary_imagenet_dataset_directory'
DataGenerator.build_imagenet_tf_record(
imagenet_original_data_path,
'validation',
output_directory=output_directory)
ds, img_shape, num_examples, num_classes = DataGenerator.get_imagenet(
output_directory, 'validation')
self.assertTrue(isinstance(ds, tf.data.Dataset))
self.assertEqual(img_shape, (224, 224, 3))
self.assertEqual(num_examples, 50000)
self.assertEqual(num_classes, 1000)
shutil.rmtree(output_directory)
