File size: 3,750 Bytes
d453003 89a037c d453003 89a037c d453003 |
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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import streamlit as st
import cv2
import tempfile
import numpy as np
import torch
from collections import deque
from transformers import AutoFeatureExtractor, AutoModelForVideoClassification
from streamlit_webrtc import webrtc_streamer, VideoTransformerBase, RTCConfiguration, WebRtcMode
# Constants
NUM_FRAMES = 16
MODEL_NAME = "jatinmehra/Accident-Detection-using-Dashcam"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@st.cache_resource
def load_model_and_extractor():
extractor = AutoFeatureExtractor.from_pretrained("facebook/timesformer-base-finetuned-k400")
model = AutoModelForVideoClassification.from_pretrained(
MODEL_NAME,
num_labels=2,
ignore_mismatched_sizes=True
).to(DEVICE)
model.eval()
return extractor, model
extractor, model = load_model_and_extractor()
st.title("Dashcam Accident Predictor")
st.write("**higher score = higher accident probability**")
# Function to run inference on a saved video file
def run_inference_on_video(video_path):
cap = cv2.VideoCapture(video_path)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
fps = cap.get(cv2.CAP_PROP_FPS) or 30
if total_frames <= 0:
st.error("Failed to read video frames.")
return None
# Uniform sampling
indices = np.linspace(0, total_frames-1, NUM_FRAMES, dtype=int)
frames = []
for idx in indices:
cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
ret, frame = cap.read()
if not ret:
frames.append(np.zeros((224,224,3), dtype=np.uint8))
else:
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
resized = cv2.resize(rgb, (224,224))
frames.append(resized)
cap.release()
# Preprocess and predict
inputs = extractor(frames, return_tensors="pt")
pixel_values = inputs['pixel_values'].to(DEVICE)
with torch.no_grad():
outputs = model(pixel_values=pixel_values).logits
prob = torch.softmax(outputs, dim=1)[0,1].item()
return prob
# UI Selection
source = st.radio("Choose input source", ("Upload Video", "Webcam"))
if source == "Upload Video":
uploaded_file = st.file_uploader("Upload a video file", type=["mp4", "avi", "mov"])
if uploaded_file is not None:
tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
tfile.write(uploaded_file.read())
st.video(uploaded_file)
st.write("Running inference...")
score = run_inference_on_video(tfile.name)
if score is not None:
st.success(f"Accident probability: {score:.2f}")
else:
# Webcam stream processing
class AcciTransformer(VideoTransformerBase):
def __init__(self):
self.buffer = deque(maxlen=NUM_FRAMES)
def transform(self, frame):
img = frame.to_ndarray(format="bgr24")
rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
resized = cv2.resize(rgb, (224,224))
self.buffer.append(resized)
if len(self.buffer) == NUM_FRAMES:
inputs = extractor(list(self.buffer), return_tensors="pt")
pixel_values = inputs['pixel_values'].to(DEVICE)
with torch.no_grad():
outputs = model(pixel_values=pixel_values).logits
prob = torch.softmax(outputs, dim=1)[0,1].item()
cv2.putText(img, f"Prob: {prob:.2f}", (10,30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2)
return img
webrtc_streamer(
key="dashcam-webcam",
mode=WebRtcMode.RECVONLY,
rtc_configuration=RTCConfiguration({
"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]
}),
video_transformer_factory=AcciTransformer
) |