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()