Update app.py

app.py

@@ -1,66 +1,135 @@
 import streamlit as st
-import numpy as np
-import torch
 from ultralytics import YOLO
-import easyocr
+import numpy as np
+import cv2
 from PIL import Image
 
-# Title of the app
-st.title("License Plate Recognition 🚗")
-
-# Load the YOLO model for license plate detection
-@st.cache_resource
-def load_yolo_model():
-    model_path = "best.pt"  # Replace with your model file
-    model = YOLO(model_path)
-    if torch.cuda.is_available():
-        model.to("cuda")
-    return model
-
-# Load EasyOCR reader
-@st.cache_resource
-def load_easyocr_reader():
-    return easyocr.Reader(['en'], gpu=torch.cuda.is_available())
-
-# Initialize models
-yolo_model = load_yolo_model()
-ocr_reader = load_easyocr_reader()
-
-# Function to process the uploaded image
-def process_image(image, confidence_threshold=0.5):
-    # Convert the PIL image to a numpy array
-    image_np = np.array(image)
-    
-    # Perform license plate detection
-    results = yolo_model(image_np, conf=confidence_threshold)
-    annotated_image = Image.fromarray(results[0].plot())
-    st.image(annotated_image, caption="Detected License Plate(s)", use_container_width=True)
-    
-    # Loop through detections
-    for result in results:
-        boxes = result.boxes.xyxy.cpu().numpy().astype(int)
-        if len(boxes) == 0:
-            st.warning("No license plate detected!")
-            return
-        for i, box in enumerate(boxes):
-            x1, y1, x2, y2 = box
-            cropped_image = image_np[y1:y2, x1:x2]
-            cropped_image_pil = Image.fromarray(cropped_image)
-            st.image(cropped_image_pil, caption=f"Cropped License Plate {i+1}", use_container_width=True)
-            
-            # Perform OCR on the cropped image
-            text_results = ocr_reader.readtext(np.array(cropped_image_pil), detail=0)
-            detected_text = " ".join(text_results)
-            st.write(f"**Extracted Text (Plate {i+1}):** {detected_text}")
-
-# Sidebar inputs
-confidence_threshold = st.sidebar.slider("Confidence Threshold", 0.0, 1.0, 0.5, 0.01)
-uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
+# 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:
-    # Open the uploaded file as a PIL image
-    image = Image.open(uploaded_file)
-    process_image(image, confidence_threshold)
+    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')
 
-st.markdown("---")
-st.info("**Note:** This application uses EasyOCR for text recognition. Results may vary depending on image quality and lighting conditions.")
+    with c3:
+        st.success(', '.join(license_plate_text))
+        st.write('License Plate Text')