forked from theAIGuysCode/tensorflow-yolov4-tflite
-
Notifications
You must be signed in to change notification settings - Fork 0
/
stack_video_streamlit.py
141 lines (108 loc) · 4.5 KB
/
stack_video_streamlit.py
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
129
130
131
132
133
134
135
136
137
138
139
140
141
import streamlit as st
import glob, os
from pprint import pprint
from moviepy.editor import VideoFileClip, clips_array, vfx
def stack(file1, file2, nameout, audio=False):
"""
Stack two videos on top of each other.
"""
clip1 = VideoFileClip(file1, audio=audio)
clip2 = VideoFileClip(file2, audio=audio)
final_clip = clips_array([
[clip1], # First row (add to this list if you want more than one video in row)
[clip2] # Second row (add to this list if you want more than one video in row)
])
final_clip.write_videofile(nameout)
def detect(input_file, output_file, default=True):
"""
Call the detect_video file.
"""
size_arg = "--size 416"
model_arg = "--model yolov4"
vid_arg = f"--video {input_file}"
out_arg = f"--output {output_file}"
if default:
weights_arg = "--weights ./checkpoints/yolov4-416"
os.system(f"python detect_video.py {weights_arg} {size_arg} {model_arg} {vid_arg} {out_arg}")
else:
weights_arg = "--weights ./checkpoints/custom-416"
os.system(f"python detect_video_custom.py {weights_arg} {size_arg} {model_arg} {vid_arg} {out_arg}")
def estimate_time(file):
"""
Get the clip's duration to estimate detection time and stack time.
"""
# Create an object by passing the location as a string
clip = VideoFileClip(file)
# Get stats
clip_duration = int(clip.duration)
clip_fps = int(clip.fps)
total_frames = clip_fps * clip_duration
detect_time = 20 + total_frames / 20
stack_time = 15 * clip_duration
return clip_duration, detect_time, stack_time
################################################################################
# Streamlit Deployment
#
# In order to run this file successfully in a streamlit environment, run the
# following command from the terminal:
# streamlit run stack_video.py --server.maxUploadSize 1024
################################################################################
# Title of app
st.title('YOLOv4 Demo')
# Sidebar
st.sidebar.subheader('Tab1')
option = st.selectbox(
'Pre-Generated Videos',
[
f"{dirpath}/{filename}"
for (dirpath, dirnames, filenames) in os.walk("detections/stacked/")
for filename in filenames
]
)
# File upload
uploaded_video = st.sidebar.file_uploader(
label="Video Upload",
type=['avi', 'mp4']
)
if uploaded_video is None:
st.write('You selected:', option)
video_file = open(f"{option}", 'rb')
else:
upload_name = uploaded_video.name
upload_name_avi = upload_name.replace('.mp4', '.avi')
upload_name_base = upload_name_avi.replace('.avi', '')
upload_folder = f"streamlit_detections/{upload_name_base}"
if os.path.exists(f'{upload_folder}'):
st.write('Video has already been analyzed. Retrieved the previously analyzed video.')
else:
os.makedirs(f'{upload_folder}')
os.makedirs(f'{upload_folder}/start')
os.makedirs(f'{upload_folder}/custom')
os.makedirs(f'{upload_folder}/default')
start_file = f'{upload_folder}/start/{upload_name}'
# Checks and deletes the output file
# You cant have a existing file or it will throw an error
if os.path.isfile(start_file):
os.remove(start_file)
# opens the file 'output.avi' which is accessable as 'out_file'
with open(start_file, "wb") as out_file: # open for [w]riting as [b]inary
out_file.write(uploaded_video.read())
input_file = f"./{upload_folder}/start/{upload_name}"
custom_output = f"./{upload_folder}/custom/{upload_name_avi}"
default_output = f"./{upload_folder}/default/{upload_name_avi}"
clip_duration, detect_time, stack_time = estimate_time(input_file)
st.write(
f'Your clip is {clip_duration} seconds long. Expect results in ' +\
f'approximately {2 * detect_time + stack_time} seconds.'
)
with st.spinner(text='Detecting with default YOLO weights...'):
detect(input_file, default_output, default=True)
with st.spinner(text='Detecting with custom YOLO weights...'):
detect(input_file, custom_output, default=False)
with st.spinner(text='Stacking the two videos for easy comparison...'):
stack(default_output, custom_output, f"{upload_folder}/{upload_name_base}.mp4")
st.balloons()
video_file = open(f"{upload_folder}/{upload_name_base}.mp4", 'rb')
uploaded_video = None
video_bytes = video_file.read()
st.video(video_bytes)