Update app.py

app.py

@@ -1,54 +1,69 @@
-import gradio as gr
-import cv2
-import numpy as np
-from PIL import Image
-from ultralytics import YOLO
-
-# Load the YOLO model
-model = YOLO('best.pt')
-
-def predict(image):
-    # Convert PIL Image to numpy array
-    img_array = np.array(image)
-    
-    # Perform inference
-    results = model(img_array)
-    
-    # Plot results
-    for r in results:
-        im_array = r.plot()
-        im = Image.fromarray(im_array[..., ::-1])  # RGB PIL image
-    
-    # Get detection counts and confidences
-    class_counts = {i: 0 for i in range(len(model.names))}
-    class_confidences = {i: [] for i in range(len(model.names))}
-    
-    for box in results[0].boxes:
-        cls = int(box.cls[0])
-        conf = float(box.conf[0])
-        class_counts[cls] += 1
-        class_confidences[cls].append(conf)
-    
-    # Create output string
-    output = "Detection Results:\n"
-    for i in range(len(model.names)):
-        count = class_counts[i]
-        avg_conf = np.mean(class_confidences[i]) if class_confidences[i] else 0
-        output += f"{model.names[i]}: {count} detections (Avg. Confidence: {avg_conf:.2f})\n"
-    
-    return im, output
-
-# Create Gradio interface
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=[
-        gr.Image(type="pil", label="Detected Image"),
-        gr.Textbox(label="Detection Results")
-    ],
-    title="Helmet and License Plate Detection",
-    description="Upload an image to detect helmets, license plates, and motorcyclists."
-)
-
-# Launch the interface
-iface.launch()
+import gradio as gr
+import cv2
+import numpy as np
+from PIL import Image
+from ultralytics import YOLO
+
+# Load the YOLO models
+model_all = YOLO('best.pt')  # Model for helmet, license plate, and motorcyclist
+model_np = YOLO('best_1.pt')  # Model for number plate detection
+
+def predict(image):
+    # Convert PIL Image to numpy array
+    img_array = np.array(image)
+
+    # Perform inference with both models
+    results_all = model_all(img_array)
+    results_np = model_np(img_array)
+
+    # Plot results from the first model (all detections)
+    for r in results_all:
+        im_array = r.plot()
+    im = Image.fromarray(im_array[..., ::-1])  # RGB PIL image
+
+    # Initialize counters and confidence lists
+    class_counts_all = {i: 0 for i in range(len(model_all.names))}
+    class_confidences_all = {i: [] for i in range(len(model_all.names))}
+    np_count = 0
+    np_confidences = []
+
+    # Process results from the first model (all detections)
+    for box in results_all[0].boxes:
+        cls = int(box.cls[0])
+        conf = float(box.conf[0])
+        class_counts_all[cls] += 1
+        class_confidences_all[cls].append(conf)
+
+    # Process results from the second model (number plate detection)
+    for box in results_np[0].boxes:
+        np_count += 1
+        np_confidences.append(float(box.conf[0]))
+
+    # Create output string
+    output = "Detection Results:\n"
+    output += "Model 1 (Helmet, License Plate, Motorcyclist):\n"
+    for i in range(len(model_all.names)):
+        count = class_counts_all[i]
+        avg_conf = np.mean(class_confidences_all[i]) if class_confidences_all[i] else 0
+        output += f"{model_all.names[i]}: {count} detections (Avg. Confidence: {avg_conf:.2f})\n"
+
+    output += "\nModel 2 (Number Plate Detection):\n"
+    avg_np_conf = np.mean(np_confidences) if np_confidences else 0
+    output += f"Number Plates: {np_count} detections (Avg. Confidence: {avg_np_conf:.2f})\n"
+
+    return im, output
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Image(type="pil", label="Detected Image"),
+        gr.Textbox(label="Detection Results")
+    ],
+    title="Helmet, License Plate, and Motorcyclist Detection",
+    description="Upload an image to detect helmets, license plates, and motorcyclists using two specialized models."
+)
+
+# Launch the interface
+iface.launch(share='True')