app.py

import base64
import gradio as gr
import random
from fastai.vision.all import *
from pathlib import Path
from PIL import Image
import io
import os
import requests

import fal_client

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

flowers_endangerment = {
    "blazing star": "Not considered endangered.",
    "bristlecone pine": "Least Concern (stable population).",
    "california bluebell": "Not listed as endangered or threatened.",
    "california buckeye": "Not endangered.",
    "california buckwheat": "Generally secure.",
    "california fuchsia": "Not endangered overall; some subspecies at risk.",
    "california checkerbloom": "Not generally endangered; some subspecies critically imperiled.",
    "california lilac": "Most species not endangered; some species are endangered.",
    "california poppy": "Generally secure; some subspecies face threats.",
    "california sagebrush": "Considered secure (G4-G5).",
    "california wild grape": "Apparently secure (G4).",
    "california wild rose": "Secure (G4).",
    "coyote mint": "Varies by species; some federally listed as endangered.",
    "elegant clarkia": "Secure (G5).",
    "baby blue eyes": "Secure.",
    "hummingbird sage": "Apparently secure (G4).",
    "delphinium": "Varies by species; some are endangered.",
    "matilija poppy": "Not currently endangered.",
    "blue-eyed grass": "Not endangered.",
    "penstemon spectabilis": "Not endangered.",
    "seaside daisy": "Not endangered.",
    "sticky monkeyflower": "Not endangered.",
    "tidy tips": "Generally not endangered; some subspecies may be at risk.",
    "wild cucumber": "Generally not endangered.",
    "douglas iris": "Not endangered.",
    "goldfields coreopsis": "Varies by species; many not endangered."
}


def get_status(flower_name):
    """Return the endangerment status of a given flower name."""
    return flowers_endangerment.get(flower_name, "Flower not found in database.")


# Templates for AI image generation
prompt_templates = [
    "A dreamy Picasso-like 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 crayon etching, 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.",
    "In a dim, graffiti-covered subway station, a vibrant spray-painted {flower} with neon petals and shimmering gold accents blooms defiantly on a grimy tiled wall, its dripping paint and curling vines injecting unexpected life into the underground gloom.",
    "Amid a bleak industrial wasteland strewn with rusted metal and trash, a single {flower} bathed in a shaft of golden light stands in stark contrast, its delicate form glowing with quiet resilience against the surrounding decay."
]

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

# Function to handle FAL queue updates
def on_queue_update(update):
    """Handle queue updates from FAL."""
    print(f"Queue update: {update}")


# Main function to process the uploaded image
def process_image(img, generate_image=True):
    print("Starting prediction...")
    predicted_class, _, probs = learn.predict(img)
    print(f"Prediction complete: {predicted_class}")

    # Format classification results for Gradio Label
    classification_results = {
        predicted_class: float(probs[learn.dls.vocab.o2i[predicted_class]])
    }

    # Get Wikipedia link
    wiki_url = search_terms_wikipedia.get(predicted_class, "No Wikipedia entry found.")

    # Get endangerment status
    endangerment_status = get_status(predicted_class)
    print(f"Status found: {endangerment_status}")

    # Generate artistic interpretation using FAL
    print("Sending request to FAL...")
    try:
        if generate_image:
            
            # Use fal-client to generate image
            result = fal_client.subscribe(
                "fal-ai/flux/dev",
                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 image from the result

            image_url = result["images"][0]["url"]
            response = requests.get(image_url)
            generated_image = Image.open(io.BytesIO(response.content))

        else:
            generated_image = None

    except Exception as e:
        print(f"Error generating image: {e}")
        generated_image = None

    print("Image retrieved and ready to return")
    return classification_results, generated_image, wiki_url, endangerment_status


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


# 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(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 Article Link", lines=1)
            endangerment_output = gr.Textbox(label="Endangerment Status", 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,
        fn=None,
        outputs=None
    )

    input_image.change(
        fn=lambda img: process_image(img, generate_image=True),
        inputs=input_image,
        outputs=[label_output, generated_image, wiki_output, endangerment_output]
    )

    input_image.clear(
        fn=clear_outputs,
        inputs=[],
        outputs=[label_output, generated_image, wiki_output, endangerment_output]
    )

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