import streamlit as st
import torch
from detect import detect
from PIL import Image
from io import *
import glob
from datetime import datetime
import os
import wget
import time

def imageInput(device, src):
    if src == 'Upload your own data.':
        image_file = st.file_uploader("Upload An Image", type=['png', 'jpeg', 'jpg'])
        col1, col2 = st.columns(2)
        if image_file is not None:
            img = Image.open(image_file)
            with col1:
                st.image(img, caption='Uploaded Image', use_column_width='always')
            ts = datetime.timestamp(datetime.now())
            imgpath = os.path.join('data', str(ts)+image_file.name)
            outputpath = os.path.join('data', os.path.basename(imgpath))
            with open(imgpath, mode="wb") as f:
                f.write(image_file.getbuffer())

            #call Model prediction--
            model = torch.hub.load('ultralytics/yolov5', 'custom', path='models/YOLOv5m.pt', force_reload=True) 
            model.cuda() if device == 'cuda' else model.cpu()
            pred = model(imgpath)
            pred.render()  # render bbox in image
            for im in pred.ims:
                im_base64 = Image.fromarray(im)
                im_base64.save(outputpath)

            #--Display predicton
            
            img_ = Image.open(outputpath)
            with col2:
                st.image(img_, caption='Model Prediction(s)', use_column_width='always')
    elif src == 'URL':
        url = st.text_input('Enter URL ...')
        if url is not None:
            image_file = wget.download(url)
            col1, col2 = st.columns(2)
            if image_file is not None:
                img = Image.open(image_file)
                with col1:
                    st.image(img, caption='Uploaded Image', use_column_width='always')
                ts = datetime.timestamp(datetime.now())
                imgpath = os.path.join('data', str(ts)+image_file.name)
                outputpath = os.path.join('data', os.path.basename(imgpath))
                with open(imgpath, mode="wb") as f:
                    f.write(image_file.getbuffer())
    
                #call Model prediction--
                model = torch.hub.load('ultralytics/yolov5', 'custom', path='models/YOLOv5m.pt', force_reload=True) 
                model.cuda() if device == 'cuda' else model.cpu()
                pred = model(imgpath)
                pred.render()  # render bbox in image
                for im in pred.ims:
                    im_base64 = Image.fromarray(im)
                    im_base64.save(outputpath)
    
                #--Display predicton
                
                img_ = Image.open(outputpath)
                with col2:
                    st.image(img_, caption='Model Prediction(s)', use_column_width='always')


def videoInput(device, src):
    uploaded_video = st.file_uploader("Upload Video", type=['mp4', 'mpeg', 'mov'])
    if uploaded_video != None:

        ts = datetime.timestamp(datetime.now())
        imgpath = os.path.join('data/uploads', str(ts)+uploaded_video.name)
        outputpath = os.path.join('data/video_output', os.path.basename(imgpath))

        with open(imgpath, mode='wb') as f:
            f.write(uploaded_video.read())  # save video to disk

        st_video = open(imgpath, 'rb')
        video_bytes = st_video.read()
        st.video(video_bytes)
        st.write("Uploaded Video")
        detect(weights="models/yoloTrained.pt", source=imgpath, device=0) if device == 'cuda' else detect(weights="models/YOLOv5m.pt", source=imgpath, device='cpu')
        st_video2 = open(outputpath, 'rb')
        video_bytes2 = st_video2.read()
        st.video(video_bytes2)
        st.write("Model Prediction")


def main():
    # -- Sidebar
    st.sidebar.title('⚙️Options')
    datasrc = st.sidebar.radio("Select input source.", ['URL', 'Upload your own data.'])  
    option = st.sidebar.radio("Select input type.", ['Image', 'Video'], disabled = True)
    if torch.cuda.is_available():
        deviceoption = st.sidebar.radio("Select compute Device.", ['cpu', 'cuda'], disabled = False, index=1)
    else:
        deviceoption = st.sidebar.radio("Select compute Device.", ['cpu', 'cuda'], disabled = True, index=0)
    # -- End of Sidebar

    st.header('Unreal Engine 5 Tank Demo')
    st.subheader('Select the options')
    if option == "Image":    
        imageInput(deviceoption, datasrc)
    elif option == "Video": 
        videoInput(deviceoption, datasrc)
    

if __name__ == '__main__':
    main()