import gradio as gr from detector import ( yolov8_detect, download_sample_images, get_ocr_status, ADVANCED_OCR_AVAILABLE, OCR_AVAILABLE ) try: from advanced_ocr import get_available_models except ImportError: def get_available_models(): return {} class UIComponents: def __init__(self): download_sample_images() self.ocr_status = get_ocr_status() self.custom_css = """ .gradio-container { max-width: 1200px !important; margin: 0 auto; } .main-header { text-align: center; margin-bottom: 2rem; padding: 1rem; } .main-title { font-size: 2rem; font-weight: 600; color: #333; margin: 0; } .subtitle { color: #666; font-size: 1rem; margin: 0.5rem 0; } .status-info { font-size: 0.9rem; color: #888; margin: 0.5rem 0; } .section-gap { margin: 1.5rem 0; } """ def toggle_sections(self, extract_text_checked, crop_checked): show_gallery = bool(extract_text_checked and crop_checked) show_ocr = bool(extract_text_checked) return ( gr.update(visible=show_gallery), gr.update(visible=show_ocr), ) def get_ocr_status_text(self): if ADVANCED_OCR_AVAILABLE: return "Advanced OCR Available" elif OCR_AVAILABLE: return "Basic OCR Available" else: return "OCR Not Available" def create_header(self): return gr.HTML(f"""

AI Helmet Detection System

Motorcyclist safety monitoring with license plate recognition

YOLOv11 • {self.get_ocr_status_text()} • Real-time Processing

""") def create_settings_panel(self): components = {} with gr.Column(scale=1): gr.Markdown("### Settings") components['input_image'] = gr.Image( type="filepath", label="Upload Image", sources=["upload", "webcam"] ) with gr.Row(): components['image_size'] = gr.Slider( minimum=320, maximum=1280, value=640, step=32, label="Image Size" ) with gr.Row(): components['conf_threshold'] = gr.Slider( minimum=0.0, maximum=1.0, value=0.4, step=0.05, label="Confidence" ) components['iou_threshold'] = gr.Slider( minimum=0.0, maximum=1.0, value=0.5, step=0.05, label="IoU Threshold" ) components['show_stats'] = gr.Checkbox( value=True, label="Show Statistics" ) components['crop_plates'] = gr.Checkbox( value=True, label="Extract License Plates" ) if self.ocr_status["any_available"]: components['extract_text'] = gr.Checkbox( value=False, label="Enable OCR" ) components['ocr_on_no_helmet'] = gr.Checkbox( value=True, label="Auto-OCR for No Helmet" ) if ADVANCED_OCR_AVAILABLE: models = get_available_models() model_choices = [("Auto (Recommended)", "auto"), ("Basic EasyOCR", "basic")] for key, info in models.items(): model_choices.append((info['name'], key)) components['selected_ocr_model'] = gr.Dropdown( choices=model_choices, value="auto", label="OCR Model" ) else: components['selected_ocr_model'] = gr.State("basic") gr.Markdown("*Note: OCR processing may increase detection time.*") else: components['extract_text'] = gr.Checkbox( value=False, label="OCR Not Available", interactive=False ) components['ocr_on_no_helmet'] = gr.Checkbox( value=False, label="Auto-OCR (Not Available)", interactive=False ) components['selected_ocr_model'] = gr.State("basic") with gr.Row(): components['submit_btn'] = gr.Button("Start Detection", variant="primary") components['clear_btn'] = gr.Button("Clear") return components def create_results_panel(self): components = {} with gr.Column(scale=2): gr.Markdown("### Results") components['output_image'] = gr.Image( type="pil", label="Detection Results" ) with gr.Row(): components['output_table'] = gr.Dataframe( headers=["Object", "Confidence", "Position", "Dimensions"], label="Detection Details", interactive=False ) components['output_stats'] = gr.Textbox( label="Statistics", interactive=False, lines=6 ) components['license_gallery'] = gr.Gallery( label="License Plates", columns=3, visible=False ) components['ocr_group'] = gr.Group(visible=False) with components['ocr_group']: components['plate_text_output'] = gr.Textbox( label="OCR Results", lines=4, interactive=False ) components['download_file'] = gr.File( label="Download Results (ZIP)", interactive=False ) return components def create_examples_tab(self, input_image, output_components): with gr.TabItem("Examples"): gr.Markdown("### Sample Images") gr.Markdown("Click any example to test the detection system:") gr.Examples( examples=[ ["sample_1.jpg"], ["sample_2.jpg"], ["sample_3.jpg"], ["sample_4.jpg"], ["sample_6.jpg"], ["sample_7.jpg"], ["sample_8.jpg"], ], inputs=input_image, outputs=[ output_components['output_image'], output_components['output_table'], output_components['output_stats'], output_components['license_gallery'], output_components['download_file'], output_components['plate_text_output'], ], fn=lambda img: yolov8_detect( img, 640, 0.4, 0.5, True, True, True, False ), cache_examples=True ) def create_info_tab(self): with gr.TabItem("Info"): gr.Markdown("### System Information") gr.Markdown(f""" **AI Model:** YOLOv11 **Classes:** Helmet, No Helmet, License Plate **OCR Status:** {self.get_ocr_status_text()} **Features:** Detection, extraction, text recognition **Privacy:** All processing is local. No data stored. **Usage:** For demonstration and research purposes only. """) def setup_event_handlers(self, settings_components, results_components): settings_components['submit_btn'].click( fn=yolov8_detect, inputs=[ settings_components['input_image'], settings_components['image_size'], settings_components['conf_threshold'], settings_components['iou_threshold'], settings_components['show_stats'], gr.State(True), settings_components['crop_plates'], settings_components['extract_text'], settings_components['ocr_on_no_helmet'], settings_components['selected_ocr_model'], ], outputs=[ results_components['output_image'], results_components['output_table'], results_components['output_stats'], results_components['license_gallery'], results_components['download_file'], results_components['plate_text_output'], ], ) settings_components['clear_btn'].click( fn=lambda: [None, None, None, None, None, None], inputs=[], outputs=[ settings_components['input_image'], results_components['output_image'], results_components['output_table'], results_components['output_stats'], results_components['license_gallery'], results_components['download_file'], results_components['plate_text_output'], ], ) settings_components['extract_text'].change( fn=self.toggle_sections, inputs=[settings_components['extract_text'], settings_components['crop_plates']], outputs=[results_components['license_gallery'], results_components['ocr_group']], ) settings_components['crop_plates'].change( fn=self.toggle_sections, inputs=[settings_components['extract_text'], settings_components['crop_plates']], outputs=[results_components['license_gallery'], results_components['ocr_group']], )