File size: 1,607 Bytes
e68321e |
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 |
import os
import os.path as osp
import cv2
from PIL import Image
import streamlit as st
def runVideo(model, video, vdo_view, warn):
video_name = osp.basename(video)
outputpath = osp.join('data/video_output', video_name)
# Create A Dir to save Video Frames
os.makedirs('data/video_frames', exist_ok=True)
frames_dir = osp.join('data/video_frames', video_name)
os.makedirs(frames_dir, exist_ok=True)
cap = cv2.VideoCapture(video)
frame_count = 0
with st.spinner(text="Predicting..."):
warn.warning(
'This is realtime prediction, If you wish to download the final prediction result wait until the process done.', icon="⚠️")
while True:
frame_count += 1
ret, frame = cap.read()
if ret == False:
break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
result = model(frame)
result.render()
image = Image.fromarray(result.ims[0])
vdo_view.image(image, caption='Current Model Prediction(s)')
image.save(osp.join(frames_dir, f'{frame_count}.jpg'))
cap.release()
# convert frames in dir to a single video file
os.system(
f'ffmpeg -framerate 30 -i {frames_dir}/%d.jpg -c:v libx264 -pix_fmt yuv420p {outputpath}')
# Clean up Frames Dir
os.system(f'rm -rf {frames_dir}')
# Display Video
output_video = open(outputpath, 'rb')
output_video_bytes = output_video.read()
st.video(output_video_bytes)
st.write("Model Prediction")
vdo_view.empty()
warn.empty()
|