Update app.py

app.py

@@ -15,10 +15,14 @@ reader = easyocr.Reader(['en'])
 # Directory to save images of non-helmet riders
 os.makedirs("non_helmet_riders", exist_ok=True)
 
-# Function to preprocess image for better OCR performance
+# Function to enhance the image for better number plate recognition
 def preprocess_image_for_ocr(image):
+    # Convert image to grayscale
     gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
+    
+    # Apply thresholding to binarize the image (white text on black background)
     _, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)
+    
     return thresh
 
 # Function to detect non-helmet riders and their license plates
@@ -26,50 +30,81 @@ def detect_non_helmet_and_plate(image):
     img_np = np.array(image)
     results = model(image)
 
+    # Default outputs
     helmet_status = "Pass"
-    license_plate_text = "No license plate detected"
+    license_plate_text = "I can't detect image"
     license_plate_image = None
-    non_helmet_detected = False
 
+    # Parse YOLO results
+    non_helmet_detected = False
     for *xyxy, conf, cls in results.xyxy[0]:
         class_id = int(cls)
-        if class_id == 0:  # Class 0 is 'person'
+        if class_id == 0:  # Class 0 is 'person' in YOLOv5s
             non_helmet_detected = True
             helmet_status = "Fail"
-            cv2.rectangle(img_np, (int(xyxy[0]), int(xyxy[1])),
-                          (int(xyxy[2]), int(xyxy[3])), (0, 0, 255), 2)
-            cv2.putText(img_np, "No Helmet",
-                        (int(xyxy[0]), int(xyxy[1]) - 10),
+            cv2.rectangle(img_np, (int(xyxy[0]), int(xyxy[1])), 
+                          (int(xyxy[2]), int(xyxy[3])), (0, 0, 255), 2)  # Red box
+            cv2.putText(img_np, "No Helmet", 
+                        (int(xyxy[0]), int(xyxy[1]) - 10), 
                         cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
 
+            # Save the image of the non-helmet rider
+            cropped_img = img_np[int(xyxy[1]):int(xyxy[3]), int(xyxy[0]):int(xyxy[2])]
+            rider_image_path = f"non_helmet_riders/rider_{np.random.randint(10000)}.jpg"
+            cv2.imwrite(rider_image_path, cropped_img)
+
+    # Detect license plate if a non-helmet rider is found
     if non_helmet_detected:
-        plate_text = reader.readtext(preprocess_image_for_ocr(image))
+        plate_text = reader.readtext(preprocess_image_for_ocr(image))  # Preprocess image before passing to OCR
         for detection in plate_text:
             text = detection[1]
-            if len(text) > 5 and text.isalnum():
+            # Filter for license plate-like text
+            if len(text) > 5 and text.isalnum():  # Assuming plates have a minimum length and alphanumeric
                 license_plate_text = text
+
+                # Create the cropped image of the plate
+                plate_img = np.array(image)[int(detection[0][0][1]):int(detection[0][2][1]),
+                                            int(detection[0][0][0]):int(detection[0][2][0])]
+                license_plate_image = Image.fromarray(plate_img)
                 break
 
+    # Convert the processed image back to PIL for Gradio display
     img_pil = Image.fromarray(img_np)
-    return img_pil, helmet_status, None, license_plate_text
+    return img_pil, helmet_status, license_plate_image, license_plate_text  # Returning the image, helmet status, plate image, and license plate number
+
+# Function to capture live video frame from webcam
+def capture_webcam_frame():
+    cap = cv2.VideoCapture(0)
+    ret, frame = cap.read()
+    cap.release()
+    if ret:
+        img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        return detect_non_helmet_and_plate(img)
+    return None, "Error", "I can't detect image", "I can't detect image"
 
-# Gradio interface
-def interface_fn(image):
-    return detect_non_helmet_and_plate(image)
+# Set up Gradio interface with both upload and webcam inputs
+def interface_fn(image, capture_from_webcam):
+    if capture_from_webcam:
+        return capture_webcam_frame()
+    else:
+        return detect_non_helmet_and_plate(image)
 
+# Set up Gradio interface
 interface = gr.Interface(
     fn=interface_fn,
-    inputs=gr.Image(type="pil", label="Upload Image"),
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Checkbox(label="Capture from Webcam")  # Use Checkbox to toggle between upload and webcam
+    ],
     outputs=[
-        gr.Image(type="pil", label="Processed Image"),
-        gr.Textbox(label="Helmet Status"),
-        gr.Image(type="pil", label="License Plate Image"),
-        gr.Textbox(label="License Plate Number")
+        gr.Image(type="pil", label="Processed Image"),  # Output: Processed Image
+        gr.Textbox(label="Helmet Status"),  # Output: Helmet Status
+        gr.Image(type="pil", label="License Plate Image"),  # Output: License Plate Image
+        gr.Textbox(label="License Plate Number")  # Output: License Plate Number
     ],
     title="Helmet and License Plate Detection",
-    description="Detect riders without helmets. If a rider is without a helmet, capture their image and license plate."
+    description="Detect riders without helmets. If a rider is without a helmet, capture their image and license plate.",
 )
 
 # Launch Gradio app
-if __name__ == "__main__":
-    interface.launch()
+interface.launch(share=True)