import streamlit as st
import streamlit as st
import argparse
import io
from io import BytesIO
from PIL import ImageOps, Image
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

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

top_image = Image.open('banner_top.png')
bottom_image = Image.open('banner_bottom.png')
main_image = Image.open('main_banner.png')

upload_path = ""
download_path = ""

st.image(main_image,use_column_width='auto')
st.title(' Automatic Number Plate Recognition 🚘🚙')
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"])
st.sidebar.image(bottom_image,use_column_width='auto')

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

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(img, raw_plate, digits)

if selected_type == "Upload Image":
    st.info('✨ Supports all popular image formats 📷 - PNG, JPG, BMP 😉')
    uploaded_file = st.file_uploader("Upload Image of car's number plate 🚓", type=["png","jpg","bmp","jpeg"])

    if uploaded_file is not None:
        
        image = Image.open(uploaded_file).convert('RGB')
        st.image(image, use_column_width=True)

        with st.spinner(f"Working... 💫"):
            result = FINAL(image)

        st.image(result)    
        
        #with st.spinner(f"Working... 💫"):
         #   uploaded_image = os.path.abspath(os.path.join(upload_path,uploaded_file.name))
          #  downloaded_image = os.path.abspath(os.path.join(download_path,str("output_"+uploaded_file.name)))

            # upload file
#file = st.file_uploader('', type=['jpeg', 'jpg', 'png'])

# display image
#if file is not None:
 #   image = Image.open(file).convert('RGB')
  #  st.image(image, use_column_width=True)

   #         with open(uploaded_image,'rb') as image:
    #            img_bytes = image.read()


     #       image_stream = io.BytesIO()
      #      image_stream.seek(0)
       #     file_bytes = np.asarray(bytearray(image_stream.read()), dtype=np.uint8)
        #    img = cv.imdecode(file_bytes, cv.IMREAD_COLOR)    

         #   img = Image.open(io.BytesIO(img_bytes))


          #  img = np.asarray(bytearray(img.read()), dtype=np.uint8)
            
           # img = cv2.imdecode(img, cv2.IMREAD_COLOR)

            #results = FINAL(image)  
    
            #for img in results.imgs:
                #img_base64 = Image.fromarray(img)
                #img_base64.save(downloaded_image, format="JPEG")
              

            
    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)