Spaces:
Build error
Build error
import streamlit as st | |
from PIL import Image | |
import torch | |
pip install torch torchvision torchaudio | |
from transformers import AutoFeatureExtractor, AutoModelForImageClassification | |
# Load model and extractor | |
model = AutoModelForImageClassification.from_pretrained("best.pt") | |
extractor = AutoFeatureExtractor.from_pretrained("best.pt") | |
st.title('Smoke Detection App') | |
uploaded_image = st.file_uploader("Choose an image...", type="jpg") | |
if uploaded_image is not None: | |
image = Image.open(uploaded_image) | |
st.image(image, caption='Uploaded Image.', use_column_width=True) | |
st.write("") | |
st.write("Classifying...") | |
inputs = extractor(images=image, return_tensors="pt") | |
with torch.no_grad(): | |
outputs = model(**inputs) | |
predictions = torch.argmax(outputs.logits, dim=1) | |
if predictions.item() == 1: | |
st.write("Smoke detected!") | |
else: | |
st.write("No smoke detected.") | |