ccr-colorado / app23.py
tstone87's picture
Rename app.py to app23.py
85b39e8 verified
raw
history blame contribute delete
913 Bytes
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()