cat-vs-dog-test / app.py
hjay's picture
Create app.py
375668f verified
raw
history blame contribute delete
929 Bytes
import gradio as gr
import requests
from PIL import Image
from io import BytesIO
from transformers import AutoImageProcessor, AutoModelForImageClassification
preprocessor = AutoImageProcessor.from_pretrained("hjay/autotrain-z7ygf-g8xy8")
model = AutoModelForImageClassification.from_pretrained("hjay/autotrain-z7ygf-g8xy8")
def predict(img_url):
response = requests.get(img_url)
input_img = Image.open(BytesIO(response.content))
inputs = preprocessor(images=input_img, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
return input_img, model.config.id2label[predicted_class_idx]
gradio_app = gr.Interface(
predict,
inputs="textbox",
outputs=[gr.Image(label="Processed Image"), gr.Label(label="Result", num_top_classes=2)],
title="Cat Or Dog?",
)
if __name__ == "__main__":
gradio_app.launch()