Spaces:
Running
Running
File size: 3,927 Bytes
5667a93 ad9abdd |
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 matplotlib
matplotlib.use('Agg')
import streamlit as st
import cv2
import numpy as np
from yolov5 import YOLOv5
from sort.sort import Sort
import tempfile
import shutil
from moviepy.editor import VideoFileClip, concatenate_videoclips, ImageSequenceClip
import os
# Load the pre-trained model and initialize the SORT tracker
model_path = 'yolov5s.pt' # Ensure this path points to the model file
model = YOLOv5(model_path, device='cpu')
tracker = Sort()
def process_video(uploaded_file):
# Save the uploaded file to a temporary location
temp_file_path = "temp_video.mp4"
with open(temp_file_path, "wb") as temp_file:
temp_file.write(uploaded_file.getbuffer())
# Use moviepy to read the video file
video_clip = VideoFileClip(temp_file_path)
total_frames = int(video_clip.fps * video_clip.duration)
width, height = video_clip.size
# Temporary directory to save processed video frames
temp_dir = tempfile.mkdtemp()
unique_cars = set()
progress_bar = st.progress(0)
for frame_idx, frame in enumerate(video_clip.iter_frames()):
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
progress_percentage = int((frame_idx + 1) / total_frames * 100)
progress_bar.progress(progress_percentage)
# Detection and tracking logic
results = model.predict(frame)
preds = results.pandas().xyxy[0]
detections = []
for index, row in preds.iterrows():
if row['name'] == 'car':
xmin, ymin, xmax, ymax, conf = row['xmin'], row['ymin'], row['xmax'], row['ymax'], row['confidence']
detections.append([xmin, ymin, xmax, ymax, conf])
if detections:
detections_np = np.array(detections)
trackers = tracker.update(detections_np)
for d in trackers:
unique_cars.add(int(d[4]))
xmin, ymin, xmax, ymax = map(int, d[:4])
cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (255, 0, 0), 2)
cv2.putText(frame, f'ID: {int(d[4])}', (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
cv2.putText(frame, f'Unique Cars: {len(unique_cars)}', (10, 35), cv2.FONT_HERSHEY_SIMPLEX, 1.25, (0, 255, 0), 2)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Convert back to RGB for moviepy
cv2.imwrite(f"{temp_dir}/{frame_idx:04d}.jpg", frame)
frames_files = [os.path.join(temp_dir, f"{i:04d}.jpg") for i in range(total_frames)]
clip = ImageSequenceClip(frames_files, fps=video_clip.fps)
output_video_path = 'processed_video.mp4'
clip.write_videofile(output_video_path, codec='libx264') # Use libx264 codec for compatibility
# Remove temporary directory and temporary files
shutil.rmtree(temp_dir)
return output_video_path
def main():
# Initialize session state variables if they don't exist
if 'output_video_path' not in st.session_state:
st.session_state.output_video_path = None
st.sidebar.image("logo.jpg", use_column_width=True)
uploaded_file = st.sidebar.file_uploader("Upload a video", type=['mp4'])
st.title("Car Detection and Tracking")
if uploaded_file is not None:
# Process the video only if it hasn't been processed yet or a new file is uploaded
if st.session_state.output_video_path is None or st.session_state.uploaded_file_name != uploaded_file.name:
st.session_state.uploaded_file_name = uploaded_file.name
st.session_state.output_video_path = process_video(uploaded_file)
# Display the processed video
st.video(st.session_state.output_video_path)
# Provide a download link for the processed video
with open(st.session_state.output_video_path, "rb") as file:
st.download_button("Download Processed Video", file, file_name="processed_video.mp4")
if __name__ == "__main__":
main()
|