Spaces:
Running
Running
import os | |
import sys | |
import tempfile | |
import cv2 | |
import numpy as np | |
import streamlit as st | |
# Append src directory to import custom modules | |
sys.path.append(os.path.join(os.path.dirname(__file__), 'src')) | |
# Environment fixes for Hugging Face Spaces | |
os.environ["STREAMLIT_BROWSER_GATHER_USAGE_STATS"] = "false" | |
os.environ["XDG_CONFIG_HOME"] = "/tmp" | |
os.environ["HF_HUB_CACHE"] = "/tmp/huggingface" | |
# Local utility imports | |
from preprocess import preprocess_frame | |
from predict import run_prediction, load_trained_model | |
# Streamlit setup | |
st.set_page_config(layout="wide") | |
st.title("π Violence Detection in Video") | |
st.markdown("Upload a video and let the model detect violent scenes in real-time.") | |
# Load model with proper feedback | |
st.info("π Loading model...") | |
try: | |
model = load_trained_model() | |
st.success("β Model loaded successfully.") | |
except Exception as e: | |
st.error(f"β Model loading failed: {e}") | |
st.stop() | |
# Upload video section | |
uploaded_file = st.file_uploader("π€ Upload a video", type=["mp4", "avi", "mpeg", "mov", "mpg"]) | |
if uploaded_file is not None: | |
try: | |
st.info("π₯ Reading and saving uploaded file...") | |
# Read file into memory first to avoid Hugging Face frontend 403 issues | |
file_bytes = uploaded_file.read() | |
if not file_bytes: | |
st.error("β Uploaded file could not be read.") | |
st.stop() | |
with tempfile.NamedTemporaryFile(delete=False, dir='/tmp', suffix='.mp4') as tfile: | |
tfile.write(file_bytes) | |
video_path = tfile.name | |
st.success("β Video saved. Loading...") | |
cap = cv2.VideoCapture(video_path) | |
if not cap.isOpened(): | |
st.error("β Could not open video. Please try another format.") | |
st.stop() | |
stframe = st.empty() | |
frame_count = 0 | |
while cap.isOpened(): | |
ret, frame = cap.read() | |
if not ret: | |
break | |
frame_count += 1 | |
processed = preprocess_frame(frame) | |
pred = run_prediction(model, processed) | |
label = "Violent" if pred <= 0.5 else "Non-Violent" | |
color = (0, 0, 255) if label == "Violent" else (0, 255, 0) | |
cv2.putText(frame, f'{label} ({pred:.2f})', (10, 30), | |
cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2) | |
stframe.image(frame, channels="BGR") | |
cap.release() | |
st.success(f"β Done! {frame_count} frames processed.") | |
except Exception as e: | |
st.error(f"β An unexpected error occurred:\n\n{e}") | |