change to small dataset

app.py

@@ -56,12 +56,12 @@ if ims is not None:
         #     for r in retrieved_examples["image"]:
         #         st.image(r)
     
-if any(picks):
-    # st.write("Nearest butterflies:")
-    for i,pick in enumerate(picks):
-        if pick:
-            scores, retrieved_examples=dataset.get_nearest_examples('beit_embeddings', embed(ims[i]), k=5)
-            for r in retrieved_examples["image"]:
-                cols[i].image(r)
+    if any(picks):
+        # st.write("Nearest butterflies:")
+        for i,pick in enumerate(picks):
+            if pick:
+                scores, retrieved_examples=dataset.get_nearest_examples('beit_embeddings', embed(ims[i]), k=5)
+                for r in retrieved_examples["image"]:
+                    cols[i].image(r)
 
 

demo.py

@@ -2,12 +2,9 @@ import torch
 from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
 from datasets import load_dataset
 
-def get_train_data(dataset_name="ceyda/smithsonian_butterflies_transparent_cropped",data_limit=1000):
+def get_train_data(dataset_name="huggan/smithsonian_butterflies_subset"):
     dataset=load_dataset(dataset_name)
     dataset=dataset.sort("sim_score")
-    score_thresh = dataset["train"][data_limit]['sim_score']
-    dataset = dataset.filter(lambda x: x['sim_score'] < score_thresh)
-    dataset = dataset.map(lambda x: {'image' : x['image'].convert("RGB")})
     return dataset["train"]
 
 from transformers import BeitFeatureExtractor, BeitForImageClassification