Spaces:
Sleeping
Sleeping
| import os | |
| import io | |
| import time | |
| import cv2 | |
| import numpy as np | |
| from flask import Flask, render_template, Response, request, send_file, jsonify | |
| from ultralytics import YOLO | |
| from PIL import Image | |
| app = Flask(__name__, template_folder="templates") | |
| # ---------- CONFIG ---------- | |
| MODEL_PATH = "./bestbest.pt" # ensure model is present | |
| IMGSZ = 540 | |
| WARNING_PX = 100 | |
| BOX_W = 100 | |
| BOX_H = 100 | |
| # ---------------------------- | |
| # 1. SETUP FACE DETECTION (Standard OpenCV Haarcascade) | |
| # This loads the face detection model built into OpenCV | |
| face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml') | |
| # Load Hand model once | |
| model = None | |
| if os.path.exists(MODEL_PATH): | |
| try: | |
| model = YOLO(MODEL_PATH) | |
| print("Model loaded.") | |
| except Exception as e: | |
| print("Model load failed:", e) | |
| model = None | |
| else: | |
| print(f"WARNING: model not found at {MODEL_PATH}") | |
| # --- helper functions (exact logic) --- | |
| def get_box_coords(frame_w, frame_h, box_w=BOX_W, box_h=BOX_H): | |
| cx, cy = frame_w // 2, frame_h // 2 | |
| x1 = cx - box_w // 2 | |
| y1 = cy - box_h // 2 | |
| x2 = cx + box_w // 2 | |
| y2 = cy + box_h // 2 | |
| return int(x1), int(y1), int(x2), int(y2) | |
| def rects_intersect(rectA, rectB): | |
| ax1, ay1, ax2, ay2 = rectA | |
| bx1, by1, bx2, by2 = rectB | |
| inter_x1 = max(ax1, bx1) | |
| inter_y1 = max(ay1, by1) | |
| inter_x2 = min(ax2, bx2) | |
| inter_y2 = min(ay2, by2) | |
| iw = inter_x2 - inter_x1 | |
| ih = inter_y2 - inter_y1 | |
| return (iw > 0) and (ih > 0) | |
| def rect_distance(rectA, rectB): | |
| ax1, ay1, ax2, ay2 = rectA | |
| bx1, by1, bx2, by2 = rectB | |
| dx = 0 | |
| if ax2 < bx1: dx = bx1 - ax2 | |
| elif bx2 < ax1: dx = ax1 - bx2 | |
| dy = 0 | |
| if ay2 < by1: dy = by1 - ay2 | |
| elif by2 < ay1: dy = ay1 - by2 | |
| return int(np.hypot(dx, dy)) | |
| def annotate_frame_bgr(frame_bgr): | |
| """Apply your detection + drawing logic on a BGR OpenCV frame, return annotated BGR""" | |
| h, w = frame_bgr.shape[:2] | |
| x1, y1, x2, y2 = get_box_coords(w, h, BOX_W, BOX_H) | |
| cv2.rectangle(frame_bgr, (x1, y1), (x2, y2), (220,220,220), 2) | |
| cv2.putText(frame_bgr, "DANGER ZONE", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (220,220,220), 2) | |
| # --- STEP 1: Detect Hand (Your original logic) --- | |
| hand_bbox = None | |
| if model is not None: | |
| try: | |
| results = model.predict(frame_bgr, imgsz=IMGSZ, conf=0.25, max_det=1, verbose=False) | |
| if len(results) > 0 and len(results[0].boxes) > 0: | |
| b = results[0].boxes[0] | |
| xyxy = b.xyxy.cpu().numpy().astype(int)[0] | |
| hand_bbox = (int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])) | |
| except Exception as e: | |
| print("Inference error:", e) | |
| hand_bbox = None | |
| # --- STEP 2: Detect Faces & Filter --- | |
| # Convert to grayscale for face detection (required by haar cascade) | |
| gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY) | |
| faces = face_cascade.detectMultiScale(gray, 1.1, 4) | |
| # Check if the detected hand overlaps with ANY detected face | |
| if hand_bbox is not None: | |
| for (fx, fy, fw, fh) in faces: | |
| # Create face rectangle (x1, y1, x2, y2) | |
| face_rect = (fx, fy, fx + fw, fy + fh) | |
| # If hand overlaps face, HIDE HAND (set to None) | |
| if rects_intersect(hand_bbox, face_rect): | |
| print("Ignored hand detection (Overlap with Face)") | |
| hand_bbox = None | |
| break # Stop checking other faces, we already hid the hand | |
| # --- STEP 3: Draw Logic (Only draws if hand_bbox is still valid) --- | |
| state = "NO HAND" | |
| color = (100,100,100) | |
| dist = 0 | |
| if hand_bbox is not None: | |
| danger_rect = (x1, y1, x2, y2) | |
| if rects_intersect(hand_bbox, danger_rect): | |
| state = "DANGER"; color = (0,0,255); dist = 0 | |
| else: | |
| dist = rect_distance(hand_bbox, danger_rect) | |
| if dist <= WARNING_PX: | |
| state = "WARNING"; color = (0,165,255) | |
| else: | |
| state = "SAFE"; color = (0,255,0) | |
| cv2.rectangle(frame_bgr, (hand_bbox[0], hand_bbox[1]), (hand_bbox[2], hand_bbox[3]), color, 2) | |
| cv2.putText(frame_bgr, f"Dist: {dist}px", (10,110), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2) | |
| cv2.putText(frame_bgr, state, (10,70), cv2.FONT_HERSHEY_DUPLEX, 1.5, color, 3) | |
| if state == "DANGER": | |
| cv2.putText(frame_bgr, "HAND TOO CLOSE!", (w//2 - 200, h//2), cv2.FONT_HERSHEY_DUPLEX, 1.5, (0,0,255), 4) | |
| return frame_bgr | |
| # ---- Routes ---- | |
| def index(): | |
| return render_template("index.html") | |
| def predict(): | |
| """ | |
| Receives a single JPEG (multipart form under key 'frame'). | |
| Returns processed JPEG image as response (image/jpeg). | |
| """ | |
| if "frame" not in request.files: | |
| return jsonify({"error":"no frame"}), 400 | |
| file = request.files["frame"] | |
| in_bytes = file.read() | |
| # read into OpenCV | |
| nparr = np.frombuffer(in_bytes, np.uint8) | |
| img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) # BGR | |
| if img is None: | |
| return jsonify({"error":"invalid image"}), 400 | |
| # annotate | |
| out_bgr = annotate_frame_bgr(img) | |
| # encode to jpeg | |
| ret, buf = cv2.imencode(".jpg", out_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 85]) | |
| if not ret: | |
| return jsonify({"error":"encode failed"}), 500 | |
| return Response(buf.tobytes(), mimetype="image/jpeg") | |
| if __name__ == "__main__": | |
| # Run with flask for debugging | |
| app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)), debug=False) |