Spaces:
Running
Running
File size: 5,447 Bytes
32c3c28 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | from flask import Flask, Response, jsonify, request
from flask_cors import CORS
import cv2
import json
import time
import os
import base64
import numpy as np
from datetime import datetime
from ppe_detection import load_model, process_frame
app = Flask(__name__)
CORS(app) # Allow cross-origin requests from Vercel frontend
# Global variables
model = None
detection_active = False
detection_results = []
violation_count = 0
helmet_count = 0
vest_count = 0
glove_count = 0
def init_model():
"""Load the YOLOv8 model globally"""
global model
if model is None:
print("Loading YOLOv8 model...")
model = load_model()
if model is None:
print("Failed to load model!")
def update_counters(detections):
"""Update global counters based on detection results"""
global violation_count, helmet_count, vest_count, glove_count
current_violation_count = 0
current_helmet_count = 0
current_vest_count = 0
current_glove_count = 0
for detection in detections:
type_name = detection.get("type", "")
if type_name.startswith("NO-") and detection.get("detected", False):
current_violation_count += 1
elif (type_name == "Hardhat" or type_name == "helmet") and detection.get("detected", False):
current_helmet_count += 1
elif (type_name == "Safety Vest" or type_name == "vest") and detection.get("detected", False):
current_vest_count += 1
elif (type_name == "Gloves" or type_name == "hand gloves") and detection.get("detected", False):
current_glove_count += 1
# Update global counters
violation_count = current_violation_count
helmet_count = current_helmet_count
vest_count = current_vest_count
glove_count = current_glove_count
@app.route("/")
def index():
"""Home page route"""
return "PPE Detection Cloud API Server is running!"
@app.route("/api/start", methods=["POST"])
def start_detection():
"""Start PPE detection"""
global detection_active
# Initialize the model on first start
init_model()
detection_active = True
return jsonify({"success": True, "message": "Detection started"})
@app.route("/api/stop", methods=["POST"])
def stop_detection():
"""Stop PPE detection"""
global detection_active
detection_active = False
return jsonify({"success": True, "message": "Detection stopped"})
@app.route("/api/status")
def get_status():
"""Get detection status and counters"""
global detection_active, violation_count, helmet_count, vest_count, glove_count
return jsonify({
"active": detection_active,
"violations": violation_count,
"helmets": helmet_count,
"vests": vest_count,
"gloves": glove_count
})
@app.route("/api/results")
def get_results():
"""Get recent detection results"""
global detection_results
return jsonify({"results": detection_results})
@app.route("/api/socket", methods=["POST", "OPTIONS"])
def process_socket_frame():
"""Process incoming base64 video frames from the frontend"""
if request.method == "OPTIONS":
return jsonify({"success": True})
global model, detection_active, detection_results
if not detection_active:
return jsonify({"success": False, "message": "Detection is not active"})
if model is None:
init_model()
if model is None:
return jsonify({"success": False, "message": "Model failed to load"})
try:
data = request.json
frame_data = data.get("frame", "")
# Remove base64 prefix if present
if "," in frame_data:
frame_data = frame_data.split(",")[1]
# Decode base64 to OpenCV image
img_bytes = base64.b64decode(frame_data)
np_arr = np.frombuffer(img_bytes, np.uint8)
frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
if frame is None:
return jsonify({"success": False, "message": "Failed to decode frame"})
# Run YOLO inference
_, detections = process_frame(frame, model, 0)
timestamp = datetime.now().strftime("%H:%M:%S")
# Update counters
update_counters(detections)
# Update results log
result_entry = {
"timestamp": timestamp,
"detections": detections
}
detection_results.append(result_entry)
# Keep only the last 50 results
if len(detection_results) > 50:
detection_results = detection_results[-50:]
return jsonify({
"success": True,
"processed": True,
"timestamp": timestamp,
"message": "Frame processed successfully",
"detections": detections
})
except Exception as e:
print(f"Error processing frame: {e}")
return jsonify({"success": False, "message": str(e)})
if __name__ == "__main__":
print("Starting Cloud PPE Detection API Server...")
# Render binds to the PORT environment variable
port = int(os.environ.get("PORT", 5000))
app.run(host="0.0.0.0", port=port, debug=False, threaded=True) |