Create ui.py

ui.py

@@ -0,0 +1,293 @@
+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"""
+            <div class="main-header">
+                <h1 class="main-title">AI Helmet Detection System</h1>
+                <p class="subtitle">Motorcyclist safety monitoring with license plate recognition</p>
+                <p class="status-info">YOLOv11 • {self.get_ocr_status_text()} • Real-time Processing</p>
+            </div>
+        """)
+
+    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']],
+        )