Update app.py

app.py

@@ -3,6 +3,7 @@ import cv2
 import numpy as np
 import easyocr
 from ultralytics import YOLO
+from PIL import Image
 
 # Title of the app
 st.title("License Plate Recognition System🚗")
@@ -28,101 +29,58 @@ def process_image(image, confidence_threshold=0.5):
     # Perform license plate detection
     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
+    license_plate_text = []
+    cropped_images = []
+
     for result in results:
         boxes = result.boxes.xyxy.cpu().numpy().astype(int)
         if len(boxes) == 0:
             st.warning("No license plate detected!")
-            return
+            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)
-            
+            cropped_images.append(cropped_image_rgb)
+
             # 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}")
-            st.write(f"**Confidence Score:** {result.boxes.conf.cpu().numpy()[i]:.2f}")
-
-# Function to process video and detect license plates
-def process_video(video_path, confidence_threshold=0.5, output_path="output_video.mp4"):
-    # Open the video file
-    cap = cv2.VideoCapture(video_path)
-
-    # Get video frame dimensions
-    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
-    # Create VideoWriter object to save the output video
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Codec for mp4
-    out = cv2.VideoWriter(output_path, fourcc, 20.0, (frame_width, frame_height))  # 20 FPS
-
-    if not cap.isOpened():
-        st.error("Error opening video stream or file")
-        return
-
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
+            license_plate_text.append(detected_text)
 
-        results = yolo_model(frame, conf=confidence_threshold)
-        annotated_frame = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
-
-        # Loop through detections and perform OCR
-        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)
-
-                # Optionally add detected text on the annotated frame
-                cv2.putText(annotated_frame, detected_text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
-
-        # Write the annotated frame to the output video
-        out.write(cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR))
-
-    cap.release()
-    out.release()
-
-    st.success(f"Video processing complete. Output video saved to {output_path}")
-
-    # Provide a download link for the processed video
-    with open(output_path, "rb") as f:
-        st.download_button(label="Download Processed Video", data=f, file_name=output_path)
+    return license_plate_text, cropped_images
 
 # Sidebar input for file upload
-uploaded_file = st.file_uploader("Upload an Image or Video", type=["mp4", "avi", "mov", "jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    # Check if it's an image or video
-    file_type = uploaded_file.type
-
-    if file_type.startswith("image"):
-        # Read and process the image
-        image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
-        confidence_threshold = st.sidebar.slider("Confidence Threshold", 0.0, 1.0, 0.5, 0.01)
-        process_image(image, confidence_threshold)
-
-    elif file_type.startswith("video"):
-        # Save the uploaded video to a temporary file
-        video_bytes = uploaded_file.read()
-        video_path = "/tmp/uploaded_video.mp4"
-        with open(video_path, "wb") as f:
-            f.write(video_bytes)
-
-        # Process the video and save the output
-        output_path = "/tmp/output_video.mp4"
-        process_video(video_path, confidence_threshold=0.5, output_path=output_path)
-
+    # Read and process the image
+    image = np.array(Image.open(uploaded_file))
+    confidence_threshold = st.sidebar.slider("Confidence Threshold", 0.0, 1.0, 0.5, 0.01)
+
+    # Create three columns to display the images side by side
+    c1, c2, c3 = st.columns(3)
+
+    with c1:
+        st.image(image, caption='Uploaded Image', use_container_width=True)
+
+    license_plate_text, cropped_images = process_image(image, confidence_threshold)
+
+    with c2:
+        if cropped_images:
+            for i, cropped_image in enumerate(cropped_images):
+                st.image(cropped_image, caption=f'Cropped License Plate {i+1}', use_container_width=True)
+        else:
+            st.write('No License Plate Detected')
+
+    with c3:
+        if license_plate_text:
+            st.success(', '.join(license_plate_text))
+            st.write('License Plate Text')
+        else:
+            st.write('No text detected')
+    
 st.markdown("---")
-st.info("This application uses Fine Tuned YOLOv8 for detection and EasyOCR for text recognition.")
+st.info("This application uses Fine Tuned YOLOv8 for detection and EasyOCR for text recognition.")