app.py

from flask import Flask, request, send_file, jsonify
import cv2
import numpy as np
import os
import uuid

app = Flask(__name__)

# Paths
YOLO_CFG = "yolo/yolov3.cfg"
YOLO_WEIGHTS = "yolo/yolov3.weights"
CLASSES_FILE = "yolo/coco.names"

CONF_THRESHOLD = 0.5
NMS_THRESHOLD = 0.4

# Load classes
with open(CLASSES_FILE, "r") as f:
    classes = [line.strip() for line in f.readlines()]

# Load YOLO
net = cv2.dnn.readNet(YOLO_WEIGHTS, YOLO_CFG)
layer_names = net.getLayerNames()
output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]

@app.route("/predict", methods=["POST"])
def predict():
    if "image" not in request.files:
        return jsonify({"error": "No image provided"}), 400

    file = request.files["image"]
    img = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_COLOR)
    height, width, _ = img.shape

    # YOLO inference
    blob = cv2.dnn.blobFromImage(img, 1/255.0, (416, 416), swapRB=True, crop=False)
    net.setInput(blob)
    outputs = net.forward(output_layers)

    boxes = []
    confidences = []
    class_ids = []

    for output in outputs:
        for detection in output:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]

            if confidence > CONF_THRESHOLD:
                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)

    # 🔥 NMS
    indexes = cv2.dnn.NMSBoxes(
        boxes, confidences, CONF_THRESHOLD, NMS_THRESHOLD
    )

    for i in indexes.flatten():
        x, y, w, h = boxes[i]
        label = f"{classes[class_ids[i]]}: {confidences[i]:.2f}"
        color = (0, 255, 0)

        cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
        cv2.putText(
            img, label, (x, y - 10),
            cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2
        )

    # Save annotated image
    os.makedirs("outputs", exist_ok=True)
    filename = f"outputs/result_{uuid.uuid4().hex}.png"
    cv2.imwrite(filename, img)

    return send_file(filename, mimetype="image/png")

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)