import torch
from ultralytics import YOLO
from ultralytics.nn.tasks import DetectionModel
import torch.serialization

import cv2
import numpy as np
import gradio as gr
import easyocr
from transformers import T5ForConditionalGeneration, T5Tokenizer

# ✅ 1. Allow YOLO model class from Ultralytics to be loaded safely
# torch.serialization.add_safe_globals([DetectionModel])

# ✅ 2. Load YOLOv8 fine-tuned model (make sure best.pt is in same directory)
model = YOLO("best.pt")

# ✅ 3. Load EasyOCR Reader
reader = easyocr.Reader(['en'], gpu=False)

# ✅ 4. Load T5 GenAI Model for OCR correction
t5_tokenizer = T5Tokenizer.from_pretrained("vennify/t5-base-grammar-correction")
t5_model = T5ForConditionalGeneration.from_pretrained("vennify/t5-base-grammar-correction")


# === Pipeline Function ===
def detect_and_read(image):
    # Step 1: Run YOLOv8 detection
    results = model(image)[0]
    annotated_frame = image.copy()

    final_texts = []

    for box in results.boxes:
        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
        plate_crop = image[y1:y2, x1:x2]

        # Step 2: OCR on the detected plate
        ocr_result = reader.readtext(plate_crop, detail=0)
        ocr_text = " ".join(ocr_result)

        # Step 3: GenAI Correction with T5
        input_text = "correct: " + ocr_text
        input_ids = t5_tokenizer.encode(input_text, return_tensors="pt")
        outputs = t5_model.generate(input_ids, max_length=32, num_beams=4, early_stopping=True)
        corrected_text = t5_tokenizer.decode(outputs[0], skip_special_tokens=True)

        final_texts.append(corrected_text)

        # Draw box and text
        cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(annotated_frame, corrected_text, (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

    return annotated_frame, final_texts


# === Gradio Interface ===
interface = gr.Interface(
    fn=detect_and_read,
    inputs=gr.Image(type="numpy", label="Upload Car Image"),
    outputs=[
        gr.Image(type="numpy", label="Detected + Annotated Image"),
        gr.Textbox(label="Corrected Plate Text")
    ],
    title="License Plate Detection with YOLOv8 + EasyOCR + GenAI"
)

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