Spaces:
Sleeping
Sleeping
File size: 2,303 Bytes
3186ef0 c071f67 3186ef0 a0e83fb 3186ef0 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | import streamlit as st
import tensorflow as tf
import numpy as np
import cv2
import os
from tensorflow.keras.models import load_model
# Load the trained model
@st.cache_resource
def load_trained_model():
model = load_model("model.h5") # Updated model name
return model
# Preprocess a video to extract frames and prepare for prediction
def preprocess_video(video_path, img_height=224, img_width=224):
cap = cv2.VideoCapture(video_path)
frames = []
while True:
ret, frame = cap.read()
if not ret:
break
# Resize frame to match model input size
frame = cv2.resize(frame, (img_height, img_width))
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Convert BGR to RGB
frames.append(frame)
cap.release()
frames = np.array(frames, dtype=np.float32) / 255.0 # Normalize
return frames
# Predict from video frames
def predict_video(model, frames):
predictions = model.predict(frames, batch_size=32) # Batch size for efficient inference
avg_prediction = np.mean(predictions, axis=0) # Average over frames
class_idx = np.argmax(avg_prediction) # Get the class index
return class_idx, avg_prediction
# App UI
st.title("Driver Distraction Detection")
st.write("Team 18 Video Project: Sayandip Bhattacharyya, Purnendu Rudrapal, Sridatta Das, Sidhartha Karjee")
# File uploader
uploaded_file = st.file_uploader("Upload a video", type=["mp4", "avi", "mov"])
# Classes (Ensure these match the class names from training)
class_names = ["Microsleep", "Yawning"]
if uploaded_file:
# Save uploaded video temporarily
temp_video_path = "temp_video.mp4"
with open(temp_video_path, "wb") as f:
f.write(uploaded_file.read())
st.video(temp_video_path)
st.write("Processing video...")
# Load model
model = load_trained_model()
# Preprocess video
frames = preprocess_video(temp_video_path)
st.write(f"Extracted {len(frames)} frames for prediction.")
# Make prediction
class_idx, avg_prediction = predict_video(model, frames)
# Display result
st.write(f"**Prediction:** {class_names[class_idx]}")
st.write(f"Confidence: {avg_prediction[class_idx] * 100:.2f}%")
# Cleanup temporary file
os.remove(temp_video_path)
|