app.py

from fastai.vision.all import *
import gradio as gr
import fal_client
from PIL import Image
import io
import random
import requests

# 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",
    "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"
}

# Templates for AI image generation
prompt_templates = [
    "A delicate watercolor painting of a {flower} with morning dew drops, soft pastel colors blending together as the sunrise creates gentle shadows.",
    "A whimsical nature journal sketch of a {flower} surrounded by butterflies and bees, with loose watercolor washes in fresh spring colors.",
    "An impressionistic scene of a {flower} swaying in the morning breeze, painted in loose brushstrokes with soft blues and golden morning light.",
    "A botanical illustration of a {flower} with detailed pencil lines and gentle watercolor washes, surrounded by notes about its features and habitat.",
    "A dreamy plein air painting of a {flower} along a morning hiking trail, with misty mountains in the background and soft morning colors."
]

# Example images for the interface
example_images = [
    'https://www.deserthorizonnursery.com/wp-content/uploads/2024/03/Brittlebush-Encelia-Farinosa-desert-horizon-nursery.jpg',
    'https://cdn.mos.cms.futurecdn.net/VJE7gSuQ9KWbkqEsWgX5zS.jpg',
    'https://www.parksconservancy.org/sites/default/files/styles/basic/public/A_GEN_131213_WTE_109.jpg?itok=9SDtr4b2',
    'https://silverfallsseed.com/wp-content/uploads/2016/01/tidy-tips-_-9.jpg',
    'https://valleywaternews.org/wp-content/uploads/2016/06/ceanothus-clusters.jpg?w=1440',
    'https://cdn11.bigcommerce.com/s-1b9100svju/images/stencil/1280x1280/products/2044/1440/DETA-635__75522.1664817787.jpg?c=1'
]

# 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"])

# 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)))
    
    # 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={
            "prompt": random.choice(prompt_templates).format(flower=predicted_class),
            "image_size": "portrait_4_3"
        },
        with_logs=True,
        on_queue_update=on_queue_update,
    )
    
    # 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

# Function to clear all outputs
def clear_outputs():
    return {
        label_output: None,
        generated_image: None,
        wiki_output: None
    }

# Load the AI model
learn = load_learner('export.pkl')

# Create the web interface
with gr.Blocks() as demo:
    # Input section
    with gr.Row():
        input_image = gr.Image(height=192, width=192, 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 Article Link", lines=1)
        generated_image = gr.Image(label="AI Generated Interpretation")
    
    # Add example images
    gr.Examples(
        examples=example_images,
        inputs=input_image,
        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]
    )

# Start the application
demo.launch(inline=False)