|
|
import gradio as gr |
|
|
from PIL import Image |
|
|
import requests |
|
|
from io import BytesIO |
|
|
import os |
|
|
from transformers import pipeline |
|
|
|
|
|
|
|
|
classifier = pipeline("image-classification", model="Docty/solacies") |
|
|
|
|
|
def classify_image(input_img=None, img_url=None): |
|
|
""" |
|
|
Accepts either an uploaded image (input_img) or an image URL (img_url). |
|
|
""" |
|
|
if img_url: |
|
|
try: |
|
|
response = requests.get(img_url) |
|
|
response.raise_for_status() |
|
|
img = Image.open(BytesIO(response.content)).convert("RGB") |
|
|
except Exception as e: |
|
|
return {"Error": f"Failed to load image from URL: {e}"} |
|
|
elif input_img: |
|
|
if not isinstance(input_img, Image.Image): |
|
|
img = Image.fromarray(input_img) |
|
|
else: |
|
|
img = input_img |
|
|
else: |
|
|
return {"Error": "No image provided."} |
|
|
|
|
|
results = classifier(img) |
|
|
return {res["label"]: float(res["score"]) for res in results} |
|
|
|
|
|
|
|
|
theme = gr.themes.Soft( |
|
|
primary_hue="blue", |
|
|
secondary_hue="lime", |
|
|
neutral_hue="slate" |
|
|
) |
|
|
|
|
|
with gr.Blocks(theme=theme) as demo: |
|
|
gr.Markdown("## Image Classifier") |
|
|
gr.Markdown("Upload an image **or** enter an image URL to classify it using the model.") |
|
|
|
|
|
with gr.Row(): |
|
|
image_input = gr.Image(type="pil", label="Upload Image") |
|
|
url_input = gr.Textbox(label="Image URL (optional)", placeholder="Paste image URL here...") |
|
|
label_output = gr.Label(num_top_classes=3, label="Predictions") |
|
|
|
|
|
classify_btn = gr.Button("Classify Image", variant="primary") |
|
|
|
|
|
gr.Examples( |
|
|
examples=[f'./images_samples/{i}' for i in os.listdir('./images_samples')], |
|
|
inputs=image_input, |
|
|
outputs=label_output, |
|
|
fn=classify_image, |
|
|
cache_examples=False |
|
|
) |
|
|
|
|
|
classify_btn.click( |
|
|
fn=classify_image, |
|
|
inputs=[image_input, url_input], |
|
|
outputs=label_output |
|
|
) |
|
|
|
|
|
demo.launch(share=True) |
|
|
|
