Update app.py

app.py

@@ -4,69 +4,165 @@ import cv2
 import numpy as np
 from PIL import Image
 
+# Initialize Roboflow client
 CLIENT = InferenceHTTPClient(
     api_url="https://serverless.roboflow.com",
     api_key="DIAhXQf6AUsyM1PRfdFa"
 )
 
-def draw_detections(image, predictions):
+def draw_detections(image, predictions, color=(0, 255, 0), prefix=""):
+    """Draw bounding boxes and labels on the image"""
     img_array = np.array(image)
+    
     if 'predictions' in predictions:
         for pred in predictions['predictions']:
-            x, y = int(pred['x']), int(pred['y'])
-            w, h = int(pred['width']), int(pred['height'])
-            x1, y1 = int(x - w/2), int(y - h/2)
-            x2, y2 = int(x + w/2), int(y + h/2)
-            cv2.rectangle(img_array, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            label = f"{pred['class']}: {pred['confidence']:.2f}"
-            cv2.putText(img_array, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
-    return Image.fromarray(img_array)
+            x = int(pred['x'])
+            y = int(pred['y'])
+            width = int(pred['width'])
+            height = int(pred['height'])
+            confidence = pred['confidence']
+            class_name = pred['class']
+            
+            # Calculate bounding box coordinates
+            x1 = int(x - width / 2)
+            y1 = int(y - height / 2)
+            x2 = int(x + width / 2)
+            y2 = int(y + height / 2)
+            
+            # Draw rectangle
+            cv2.rectangle(img_array, (x1, y1), (x2, y2), color, 2)
+            
+            # Draw label with prefix
+            label = f"{prefix}{class_name}: {confidence:.2f}"
+            cv2.putText(img_array, label, (x1, y1 - 10), 
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+    
+    return img_array
+
+def detect_both(image):
+    """Detect both road lanes and road signs in one image"""
+    if image is None:
+        return None, "Please upload an image first."
+    
+    try:
+        # Detect road lanes
+        lanes_result = CLIENT.infer(image, model_id="road-lean/1")
+        
+        # Detect road signs
+        signs_result = CLIENT.infer(image, model_id="road-sign-peqgi/1")
+        
+        # Draw both detections on the same image
+        img_array = np.array(image)
+        
+        # Draw lanes in green
+        img_array = draw_detections(Image.fromarray(img_array), lanes_result, 
+                                    color=(0, 255, 0), prefix="Lane: ")
+        
+        # Draw signs in blue
+        img_array = draw_detections(Image.fromarray(img_array), signs_result, 
+                                    color=(255, 0, 0), prefix="Sign: ")
+        
+        output_image = Image.fromarray(img_array)
+        
+        # Format results text
+        num_lanes = len(lanes_result.get('predictions', []))
+        num_signs = len(signs_result.get('predictions', []))
+        
+        results_text = f"🛣️ Road Lanes Detected: {num_lanes}\n"
+        if num_lanes > 0:
+            for i, pred in enumerate(lanes_result.get('predictions', []), 1):
+                results_text += f"  {i}. {pred['class']}: {pred['confidence']:.2%}\n"
+        else:
+            results_text += "  No lanes detected\n"
+        
+        results_text += f"\n🚦 Road Signs Detected: {num_signs}\n"
+        if num_signs > 0:
+            for i, pred in enumerate(signs_result.get('predictions', []), 1):
+                results_text += f"  {i}. {pred['class']}: {pred['confidence']:.2%}\n"
+        else:
+            results_text += "  No signs detected\n"
+        
+        results_text += f"\n📊 Total Detections: {num_lanes + num_signs}"
+        results_text += "\n\n🟢 Green boxes = Road Lanes\n🔵 Blue boxes = Road Signs"
+        
+        return output_image, results_text
+    except Exception as e:
+        return image, f"Error: {str(e)}"
 
 def detect_road_lanes(image):
+    """Detect road lanes only"""
     if image is None:
-        return None, "Please upload an image."
+        return None, "Please upload an image first."
+    
     try:
         result = CLIENT.infer(image, model_id="road-lean/1")
-        output = draw_detections(image, result)
-        n = len(result.get('predictions', []))
-        text = f"Detections: {n}\n\n"
-        text += "No lanes detected." if n == 0 else "\n".join(
-            f"{i}. {p['class']}: {p['confidence']:.2%}" 
-            for i, p in enumerate(result.get('predictions', []), 1)
-        )
-        return output, text
+        img_array = draw_detections(image, result, color=(0, 255, 0))
+        output_image = Image.fromarray(img_array)
+        
+        num_detections = len(result.get('predictions', []))
+        results_text = f"Detections: {num_detections}\n\n"
+        
+        if num_detections == 0:
+            results_text += "No road lanes detected."
+        else:
+            for i, pred in enumerate(result.get('predictions', []), 1):
+                results_text += f"{i}. {pred['class']}: {pred['confidence']:.2%}\n"
+        
+        return output_image, results_text
     except Exception as e:
-        return image, f"Error: {e}"
+        return image, f"Error: {str(e)}"
 
 def detect_road_signs(image):
+    """Detect road signs only"""
     if image is None:
-        return None, "Please upload an image."
+        return None, "Please upload an image first."
+    
     try:
         result = CLIENT.infer(image, model_id="road-sign-peqgi/1")
-        output = draw_detections(image, result)
-        n = len(result.get('predictions', []))
-        text = f"Detections: {n}\n\n"
-        text += "No signs detected." if n == 0 else "\n".join(
-            f"{i}. {p['class']}: {p['confidence']:.2%}" 
-            for i, p in enumerate(result.get('predictions', []), 1)
-        )
-        return output, text
+        img_array = draw_detections(image, result, color=(255, 0, 0))
+        output_image = Image.fromarray(img_array)
+        
+        num_detections = len(result.get('predictions', []))
+        results_text = f"Detections: {num_detections}\n\n"
+        
+        if num_detections == 0:
+            results_text += "No road signs detected."
+        else:
+            for i, pred in enumerate(result.get('predictions', []), 1):
+                results_text += f"{i}. {pred['class']}: {pred['confidence']:.2%}\n"
+        
+        return output_image, results_text
     except Exception as e:
-        return image, f"Error: {e}"
+        return image, f"Error: {str(e)}"
 
+# Create interface
 with gr.Blocks() as demo:
     gr.Markdown("# 🚗 Road Detection System")
-    with gr.Tab("Road Lanes"):
+    gr.Markdown("Upload an image to detect road lanes and road signs using AI.")
+    
+    with gr.Tab("🎯 Detect Both"):
+        gr.Markdown("### Detect both road lanes and road signs in one image")
+        with gr.Row():
+            both_input = gr.Image(type="pil", label="Upload Image")
+            both_output = gr.Image(label="Detection Results")
+        both_text = gr.Textbox(label="Details", lines=10)
+        both_btn = gr.Button("🔍 Detect All", variant="primary", size="lg")
+        both_btn.click(detect_both, inputs=both_input, outputs=[both_output, both_text])
+    
+    with gr.Tab("🛣️ Road Lanes Only"):
         with gr.Row():
-            lane_in = gr.Image(type="pil")
-            lane_out = gr.Image()
-        lane_text = gr.Textbox(lines=5)
-        gr.Button("Detect").click(detect_road_lanes, lane_in, [lane_out, lane_text])
-    with gr.Tab("Road Signs"):
+            lane_input = gr.Image(type="pil", label="Upload Image")
+            lane_output = gr.Image(label="Detection Results")
+        lane_text = gr.Textbox(label="Details", lines=5)
+        lane_btn = gr.Button("Detect Road Lanes", variant="primary")
+        lane_btn.click(detect_road_lanes, inputs=lane_input, outputs=[lane_output, lane_text])
+    
+    with gr.Tab("🚦 Road Signs Only"):
         with gr.Row():
-            sign_in = gr.Image(type="pil")
-            sign_out = gr.Image()
-        sign_text = gr.Textbox(lines=5)
-        gr.Button("Detect").click(detect_road_signs, sign_in, [sign_out, sign_text])
+            sign_input = gr.Image(type="pil", label="Upload Image")
+            sign_output = gr.Image(label="Detection Results")
+        sign_text = gr.Textbox(label="Details", lines=5)
+        sign_btn = gr.Button("Detect Road Signs", variant="primary")
+        sign_btn.click(detect_road_signs, inputs=sign_input, outputs=[sign_output, sign_text])
 
 demo.launch()