Update app.py

app.py

@@ -15,58 +15,61 @@ reader = easyocr.Reader(['en'])
 # Directory to save images of non-helmet riders
 os.makedirs("non_helmet_riders", exist_ok=True)
 
+# Function to enhance the image for better number plate recognition
+def preprocess_image_for_ocr(image):
+    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
+    _, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)
+    return thresh
+
 # Function to detect non-helmet riders and their license plates
 def detect_non_helmet_and_plate(image):
     img_np = np.array(image)
     results = model(image)
 
-    # Default outputs
     helmet_status = "Pass"
-    license_plate_text = "Pass"
+    license_plate_text = "I can't detect image"
+    license_plate_image = None
 
-    # 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' 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)  # Red box
-            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)
+            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(img_np)
+        plate_text = reader.readtext(preprocess_image_for_ocr(image))
         for detection in plate_text:
             text = detection[1]
-            if len(text) > 5:  # Assuming plates have a minimum length
+            if len(text) > 5 and text.isalnum():
                 license_plate_text = text
                 break
 
-    # Convert the processed image back to PIL for Gradio display
     img_pil = Image.fromarray(img_np)
-    return img_pil, helmet_status, license_plate_text
+    return img_pil, helmet_status, None, license_plate_text
+
+# Gradio interface
+def interface_fn(image):
+    return detect_non_helmet_and_plate(image)
 
-# Set up Gradio interface
 interface = gr.Interface(
-    fn=detect_non_helmet_and_plate,
+    fn=interface_fn,
     inputs=gr.Image(type="pil", label="Upload Image"),
     outputs=[
         gr.Image(type="pil", label="Processed Image"),
         gr.Textbox(label="Helmet Status"),
-        gr.Textbox(label="License Plate"),
+        gr.Image(type="pil", label="License Plate Image"),
+        gr.Textbox(label="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
-interface.launch(share=True)
+if __name__ == "__main__":
+    interface.launch()