smoke / app.py
Harika12323's picture
init
17d51ee verified
raw
history blame contribute delete
No virus
916 Bytes
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.")