import streamlit as st
import cv2
import numpy as np
from PIL import Image
import imutils
import easyocr
import os
import pathlib
import platform
from xyxy_converter import yolov5_to_image_coordinates
import shutil

system_platform = platform.system()
if system_platform == 'Windows': pathlib.PosixPath = pathlib.WindowsPath

CUR_DIR = os.getcwd()
YOLO_PATH = f"{CUR_DIR}/yolov5"
MODEL_PATH = "runs/train/exp/weights/best.pt"

def main():
    st.title("Odometer value extractor with Streamlit")

    # Use st.camera to capture images from the user's camera
    img_file_buffer = st.camera_input(label='Please, take a photo of odometer', key="odometer")

    # Check if an image is captured
    if img_file_buffer is not None:
        # Convert the image to a NumPy array
        image = Image.open(img_file_buffer)
        image_np = np.array(image)
        resized_image = cv2.resize(image_np, (640, 640))
        resized_image = resized_image.astype(np.uint8)
        cv2.imwrite('odometer_image.jpg', resized_image)

        # detect(
        #     weights='yolov5\runs\train\exp\weights\best.pt',
        #     source='odometer_image.jpg',
        #     img=640,
        #     conf=0.4,
        #     name='temp_exp',
        #     hide_labels=True,
        #     hide_conf=True,
        #     save_txt=True,
        #     exist_ok=True
        # )

        # os.system('wandb disabled')

        os.chdir(YOLO_PATH)

        # try:
        #     shutil.rmtree('runs/detect/temp_exp')
        # except:
        #     pass

        image_path = "../odometer_image.jpg"
        # command = f"python detect.py --weights {MODEL_PATH} --source {image_path} --img 640 --conf 0.4 --name 'temp_exp' --hide-labels --hide-conf --save-txt --exist-ok"
        command = f'''
        python detect.py \
            --weights {MODEL_PATH} \
            --source {image_path} \
            --img 640 \
            --conf 0.4 \
            --name temp_exp \
            --hide-labels \
            --hide-conf \
            --save-txt \
            --exist-ok \
            --save-conf
        '''

        # Run the command
        os.system(command)

        # st.write('The detection is completed!!!')

        os.chdir(CUR_DIR)

        # st.write(os.path.exists('yolov5/runs/detect/temp_exp'))

        if os.path.exists('yolov5/runs/detect/temp_exp'):
            processed_image = cv2.imread('yolov5/runs/detect/temp_exp/odometer_image.jpg')
            # st.write('Image boxed and loaded')
            text_files = os.listdir('yolov5/runs/detect/temp_exp/labels')
            original_img = cv2.imread('odometer_image.jpg')
            gray = cv2.cvtColor(original_img, cv2.COLOR_BGR2GRAY)

            if len(text_files) == 0:
                display_text = "An odometer is not detected in the image!!!"
            else:
                text_file_path = f'yolov5/runs/detect/temp_exp/labels/{text_files[0]}'
                x1, y1, x2, y2 = yolov5_to_image_coordinates(text_file_path)
                # st.write(x1, y1, x2, y2)
                cropped_image = gray[x1:x2, y1:y2]

                reader = easyocr.Reader(['en'])
                result = reader.readtext(cropped_image)

                if len(result) != 0:
                    odometer_value = sorted(result, key=lambda x: x[2], reverse=True)[0][1]
                    display_text = f"Odometer value: {odometer_value}"
                else:
                    odometer_value = 'not detected'
                    display_text = f"The odometer value is {odometer_value}!!!"
        else:
            display_text = "An odometer is not detected in the image!!!"
            processed_image = cv2.imread('odometer_image.jpg')

        try:
            shutil.rmtree('odometer_image.jpg')
            shutil.rmtree('runs/detect/temp_exp')
        except:
            pass

        # Resize or preprocess the image as needed for your model
        # For example, resizing to a specific input size
        # processed_image = cv2.resize(image_np, (224, 224))
        
        # Process the image using your deep learning model
        # processed_image = process_image(image_np)

        # Display the processed image
        st.image(processed_image, caption=f"{display_text}", use_column_width=True)

        st.session_state.pop("odometer")

if __name__ == "__main__":
    main()