import streamlit as st
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
from tensorflow.keras.models import load_model

# Title for the Streamlit App
st.title("Nepal Vehicle License Plate and Character Recognition")

# Description
st.write("Upload an image to detect license plates, segment characters, and recognize each character using advanced YOLO and CNN models.")

# Download YOLO and CNN model weights from Hugging Face
@st.cache_resource
def load_models():
    # Full license plate detection model
    full_plate_model_path = hf_hub_download(
        repo_id="krishnamishra8848/Nepal-Vehicle-License-Plate-Detection", filename="last.pt"
    )
    full_plate_model = YOLO(full_plate_model_path)

    # Character detection model
    character_model_path = hf_hub_download(
        repo_id="krishnamishra8848/Nepal_Vehicle_License_Plates_Detection_Version3", filename="best.pt"
    )
    character_model = YOLO(character_model_path)

    # Character recognition model
    recognition_model_path = hf_hub_download(
        repo_id="krishnamishra8848/Nepal_Vehicle_License_Plates_Character_Recognisation", filename="model.h5"
    )
    recognition_model = load_model(recognition_model_path)

    return full_plate_model, character_model, recognition_model

# Load models
full_plate_model, character_model, recognition_model = load_models()

# Function to detect license plates
def detect_license_plate(image):
    img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    results = full_plate_model(img_bgr)

    detected_image = img_bgr.copy()
    cropped_plates = []
    for result in results:
        if hasattr(result, 'boxes') and result.boxes is not None:
            for box in result.boxes.xyxy:
                x1, y1, x2, y2 = map(int, box)
                cv2.rectangle(detected_image, (x1, y1), (x2, y2), (255, 0, 0), 2)  # Draw bounding box
                cropped_plate = img_bgr[y1:y2, x1:x2]
                cropped_plates.append(cropped_plate)

    return detected_image, img_bgr, cropped_plates

# Function to detect characters
def detect_characters(image):
    results = character_model(image)
    detected_image = image.copy()
    character_crops = []
    for result in results:
        if hasattr(result, 'boxes') and result.boxes is not None:
            for box in result.boxes.xyxy:
                x1, y1, x2, y2 = map(int, box)
                cv2.rectangle(detected_image, (x1, y1), (x2, y2), (0, 255, 0), 2)  # Draw bounding box
                character_crops.append(image[y1:y2, x1:x2])
    return detected_image, character_crops

# Function to recognize characters
def recognize_characters(character_crops):
    class_labels = [
        'क', 'को', 'ख', 'ग', 'च', 'ज', 'झ', 'ञ', 'डि', 'त', 'ना', 'प',
        'प्र', 'ब', 'बा', 'भे', 'म', 'मे', 'य', 'लु', 'सी', 'सु', 'से', 'ह',
        '०', '१', '२', '३', '४', '५', '६', '७', '८', '९'
    ]
    recognized_characters = []
    for char_crop in character_crops:
        # Preprocess the cropped character for recognition model
        resized = cv2.resize(char_crop, (64, 64))
        normalized = resized / 255.0
        reshaped = np.expand_dims(normalized, axis=0)  # Add batch dimension

        # Predict the character
        prediction = recognition_model.predict(reshaped)
        predicted_class = class_labels[np.argmax(prediction)]
        recognized_characters.append(predicted_class)
    return recognized_characters

# Upload an image file
uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Load image
    image = Image.open(uploaded_file)

    # Detect license plates
    with st.spinner("Processing image..."):
        detected_image, original_image, cropped_plates = detect_license_plate(image)

        # Show the detected license plate image with bounding boxes
        st.image(cv2.cvtColor(detected_image, cv2.COLOR_BGR2RGB), caption="Detected License Plates", use_container_width=True)

        # Initialize variables for comparison
        max_characters = []
        best_recognition_method = ""

        # Process original image for character detection
        char_detected_img, character_crops_original = detect_characters(original_image)

        # Recognize characters from original image
        recognized_original = recognize_characters(character_crops_original)

        if recognized_original:
            if len(recognized_original) > len(max_characters):
                max_characters = recognized_original
                best_recognition_method = "Original Image"

        # Process cropped plates for character detection
        for idx, cropped_plate in enumerate(cropped_plates, 1):
            st.write(f"Processing Cropped Plate {idx}:")
            cropped_char_detected_img, character_crops_cropped = detect_characters(cropped_plate)
            recognized_cropped = recognize_characters(character_crops_cropped)

            if recognized_cropped:
                if len(recognized_cropped) > len(max_characters):
                    max_characters = recognized_cropped
                    best_recognition_method = f"Cropped Plate {idx}"

        # Show the result
        st.image(cv2.cvtColor(char_detected_img, cv2.COLOR_BGR2RGB), caption="Detected Characters from Original Image", use_container_width=True)

        if best_recognition_method:
            st.write(f"Best Recognition Method: {best_recognition_method}")
            st.write("Recognized Characters:", "".join(max_characters))
        else:
            st.write("No characters detected in any method.")

    st.success("Processing complete!")