Spaces:
Sleeping
Sleeping
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() | |