import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# Load the model and processor
processor = AutoImageProcessor.from_pretrained("prithivMLmods/Fire-Detection-Engine")
model = AutoModelForImageClassification.from_pretrained("prithivMLmods/Fire-Detection-Engine")

def predict(image):
    # Convert image to expected format
    image = Image.fromarray(image)
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = logits.argmax(-1).item()
    
    return f"Predicted class: {predicted_class}"

# Create Gradio app
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Textbox(),
    title="Fire Detection Engine",
    description="Upload an image to check for fire."
)

iface.launch()