Update app.py

app.py

@@ -7,10 +7,7 @@ import random
 import requests
 from pathlib import Path
 
-# Load model
-learn = load_learner('resnet50_30_categories.pkl')
-
-# Wikipedia links dictionary
+# Dictionary of plant names and their Wikipedia links
 search_terms_wikipedia = {
     "blazing star": "https://en.wikipedia.org/wiki/Mentzelia",
     "bristlecone pine": "https://en.wikipedia.org/wiki/Pinus_longaeva",
@@ -40,37 +37,41 @@ search_terms_wikipedia = {
     "goldfields coreopsis": "https://en.wikipedia.org/wiki/Coreopsis"
 }
 
-# Prompt templates for AI generation
+# Templates for AI image generation
 prompt_templates = [
-    "A dreamy watercolor scene of a {flower} on a misty morning trail...",
-    "A loose, expressive watercolor sketch of a {flower} in a wild meadow...",
-    "An artist's nature journal page featuring a detailed {flower} study...",
-    "A vibrant plein air painting of a {flower} patch along a coastal trail...",
-    "A whimsical mixed-media scene of a {flower} garden at sunrise..."
+    "A dreamy watercolor scene of a {flower} on a misty morning trail, with golden sunbeams filtering through towering redwoods, and a curious hummingbird hovering nearby.",
+    "A loose, expressive watercolor sketch of a {flower} in a wild meadow, surrounded by dancing butterflies and morning dew drops sparkling like diamonds in the dawn light.",
+    "An artist's nature journal page featuring a detailed {flower} study, with delicate ink lines and soft watercolor washes, complete with small sketches of bees and field notes in the margins.",
+    "A vibrant plein air painting of a {flower} patch along a coastal hiking trail, with crashing waves and rugged cliffs in the background, painted in bold, energetic brushstrokes.",
+    "A whimsical mixed-media scene of a {flower} garden at sunrise, combining loose watercolor washes with detailed botanical illustrations, featuring hidden wildlife and morning fog rolling through the valley."
 ]
 
-# Local example image paths
+# Example images (using local paths)
 example_images = [
-    str(Path("example_images/example_1.jpg")),
-    str(Path("example_images/example_2.jpg")),
-    str(Path("example_images/example_3.jpg")),
-    str(Path("example_images/example_4.jpg")),
-    str(Path("example_images/example_5.jpg")),
+    str(Path('example_images/example_1.jpg')),
+    str(Path('example_images/example_2.jpg')),
+    str(Path('example_images/example_3.jpg')),
+    str(Path('example_images/example_4.jpg')),
+    str(Path('example_images/example_5.jpg'))
+    
 ]
 
-# Logging for FAL client
+# Function to handle AI generation progress updates
 def on_queue_update(update):
     if isinstance(update, fal_client.InProgress):
         for log in update.logs:
-            print(log["message"])
+           print(log["message"])
 
-# Process image and return classification + AI-generated artwork + Wiki URL
+# Main function to process the uploaded image
 def process_image(img):
+    # Classify the image
     predicted_class, _, probs = learn.predict(img)
     classification_results = dict(zip(learn.dls.vocab, map(float, probs)))
-    wiki_url = search_terms_wikipedia.get(predicted_class.lower(), "No Wikipedia entry found.")
-
-    # Generate image via FAL API
+    
+    # Get Wikipedia link
+    wiki_url = search_terms_wikipedia.get(predicted_class, "No Wikipedia entry found.")
+    
+    # Generate artistic interpretation by calling the Flux API
     result = fal_client.subscribe(
         "fal-ai/flux/schnell",
         arguments={
@@ -80,35 +81,59 @@ def process_image(img):
         with_logs=True,
         on_queue_update=on_queue_update,
     )
-
-    image_url = result["images"][0]["url"]
+    
+    # Get the generated image
+    image_url = result['images'][0]['url']
     response = requests.get(image_url)
     generated_image = Image.open(io.BytesIO(response.content))
+    
+    return classification_results, generated_image, wiki_url
 
-    return classification_results, generated_image, str(wiki_url)
+# Function to clear all outputs
+def clear_outputs():
+    return {
+        label_output: None,
+        generated_image: None,
+        wiki_output: None
+    }
 
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# 🌼 California Native Plant Classifier & AI Art Generator")
+# Load the AI model
+learn = load_learner('resnet50_30_categories.pkl')
 
+# Create the web interface
+with gr.Blocks() as demo:
+    # Input section
     with gr.Row():
-        input_image = gr.Image(type="pil", label="Upload a Photo", height=250)
-
+        input_image = gr.Image(height=230, width=230, label="Upload Image for Classification", type="pil")
+    
+    # Output section
     with gr.Row():
         with gr.Column():
             label_output = gr.Label(label="Classification Results")
-            wiki_output = gr.Textbox(label="Wikipedia Link")
-        generated_image = gr.Image(label="AI-Generated Artistic Interpretation")
-
-    # Submit button to trigger image processing
-    submit_btn = gr.Button("Submit")
-    submit_btn.click(fn=process_image, inputs=input_image, outputs=[label_output, generated_image, wiki_output])
-
-    # Examples
+            wiki_output = gr.Textbox(label="Wikipedia Article Link", lines=1)
+        generated_image = gr.Image(label="AI Generated Interpretation")
+    
+    # Add example images using local paths
     gr.Examples(
         examples=example_images,
         inputs=input_image,
-        examples_per_page=6
+        examples_per_page=6,
+        fn=process_image,
+        outputs=[label_output, generated_image, wiki_output]
+    )
+    
+    # Set up what happens when an image is uploaded or removed
+    input_image.change(
+        fn=process_image,
+        inputs=input_image,
+        outputs=[label_output, generated_image, wiki_output]
+    )
+    
+    input_image.clear(
+        fn=clear_outputs,
+        inputs=[],
+        outputs=[label_output, generated_image, wiki_output]
     )
 
-demo.launch()
+# Start the application
+demo.launch(inline=False)