CalamityJon's picture
Rename animal_image_app.py to app.py
66b8d3c verified
import gradio as gr
from transformers import pipeline
from PIL import Image
# Load image classification model
# Using a pre-trained model that can classify various animals and objects
classifier = pipeline("image-classification", model="google/vit-base-patch16-224")
def classify_image(image):
"""
Classify an uploaded animal image and return top predictions with progress bars
"""
if image is None:
return "<div style='font-family: Arial, sans-serif; padding: 20px; text-align: center; color: #666;'>Please upload an image.</div>"
# Classify the image
results = classifier(image)
# Format results with HTML progress bars - show top 5 predictions
html_results = "<div style='font-family: Arial, sans-serif;'>"
html_results += "<h3 style='margin-top: 0;'>Top Predictions:</h3>"
for i, result in enumerate(results[:5], 1):
label = result['label']
score = result['score'] * 100
score_int = int(score)
# Create progress bar with color gradient (green for high, yellow for medium, red for low)
if score_int >= 70:
bar_color = "#4CAF50" # Green
elif score_int >= 40:
bar_color = "#FF9800" # Orange
else:
bar_color = "#F44336" # Red
html_results += f"""
<div style='margin-bottom: 15px;'>
<div style='display: flex; justify-content: space-between; margin-bottom: 5px;'>
<span style='font-weight: bold;'>{i}. {label}</span>
<span style='font-weight: bold; color: #333;'>{score:.2f}%</span>
</div>
<div style='background-color: #e0e0e0; border-radius: 10px; height: 25px; overflow: hidden;'>
<div style='background-color: {bar_color}; height: 100%; width: {score:.2f}%; transition: width 0.3s ease; border-radius: 10px;'></div>
</div>
</div>
"""
html_results += "</div>"
return html_results
# Create the Gradio interface
with gr.Blocks(title="Animal Image Classifier") as demo:
gr.Markdown("# Animal Image Classifier")
gr.Markdown("Upload an animal photo to classify it using AI!")
with gr.Row():
with gr.Column():
# Image input
input_image = gr.Image(
type="pil",
label="Upload Animal Photo"
)
# Classify button
classify_btn = gr.Button("Classify Image", variant="primary", size="lg")
clear_btn = gr.Button("Clear", variant="secondary")
with gr.Column():
# Output for classification results with HTML progress bars
output_html = gr.HTML(
label="Classification Results"
)
# Example images at the bottom
gr.Markdown("### Example Images")
gr.Markdown("Try these example images:")
example_images = [
"cat.png",
"frog.png",
"hippo.png",
"jaguar.png",
"sloth.png",
"toucan.png",
"turtle.png"
]
# Create example gallery - images are in the same directory as this script
import os
script_dir = os.path.dirname(os.path.abspath(__file__))
example_paths = [[os.path.join(script_dir, img)] for img in example_images]
gr.Examples(
examples=example_paths,
inputs=input_image,
label="Click on an example image to load it"
)
# Define button actions
classify_btn.click(
fn=classify_image,
inputs=input_image,
outputs=output_html
)
clear_btn.click(
fn=lambda: (None, "<div></div>"),
inputs=None,
outputs=[input_image, output_html]
)
# Launch the app
if __name__ == "__main__":
demo.launch()