Ely-testa's picture
Update app.py
ddf33a1 verified
raw
history blame
3.61 kB
from fastai.vision.all import *
import gradio as gr
import fal_client
from PIL import Image
import io
import random
import requests
from pathlib import Path
# Load your model
learn = load_learner('resnet50_30_categories.pkl')
# Wikipedia links
search_terms_wikipedia = {
"blazing star": "https://en.wikipedia.org/wiki/Mentzelia",
"bristlecone pine": "https://en.wikipedia.org/wiki/Pinus_longaeva",
# ... (same as before)
"goldfields coreopsis": "https://en.wikipedia.org/wiki/Coreopsis"
}
# Prompt templates for art generation
prompt_templates = [
"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."
]
# Example image paths (replace with actual paths on your system or Hugging Face space)
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'))
]
def on_queue_update(update):
if isinstance(update, fal_client.InProgress):
for log in update.logs:
print(log["message"])
else:
print("Received non-InProgress update:", update)
# Processing function
def process_image(img):
predicted_class, _, probs = learn.predict(img)
classification_results = dict(zip(learn.dls.vocab, map(float, probs)))
# Wikipedia
wiki_url = search_terms_wikipedia.get(predicted_class.lower(), "No Wikipedia entry found.")
# Generate image via FAL
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,
)
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
# Interface
with gr.Blocks() as demo:
gr.Markdown("# 🌼 Wildflower Classifier & Artistic Generator")
with gr.Row():
input_image = gr.Image(height=230, width=230, label="Upload an image", type="pil")
with gr.Row():
with gr.Column():
label_output = gr.Label(label="Prediction")
wiki_output = gr.Textbox(label="Wikipedia Link")
generated_image = gr.Image(label="AI Artistic Interpretation")
gr.Examples(
examples=example_images,
inputs=input_image,
examples_per_page=6
)
input_image.upload(
fn=process_image,
inputs=input_image,
outputs=[label_output, generated_image, wiki_output]
)
demo.launch()