import streamlit as st
import cv2
import numpy as np
import easyocr
from ultralytics import YOLO

# Title of the app with enhanced styling
st.markdown("""
    <style>
        .stApp {
            background-image: url("https://i.postimg.cc/cLpB502d/black-car.png");
            background-size: cover;
            background-position: center center;
        }
        
        h1 {
            color: white;
            font-size: 3em;
            text-shadow: 3px 3px 5px rgba(0, 0, 0, 0.7);
            text-align: center;
        }

        .sidebar .sidebar-content {
            background-color: rgba(0, 0, 0, 0.6);
            color: white;
        }

        .stButton>button {
            background-color: #4CAF50;
            color: white;
            font-size: 18px;
            padding: 10px 24px;
            border-radius: 10px;
            border: none;
            cursor: pointer;
        }

        .stButton>button:hover {
            background-color: #45a049;
        }

        .stTextInput>div>div>input {
            background-color: #ffffff;
            color: black;
        }

        .stTextInput>div>label {
            color: white;
        }

        .stSlider>div>label {
            color: white;
        }

        .stMarkdown {
            color: white;
        }
        
        .stInfo {
            background-color: rgba(0, 0, 0, 0.6);
            color: white;
        }
    </style>
""", unsafe_allow_html=True)

# Title of the app
st.title("License Plate Recognition System")

# Load the YOLO model for license plate detection
@st.cache_resource
def load_yolo_model():
    model_path = "best.pt"  # Replace with your model file
    model = YOLO(model_path)
    return model

# Load EasyOCR reader
@st.cache_resource
def load_easyocr_reader():
    return easyocr.Reader(['en'], gpu=False)

# Initialize models
yolo_model = load_yolo_model()
ocr_reader = load_easyocr_reader()

# Function to process image and detect license plates
def process_image(image, confidence_threshold=0.5):
    # Perform license plate detection
    results = yolo_model(image, conf=confidence_threshold)
    annotated_image = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
    
    # Create columns to display the uploaded image, cropped image, and extracted text side by side
    c1, c2 = st.columns(2)
    
    with c1:
        st.image(annotated_image, caption="Detected License Plate(s)", use_container_width=True)

    # Loop through detections and perform OCR
    for result in results:
        boxes = result.boxes.xyxy.cpu().numpy().astype(int)
        if len(boxes) == 0:
            st.warning("No license plate detected!")
            return
        
        for i, box in enumerate(boxes):
            x1, y1, x2, y2 = box
            cropped_image = image[y1:y2, x1:x2]
            cropped_image_rgb = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2RGB)

            with c2:
                st.image(cropped_image_rgb, caption=f"Cropped License Plate {i+1}", use_container_width=True)

            # Perform OCR on the cropped image
            text_results = ocr_reader.readtext(cropped_image_rgb, detail=0)
            detected_text = " ".join(text_results)
            
            with c2:
                st.write(f"**Extracted Text (Plate {i+1}):** {detected_text}")
                st.write(f"**Confidence Score:** {result.boxes.conf.cpu().numpy()[i]:.2f}")

# Sidebar input for file upload
uploaded_file = st.file_uploader("Upload an Image or Video", type=["mp4", "avi", "mov", "jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Check if it's an image or video
    file_type = uploaded_file.type

    if file_type.startswith("image"):
        # Read and process the image
        image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
        confidence_threshold = st.sidebar.slider("Confidence Threshold", 0.0, 1.0, 0.5, 0.01)
        process_image(image, confidence_threshold)

st.markdown("---")
st.info("Please use images as input, because processing videos with the CPU may result in longer times.")
st.info("This application uses Fine Tuned YOLOv8 for detection and EasyOCR for text recognition.")