Upload app.py

app.py

@@ -21,6 +21,7 @@ main_language_values = sorted([[name, code] for code, name in language_names.ite
 
 babel_imagenet = json.load(open("data/babel_imagenet-298.json", encoding="utf-8"))
 babelnet_images = json.load(open("data/images.json", encoding="utf-8"))
+max_image_choices = 10   # Currently up to 30 images but relevance degrades quickly in my experience. Limiting to 10
 no_image_idxs = [i for i, imgs in enumerate(babelnet_images) if len(imgs) == 0]
 IMG_HEIGHT, IMG_WIDTH = 512, 512
 
@@ -95,7 +96,7 @@ def prepare(raw_idx, lang, text_embeddings, class_order, randomize_images):
 
     img_idx = 0
     if randomize_images:
-        img_idx = np.random.choice(len(babelnet_images[class_idx]))
+        img_idx = np.random.choice(min(len(babelnet_images[class_idx]), max_image_choices))
     img_url = babelnet_images[class_idx][img_idx]["url"]
     class_labels = babel_imagenet[lang][1] if lang != "EN" else openai_en_classes
 
@@ -144,7 +145,7 @@ def reroll(raw_idx, lang, text_embeddings, class_order, randomize_images):
 
     img_idx = 0
     if randomize_images:
-        img_idx = np.random.choice(len(babelnet_images[class_idx]))
+        img_idx = np.random.choice(min(len(babelnet_images[class_idx]), max_image_choices))
     img_url = babelnet_images[class_idx][img_idx]["url"]
     class_labels = babel_imagenet[lang][1] if lang != "EN" else openai_en_classes