import torch
import easyocr
import gradio as gr
from PIL import Image
import cv2
import numpy as np
import os

# Load YOLOv5 pre-trained model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # Use YOLOv5s for speed

# Initialize EasyOCR for license plate recognition
reader = easyocr.Reader(['en'])

# Directory to save images of non-helmet riders
os.makedirs("non_helmet_riders", exist_ok=True)

# Function to enhance the image for better number plate recognition
def preprocess_image_for_ocr(image):
    # Convert image to grayscale
    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    
    # Apply thresholding to binarize the image (white text on black background)
    _, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)
    
    return thresh

# Function to detect non-helmet riders and their license plates
def detect_non_helmet_and_plate(image):
    img_np = np.array(image)
    results = model(image)

    # Default outputs
    helmet_status = "Pass"
    license_plate_text = "I can't detect image"
    license_plate_image = None

    # Parse YOLO results
    non_helmet_detected = False
    for *xyxy, conf, cls in results.xyxy[0]:
        class_id = int(cls)
        if class_id == 0:  # Class 0 is 'person' in YOLOv5s
            non_helmet_detected = True
            helmet_status = "Fail"
            cv2.rectangle(img_np, (int(xyxy[0]), int(xyxy[1])), 
                          (int(xyxy[2]), int(xyxy[3])), (0, 0, 255), 2)  # Red box
            cv2.putText(img_np, "No Helmet", 
                        (int(xyxy[0]), int(xyxy[1]) - 10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

            # Save the image of the non-helmet rider
            cropped_img = img_np[int(xyxy[1]):int(xyxy[3]), int(xyxy[0]):int(xyxy[2])]
            rider_image_path = f"non_helmet_riders/rider_{np.random.randint(10000)}.jpg"
            cv2.imwrite(rider_image_path, cropped_img)

    # Detect license plate if a non-helmet rider is found
    if non_helmet_detected:
        plate_text = reader.readtext(preprocess_image_for_ocr(image))  # Preprocess image before passing to OCR
        for detection in plate_text:
            text = detection[1]
            # Filter for license plate-like text
            if len(text) > 5 and text.isalnum():  # Assuming plates have a minimum length and alphanumeric
                license_plate_text = text

                # Create the cropped image of the plate
                plate_img = np.array(image)[int(detection[0][0][1]):int(detection[0][2][1]),
                                            int(detection[0][0][0]):int(detection[0][2][0])]
                license_plate_image = Image.fromarray(plate_img)
                break

    # Convert the processed image back to PIL for Gradio display
    img_pil = Image.fromarray(img_np)
    return img_pil, helmet_status, license_plate_image, license_plate_text  # Returning the image, helmet status, plate image, and license plate number

# Function to capture live video frame from webcam
def capture_webcam_frame():
    cap = cv2.VideoCapture(0)
    ret, frame = cap.read()
    cap.release()
    if ret:
        img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        return detect_non_helmet_and_plate(img)
    return None, "Error", "I can't detect image", "I can't detect image"

# Set up Gradio interface with both upload and webcam inputs
def interface_fn(image, capture_from_webcam):
    if capture_from_webcam:
        return capture_webcam_frame()
    else:
        return detect_non_helmet_and_plate(image)

# Set up Gradio interface
interface = gr.Interface(
    fn=interface_fn,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Checkbox(label="Capture from Webcam")  # Use Checkbox to toggle between upload and webcam
    ],
    outputs=[
        gr.Image(type="pil", label="Processed Image"),  # Output: Processed Image
        gr.Textbox(label="Helmet Status"),  # Output: Helmet Status
        gr.Image(type="pil", label="License Plate Image"),  # Output: License Plate Image
        gr.Textbox(label="License Plate Number")  # Output: License Plate Number
    ],
    title="Helmet and License Plate Detection",
    description="Detect riders without helmets. If a rider is without a helmet, capture their image and license plate.",
)

# Launch Gradio app
interface.launch(share=True)