Update app.py

app.py

@@ -1,41 +1,79 @@
-import os
-from ultralytics import YOLO
-from PIL import Image, ImageDraw
-from flask import Flask, request, jsonify
+import cv2
+import numpy as np
+import datetime
+import gradio as gr
 
-# Initialize Flask app
-app = Flask(__name__)
+# Ensure these files are available in the Hugging Face Space working directory
+net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")
+layer_names = net.getLayerNames()
+output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
+classes = []
+with open("coco.names", "r") as f:
+    classes = [line.strip() for line in f.readlines()]
 
-# Load YOLO model for text detection (Optimized for CPU)
-YOLO_MODEL_PATH = "best.pt"
-detection_model = YOLO(YOLO_MODEL_PATH, task='detect').to("cpu")  # Force CPU usage
+def detect_objects(image):
+    height, width, channels = image.shape
+    blob = cv2.dnn.blobFromImage(image, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
+    net.setInput(blob)
+    outs = net.forward(output_layers)
+    
+    class_ids = []
+    confidences = []
+    boxes = []
+    
+    for out in outs:
+        for detection in out:
+            scores = detection[5:]
+            class_id = np.argmax(scores)
+            confidence = scores[class_id]
+            if confidence > 0.5:
+                center_x = int(detection[0] * width)
+                center_y = int(detection[1] * height)
+                w = int(detection[2] * width)
+                h = int(detection[3] * height)
+                
+                x = int(center_x - w / 2)
+                y = int(center_y - h / 2)
+                
+                boxes.append([x, y, w, h])
+                confidences.append(float(confidence))
+                class_ids.append(class_id)
+    
+    indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
+    
+    return [(boxes[i], class_ids[i], confidences[i]) for i in range(len(boxes)) if i in indexes]
 
-@app.route("/detect", methods=["POST"])
-def detect_text():
-    if "image" not in request.files:
-        return jsonify({"error": "No image provided"}), 400
+def process_image(image):
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    detections = detect_objects(image)
+    for (box, class_id, confidence) in detections:
+        x, y, w, h = box
+        label = str(classes[class_id])
+        color = (0, 255, 0) if label == "person" else (0, 0, 255)
+        cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
+        cv2.putText(image, f'{label} {confidence:.2f}', (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
-    image_file = request.files["image"]
-    image = Image.open(image_file).convert("RGB")
+def capture_and_process():
+    cap = cv2.VideoCapture(0)
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        processed_frame = process_image(frame)
+        yield processed_frame
+    cap.release()
 
-    # Run YOLO detection (Optimized for CPU)
-    results = detection_model.predict(image, conf=0.3, iou=0.4, device="cpu")
-    detected_boxes = results[0].boxes.xyxy.tolist()
-    detected_boxes = [list(map(int, box)) for box in detected_boxes]
+# Define Gradio interface
+with gr.Blocks() as iface:
+    gr.Markdown("# YOLO Object Detection")
+    gr.Markdown("## Real-time object detection using YOLO")
 
-    # Draw bounding boxes on the image
-    image_with_boxes = image.copy()
-    draw = ImageDraw.Draw(image_with_boxes)
+    with gr.Tab("Upload Image"):
+        gr.Markdown("Upload an image and the YOLO model will detect objects in the image, highlighting humans.")
+        image_input = gr.Image(type="numpy", label="Upload an image")
+        image_output = gr.Image(type="numpy", label="Detected objects")
+        image_input.upload(process_image, inputs=image_input, outputs=image_output)
 
-    for box in detected_boxes:
-        x1, y1, x2, y2 = box
-        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
-
-    # Save the processed image with bounding boxes
-    output_path = "detected_image.jpg"
-    image_with_boxes.save(output_path)
-
-    return jsonify({"boxes": detected_boxes, "output_image": output_path})
-
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=5000, debug=True)
+# Launch Gradio interface
+iface.launch()