DeepActionPotential/DrowSeeAi

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