import streamlit as st
import argparse
import io
from io import BytesIO
from PIL import ImageOps, Image
import warnings 
import os                        # To work with operation system commands
import cv2                       # To process images
import random                    # to generate random choices
import warnings                  # To avoid python warnings
import numpy as np               # To work with arrays
import pandas as pd              # To work with DataFrames
import seaborn as sns            # To visualizations
from tqdm import tqdm            # To represent progress bars
from ultralytics import YOLO     # To Create Yolo model
from termcolor import colored    # To colorfull outputs
import matplotlib.pyplot as plt  # To visualizations

warnings.filterwarnings('ignore')

st.set_page_config(
    page_title="Auto NPR",
    page_icon="✨",
    layout="centered",
    initial_sidebar_state="expanded",
)

top_image = Image.open('favicon.ico')
main_image = Image.open('main_banner.png')

upload_path = ""
download_path = ""

st.image(main_image,use_column_width='auto')
st.sidebar.image(top_image,use_column_width='auto')
st.sidebar.header('Input')
selected_type = st.sidebar.selectbox('Please select an activity type', ["Upload Image", "Live Video Feed"])

#load weights
best_model_plate = YOLO('best.pt')
best_model_digits = YOLO('best2.pt')

def Detect_Plate(img):
    '''
    Predict on a car image and return result folder and predicted image path. 
    '''
    result1 = best_model_plate.predict([img], save=True, iou=0.7)
    
    # Location of saved predicted images
    result_path1 = result1[0].save_dir
    pred1_path = os.path.join(result_path1, os.listdir(result_path1)[-1])
    return result1, pred1_path

def Detect_Digits(img2) :
    '''
    Predict on a croped plate and return result folder and predicted image.
    '''
    result2 = best_model_digits.predict([img2], save=True, iou=0.7, show_conf=False)
    
    # Location of saved predicted images
    result_path = result2[0].save_dir
    pred2_path = os.path.join(result_path, os.listdir(result_path)[0])
    return result2, pred2_path

def Sort_Digits(result2) :
    '''
    Sort detected labels base on their X-cordinate,
    --- Sort from lef to right ---
    '''
    list_of_items = []
    for i in range(len(result2[0])) :
        # Class labels decoded
        class_label = classes_dict[str(int(result2[0][i].boxes.cls))]
        # X cordinate to find the situation on X-axis
        cord_x = int(result2[0][i].boxes.xyxy[0][0])
        list_of_items.append((cord_x, class_label))

    list_of_items.sort()
    #list_of_items

    digits = []
    for digit in list_of_items :
        digits.append(digit[1])
    if len(digits) == 7 : digits.append('0')
    return digits


raw_plate = ('raw_plate.png')

classes_dict =  {
            '0':'0',
            '1':'1',
            '2':'2',
            '3':'3',
            '4':'4',
            '5':'5',
            '6':'6',
            '7':'7', 
            '8':'8',
            '9':'9',
            '10':'B',
            '11':'C',
            '12':'D',
            '13':'G', 
            '14':'H',
            '15':'J',
            '16':'L',
            '17':'M', 
            '18':'N', 
            '19':'S', 
            '20':'T', 
            '21':'V', 
            '22':'Y'
            }

def Plot_Result(raw_plate, digits) :
    
    raw_plate = cv2.imread(raw_plate)
    raw_plate = cv2.cvtColor(raw_plate, cv2.COLOR_BGR2RGB)

    # Cordinates of free locations in raw plate
    p1 = (18, 150)
    p2 = (143, 150)
    p3 = (268, 150)
    p4 = (393, 150)
    p5 = (518, 150)
    p6 = (643, 150)
    p7 = (778, 180)
    p8 = (878, 180)

    # Fill the form
    cv2.putText(raw_plate, digits[0], p1, cv2.FONT_HERSHEY_COMPLEX, 5, (0, 0, 0), 5, cv2.LINE_AA);
    cv2.putText(raw_plate, digits[1], p2, cv2.FONT_HERSHEY_COMPLEX, 5, (0, 0, 0), 5, cv2.LINE_AA);
    cv2.putText(raw_plate, digits[2], p3, cv2.FONT_HERSHEY_COMPLEX, 5, (0, 0, 0), 5, cv2.LINE_AA);
    cv2.putText(raw_plate, digits[3], p4, cv2.FONT_HERSHEY_COMPLEX, 5, (0, 0, 0), 5, cv2.LINE_AA);
    cv2.putText(raw_plate, digits[4], p5, cv2.FONT_HERSHEY_COMPLEX, 5, (0, 0, 0), 5, cv2.LINE_AA);
    cv2.putText(raw_plate, digits[5], p6, cv2.FONT_HERSHEY_COMPLEX, 5, (0, 0, 0), 5, cv2.LINE_AA);
    cv2.putText(raw_plate, digits[6], p7, cv2.FONT_HERSHEY_COMPLEX, 5, (0, 0, 0), 5, cv2.LINE_AA);
    cv2.putText(raw_plate, digits[7], p8, cv2.FONT_HERSHEY_COMPLEX, 5, (0, 0, 0), 5, cv2.LINE_AA);


    filename = 'output_image.jpg'
    cv2.imwrite(filename, raw_plate)

    out_img = cv2.imread("output_image.jpg")
    st.image(out_img, caption='final image ✅')

    
def FINAL(img) :
    '''
    A pipeline for all parts of phase 3.
    start with a car image.
    result is digits and char on car plate.
    '''
    # Read car image ( STEP-1 )
    #img = cv2.imread(img)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # First prediction -> Detect car-plate ( STEP-2 )
    result1, _ = Detect_Plate(img)

    # Plate Cordinates ( STEP-3 )
    pts = result1[0].boxes.xyxy.tolist()[0]

    # Crop plate
    img2 = img[round(pts[1]):round(pts[3]), round(pts[0]):round(pts[2])]

    # Resize plate to feed to second model ( STEP-4 )
    img2 = cv2.resize(img2, (120, 70))

    # Second prediction -> Detect digits in plate
    result2, _ = Detect_Digits(img2)

    # Sort detected digits ( STEP-5 )
    digits = Sort_Digits(result2)

    Plot_Result(raw_plate, digits)



if selected_type == "Upload Image":
    uploaded_file = st.file_uploader("Upload Image of car's plate 🚘", type=["png","jpg","bmp","jpeg"])

    if uploaded_file is not None:
        image = Image.open(uploaded_file).convert('RGB')
        st.image(image, width=400, caption='input image 📷')

        image.save('input_image.jpg')

        img = cv2.imread('input_image.jpg')
            
        st.markdown("------")
        
        FINAL(img)
        
              

            
    else:
        st.warning('⚠ Please upload your Image 😯')

st.markdown("<br><hr><center>Made with ❤️ by <a href='https://omidsakaki.ir/'><strong>omid sakaki ghazvini</strong></a></center><hr>", unsafe_allow_html=True)