Spaces:
Runtime error
Runtime error
# Bissmillah | |
import streamlit as st | |
import os | |
from os import listdir | |
import wget | |
from PIL import Image | |
import io | |
import numpy as np | |
import cv2 | |
import itertools | |
import sys | |
def load_model(): | |
wpath = 'test_detection/yolov5/weights/crowdhuman_yolov5m.pt' | |
if not os.path.exists(wpath): | |
st.write('path didnt exist, so creation ! ') | |
#os.system("python pip uninstall opencv-python") | |
os.system("python -m pip install numpy torch pandas Pillow opencv-python-headless PyYAML>=5.3.1 torchvision>=0.8.1 matplotlib seaborn>=0.11.0 easydict") | |
with st.spinner('Downloading model weights for crowdhuman_yolov5m'): | |
#os.system('wget -O yolov5/weights/crowdhuman_yolov5m.pt https://github.com/mikel-brostrom/Yolov5_DeepSort_Pytorch/releases/download/v.2.0/crowdhuman_yolov5m.pt') | |
os.system('wget -nc https://github.com/mikel-brostrom/Yolov5_DeepSort_Pytorch/releases/download/v.2.0/crowdhuman_yolov5m.pt -O test_detection/yolov5/weights/crowdhuman_yolov5m.pt') | |
#st.write('in function load_model', os.listdir('yolov5/weights/')) | |
else: | |
st.write('path alredy exist, so no creation ! ') | |
print("Model is here.") | |
#=================================================================================== | |
# Ft saving uploaded video to directory | |
def save_uploaded_vid(uploadedfile): | |
with open(os.path.join("data", uploadedfile.name),"wb") as f: | |
f.write(uploadedfile.getbuffer()) | |
return st.success("Video saved in data dir ") | |
#@st.cache(ttl=3600, max_entries=10) | |
def load_output_video(vid): | |
if isinstance(vid, str): | |
video = open(vid, 'rb') | |
else: | |
video = vid.read() | |
vname = vid.name | |
save_uploaded_vid(vid) | |
return video | |
def starter(): | |
st.image('test_detection/data/LOGOGlob.png', width = 400) | |
st.title("Test of Person detection") | |
st.text("") | |
st.text("") | |
st.success("Welcome! Please upload a video!") | |
args = { 'HirakAlger' : '112vHirakAlger_09042021_s.mp4' } | |
vid_upload = st.file_uploader(label= 'Upload Video', type = ['mp4', 'avi']) | |
vid_open = "test_detection/data/"+args['HirakAlger'] if vid_upload is None else vid_upload | |
vname = args['HirakAlger'] if vid_upload is None else vid_upload.name | |
video = load_output_video(vid_open) | |
st.video(video) | |
st.write('in function : vname = ', vname) | |
st.write('in function ', os.listdir('test_detection/data/')) | |
st.write('in function ', os.listdir('test_detection/yolov5/weights/')) | |
vidcap = cv2.VideoCapture( "test_detection/data/"+vname) | |
#frames = cv.get_frames("data/"+vname) | |
success, frame0 = vidcap.read() | |
frame0 = cv2.cvtColor(frame0, cv2.COLOR_BGR2RGB) | |
st.write('shape of frame 01 : ', frame0.shape) | |
return vname, frame0 | |
#=================================================================================== | |
def prediction(vname): | |
vpath='test_detection/data/'+vname | |
wpath = 'test_detection/yolov5/weights/crowdhuman_yolov5m.pt' | |
if os.path.exists(wpath): | |
st.write('start prediction') | |
os.system("python test_detection/track.py --yolo_weights test_detection/yolov5/weights/crowdhuman_yolov5m.pt --img 352 --save-vid --save-txt --classes 1 --conf-thres 0.4 --source " + vpath) | |
os.system("ffmpeg -i test_detection/inference/output/"+vname + " -vcodec libx264 -y test_detection/inference/output/output_video.mp4") | |
path = 'test_detection/inference/output/output_video.mp4' | |
if os.path.exists(path): | |
video_file = open(path, 'rb') | |
video_bytes = video_file.read() | |
st.video(video_bytes) | |
#=================================================================================== | |
def main(): | |
vname, frame0 = starter() | |
st.write('vname befor prediction ',vname) | |
if st.button('Heads detection!'): | |
prediction(vname) | |
st.success("Click again to retry or try a different video by uploading") | |
return | |
if __name__ == '__main__': | |
os.system('git clone --recurse-submodules https://github.com/nnassime/test_detection.git') | |
os.system("sys.path.insert(0, './test_detection/yolov5/models/')") | |
os.system("sys.path.insert(0, './test_detection/yolov5/')") | |
sys.path.insert(0, './yolov5') | |
sys.path.insert(0, 'home/user/app/test_detection/yolov5') | |
os.system("sys.path.insert(0, 'home/user/app/test_detection/yolov5')") | |
st.write(os.listdir('test_detection/yolov5/')) | |
#os.system('python test_detection/__streamlit_app.py') | |
load_model() | |
st.write("bismillah") | |
print("bismillah") | |
main() | |
st.write('out function ', os.listdir('test_detection/data/')) | |
st.write('out function ', os.listdir('test_detection/yolov5/weights/')) | |