app.py

import streamlit as st
from ultralytics import YOLO
import numpy as np
import cv2
from PIL import Image

# Model labels
model1Labels = {0: 'single_number_plate', 1: 'double_number_plate'}

model2Labels = {
    0: '0', 1: '1', 2: '2', 3: '3', 4: '4', 5: '5', 6: '6', 7: '7', 8: '8', 9: '9', 10: 'A', 11: 'B', 12: 'C',
    13: 'D', 14: 'E', 15: 'F', 16: 'G', 17: 'H', 18: 'I', 19: 'J', 20: 'K', 21: 'L', 22: 'M', 23: 'N', 24: 'O',
    25: 'P', 26: 'Q', 27: 'R', 28: 'S', 29: 'T', 30: 'U', 31: 'V', 32: 'W', 33: 'X', 34: 'Y', 35: 'Z'
}

# Load models
model = YOLO("models/LP-detection.pt")
model2 = YOLO("models/Charcter-LP.pt")

def prediction(image):
    result = model.predict(source=image, conf=0.5)
    boxes = result[0].boxes
    height = boxes.xywh
    crd = boxes.data

    n = len(crd)
    lp_number = []
    img_lp_final = None  

    for i in range(n):
        ht = int(height[i][3])
        c = int(crd[i][5])

        xmin = int(crd[i][0])
        ymin = int(crd[i][1])
        xmax = int(crd[i][2])
        ymax = int(crd[i][3])

        img_lp = image[ymin:ymax, xmin:xmax]
        img_lp_final = img_lp.copy()  # Store the cropped image for display
        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)

        h = np.median(ht)

        # Second Model Prediction
        result2 = model2.predict(source=img_lp, conf=0.25)
        boxes_ocr = result2[0].boxes
        data2 = boxes_ocr.data

        n2 = len(data2)
        xaxis0, xaxis11, xaxis12 = [], [], []
        label0, label11, label12 = [], [], []
        numberPlate = ""

        if c == 0:  # Single line license plate
            for i in range(n2):
                x = int(data2[i][2])
                xaxis0.append(x)
                l = int(data2[i][5])
                label0.append(l)

            # Sort characters by x-axis for single line
            sorted_labels = [label0[i] for i in np.argsort(xaxis0)]
            numberPlate = ''.join([model2Labels.get(l) for l in sorted_labels])
            lp_number.append(numberPlate)

        elif c == 1:  # Double line license plate
            for i in range(n2):
                x = int(data2[i][0])
                y = int(data2[i][3])
                l = int(data2[i][5])
                if y < (h / 2):
                    xaxis11.append(x)
                    label11.append(l)
                else:
                    xaxis12.append(x)
                    label12.append(l)

            # Sort characters by x-axis for double line (upper and lower separately)
            sorted_labels11 = [label11[i] for i in np.argsort(xaxis11)]
            sorted_labels12 = [label12[i] for i in np.argsort(xaxis12)]
            numberPlate = ''.join([model2Labels.get(l) for l in sorted_labels11 + sorted_labels12])
            lp_number.append(numberPlate)

    return lp_number, img_lp_final  

st.title('License Plate Recognition 🚗')
st.header('Upload an image of a license plate to get the License number.')

# Define example images (update with actual paths)
example_images = {
    "Car ": "test/audiR8V10.jpg",
    "Car 2": "test/c7.jpg",
    "Car 3": "test/c4.jpg",
    "CCTV B/W": "test/cctv img plate.jpg",
    "Bike": "test/BikeNumberPlate.jpg",
    "Bus": "test/bus.jpg",
}

# File uploader
uploaded_file = st.file_uploader("Choose an image...", type="jpg")



c1, c2 = st.columns(2)
for name, path in example_images.items():
    with c1:
        example_img = Image.open(path)

image = None
if uploaded_file is not None:
    image = np.array(Image.open(uploaded_file))
else:
    st.header("Or choose an example image from below dropdown:")
    selected_example = st.selectbox("", list(example_images.keys()))
    if selected_example:
        image = np.array(Image.open(example_images[selected_example]))

if image is not None:
    c1, c2, c3 = st.columns(3)

    with c1:
        st.image(image, caption='Uploaded Image', use_column_width=True)

    license_plate_text, img_lp = prediction(image)

    with c2:
        if img_lp is not None:
            st.image(img_lp, caption='Cropped License Plate', use_column_width=True)
        else:
            st.write('No License Plate Detected')

    with c3:
        st.success(', '.join(license_plate_text))
        st.write('License Plate Text')