Brain_Tumor_Yolo / video_view.py
davron04's picture
Update video_view.py
f0ac3f9 verified
import streamlit as st
import cv2
import tempfile
import time
import os
class Video_View:
def __init__(self, app, model):
self.app = app
self.model = model
def toggle_video_processing(self):
"""Toggle video processing state."""
st.session_state.video_processed = False
def show(self):
# Top navigation
col1_back, col2_back = st.columns([0.2, 0.8])
with col1_back:
if st.button("Back", key='video_back', icon=':material/arrow_back:', type='primary'):
self.app.change_page("Main")
st.markdown("<h1 style='text-align: center;'>🧠 Video Detection</h1>", unsafe_allow_html=True)
st.divider()
uploaded_file = st.file_uploader("Upload a video", type=["mp4", "avi", "mov"],on_change=self.toggle_video_processing)
if not st.session_state.video_processed:
if uploaded_file is not None:
# Save to temp file
tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
tfile.write(uploaded_file.read())
cap = cv2.VideoCapture(tfile.name)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # Total number of frames
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
out_path = os.path.join(tempfile.gettempdir(), "predicted_video.mp4")
out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'XVID'), fps, (width, height))
frame_count = 0 # To track the number of processed frames
with st.spinner("Processing video... ⏳"):
progress_bar = st.progress(0) # Create a progress bar
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# Run YOLO model on the frame
results = self.model(frame)[0]
for result in results.boxes.data.tolist():
x1, y1, x2, y2, score, _ = result
color = (0, 0, 255) if score > 0.5 else (0, 255, 0)
label = f"{score:.2f}"
cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), color, 1)
cv2.putText(frame, label, (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1)
out.write(frame)
# Update the progress bar
frame_count += 1
progress_percentage = frame_count / total_frames # Calculate the percentage
progress_bar.progress(progress_percentage) # Update the progress bar
cap.release()
out.release()
cv2.destroyAllWindows()
st.session_state.video_processed = True
st.success("βœ… Detection complete!")
# Read and display the video
with open(out_path, 'rb') as video_file:
video_bytes = video_file.read()
st.video(uploaded_file, loop=True, autoplay=True, muted=False)
st.download_button("πŸ“₯ Download Predicted Video", video_bytes, file_name="predicted_video.mp4",
mime="video/mp4")