Update app.py

app.py

@@ -1,154 +1,178 @@
 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
-from ultralytics import YOLO
-from PIL import Image
-import requests
-import json
+import time
 
-# Download sample images
+# Download sample images (optional)
 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')
 
+# Load model (cached for performance)
 model = YOLO("best.pt")
+class_names = {0: 'With Helmet', 1: 'Without Helmet', 2: 'License Plate'}
 
-GEMINI_API_KEY = "AIzaSyCBs4TumAonKI0AodIzbl4b8Vmu9eM_r9I"
-GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
-
-def get_gemini_response(prompt: str) -> str:
-    headers = {
-        'Content-Type': 'application/json',
-    }
-    
-    data = {
-        "contents": [{"parts": [{"text": prompt}]}]
-    }
-    
-    try:
-        response = requests.post(
-            f"{GEMINI_API_URL}?key={GEMINI_API_KEY}",
-            headers=headers,
-            json=data
-        )
-        response.raise_for_status()
-        result = response.json()
-        return result['candidates'][0]['content']['parts'][0]['text']
-    except Exception as e:
-        print(f"Error calling Gemini API: {str(e)}")
-        return "Unable to get analysis at the moment."
-
-def get_safety_analysis(stats: dict) -> str:
-    prompt = f"""
-    You are a traffic safety analyst. Analyze the following statistics and provide a brief safety report:
-    
-    Detection Results:
-    - Total Detections: {stats.get('total_detections', 0)}
-    - Riders with Helmet: {stats.get('with_helmet', 0)}
-    - Riders without Helmet: {stats.get('without_helmet', 0)}
-    - Helmet Compliance Rate: {stats.get('helmet_compliance', 0)}%
-    - License Plates Detected: {stats.get('license_plates', 0)}
-    
-    Provide a 3-4 sentence safety analysis focusing on helmet compliance and potential safety concerns.
-    """
-    return get_gemini_response(prompt)
-
-def yoloV8_func(image=None, image_size=640, conf_threshold=0.4, iou_threshold=0.5):
-    print(f"Received image_size: {image_size}")
-    
+def yoloV8_func(
+    image=None, 
+    image_size=640, 
+    conf_threshold=0.4, 
+    iou_threshold=0.5,
+    show_stats=True,
+    show_confidence=True
+):
+    # Handle NoneType for image_size
     if image_size is None:
         image_size = 640
     
+    # Ensure image_size is an integer
     if not isinstance(image_size, int):
         image_size = int(image_size)
     
+    # Construct imgsz as a list of two integers [width, height]
     imgsz = [image_size, image_size]
 
-    results = model.predict(
-        source=image,
-        conf=conf_threshold,
-        iou=iou_threshold,
-        imgsz=imgsz,
-        verbose=False
-    )
+    # Make predictions
+    results = model.predict(image, conf=conf_threshold, iou=iou_threshold, imgsz=imgsz)
+
+    # Get the output image with bounding boxes
+    annotated_image = results[0].plot()  # This returns a PIL image
+    
+    # Convert to PIL if it's a numpy array
+    if isinstance(annotated_image, np.ndarray):
+        annotated_image = Image.fromarray(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))
     
-    boxes = results[0].boxes.xyxy.cpu().numpy()
-    class_ids = results[0].boxes.cls.cpu().numpy().astype(int)
+    # Extract detection information
+    boxes = results[0].boxes
+    detections = []
     
-    total_riders = int(sum((class_ids == 0) | (class_ids == 1)))
-    helmet_compliance = 0 if total_riders == 0 else int(sum(class_ids == 0) / total_riders * 100)
+    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)}"
+            })
     
-    stats = {
-        'total_detections': len(boxes),
-        'with_helmet': int(sum(class_ids == 0)),
-        'without_helmet': int(sum(class_ids == 1)),
-        'license_plates': int(sum(class_ids == 2)),
-        'helmet_compliance': helmet_compliance,
-        'total_riders': total_riders
-    }
+    # Create stats text
+    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"
     
-    safety_analysis = get_safety_analysis(stats)
-    print("\nSafety Analysis:", safety_analysis)
+    # Add stats to image if requested
+    if show_stats and stats_text:
+        draw = ImageDraw.Draw(annotated_image)
+        try:
+            font = ImageFont.truetype("arial.ttf", 20)
+        except:
+            font = ImageFont.load_default()
+        
+        # Add semi-transparent background for text
+        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))
+        
+        # Add text
+        draw.text((15, 15), stats_text, font=font, fill=(255, 255, 255))
     
-    annotated_image = results[0].plot()
+    # Create a detection table for display
+    detection_table = pd.DataFrame(detections) if detections else pd.DataFrame(columns=["Object", "Confidence", "Position", "Dimensions"])
     
-    return annotated_image
+    return annotated_image, detection_table, stats_text
 
-with gr.Blocks() as demo:
-    with gr.Row():
-        image_input = gr.Image(type="filepath", label="Input Image")
-        output_image = gr.Image(type="pil", label="Output Image")
+# Define custom CSS for styling
+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; }
+"""
+
+# Set up Gradio interface with Blocks for more control
+with gr.Blocks(css=custom_css, title="YOLOv11 Motorcyclist Helmet Detection") as demo:
+    gr.HTML("<h1 id='title'>YOLOv11 Motorcyclist Helmet Detection</h1>")
+    gr.HTML("""
+    <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>
+    </div>
+    """)
     
     with gr.Row():
-        image_size = gr.Slider(
-            minimum=320, 
-            maximum=1280, 
-            value=640, 
-            step=32, 
-            label="Image Size",
-            interactive=True
-        )
-        conf_threshold = gr.Slider(
-            minimum=0.1, 
-            maximum=1.0, 
-            value=0.4, 
-            step=0.05, 
-            label="Confidence Threshold",
-            interactive=True
-        )
-        iou_threshold = gr.Slider(
-            minimum=0.1, 
-            maximum=1.0, 
-            value=0.5, 
-            step=0.05, 
-            label="IOU Threshold",
-            interactive=True
-        )
+        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")
+            
+            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="Output Image")
+            output_table = gr.Dataframe(
+                headers=["Object", "Confidence", "Position", "Dimensions"],
+                label="Detection Details",
+                interactive=False
+            )
+            output_stats = gr.Textbox(label="Detection Summary", interactive=False)
     
-    process_btn = gr.Button("Process Image")
-    process_btn.click(
+    # Examples
+    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],
         fn=yoloV8_func,
-        inputs=[image_input, image_size, conf_threshold, iou_threshold],
-        outputs=output_image
+        cache_examples=True,
+    )
+    
+    # Footer
+    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>
+    </div>
+    """)
+    
+    # Button actions
+    submit_btn.click(
+        fn=yoloV8_func,
+        inputs=[input_image, image_size, conf_threshold, iou_threshold, show_stats],
+        outputs=[output_image, output_table, output_stats]
+    )
+    
+    clear_btn.click(
+        fn=lambda: [None, None, None],
+        inputs=[],
+        outputs=[input_image, output_image, output_table, output_stats]
     )
-
-outputs = gr.Image(type="pil", label="Output Image")
-
-title = "YOLOv11 Motorcyclist Helmet Detection"
-description = """
-    This application uses YOLOv11 to detect Motorcyclists with and without Helmets in images. 
-    Upload an image, adjust the confidence and IOU thresholds, and view the detection results. 
-    You can customize the model's performance to fit your needs.
-"""
-article = """
-    <h2>How It Works:</h2>
-    <p>This model detects Motorcyclists with and without Helmets in images and highlights them with bounding boxes. 
-    Adjust the confidence threshold to control detection accuracy and the IOU threshold for overlap sensitivity.</p>
-    <p>Upload your images and try it out!</p>
-"""
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch(debug=True, share=True)