Update app.py

app.py

@@ -29,24 +29,27 @@ 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)
-
-    # Loop through detections and perform OCR
+    
+    # Prepare the cropped images and detected text for each plate
     license_plate_text = []
     cropped_images = []
-
+    
+    # Loop through detections and perform OCR
     for result in results:
         boxes = result.boxes.xyxy.cpu().numpy().astype(int)
         confidences = result.boxes.conf.cpu().numpy()
-
+        
         if len(boxes) == 0:
             st.warning("No license plate detected!")
             return [], [], None
+        
         for i, (box, conf) in enumerate(zip(boxes, confidences)):
             x1, y1, x2, y2 = box
             # Draw bounding box on the annotated image
             cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (255, 0, 0), 2)
             cv2.putText(annotated_image, f"{conf:.2f}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
 
+            # Crop the license plate from the image
             cropped_image = image[y1:y2, x1:x2]
             cropped_image_rgb = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2RGB)
             cropped_images.append(cropped_image_rgb)
@@ -55,11 +58,11 @@ def process_image(image, confidence_threshold=0.5):
             text_results = ocr_reader.readtext(cropped_image_rgb, detail=0)
             detected_text = " ".join(text_results)
             license_plate_text.append(detected_text)
-
+    
     return license_plate_text, cropped_images, annotated_image
 
 # Sidebar input for file upload
-uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("Upload an Image or Video", type=["mp4", "avi", "mov", "jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
     # Read and process the image
@@ -70,11 +73,13 @@ if uploaded_file is not None:
     c1, c2, c3 = st.columns(3)
 
     with c1:
-        st.image(image, caption='Uploaded Image with Bounding Boxes', use_container_width=True)
-
-    license_plate_text, cropped_images, annotated_image = process_image(image, confidence_threshold)
+        # Display uploaded image with bounding boxes
+        license_plate_text, cropped_images, annotated_image = process_image(image, confidence_threshold)
+        if annotated_image is not None:
+            st.image(annotated_image, caption='Uploaded Image with Bounding Boxes', use_container_width=True)
 
     with c2:
+        # Display cropped license plates
         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)
@@ -82,14 +87,12 @@ if uploaded_file is not None:
             st.write('No License Plate Detected')
 
     with c3:
+        # Display the extracted text for license plates
         if license_plate_text:
             st.success(', '.join(license_plate_text))
             st.write('License Plate Text')
         else:
             st.write('No text detected')
 
-    # Display the annotated image with bounding boxes
-    st.image(annotated_image, caption='Annotated Image with Bounding Boxes', use_container_width=True)
-
 st.markdown("---")
 st.info("This application uses Fine Tuned YOLOv8 for detection and EasyOCR for text recognition.")