Skip to content

Commit b5f29cc

Browse files
pmeierdatumbox
andauthored
add tests for SBU dataset (#3464)
Co-authored-by: Vasilis Vryniotis <[email protected]>
1 parent 1e20b9d commit b5f29cc

File tree

1 file changed

+33
-0
lines changed

1 file changed

+33
-0
lines changed

test/test_datasets.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1122,5 +1122,38 @@ def _create_alphabet_folder(self, root, name):
11221122
return num_images_total
11231123

11241124

1125+
class SBUTestCase(datasets_utils.ImageDatasetTestCase):
1126+
DATASET_CLASS = datasets.SBU
1127+
FEATURE_TYPES = (PIL.Image.Image, str)
1128+
1129+
def inject_fake_data(self, tmpdir, config):
1130+
num_images = 3
1131+
1132+
dataset_folder = pathlib.Path(tmpdir) / "dataset"
1133+
images = datasets_utils.create_image_folder(tmpdir, "dataset", self._create_file_name, num_images)
1134+
1135+
self._create_urls_txt(dataset_folder, images)
1136+
self._create_captions_txt(dataset_folder, num_images)
1137+
1138+
return num_images
1139+
1140+
def _create_file_name(self, idx):
1141+
part1 = datasets_utils.create_random_string(10, string.digits)
1142+
part2 = datasets_utils.create_random_string(10, string.ascii_lowercase, string.digits[:6])
1143+
return f"{part1}_{part2}.jpg"
1144+
1145+
def _create_urls_txt(self, root, images):
1146+
with open(root / "SBU_captioned_photo_dataset_urls.txt", "w") as fh:
1147+
for image in images:
1148+
fh.write(
1149+
f"http://static.flickr.com/{datasets_utils.create_random_string(4, string.digits)}/{image.name}\n"
1150+
)
1151+
1152+
def _create_captions_txt(self, root, num_images):
1153+
with open(root / "SBU_captioned_photo_dataset_captions.txt", "w") as fh:
1154+
for _ in range(num_images):
1155+
fh.write(f"{datasets_utils.create_random_string(10)}\n")
1156+
1157+
11251158
if __name__ == "__main__":
11261159
unittest.main()

0 commit comments

Comments
 (0)