Update app.py

app.py

@@ -1,41 +1,52 @@
-import gradio as gr
 import cv2
-import numpy as np
+import gradio as gr
+from ultralytics import YOLO
 from PIL import Image
+import numpy as np
+
+# Load the YOLOv5 model (pre-trained on COCO dataset)
+model = YOLO('yolov8n.pt')  # You can replace with your custom model if available
 
-def process_image(image):
+def detect_objects(image):
     """
-    Process an uploaded or captured image: convert it to grayscale.
+    Detect suspicious objects in the image using YOLO.
     """
-    if isinstance(image, np.ndarray):
-        # Convert to grayscale if it's a NumPy array (e.g., from webcam or upload)
-        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-        processed_image = cv2.cvtColor(gray_image, cv2.COLOR_GRAY2RGB)
-        return processed_image
-    elif isinstance(image, Image.Image):
-        # Handle PIL image
-        image = np.array(image)
-        gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
-        processed_image = cv2.cvtColor(gray_image, cv2.COLOR_GRAY2RGB)
-        return Image.fromarray(processed_image)
-    else:
-        return "Invalid image format"
-
-# Create the Gradio interface
+    # Convert PIL image to OpenCV format (numpy array)
+    image_np = np.array(image)
+    results = model(image_np)  # Perform detection
+
+    # Draw bounding boxes on the detected objects
+    for result in results:
+        boxes = result.boxes  # Bounding boxes
+        for box in boxes:
+            x1, y1, x2, y2 = map(int, box.xyxy[0])  # Extract box coordinates
+            label = box.cls[0]  # Class label
+            confidence = box.conf[0]  # Confidence score
+
+            # Draw rectangle and label on the image
+            cv2.rectangle(image_np, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            text = f"{model.names[int(label)]} ({confidence:.2f})"
+            cv2.putText(image_np, text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+
+    # Convert back to PIL image
+    processed_image = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
+    return processed_image
+
+# Create a Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Image Capture and Processing App")
-    gr.Markdown("Upload an image or capture one, and convert it to grayscale!")
+    gr.Markdown("# Suspicious Object Detection")
+    gr.Markdown("Upload an image or use your webcam to capture one. The app will detect objects using YOLOv5.")
 
     # Input section
     image_input = gr.Image(type="pil", label="Upload or Capture Image")
 
     # Output section
-    output_image = gr.Image(type="pil", label="Processed Image")
+    output_image = gr.Image(type="pil", label="Processed Image with Annotations")
 
-    # Button to process the image
-    process_button = gr.Button("Process Image")
+    # Button for detection
+    detect_button = gr.Button("Detect Suspicious Objects")
 
-    # Link button to the function
-    process_button.click(process_image, inputs=[image_input], outputs=output_image)
+    # Link the button to the detection function
+    detect_button.click(detect_objects, inputs=[image_input], outputs=[output_image])
 
 demo.launch()