from typing import get_args

import cv2
import numpy as np
import streamlit as st
from PIL import Image

from open_image_models import LicensePlateDetector
from open_image_models.detection.core.hub import PlateDetectorModel

# Define the available models
detection_models = get_args(PlateDetectorModel)

# Streamlit interface
st.title("🦀 Open Image Models: Pre-trained Models for Object Detection")
st.write("Leverage fast and efficient pre-trained ONNX models for various object detection tasks, starting with license plate detection.")
st.markdown("---")

# Model selection dropdown (specific to license plate detection in this example)
selected_model = st.selectbox("🔍 Select a License Plate Detection Model", detection_models)

# File uploader for images
uploaded_file = st.file_uploader("📂 Upload an image...", type=["jpg", "png", "jpeg", "webp"])

if uploaded_file is not None:
    # Load the image using PIL
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image', use_column_width=True)

    st.write("")
    st.write("🔍 **Detecting license plates...**")

    # Convert the PIL image to an OpenCV format (NumPy array)
    image_np = np.array(image)
    image_cv2 = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)

    # Initialize the License Plate Detector
    lp_detector = LicensePlateDetector(detection_model=selected_model)

    # Perform license plate detection
    detections = lp_detector.predict(image_cv2)

    # Streamlit display for detections
    if detections:
        st.success(f"✅ {len(detections)} License Plates Detected!")

        # Use an expander to show details in a more organized way
        with st.expander("See detected plates details"):
            for i, detection in enumerate(detections):
                # Access attributes of the DetectionResult class
                bbox = detection.bounding_box
                st.markdown(f"""
                **Plate {i+1}:**
                - **Label:** {detection.label}
                - **Confidence:** {detection.confidence:.2f}
                - **Bounding Box:** (x1: {bbox.x1}, y1: {bbox.y1}, x2: {bbox.x2}, y2: {bbox.y2})
                """)
    else:
        st.warning("⚠️ No license plates detected!")

    # Annotate and display the image with detected plates
    annotated_image = lp_detector.display_predictions(image_cv2)

    # Convert the annotated image from BGR to RGB for Streamlit display
    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
    st.image(annotated_image_rgb, caption='Annotated Image with Detections', use_column_width=True)

# Add some additional style or layout to make the app more attractive
st.markdown("""
<style>
    .stButton>button {
        font-size: 16px;
        background-color: #4CAF50;
        color: white;
        border-radius: 8px;
    }
    .stImage img {
        border-radius: 10px;
        padding: 10px;
    }
</style>
""", unsafe_allow_html=True)