import streamlit as st
from streamlit_webrtc import webrtc_streamer, VideoProcessorBase
import av
from transformers import DetrImageProcessor, DetrForObjectDetection, TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image, ImageDraw
import numpy as np
import torch

# Step 1: Load Models
# DETR for object detection
detr_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
detr_model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")

# TrOCR for text recognition
trocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
trocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-stage1")

# Authorized car database for verification
authorized_cars = {"KA01AB1234", "MH12XY5678", "DL8CAF9090", "CH01AG2863"}  # Example data


# Step 2: Define Helper Functions
def detect_license_plate(frame):
    """
    Detect license plates in the frame using DETR.
    """
    pil_image = Image.fromarray(frame)
    inputs = detr_processor(images=pil_image, return_tensors="pt")
    outputs = detr_model(**inputs)
    
    # Get bounding boxes
    target_sizes = torch.tensor([pil_image.size[::-1]])
    results = detr_processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)
    return results[0]["boxes"], pil_image


def recognize_text_from_plate(cropped_plate):
    """
    Recognize text from the cropped license plate image using TrOCR.
    """
    inputs = trocr_processor(images=cropped_plate, return_tensors="pt")
    outputs = trocr_model.generate(**inputs)
    return trocr_processor.batch_decode(outputs, skip_special_tokens=True)[0]


def verify_plate(plate_text):
    """
    Check if the recognized plate text exists in the authorized cars database.
    """
    if plate_text in authorized_cars:
        return f"✅ Access Granted: {plate_text}"
    else:
        return f"❌ Access Denied: {plate_text}"


# Step 3: Custom Video Processor for WebRTC
class LicensePlateProcessor(VideoProcessorBase):
    """
    Custom video processor to handle video frames in real-time.
    """
    def recv(self, frame: av.VideoFrame):
        frame = frame.to_ndarray(format="bgr24")  # Convert frame to NumPy array
        boxes, pil_image = detect_license_plate(frame)
        draw = ImageDraw.Draw(pil_image)

        recognized_plates = []
        for box in boxes:
            # Crop detected license plate
            cropped_plate = pil_image.crop((box[0], box[1], box[2], box[3]))
            plate_text = recognize_text_from_plate(cropped_plate)
            recognized_plates.append(plate_text)

            # Draw bounding box and label on the image
            draw.rectangle(box.tolist(), outline="red", width=3)
            draw.text((box[0], box[1]), plate_text, fill="red")

        # Convert back to OpenCV format
        processed_frame = np.array(pil_image)

        # Log results in Streamlit UI
        for plate_text in recognized_plates:
            st.write(verify_plate(plate_text))

        return av.VideoFrame.from_ndarray(processed_frame, format="bgr24")


# Step 4: Streamlit Interface
st.title("Real-Time Car Number Plate Recognition")
st.write("This app uses Hugging Face Transformers and WebRTC for real-time processing.")

# Start WebRTC Streamer
webrtc_streamer(
    key="plate-recognition",
    video_processor_factory=LicensePlateProcessor,
    rtc_configuration={
        # Required to ensure WebRTC works across networks
        "iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]
    }
)