import gradio as gr
import numpy as np
from PIL import Image, ImageEnhance
from ultralytics import YOLO
import cv2

# Load YOLO model
model_path = "./best.pt"
modelY = YOLO(model_path)
modelY.to('cpu')

# Preprocessing function
def preprocessing(image):
    if image.mode != 'RGB':
        image = image.convert('RGB')
    image = ImageEnhance.Sharpness(image).enhance(2.0)
    image = ImageEnhance.Contrast(image).enhance(1.5)
    image = ImageEnhance.Brightness(image).enhance(0.8)
    width = 448
    aspect_ratio = image.height / image.width
    height = int(width * aspect_ratio)
    return image.resize((width, height))

# YOLO document detection and cropping
def detect_and_crop_document(image):
    image_np = np.array(image)
    results = modelY(image_np, conf=0.80, device='cpu')
    cropped_images = []
    predictions = []
    
    for result in results:
        for box in result.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = int(box.conf[0] * 100)  # Convert confidence to percentage
            cls = int(box.cls[0])
            class_name = modelY.names[cls].capitalize()  # Capitalize class names
            cropped_image_np = image_np[y1:y2, x1:x2]
            cropped_image = Image.fromarray(cropped_image_np)
            cropped_images.append(cropped_image)
            predictions.append(f"Detected: STNK {class_name} -- (Confidence: {conf}%)")
    
    if not cropped_images:
        return None, "No document detected"
    return cropped_images, predictions

# Gradio interface
def process_image(image):
    preprocessed_image = preprocessing(image)
    cropped_images, predictions = detect_and_crop_document(preprocessed_image)
    
    if cropped_images:
        return cropped_images, '\n'.join(predictions)
    return None, "No document detected"

with gr.Blocks(css=".gr-button {background-color: #4caf50; color: white; font-size: 16px; padding: 10px 20px; border-radius: 8px;}") as demo:
    gr.Markdown(
        """
        <h1 style="text-align: center; color: #4caf50;">📜 License Registration Classification</h1>
        <p style="text-align: center; font-size: 18px;">Upload an image and let the YOLO model detect and crop license documents automatically.</p>
        """
    )
    with gr.Row():
        with gr.Column(scale=1, min_width=300):
            input_image = gr.Image(type="pil", label="Upload License Image", interactive=True)
            with gr.Row():
                clear_btn = gr.Button("Clear")
                submit_btn = gr.Button("Detect Document")
        with gr.Column(scale=2):
            output_image = gr.Gallery(label="Cropped Documents", interactive=False)
            output_text = gr.Textbox(label="Detection Result", interactive=False)

    submit_btn.click(process_image, inputs=input_image, outputs=[output_image, output_text])
    clear_btn.click(lambda: (None, ""), outputs=[output_image, output_text])

demo.launch()