import streamlit as st
from PIL import Image
import numpy as np
from ultralytics import YOLO
from paddleocr import PaddleOCR
from streamlit_webrtc import webrtc_streamer, VideoProcessorBase
import av

# Load the license plate segmentation model (YOLOv8)
segmentation_model = YOLO('plate_segment_best.pt')

class LicensePlateProcessor(VideoProcessorBase):
    def __init__(self):
        self.ocr = PaddleOCR(use_angle_cls=True, lang='en')

    def process_image(self, image):
        # Perform license plate segmentation using the YOLOv8 model
        results = segmentation_model(image)

        # Extract the bounding box with the highest confidence score
        boxes = results[0].boxes.xyxy.numpy()
        if len(boxes) == 0:
            return None

        # Assuming the first box is the license plate (you can adjust based on your model's output)
        x1, y1, x2, y2 = boxes[0]
        segmented_plate = image[int(y1):int(y2), int(x1):int(x2)]

        # Use PaddleOCR to extract text from the segmented plate area
        ocr_results = self.ocr.ocr(segmented_plate, cls=True)

        # Extract the plate number from OCR results
        plate_number = self.extract_plate_number(ocr_results)

        return plate_number

    def extract_plate_number(self, ocr_results):
        plate_number = ""
        for line in ocr_results:
            for word in line:
                plate_number += word[1][0] + " "
        return plate_number.strip()

    def recv(self, frame):
        img = frame.to_ndarray(format="bgr24")
        plate_number = self.process_image(img)

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

def main():
    st.title('License Plate Reader')

    # Input suspected criminal plate numbers
    suspected_numbers = st.text_area('Enter suspected criminal plate numbers (one per line):')
    suspected_plates = [num.strip() for num in suspected_numbers.split('\n') if num.strip()]

    # File uploader for image input
    uploaded_file = st.file_uploader('Choose an image...', type=['jpg', 'png'])
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)

        # Convert PIL image to numpy array
        image_np = np.array(image)

        # Process image to extract license plate number
        plate_number = LicensePlateProcessor().process_image(image_np)

        if plate_number:
            st.write('Detected License Plate Number:', plate_number)
            if suspected_plates and plate_number in suspected_plates:
                st.write('**Criminal Car Detected!**')
            else:
                st.write('Car is not in the suspected criminal list.')
        else:
            st.write('License Plate Number could not be extracted.')

    # Webcam input for live video stream
    ctx = webrtc_streamer(
        key="license-plate-reader",
        video_processor_factory=LicensePlateProcessor,
        rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]},
        media_stream_constraints={"video": True, "audio": False},
        async_processing=True,
    )

    if ctx.video_processor:
        plate_number = ctx.video_processor.recv(ctx.video_processor)

        if plate_number:
            st.write('Detected License Plate Number:', plate_number)
            if suspected_plates and plate_number in suspected_plates:
                st.write('**Criminal Car Detected!**')
            else:
                st.write('Car is not in the suspected criminal list.')
        else:
            st.write('License Plate Number could not be extracted.')

if __name__ == '__main__':
    main()