import cv2
import pytesseract
from PIL import Image
from ultralytics import YOLO
import supervision as sv
import gradio as gr
import os

# Set Tesseract path if needed
pytesseract.pytesseract_cmd = "/usr/bin/tesseract"

# Load the YOLO model (replace with the actual model path)
model_yolo = YOLO('./saved_model.pt')

def process_video(video_path):
    """
    Process a video frame-by-frame, detect the first license plate, and save the cropped plate.
    Returns the cropped plate path and text.
    """
    # Open the video file
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return "Error: Unable to open video file.", None

    frame_count = 0

    # Ensure the cropped plates directory exists
    cropped_dir = "cropped_plates"
    os.makedirs(cropped_dir, exist_ok=True)

    cropped_plate_path = None
    number_plate_text = None

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

        frame_count += 1
        print(f"Processing frame {frame_count}...")

        # Run YOLO predictions
        results = model_yolo.predict(source=frame, save=False)
        detections = sv.Detections(
            xyxy=results[0].boxes.xyxy.cpu().numpy(),
            confidence=results[0].boxes.conf.cpu().numpy(),
            class_id=results[0].boxes.cls.cpu().numpy().astype(int)
        )

        # Process the first detected license plate
        if len(detections.xyxy) > 0:
            x1, y1, x2, y2 = map(int, detections.xyxy[0])  # Take the first detection
            cropped_frame = frame[y1:y2, x1:x2]
            pil_image = Image.fromarray(cv2.cvtColor(cropped_frame, cv2.COLOR_BGR2RGB))

            try:
                # Extract text using Tesseract OCR
                number_plate_text = pytesseract.image_to_string(
                    pil_image, config="--psm 7 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
                ).strip()

                if number_plate_text:
                    # Save the cropped plate image
                    cropped_plate_path = os.path.join(cropped_dir, "first_detected_plate.jpg")
                    cv2.imwrite(cropped_plate_path, cropped_frame)

                    print(f"Detected Number Plate: {number_plate_text}")
                    break  # Exit loop after detecting the first license plate
            except Exception as e:
                print(f"Error during OCR: {e}")

    cap.release()

    if cropped_plate_path is None or number_plate_text is None:
        print("No license plate found in the video.")

    return cropped_plate_path, number_plate_text

def gradio_video_processor(video):
    """
    Gradio wrapper to process video and return the first detected license plate and its text.
    """
    video_path = video  # Use the file path provided by Gradio
    cropped_plate_path, plate_text = process_video(video_path)

    if cropped_plate_path and plate_text:
        html_result = f"<h3>Detected License Plate</h3>"
        html_result += f"<p></p><img src='{cropped_plate_path}' width='300'><br>"
        return cropped_plate_path, html_result
    else:
        return None, "<h3>No license plate detected in the video.</h3>"

# Gradio interface
iface = gr.Interface(
    fn=gradio_video_processor,
    inputs=gr.Video(label="Upload a Video to detect license plate"),
    outputs=[
        gr.Image(label="Cropped Plate Image"),  # Shows the cropped plate
        gr.HTML(label="License Plate Detected")  # Shows the extracted text
    ],
    title="License Plate Detection",
    description="Upload a video to detect the first license plate and extract its text."
)

if __name__ == "__main__":
    iface.launch()