fixed prediction, removed example images

app.py

@@ -3,7 +3,6 @@ import gradio as gr
 import fal_client
 from PIL import Image
 import io
-import base64
 import random
 import requests
 
@@ -52,11 +51,10 @@ def on_queue_update(update):
 
 def process_image(img):
     # First do the classification
-    pred, idx, probs = learn.predict(img)
-    classification_results = dict(zip(search_terms_wikipedia.keys(), map(float, probs)))
+    predicted_class, idx, probs = learn.predict(img)
+    classification_results = dict(zip(learn.dls.vocab, map(float, probs)))
     
     # Get Wikipedia URL for the predicted class
-    predicted_class = max(classification_results.items(), key=lambda x: x[1])[0]
     wiki_url = search_terms_wikipedia.get(predicted_class, "No Wikipedia entry found.")
     
     # Generate FLUX image
@@ -64,7 +62,7 @@ def process_image(img):
         "fal-ai/flux/schnell",
         arguments={
             "prompt": random.choice(prompt_templates).format(flower=predicted_class),
-            "image_size": "square"
+            "image_size": "portrait_4_3"
         },
         with_logs=True,
         on_queue_update=on_queue_update,
@@ -101,7 +99,7 @@ with gr.Blocks() as demo:
     gr.Examples(
         examples=examples,
         inputs=input_image,
-        examples_per_page=5,
+        examples_per_page=6,
         fn=process_image,
         outputs=[label_output, generated_image, wiki_output]
     )