Spaces:
Sleeping
Sleeping
File size: 3,651 Bytes
1acdef8 |
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 |
import streamlit as st
import os
import cv2
import time
import tempfile
from ultralytics import YOLO
from huggingface_hub import hf_hub_url, cached_download
@st.cache_resource
def load_model():
repo_id = 'navoditamathur/Soccer_yolo'
model_filename = 'soccer_ball.pt'
# Create a URL for the model file on the Hugging Face Hub
model_url = hf_hub_url(repo_id, model_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(*'mp4v')
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
progress_text = "Please wait..."
progress_bar = st.progress(0)
progress_bar.text(progress_text)
status_text = st.empty()
time_text = st.empty()
start_time = time.time()
for i in range(total_frames):
ret, frame = cap.read()
if not ret:
break
boxes = model(frame)
annotated_frame = boxes[0].plot()
out.write(annotated_frame)
progress = (i + 1) / total_frames
progress_bar.progress(progress)
elasped_time = time.time() - start_time
time_per_frame = elasped_time / (i + 1)
remaining_time = (total_frames - (i + 1)) * time_per_frame
status_text.text(f"Processing frame {i + 1} of {total_frames}")
time_text.text(f"Time remaining: {remaining_time:.2f} seconds")
cap.release()
out.release()
status_text.text("Video processing completed.")
progress_bar.empty()
time_text.empty()
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 = {
"Soccer Video": "preset_videos/soccer.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")
output_path = temp_file.name
process_video(video_path, output_path)
with open(output_path, 'rb') as video_file:
video_bytes = video_file.read()
st.header("Detected Video")
# Debugging: Display video size
st.write(f"Processed video size: {len(video_bytes)} bytes")
if len(video_bytes) > 0:
st.video(video_bytes)
# Generate a download button
btn = st.download_button(
label="Download Processed Video",
data=video_bytes,
file_name="processed_video.mp4",
mime="video/mp4"
)
|