Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch | |
| from ultralytics import YOLO | |
| from PIL import Image, ImageDraw, ImageFont | |
| import numpy as np | |
| import pandas as pd | |
| import os | |
| import cv2 | |
| import time | |
| import zipfile | |
| import io | |
| from datetime import datetime | |
| try: | |
| from license_plate_ocr import extract_license_plate_text | |
| OCR_AVAILABLE = True | |
| print("OCR module loaded successfully") | |
| except ImportError as e: | |
| print(f"OCR module not available: {e}") | |
| OCR_AVAILABLE = False | |
| torch.hub.download_url_to_file('https://github.com/Janno1402/Helmet-License-Plate-Detection/blob/main/Sample-Image-1.jpg?raw=true', 'sample_1.jpg') | |
| torch.hub.download_url_to_file('https://github.com/Janno1402/Helmet-License-Plate-Detection/blob/main/Sample-Image-2.jpg?raw=true', 'sample_2.jpg') | |
| torch.hub.download_url_to_file('https://github.com/Janno1402/Helmet-License-Plate-Detection/blob/main/Sample-Image-3.jpg?raw=true', 'sample_3.jpg') | |
| torch.hub.download_url_to_file('https://github.com/Janno1402/Helmet-License-Plate-Detection/blob/main/Sample-Image-4.jpg?raw=true', 'sample_4.jpg') | |
| torch.hub.download_url_to_file('https://github.com/Janno1402/Helmet-License-Plate-Detection/blob/main/Sample-Image-5.jpg?raw=true', 'sample_5.jpg') | |
| model = YOLO("best.pt") | |
| class_names = {0: 'With Helmet', 1: 'Without Helmet', 2: 'License Plate'} | |
| def crop_license_plates(image, detections, extract_text=False): | |
| cropped_plates = [] | |
| try: | |
| if isinstance(image, str): | |
| if not os.path.exists(image): | |
| print(f"Error: Image file not found: {image}") | |
| return cropped_plates | |
| image = Image.open(image) | |
| elif isinstance(image, np.ndarray): | |
| image = Image.fromarray(image) | |
| elif not isinstance(image, Image.Image): | |
| print(f"Error: Unsupported image type: {type(image)}") | |
| return cropped_plates | |
| if image.size[0] == 0 or image.size[1] == 0: | |
| print("Error: Image has zero dimensions") | |
| return cropped_plates | |
| except Exception as e: | |
| print(f"Error loading image: {e}") | |
| return cropped_plates | |
| for i, detection in enumerate(detections): | |
| try: | |
| if detection['Object'] != 'License Plate': | |
| continue | |
| pos_str = detection['Position'].strip('()') | |
| if ',' not in pos_str: | |
| print(f"Error: Invalid position format for detection {i}: {detection['Position']}") | |
| continue | |
| x1, y1 = map(int, pos_str.split(', ')) | |
| dims_str = detection['Dimensions'] | |
| if 'x' not in dims_str: | |
| print(f"Error: Invalid dimensions format for detection {i}: {detection['Dimensions']}") | |
| continue | |
| width, height = map(int, dims_str.split('x')) | |
| if width <= 0 or height <= 0: | |
| print(f"Error: Invalid dimensions for detection {i}: {width}x{height}") | |
| continue | |
| x2, y2 = x1 + width, y1 + height | |
| if x1 < 0 or y1 < 0 or x2 > image.width or y2 > image.height: | |
| print(f"Warning: Bounding box extends beyond image boundaries for detection {i}") | |
| x1 = max(0, x1) | |
| y1 = max(0, y1) | |
| x2 = min(image.width, x2) | |
| y2 = min(image.height, y2) | |
| if x2 <= x1 or y2 <= y1: | |
| print(f"Error: Invalid crop coordinates for detection {i}: ({x1},{y1}) to ({x2},{y2})") | |
| continue | |
| cropped_plate = image.crop((x1, y1, x2, y2)) | |
| if cropped_plate.size[0] == 0 or cropped_plate.size[1] == 0: | |
| print(f"Error: Cropped image has zero dimensions for detection {i}") | |
| continue | |
| plate_data = { | |
| 'image': cropped_plate, | |
| 'confidence': detection['Confidence'], | |
| 'position': detection['Position'], | |
| 'crop_coords': f"({x1},{y1}) to ({x2},{y2})", | |
| 'text': 'Processing...' | |
| } | |
| if extract_text and OCR_AVAILABLE: | |
| try: | |
| print(f"Extracting text from license plate {i+1}...") | |
| plate_text = extract_license_plate_text(cropped_plate) | |
| if plate_text and plate_text.strip() and not plate_text.startswith('Error'): | |
| plate_data['text'] = plate_text.strip() | |
| print(f"Extracted text: {plate_text.strip()}") | |
| else: | |
| plate_data['text'] = 'No text detected' | |
| print(f"No text found in plate {i+1}") | |
| except Exception as e: | |
| print(f"OCR extraction failed for plate {i+1}: {e}") | |
| plate_data['text'] = f'OCR Failed: {str(e)}' | |
| elif extract_text and not OCR_AVAILABLE: | |
| plate_data['text'] = 'OCR not available' | |
| else: | |
| plate_data['text'] = 'OCR disabled' | |
| cropped_plates.append(plate_data) | |
| except ValueError as e: | |
| print(f"Error parsing coordinates for detection {i}: {e}") | |
| continue | |
| except Exception as e: | |
| print(f"Error cropping license plate {i}: {e}") | |
| continue | |
| return cropped_plates | |
| def create_download_files(annotated_image, cropped_plates, detections): | |
| try: | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| os.makedirs("temp", exist_ok=True) | |
| annotated_path = f"temp/annotated_image_{timestamp}.jpg" | |
| try: | |
| annotated_image.save(annotated_path, quality=95) | |
| except Exception as e: | |
| print(f"Error saving annotated image: {e}") | |
| return None, None, [] | |
| plate_paths = [] | |
| for i, plate_data in enumerate(cropped_plates): | |
| try: | |
| plate_path = f"temp/license_plate_{i+1}_{timestamp}.jpg" | |
| plate_data['image'].save(plate_path, quality=95) | |
| plate_paths.append(plate_path) | |
| except Exception as e: | |
| print(f"Error saving license plate {i+1}: {e}") | |
| continue | |
| report_data = [] | |
| for detection in detections: | |
| report_data.append(detection) | |
| for i, plate_data in enumerate(cropped_plates): | |
| report_data.append({ | |
| 'Object': f'License Plate {i+1} - Text', | |
| 'Confidence': plate_data['confidence'], | |
| 'Position': plate_data['position'], | |
| 'Dimensions': 'Extracted Text', | |
| 'Text': plate_data.get('text', 'N/A') | |
| }) | |
| report_path = f"temp/detection_report_{timestamp}.csv" | |
| if report_data: | |
| try: | |
| df = pd.DataFrame(report_data) | |
| df.to_csv(report_path, index=False) | |
| except Exception as e: | |
| print(f"Error creating detection report: {e}") | |
| report_path = None | |
| zip_path = f"temp/detection_results_{timestamp}.zip" | |
| try: | |
| with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: | |
| if os.path.exists(annotated_path): | |
| zipf.write(annotated_path, f"annotated_image_{timestamp}.jpg") | |
| for plate_path in plate_paths: | |
| if os.path.exists(plate_path): | |
| zipf.write(plate_path, os.path.basename(plate_path)) | |
| if report_path and os.path.exists(report_path): | |
| zipf.write(report_path, f"detection_report_{timestamp}.csv") | |
| except Exception as e: | |
| print(f"Error creating ZIP file: {e}") | |
| return None, annotated_path, plate_paths | |
| return zip_path, annotated_path, plate_paths | |
| except Exception as e: | |
| print(f"Error in create_download_files: {e}") | |
| return None, None, [] | |
| def yoloV8_func( | |
| image=None, | |
| image_size=640, | |
| conf_threshold=0.4, | |
| iou_threshold=0.5, | |
| show_stats=True, | |
| show_confidence=True, | |
| crop_plates=True, | |
| extract_text=False, | |
| ocr_on_no_helmet=False | |
| ): | |
| if image_size is None: | |
| image_size = 640 | |
| if not isinstance(image_size, int): | |
| image_size = int(image_size) | |
| imgsz = [image_size, image_size] | |
| results = model.predict(image, conf=conf_threshold, iou=iou_threshold, imgsz=imgsz) | |
| annotated_image = results[0].plot() | |
| if isinstance(annotated_image, np.ndarray): | |
| annotated_image = Image.fromarray(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)) | |
| boxes = results[0].boxes | |
| detections = [] | |
| if boxes is not None and len(boxes) > 0: | |
| for i, (box, cls, conf) in enumerate(zip(boxes.xyxy, boxes.cls, boxes.conf)): | |
| x1, y1, x2, y2 = box.tolist() | |
| class_id = int(cls) | |
| confidence = float(conf) | |
| label = class_names.get(class_id, f"Class {class_id}") | |
| detections.append({ | |
| "Object": label, | |
| "Confidence": f"{confidence:.2f}", | |
| "Position": f"({int(x1)}, {int(y1)})", | |
| "Dimensions": f"{int(x2-x1)}x{int(y2-y1)}" | |
| }) | |
| cropped_plates = [] | |
| license_plate_gallery = [] | |
| plate_texts = [] | |
| download_files = None | |
| has_no_helmet = any(detection['Object'] == 'Without Helmet' for detection in detections) | |
| should_extract_text = extract_text or (ocr_on_no_helmet and has_no_helmet) | |
| if crop_plates and detections: | |
| try: | |
| license_plate_count = len([d for d in detections if d['Object'] == 'License Plate']) | |
| print(f"Processing {license_plate_count} license plates...") | |
| if ocr_on_no_helmet and has_no_helmet: | |
| print("⚠️ No helmet detected - OCR will be performed on license plates") | |
| cropped_plates = crop_license_plates(image, detections, should_extract_text) | |
| print(f"Successfully cropped {len(cropped_plates)} license plates") | |
| license_plate_gallery = [plate_data['image'] for plate_data in cropped_plates] | |
| if should_extract_text and OCR_AVAILABLE: | |
| print("Extracting text from license plates...") | |
| plate_texts = [] | |
| for i, plate_data in enumerate(cropped_plates): | |
| text = plate_data.get('text', 'No text detected') | |
| print(f"Plate {i+1} text: {text}") | |
| if ocr_on_no_helmet and has_no_helmet: | |
| plate_texts.append(f"🚨 No Helmet Violation - Plate {i+1}: {text}") | |
| else: | |
| plate_texts.append(f"Plate {i+1}: {text}") | |
| elif should_extract_text and not OCR_AVAILABLE: | |
| plate_texts = ["OCR not available - install requirements: pip install transformers easyocr"] | |
| elif not should_extract_text: | |
| if ocr_on_no_helmet and not has_no_helmet: | |
| plate_texts = [f"Plate {i+1}: OCR only on no-helmet violations" for i in range(len(cropped_plates))] | |
| else: | |
| plate_texts = [f"Plate {i+1}: Text extraction disabled" for i in range(len(cropped_plates))] | |
| if cropped_plates or detections: | |
| download_files, _, _ = create_download_files(annotated_image, cropped_plates, detections) | |
| if download_files is None: | |
| print("Warning: Could not create download files") | |
| except Exception as e: | |
| print(f"Error in license plate processing: {e}") | |
| cropped_plates = [] | |
| license_plate_gallery = [] | |
| plate_texts = ["Error processing license plates"] | |
| download_files = None | |
| stats_text = "" | |
| if show_stats and detections: | |
| df = pd.DataFrame(detections) | |
| counts = df['Object'].value_counts().to_dict() | |
| stats_text = "Detection Summary:\n" | |
| for obj, count in counts.items(): | |
| stats_text += f"- {obj}: {count}\n" | |
| if cropped_plates: | |
| stats_text += f"\nLicense Plates Cropped: {len(cropped_plates)}\n" | |
| if has_no_helmet: | |
| stats_text += "⚠️ HELMET VIOLATION DETECTED!\n" | |
| if should_extract_text and OCR_AVAILABLE: | |
| stats_text += "Extracted Text:\n" | |
| for i, plate_data in enumerate(cropped_plates): | |
| text = plate_data.get('text', 'No text') | |
| if has_no_helmet and ocr_on_no_helmet: | |
| stats_text += f"🚨 Violation - Plate {i+1}: {text}\n" | |
| else: | |
| stats_text += f"- Plate {i+1}: {text}\n" | |
| if show_stats and stats_text: | |
| draw = ImageDraw.Draw(annotated_image) | |
| try: | |
| font = ImageFont.truetype("arial.ttf", 20) | |
| except: | |
| font = ImageFont.load_default() | |
| text_bbox = draw.textbbox((0, 0), stats_text, font=font) | |
| text_width = text_bbox[2] - text_bbox[0] | |
| text_height = text_bbox[3] - text_bbox[1] | |
| draw.rectangle([10, 10, 20 + text_width, 20 + text_height], fill=(0, 0, 0, 128)) | |
| draw.text((15, 15), stats_text, font=font, fill=(255, 255, 255)) | |
| detection_table = pd.DataFrame(detections) if detections else pd.DataFrame(columns=["Object", "Confidence", "Position", "Dimensions"]) | |
| plate_text_output = "\n".join(plate_texts) if plate_texts else "No license plates detected or OCR disabled" | |
| return annotated_image, detection_table, stats_text, license_plate_gallery, download_files, plate_text_output | |
| custom_css = """ | |
| #title { text-align: center; } | |
| #description { text-align: center; } | |
| .footer { | |
| text-align: center; | |
| margin-top: 20px; | |
| color: #666; | |
| } | |
| .important { font-weight: bold; color: red; } | |
| .download-section { | |
| background-color: #f0f0f0; | |
| padding: 15px; | |
| border-radius: 8px; | |
| margin-top: 10px; | |
| } | |
| .ocr-section { | |
| background-color: #e8f4fd; | |
| padding: 15px; | |
| border-radius: 8px; | |
| margin-top: 10px; | |
| } | |
| """ | |
| with gr.Blocks(css=custom_css, title="YOLOv11 Motorcyclist Helmet Detection") as demo: | |
| gr.HTML("<h1 id='title'>YOLOv11 Motorcyclist Helmet Detection with Optional OCR</h1>") | |
| gr.HTML(f""" | |
| <div id='description'> | |
| <p>This application uses YOLOv11 to detect Motorcyclists with and without Helmets in images.</p> | |
| <p>Upload an image, adjust the parameters, and view the detection results with detailed statistics.</p> | |
| <p><strong>Features:</strong> License plate cropping and optional text recognition!</p> | |
| <p><strong>OCR Status:</strong> {'✅ Available' if OCR_AVAILABLE else '❌ Not Available (install requirements)'}</p> | |
| </div> | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Input Parameters") | |
| input_image = gr.Image(type="filepath", label="Input Image", sources=["upload", "webcam"]) | |
| with gr.Row(): | |
| image_size = gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size") | |
| conf_threshold = gr.Slider(minimum=0.0, maximum=1.0, value=0.4, step=0.05, label="Confidence Threshold") | |
| with gr.Row(): | |
| iou_threshold = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, step=0.05, label="IOU Threshold") | |
| show_stats = gr.Checkbox(value=True, label="Show Statistics on Image") | |
| gr.Markdown("### License Plate Options") | |
| crop_plates = gr.Checkbox(value=True, label="Crop License Plates") | |
| if OCR_AVAILABLE: | |
| extract_text = gr.Checkbox(value=False, label="Extract Text from License Plates (OCR)") | |
| ocr_on_no_helmet = gr.Checkbox(value=True, label="🚨 Auto-OCR when No Helmet Detected") | |
| gr.Markdown("*Note: OCR processing may take additional time*") | |
| else: | |
| extract_text = gr.Checkbox(value=False, label="Extract Text (OCR Not Available)", interactive=False) | |
| ocr_on_no_helmet = gr.Checkbox(value=False, label="🚨 Auto-OCR when No Helmet (Not Available)", interactive=False) | |
| gr.Markdown("*Install requirements: `pip install torch transformers easyocr opencv-python`*") | |
| submit_btn = gr.Button("Detect Objects", variant="primary") | |
| clear_btn = gr.Button("Clear") | |
| with gr.Column(scale=2): | |
| gr.Markdown("### Output Results") | |
| output_image = gr.Image(type="pil", label="Annotated Image") | |
| output_table = gr.Dataframe( | |
| headers=["Object", "Confidence", "Position", "Dimensions"], | |
| label="Detection Details", | |
| interactive=False | |
| ) | |
| output_stats = gr.Textbox(label="Detection Summary", interactive=False) | |
| gr.Markdown("### Cropped License Plates") | |
| license_gallery = gr.Gallery( | |
| label="Extracted License Plates", | |
| show_label=True, | |
| elem_id="license_gallery", | |
| columns=3, | |
| rows=2, | |
| object_fit="contain", | |
| height="auto" | |
| ) | |
| with gr.Group(elem_classes="ocr-section"): | |
| gr.Markdown("### License Plate Text Recognition") | |
| plate_text_output = gr.Textbox( | |
| label="Extracted Text", | |
| placeholder="License plate text will appear here when OCR is enabled", | |
| lines=3, | |
| interactive=False | |
| ) | |
| gr.Markdown("### Download Results") | |
| with gr.Group(elem_classes="download-section"): | |
| download_file = gr.File( | |
| label="Download Complete Results (ZIP)", | |
| interactive=False, | |
| visible=True | |
| ) | |
| gr.Markdown("*The ZIP file contains: annotated image, cropped license plates, and detection report with OCR results*") | |
| gr.Markdown("### Example Images") | |
| gr.Examples( | |
| examples=[["sample_1.jpg"], ["sample_2.jpg"], ["sample_3.jpg"], ["sample_4.jpg"], ["sample_5.jpg"]], | |
| inputs=input_image, | |
| outputs=[output_image, output_table, output_stats, license_gallery, download_file, plate_text_output], | |
| fn=lambda img: yoloV8_func(img, 640, 0.4, 0.5, True, True, True, False), | |
| cache_examples=True, | |
| ) | |
| gr.HTML(""" | |
| <div class='footer'> | |
| <p>Built with Gradio and Ultralytics YOLO</p> | |
| <p>Note: This is a demonstration application. Detection accuracy may vary based on image quality and conditions.</p> | |
| <p><strong>License Plate Privacy:</strong> Extracted license plates and text are for demonstration purposes only.</p> | |
| <p><strong>Requirements for OCR:</strong> torch, transformers, easyocr, opencv-python</p> | |
| </div> | |
| """) | |
| submit_btn.click( | |
| fn=yoloV8_func, | |
| inputs=[input_image, image_size, conf_threshold, iou_threshold, show_stats, gr.State(True), crop_plates, extract_text, ocr_on_no_helmet], | |
| outputs=[output_image, output_table, output_stats, license_gallery, download_file, plate_text_output] | |
| ) | |
| clear_btn.click( | |
| fn=lambda: [None, None, None, None, None, None], | |
| inputs=[], | |
| outputs=[input_image, output_image, output_table, output_stats, license_gallery, download_file, plate_text_output] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(debug=True, share=True) |