implement ai generation and display it

app.py

@@ -1,34 +1,96 @@
 from fastai.vision.all import *
 import gradio as gr
 import fal_client
+from PIL import Image
+import io
+import base64
+
+search_terms_wikipedia = {
+    "blazing star": "https://en.wikipedia.org/wiki/Mentzelia",
+    "bristlecone pine": "https://en.wikipedia.org/wiki/Pinus_longaeva",
+    "california bluebell": "https://en.wikipedia.org/wiki/Phacelia_minor",
+    "california buckeye": "https://en.wikipedia.org/wiki/Aesculus_californica",
+    "california buckwheat": "https://en.wikipedia.org/wiki/Eriogonum_fasciculatum",
+    "california fuchsia": "https://en.wikipedia.org/wiki/Epilobium_canum",
+    "california checkerbloom": "https://en.wikipedia.org/wiki/Sidalcea_malviflora",
+    "california lilac": "https://en.wikipedia.org/wiki/Ceanothus",
+    "california poppy": "https://en.wikipedia.org/wiki/Eschscholzia_californica",
+    "california sagebrush": "https://en.wikipedia.org/wiki/Artemisia_californica",
+    "california wild grape": "https://en.wikipedia.org/wiki/Vitis_californica",
+    "california wild rose": "https://en.wikipedia.org/wiki/Rosa_californica",
+    "coyote mint": "https://en.wikipedia.org/wiki/Monardella",
+    "elegant clarkia": "https://en.wikipedia.org/wiki/Clarkia_unguiculata",
+    "baby blue eyes": "https://en.wikipedia.org/wiki/Nemophila_menziesii",
+    "hummingbird sage": "https://en.wikipedia.org/wiki/Salvia_spathacea",
+    "delphiniumr": "https://en.wikipedia.org/wiki/Delphinium",
+    "matilija poppy": "https://en.wikipedia.org/wiki/Romneya_coulteri",
+    "blue-eyed grass": "https://en.wikipedia.org/wiki/Sisyrinchium_bellum",
+    "penstemon spectabilis": "https://en.wikipedia.org/wiki/Penstemon_spectabilis",
+    "seaside daisy": "https://en.wikipedia.org/wiki/Erigeron_glaucus",
+    "sticky monkeyflower": "https://en.wikipedia.org/wiki/Diplacus_aurantiacus",
+    "tidy tips": "https://en.wikipedia.org/wiki/Layia_platyglossa",
+    "wild cucumber": "https://en.wikipedia.org/wiki/Marah_(plant)",
+    "douglas iris": "https://en.wikipedia.org/wiki/Iris_douglasiana",
+    "goldfields coreopsis": "https://en.wikipedia.org/wiki/Coreopsis"
+}
 
 def on_queue_update(update):
     if isinstance(update, fal_client.InProgress):
         for log in update.logs:
            print(log["message"])
 
-
-def classify_image(img):
+def process_image(img):
+    # First do the classification
     pred, idx, probs = learn.predict(img)
-    return dict(zip(categories, map(float, probs)))
-
+    classification_results = dict(zip(search_terms_wikipedia.keys(), 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
+    with gr.Status("Generating AI image based on classification..."):
+        result = fal_client.subscribe(
+            "fal-ai/flux/schnell",
+            arguments={
+                "prompt": f"A detailed, artistic interpretation of {predicted_class} flower in natural setting",
+                "image_size": "256x256"  # Low res for testing
+            },
+            with_logs=True,
+            on_queue_update=on_queue_update,
+        )
+        
+        # Convert the image data
+        image_data = base64.b64decode(result['image'])
+        generated_image = Image.open(io.BytesIO(image_data))
+        
+    return classification_results, generated_image, wiki_url
 
+# Load the learner
 learn = load_learner('export.pkl')
 
-categories = ('balsamroot', 'bladderpod', 'blazing star', 'bristlecone pine flowers', 'brittlebrush')
-
-image=gr.Image(height = 192, width = 192)
-label = gr.Label()
-examples = ['https://www.deserthorizonnursery.com/wp-content/uploads/2024/03/Brittlebush-Encelia-Farinosa-desert-horizon-nursery.jpg','https://cdn.mos.cms.futurecdn.net/VJE7gSuQ9KWbkqEsWgX5zS.jpg']
-intf = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples)
-intf.launch(inline=False)
+# Create Gradio interface
+with gr.Blocks() as demo:
+    with gr.Row():
+        input_image = gr.Image(height=192, width=192, label="Upload Image for Classification")
+    with gr.Row():
+        with gr.Column():
+            label_output = gr.Label(label="Classification Results")
+            wiki_output = gr.Textbox(label="Wikipedia Article Link", lines=1)
+        generated_image = gr.Image(label="AI Generated Interpretation")
+    
+    # Example images
+    examples = [
+        'https://www.deserthorizonnursery.com/wp-content/uploads/2024/03/Brittlebush-Encelia-Farinosa-desert-horizon-nursery.jpg',
+        'https://cdn.mos.cms.futurecdn.net/VJE7gSuQ9KWbkqEsWgX5zS.jpg'
+    ]
+    gr.Examples(examples, input_image)
+    
+    # Set up event handler
+    input_image.change(
+        fn=process_image,
+        inputs=input_image,
+        outputs=[label_output, generated_image, wiki_output]
+    )
 
-result = fal_client.subscribe(
-    "fal-ai/flux/schnell",
-    arguments={
-        "prompt": "Extreme close-up of a single tiger eye, direct frontal view. Detailed iris and pupil. Sharp focus on eye texture and color. Natural lighting to capture authentic eye shine and depth. The word \"FLUX\" is painted over it in big, white brush strokes with visible texture."
-    },
-    with_logs=True,
-    on_queue_update=on_queue_update,
-)
-print(result)
+demo.launch(inline=False)