@@ -1122,5 +1122,38 @@ def _create_alphabet_folder(self, root, name):
1122
1122
return num_images_total
1123
1123
1124
1124
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
+
1125
1158
if __name__ == "__main__" :
1126
1159
unittest .main ()
0 commit comments