ModuMLTECH commited on
Commit
f145ba4
·
verified ·
1 Parent(s): 172d5d6

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +102 -0
  2. yolov5s.pt +3 -0
app.py ADDED
@@ -0,0 +1,102 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import matplotlib
2
+ matplotlib.use('Agg')
3
+
4
+ import streamlit as st
5
+ import cv2
6
+ import numpy as np
7
+ from yolov5 import YOLOv5
8
+ from sort.sort import Sort
9
+ import tempfile
10
+ import shutil
11
+ from moviepy.editor import VideoFileClip, concatenate_videoclips, ImageSequenceClip
12
+ import os
13
+
14
+ # Load the pre-trained model and initialize the SORT tracker
15
+ model_path = 'yolov5s.pt' # Ensure this path points to the model file
16
+ model = YOLOv5(model_path, device='cpu')
17
+ tracker = Sort()
18
+
19
+ def process_video(uploaded_file):
20
+ # Save the uploaded file to a temporary location
21
+ temp_file_path = "temp_video.mp4"
22
+ with open(temp_file_path, "wb") as temp_file:
23
+ temp_file.write(uploaded_file.getbuffer())
24
+
25
+ # Use moviepy to read the video file
26
+ video_clip = VideoFileClip(temp_file_path)
27
+ total_frames = int(video_clip.fps * video_clip.duration)
28
+ width, height = video_clip.size
29
+
30
+ # Temporary directory to save processed video frames
31
+ temp_dir = tempfile.mkdtemp()
32
+
33
+ unique_cars = set()
34
+ progress_bar = st.progress(0)
35
+
36
+ for frame_idx, frame in enumerate(video_clip.iter_frames()):
37
+ frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
38
+
39
+ progress_percentage = int((frame_idx + 1) / total_frames * 100)
40
+ progress_bar.progress(progress_percentage)
41
+
42
+ # Detection and tracking logic
43
+ results = model.predict(frame)
44
+ preds = results.pandas().xyxy[0]
45
+ detections = []
46
+
47
+ for index, row in preds.iterrows():
48
+ if row['name'] == 'car':
49
+ xmin, ymin, xmax, ymax, conf = row['xmin'], row['ymin'], row['xmax'], row['ymax'], row['confidence']
50
+ detections.append([xmin, ymin, xmax, ymax, conf])
51
+
52
+ if detections:
53
+ detections_np = np.array(detections)
54
+ trackers = tracker.update(detections_np)
55
+
56
+ for d in trackers:
57
+ unique_cars.add(int(d[4]))
58
+ xmin, ymin, xmax, ymax = map(int, d[:4])
59
+ cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (255, 0, 0), 2)
60
+ cv2.putText(frame, f'ID: {int(d[4])}', (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
61
+
62
+ cv2.putText(frame, f'Unique Cars: {len(unique_cars)}', (10, 35), cv2.FONT_HERSHEY_SIMPLEX, 1.25, (0, 255, 0), 2)
63
+ frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Convert back to RGB for moviepy
64
+ cv2.imwrite(f"{temp_dir}/{frame_idx:04d}.jpg", frame)
65
+
66
+ frames_files = [os.path.join(temp_dir, f"{i:04d}.jpg") for i in range(total_frames)]
67
+ clip = ImageSequenceClip(frames_files, fps=video_clip.fps)
68
+ output_video_path = 'processed_video.mp4'
69
+ clip.write_videofile(output_video_path, codec='libx264') # Use libx264 codec for compatibility
70
+
71
+ # Remove temporary directory and temporary files
72
+ shutil.rmtree(temp_dir)
73
+
74
+ return output_video_path
75
+
76
+
77
+ def main():
78
+ # Initialize session state variables if they don't exist
79
+ if 'output_video_path' not in st.session_state:
80
+ st.session_state.output_video_path = None
81
+
82
+ st.sidebar.image("logo.jpg", use_column_width=True)
83
+ uploaded_file = st.sidebar.file_uploader("Upload a video", type=['mp4'])
84
+
85
+ st.title("Car Detection and Tracking")
86
+
87
+ if uploaded_file is not None:
88
+ # Process the video only if it hasn't been processed yet or a new file is uploaded
89
+ if st.session_state.output_video_path is None or st.session_state.uploaded_file_name != uploaded_file.name:
90
+ st.session_state.uploaded_file_name = uploaded_file.name
91
+ st.session_state.output_video_path = process_video(uploaded_file)
92
+
93
+ # Display the processed video
94
+ st.video(st.session_state.output_video_path)
95
+
96
+ # Provide a download link for the processed video
97
+ with open(st.session_state.output_video_path, "rb") as file:
98
+ st.download_button("Download Processed Video", file, file_name="processed_video.mp4")
99
+
100
+
101
+ if __name__ == "__main__":
102
+ main()
yolov5s.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8b3b748c1e592ddd8868022e8732fde20025197328490623cc16c6f24d0782ee
3
+ size 14808437