import streamlit as st
import cv2
import numpy as np
import os
import tempfile
from PIL import Image
from ultralytics import YOLO
import easyocr
import re

# ----------- Page config ----------------
st.set_page_config(
    page_title='Number plate',
    page_icon='',
    layout='wide'
)



st.title("Number plate")
st.write('Upload an image or video to detect helmet')

# -------------- load model detection -----------
@st.cache_resource
def load_model():
    return YOLO('best.pt')

@st.cache_resource
def load_ocr():
    return easyocr.Reader(['en'])

model = load_model()
ocr_model = load_ocr()



# switching tabs
x,y = st.tabs(['Image Detection','Video Detection'])

# Tabs 1 == Image detection
with x:
    st.header('Image Detection')

    img_path = st.file_uploader('Please upload an image') # upload option

    if img_path is not None:
        image = Image.open(img_path)
        image_np = np.array(image) # converting in to atrray

        # YOLO inference
        result = model(image_np,conf=0.4)
        annot_img = result[0].plot()

        # Convert BGR to RGB
        annot_img = cv2.cvtColor(annot_img, cv2.COLOR_BGR2RGB)

        # --------------------- number detection ----------------------
        detect_number = []

        for box in result[0].boxes:
            x1,y1,x2,y2 = map(int,box.xyxy[0])

            plate_crop = image_np[y1:y2,x1:x2]
            gray = cv2.cvtColor(plate_crop,cv2.COLOR_BGR2GRAY)

            ocr_result = ocr_model.readtext(gray)


            text = ""
            for t in ocr_result:
                text += t[1] + " "

            plate_text = re.sub(r'[^A-Z0-9]', '', text.upper())

            if plate_text:
                detect_number.append(plate_text)



        #st.subheader("Detected Result")
        #st.image(annot_img, use_container_width=True)
        #st.image(annot_img,width=400) # first image
        #st.image(image,width=400) # last image

        # to display side by side
        ori_img,pre_img = st.columns(2)
        with ori_img: # original image
            st.markdown('#### ***Original Image***')
            st.image(image,width=500)

        with pre_img:
            st.markdown('#### ***Detected Image***')
            st.image(annot_img,width=500)

        # Display OCR text
        if detect_number:
            st.subheader("Detected Number Plate Text")
            for num in detect_number:
                st.success(num)
        else:
            st.warning("No number plate text detected")



#------------------------------------------------------------------------
# ----------       For Video detection         -------------------------
#------------------------------------------------------------------------
# ----------       For Video detection         -------------------------
with y:
    st.header("Video Number Plate Detection")

    video_file = st.file_uploader(
        "Upload a Video",
        type=["mp4", "avi", "mov"]
    )

    if video_file is not None:
        temp_video = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
        temp_video.write(video_file.read())
        temp_video.close()

        cap = cv2.VideoCapture(temp_video.name)
        stframe = st.empty()

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            # YOLO inference
            results = model(frame, conf=0.4)
            annotated_frame = frame.copy()

            for box in results[0].boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0])

                plate_crop = frame[y1:y2, x1:x2]
                if plate_crop.size == 0:
                    continue

                gray = cv2.cvtColor(plate_crop, cv2.COLOR_BGR2GRAY)

                # OCR
                ocr_result = ocr_model.readtext(gray)
                text = " ".join([t[1] for t in ocr_result])
                plate_text = re.sub(r'[^A-Z0-9]', '', text.upper())

                # Draw bounding box
                cv2.rectangle(
                    annotated_frame,
                    (x1, y1),
                    (x2, y2),
                    (0, 255, 0),
                    2
                )

                # Draw plate text
                if plate_text:
                    cv2.putText(
                        annotated_frame,
                        plate_text,
                        (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.8,
                        (0, 255, 0),
                        2
                    )

            stframe.image(annotated_frame, channels="BGR", width=800)

        cap.release()
        os.remove(temp_video.name)
        st.success("Video processing completed")


                                        






st.markdown("""
<br>
<div style='text-align:center; padding:12px; background-color:#111111; border-radius:10px;'>
<span style='color:#AAAAAA; font-size:16px;'>
Designed & Developed by <b style='color:#CCCCCC;'>Yedeedya Injeti</b><br>
Under <b style='color:#B8860B;'>Innomatics Research Labs</b>
</span>
</div>
<br>
""", unsafe_allow_html=True)