Spaces:
Sleeping
Sleeping
File size: 1,540 Bytes
c9280e3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import streamlit as st
from ui import upload_image
from utils import load_model, predict
# -------------------------------
# 1) Set the path to your saved model file:
# Change this to the correct path where you saved your .pth/.pt
# -------------------------------
MODEL_PATH = "./models/model.pth" # ← replace with your actual path
# -------------------------------
# 2) Cache the model load so it isn't reloaded on every run:
# -------------------------------
@st.cache_resource
def get_model():
"""
Load and cache the PyTorch model so that Streamlit does not reload it on every interaction.
"""
model = load_model(MODEL_PATH)
return model
# -------------------------------
# 3) Main Streamlit UI
# -------------------------------
def main():
# apply the styles.css here
with open("./styles.css") as f:
st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
# Load the model once
model = get_model()
# Let the user upload an image via ui.upload_image()
image = upload_image()
if image is not None:
# Only show the “Predict” button if an image has been uploaded
if st.button("Predict Drowsiness"):
# Run inference
label = predict(model, image)
# Display results
if label == 1:
st.error("🚨 Drowsiness Detected (1)")
else:
st.success("✅ Not Drowsy (0)")
if __name__ == "__main__":
main()
|