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("test_detection/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= 'Welcome! Please upload a 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): | |
| with st.spinner('Running detection...'): | |
| 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 inference/output/"+vname + " -vcodec libx264 -y inference/output/output_"+vname) | |
| path = 'inference/output/output_'+vname | |
| if os.path.exists(path): | |
| video_file = open(path, 'rb') | |
| video_bytes = video_file.read() | |
| st.video(video_bytes) | |
| #=================================================================================== | |
| def extract_heads(filepath, frame0): | |
| nbperson = 0 | |
| listhead = [] | |
| if os.path.exists(filepath): | |
| #st.write("filepath : ", filepath) | |
| array_from_file = np.loadtxt(filepath, dtype=int) | |
| #st.write('np of array load : ', array_from_file.shape) | |
| array_from_file = np.delete(array_from_file,np.s_[7:10], axis=1) | |
| nbperson = np.unique(array_from_file[:,1]).shape[0] | |
| st.subheader('Display some detected heads :') | |
| st.write(' Number of heads detected : ', nbperson ) | |
| rows = 5 | |
| cols = 10 | |
| nbheads = rows*cols | |
| frame = frame0 | |
| cont = array_from_file | |
| for a in range(nbheads): | |
| numh = a | |
| head = frame[cont[numh][3]:cont[numh][3]+cont[numh][5],cont[numh][2]:cont[numh][2]+cont[numh][4],:] | |
| listhead.append(head) | |
| #st.write('Len of liste heads : ', len(listhead)) | |
| return nbperson, listhead | |
| def display_heads(nbperson, listhead): | |
| #os.system("streamlit run https://gist.githubusercontent.com/treuille/da70b4888f8b706fca7afc765751cd85/raw/0727bb67defd93774dae65a2bc6917f72e267460/image_paginator.py") | |
| image_iterator = paginator("Select a sunset page", listhead) | |
| indices_on_page, images_on_page = map(list, zip(*image_iterator)) | |
| st.image(images_on_page, width=150, caption=indices_on_page) | |
| return | |
| def paginator(label, items, items_per_page=10, on_sidebar=True): | |
| # Figure out where to display the paginator | |
| if on_sidebar: | |
| location = st.sidebar.empty() | |
| else: | |
| location = st.empty() | |
| # Display a pagination selectbox in the specified location. | |
| items = list(items) | |
| n_pages = len(items) | |
| n_pages = (len(items) - 1) // items_per_page + 1 | |
| page_format_func = lambda i: "Page %s" % i | |
| page_number = location.selectbox(label, range(n_pages), format_func=page_format_func) | |
| # Iterate over the items in the page to let the user display them. | |
| min_index = page_number * items_per_page | |
| max_index = min_index + items_per_page | |
| return itertools.islice(enumerate(items), min_index, max_index) | |
| #=================================================================================== | |
| def main(): | |
| vname, frame0 = starter() | |
| #st.write('vname befor prediction ',vname) | |
| if st.button('Heads detection!'): | |
| prediction(vname) | |
| filepath = 'inference/output/'+vname | |
| filepath = filepath[:-3]+'txt' | |
| #st.write('filepath : ',filepath) | |
| st.success("Click again to retry or try a different video by uploading") | |
| nbperson, listhead = extract_heads(filepath, frame0) | |
| display_heads(nbperson, listhead) | |
| return | |
| if __name__ == '__main__': | |
| os.system('git clone --recurse-submodules https://github.com/nnassime/test_detection.git') | |
| #st.write(os.listdir()) | |
| #st.write(os.listdir('test_detection/')) | |
| #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/')) | |