Update app.py

app.py

@@ -2,47 +2,53 @@ import streamlit as st
 import cv2
 import numpy as np
 import easyocr
+from ultralytics import YOLO
 
 # 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)
+    return model
+
 # Load EasyOCR reader
 @st.cache_resource
 def load_easyocr_reader():
     return easyocr.Reader(['en'], gpu=False)
 
+# Initialize models
+yolo_model = load_yolo_model()
 ocr_reader = load_easyocr_reader()
 
-# Load OpenCV pre-trained Haar Cascade classifier for license plates
-plate_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_russian_plate_number.xml')
-
-# Function to process image and detect license plates
-def process_image(image):
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    plates = plate_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(100, 100))
-
-    if len(plates) == 0:
-        st.warning("No license plate detected in the image!")
-        return
-
-    # Loop through detections
-    for i, (x, y, w, h) in enumerate(plates):
-        cv2.rectangle(image, (x, y), (x+w, y+h), (0, 255, 0), 2)
-        cropped_plate = image[y:y+h, x:x+w]
-        cropped_rgb = cv2.cvtColor(cropped_plate, cv2.COLOR_BGR2RGB)
-
-        # Perform OCR on the cropped image
-        text_results = ocr_reader.readtext(cropped_rgb, detail=0)
-        detected_text = " ".join(text_results)
-        
-        # Display extracted text
-        st.write(f"**Extracted Text (Plate {i+1}):** {detected_text}")
-
-    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    st.image(image_rgb, caption="Processed Image", use_container_width=True)
+# Function to process the uploaded image
+def process_image(image, confidence_threshold=0.5):
+    # Perform license plate detection using YOLO
+    results = yolo_model(image, conf=confidence_threshold)
+    annotated_image = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
+    st.image(annotated_image, caption="Detected License Plate(s)", use_container_width=True)
+    
+    # Loop through detections and perform OCR
+    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[y1:y2, x1:x2]
+            cropped_image_rgb = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2RGB)
+            st.image(cropped_image_rgb, caption=f"Cropped License Plate {i+1}", use_container_width=True)
+            
+            # Perform OCR on the cropped image
+            text_results = ocr_reader.readtext(cropped_image_rgb, detail=0)
+            detected_text = " ".join(text_results)
+            st.write(f"**Extracted Text (Plate {i+1}):** {detected_text}")
 
 # Function to process video and detect license plates
-def process_video(video_path):
+def process_video(video_path, confidence_threshold=0.5):
     cap = cv2.VideoCapture(video_path)
     
     if not cap.isOpened():
@@ -54,20 +60,24 @@ def process_video(video_path):
         if not ret:
             break
         
-        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-        plates = plate_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(100, 100))
-        
-        for i, (x, y, w, h) in enumerate(plates):
-            cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
-            cropped_plate = frame[y:y+h, x:x+w]
-            cropped_rgb = cv2.cvtColor(cropped_plate, cv2.COLOR_BGR2RGB)
-            
-            text_results = ocr_reader.readtext(cropped_rgb, detail=0)
-            detected_text = " ".join(text_results)
-            cv2.putText(frame, detected_text, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
-        
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        st.image(frame_rgb, caption="Video Frame", use_container_width=True)
+        results = yolo_model(frame, conf=confidence_threshold)
+        annotated_frame = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
+
+        st.image(annotated_frame, caption="Processed Video Frame", use_container_width=True)
+
+        for result in results:
+            boxes = result.boxes.xyxy.cpu().numpy().astype(int)
+            for i, box in enumerate(boxes):
+                x1, y1, x2, y2 = box
+                cropped_plate = frame[y1:y2, x1:x2]
+                cropped_rgb = cv2.cvtColor(cropped_plate, cv2.COLOR_BGR2RGB)
+                
+                # Perform OCR on the cropped image
+                text_results = ocr_reader.readtext(cropped_rgb, detail=0)
+                detected_text = " ".join(text_results)
+                cv2.putText(annotated_frame, detected_text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+
+        st.image(annotated_frame, caption="Video Frame with OCR", use_container_width=True)
 
     cap.release()