# To run use
# $ streamlit run yolor_streamlit_demo.py

from yolor import *

import tempfile
import os
import platform
import shutil
import time
from pathlib import Path
import cv2
import torch
import torch.backends.cudnn as cudnn

from utils.google_utils import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import (
    check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, strip_optimizer)
from utils.plots import plot_one_box
from utils.torch_utils import select_device, load_classifier, time_synchronized

from models.models import *
from utils.datasets import *
from utils.general import *

from deep_sort_pytorch.utils.parser import get_config
from deep_sort_pytorch.deep_sort import DeepSort

import streamlit as st


def main():
    
    #title
    st.title('Object Tracking Dashboard YOLOR')
    
    #side bar title
    st.sidebar.title('Settings')
    
    st.markdown(
    """
    <style>
    [data-testid="stSidebar"][aria-expanded="true"] > div:first-child {
        width: 400px;
    }
    [data-testid="stSidebar"][aria-expanded="false"] > div:first-child {
        width: 400px;
        margin-left: -400px;
    }
    </style>
    """,
    unsafe_allow_html=True,
    )

    use_webcam = st.sidebar.checkbox('Use Webcam')

    st.sidebar.markdown('---')
    confidence = st.sidebar.slider('Confidence',min_value=0.0, max_value=1.0, value = 0.25)
    st.sidebar.markdown('---')

    save_img = st.sidebar.checkbox('Save Video')
    enable_GPU = st.sidebar.checkbox('enable GPU')

    custom_classes = st.sidebar.checkbox('Use Custom Classes')
    assigned_class_id = []
    if custom_classes:
        assigned_class = st.sidebar.multiselect('Select The Custom Classes',list(names),default='person')
        for each in assigned_class:
            assigned_class_id.append(names.index(each))

    video_file_buffer = st.sidebar.file_uploader("Upload a video", type=[ "mp4", "mov",'avi','asf', 'm4v' ])

    DEMO_VIDEO = 'test.mp4'

    tfflie = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)


    ##We get our input video here

    if not video_file_buffer:
        if use_webcam:
            vid = cv2.VideoCapture(0, cv2.CAP_ARAVIS)
            tfflie.name = 0
        else:
            vid = cv2.VideoCapture(DEMO_VIDEO)
            tfflie.name = DEMO_VIDEO
            dem_vid = open(tfflie.name,'rb')
            demo_bytes = dem_vid.read()
    
            st.sidebar.text('Input Video')
            st.sidebar.video(demo_bytes)

    else:
        tfflie.write(video_file_buffer.read())
        # print("No Buffer")
        dem_vid = open(tfflie.name,'rb')
        demo_bytes = dem_vid.read()
    
        st.sidebar.text('Input Video')
        st.sidebar.video(demo_bytes)


    print(tfflie.name)
    # vid = cv2.VideoCapture(tfflie.name)
    
    stframe = st.empty()
    
    st.markdown("<hr/>", unsafe_allow_html=True)
    kpi1, kpi2, kpi3 = st.beta_columns(3) #st.columns(3)

    # stframe.image(im0,channels = 'BGR',use_column_width=True)

    with kpi1:
        st.markdown("**Frame Rate**")
        kpi1_text = st.markdown("0")

    with kpi2:
        st.markdown("**Tracked Objects**")
        kpi2_text = st.markdown("0")

    with kpi3:
        st.markdown("**Total Count**")
        kpi3_text = st.markdown("0")

    st.markdown("<hr/>", unsafe_allow_html=True)

    # call yolor 
    load_yolor_and_process_each_frame(tfflie.name, enable_GPU, confidence, assigned_class_id, kpi1_text, kpi2_text, kpi3_text, stframe)

    st.text('Video is Processed')
        
if __name__ == '__main__':
    try:
        main()
    except SystemExit:
        pass


