Spaces:
Sleeping
Sleeping
File size: 4,102 Bytes
bef9145 7423ca4 bef9145 7423ca4 bef9145 dedf935 bef9145 dedf935 bef9145 dedf935 bef9145 9c0d975 bef9145 f355b88 bef9145 | 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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import os
import time
import tempfile
import cv2
import streamlit as st
from ultralytics import YOLO
from huggingface_hub import hf_hub_url, cached_download
@st.cache_resource
def load_model():
repo_id = "BreakIntoData/cv_workshop"
filename = "soccer_ball.pt"
# Create a URL for the model file on the Hugging Face Hub
model_url = hf_hub_url(repo_id, filename)
# Download the model file from the Hub and cache it locally
cached_model_path = cached_download(model_url)
# Rename the file to have a .pt extension
new_cached_model_path = f"{cached_model_path}.pt"
os.rename(cached_model_path, new_cached_model_path)
print(f"Downloaded model to {new_cached_model_path}")
# Load the model using YOLO from the cached model file
return YOLO(new_cached_model_path)
def process_video(video_path, output_path):
cap = cv2.VideoCapture(video_path)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
fourcc = cv2.VideoWriter_fourcc(*'avc1')
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
progress_text = "Processing video... Please wait."
progress_bar = st.progress(0)
status_text = st.empty()
time_text = st.empty()
start_time = time.time()
for frame_num in range(total_frames):
ret, frame = cap.read()
if not ret:
break
results = model(frame)
annotated_frame = results[0].plot()
out.write(annotated_frame)
# Update progress
progress = (frame_num + 1) / total_frames
elapsed_time = time.time() - start_time
estimated_total_time = elapsed_time / progress
remaining_time = estimated_total_time - elapsed_time
progress_bar.progress(progress)
status_text.text(f"Processing frame {frame_num+1}/{total_frames}")
time_text.text(f"Elapsed time: {elapsed_time:.2f}s | Estimated time remaining: {remaining_time:.2f}s")
cap.release()
out.release()
progress_bar.empty()
status_text.text(f"Processed {total_frames} frames")
time_text.text(f"Total time: {time.time() - start_time:.2f}s")
model = load_model()
st.title("Soccer Ball Detection App")
# Sidebar for options
st.sidebar.header("Options")
video_option = st.sidebar.radio("Choose video source:", ("Use preset video", "Upload video"))
if video_option == "Upload video":
uploaded_file = st.sidebar.file_uploader("Choose 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())
video_path = tfile.name
else:
preset_videos = {
"Ronaldo": "preset_videos/Ronaldo.mp4",
"Sancho": "preset_videos/CityUtdR video.mp4",
"Messi": "preset_videos/Messi.mp4",
}
selected_video = st.sidebar.selectbox("Select a preset video", list(preset_videos.keys()))
video_path = preset_videos[selected_video]
if 'video_path' in locals():
st.header("Original Video")
st.video(video_path)
if st.button("Detect"):
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
with st.spinner("Processing video..."):
process_video(video_path, temp_file.name)
st.success("Video processing complete!")
st.header("Processed Video")
st.video(temp_file.name)
# Add download link
with open(temp_file.name, "rb") as file:
btn = st.download_button(
label="Download Video",
data=file,
file_name="processed_video.mp4",
mime="video/mp4"
)
# # Clean up temporary files
# os.unlink(temp_file.name)
if video_option == "Upload video":
os.unlink(video_path)
st.sidebar.markdown("---")
st.sidebar.write("Developed with ❤️ by Break Into Data")
|