# 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


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')
  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/'))