# 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

os.system("sys.path.insert(0, 'home/user/app/test_detection/yolov5')")
sys.path.insert(0, 'home/user/app/test_detection/yolov5')
  
os.system("sys.path.insert(0, 'home/user/app/test_detection/yolov5/models')")
sys.path.insert(0, 'home/user/app/test_detection/yolov5/models')
  
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')
  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/'))