Spaces:
Sleeping
Sleeping
| import torch | |
| import easyocr | |
| import gradio as gr | |
| from PIL import Image | |
| import cv2 | |
| import numpy as np | |
| import os | |
| # Load YOLOv5 pre-trained model | |
| model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # Use YOLOv5s for speed | |
| # Initialize EasyOCR for license plate recognition | |
| 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): | |
| # 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 | |
| def detect_non_helmet_and_plate(image): | |
| img_np = np.array(image) | |
| results = model(image) | |
| # Default outputs | |
| helmet_status = "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.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)) # Preprocess image before passing to OCR | |
| for detection in plate_text: | |
| text = detection[1] | |
| # 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, 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" | |
| # 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"), | |
| gr.Checkbox(label="Capture from Webcam") # Use Checkbox to toggle between upload and webcam | |
| ], | |
| outputs=[ | |
| 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.", | |
| ) | |
| # Launch Gradio app | |
| interface.launch(share=True) |