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()