Spaces:
Running
Running
Download src/web_app.py from RevDra/YOLOv12_HFD: direct link, hf CLI and curl.
- Browser
- Download file 16 kB
-
https://huggingface.co/spaces/RevDra/YOLOv12_HFD/resolve/main/src/web_app.py
- Command line
-
hf download hf://spaces/RevDra/YOLOv12_HFD/src/web_app.py
-
curl -L -o web_app.py https://huggingface.co/spaces/RevDra/YOLOv12_HFD/resolve/main/src/web_app.py
16 kB
| """ | |
| Flask Web Application for YOLOv12 Face Detection (ONNX Backend) | |
| Supports image upload, video upload, and live webcam streaming | |
| Optimized for Headless CPU Cloud environments (Hugging Face Spaces) | |
| """ | |
| import os | |
| os.environ.setdefault("OMP_NUM_THREADS", "1") | |
| os.environ.setdefault("ONNXRUNTIME_INTER_OP_NUM_THREADS", "1") | |
| os.environ.setdefault("ONNXRUNTIME_INTRA_OP_NUM_THREADS", "2") | |
| import base64 | |
| import logging | |
| from datetime import datetime | |
| from pathlib import Path | |
| from typing import Any, Dict, cast | |
| import cv2 | |
| import numpy as np | |
| from dotenv import load_dotenv | |
| from flask import Flask, jsonify, render_template, request, send_file | |
| from flask_limiter import Limiter | |
| from flask_limiter.errors import RateLimitExceeded | |
| from flask_limiter.util import get_remote_address | |
| from flask_sqlalchemy import SQLAlchemy | |
| from werkzeug.utils import secure_filename | |
| try: | |
| from turbojpeg import TurboJPEG | |
| jpeg_engine = TurboJPEG() | |
| except Exception: | |
| jpeg_engine = None | |
| def decode_image_fast(image_bytes: bytes) -> np.ndarray: | |
| """Decode images using TurboJPEG if have, else OpenCV""" | |
| if jpeg_engine is not None: | |
| try: | |
| return jpeg_engine.decode(image_bytes) | |
| except Exception: | |
| pass | |
| nparr = np.frombuffer(image_bytes, np.uint8) | |
| return cv2.imdecode(nparr, cv2.IMREAD_COLOR) | |
| def encode_jpeg_fast(image_bgr: np.ndarray, quality: int = 85) -> bytes: | |
| """Zip to TurboJPEG""" | |
| if jpeg_engine is not None: | |
| try: | |
| res = jpeg_engine.encode(image_bgr, quality=quality) | |
| if res is not None: | |
| return cast(bytes, res) | |
| except Exception: | |
| pass | |
| _, buffer = cv2.imencode(".jpg", image_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), quality]) | |
| return bytes(buffer) | |
| from face_detection_yolov12 import YOLOv12FaceDetector, detect_from_video | |
| load_dotenv() | |
| app = Flask(__name__, template_folder="../web/templates") | |
| app.config["SQLALCHEMY_DATABASE_URI"] = os.getenv("DB_URL") | |
| app.config["SECRET_KEY"] = os.getenv("SECRET_KEY", "fallback_secret_key_neu_khong_co") | |
| app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False | |
| app.config["SQLALCHEMY_ENGINE_OPTIONS"] = {"pool_recycle": 280} | |
| db = SQLAlchemy(app) | |
| PROJECT_ROOT = Path(__file__).parent.parent | |
| UPLOAD_FOLDER = PROJECT_ROOT / "data" / "uploads" | |
| MODELS_DIR = PROJECT_ROOT / "models" | |
| ALLOWED_EXTENSIONS = {"jpg", "jpeg", "png", "gif", "mp4", "avi", "mov", "mkv"} | |
| MAX_FILE_SIZE = 1024 * 1024 * 1024 # 1GB | |
| ALLOWED_MODELS = { | |
| "yolov12n-face-int8.onnx", | |
| "yolov12s-face-int8.onnx", | |
| "yolov12m-face-int8.onnx", | |
| "yolov12l-face-int8.onnx", | |
| } | |
| UPLOAD_FOLDER.mkdir(parents=True, exist_ok=True) | |
| app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER | |
| app.config["MAX_CONTENT_LENGTH"] = MAX_FILE_SIZE | |
| limiter = Limiter( | |
| key_func=get_remote_address, | |
| app=app, | |
| default_limits=["10000 per day", "1000 per hour"], | |
| storage_uri="memory://", | |
| ) | |
| class Feedback(db.Model): # type: ignore | |
| __tablename__ = "feedbacks" | |
| id = db.Column(db.Integer, primary_key=True) | |
| filename = db.Column(db.String(255)) | |
| model_name = db.Column(db.String(50)) | |
| rating = db.Column(db.Integer) | |
| comment = db.Column(db.Text) | |
| created_at = db.Column(db.DateTime, default=datetime.utcnow) | |
| with app.app_context(): | |
| try: | |
| db.create_all() | |
| print("Connect to MySQL successfully!") | |
| except Exception as e: | |
| print(f"Connection error: {e}") | |
| def get_detector(model_name): | |
| """ | |
| Get or create detector instance. | |
| The YOLOv12FaceDetector class natively handles Singleton caching internally, | |
| so we don't need a manual dictionary cache here anymore. | |
| """ | |
| safe_name = secure_filename(model_name) | |
| if safe_name.endswith(".pt"): | |
| safe_name = safe_name.replace(".pt", ".onnx") | |
| if safe_name not in ALLOWED_MODELS: | |
| logging.error(f"Attempt to load unsupported model: {safe_name}") | |
| raise ValueError(f"Unsupported model: {safe_name}") | |
| model_path = MODELS_DIR / safe_name | |
| try: | |
| final_path = model_path.resolve() | |
| safe_root = MODELS_DIR.resolve() | |
| if not str(final_path).startswith(str(safe_root)): | |
| logging.error(f"Security Alert: Symlink attack detected! {final_path}") | |
| raise ValueError("Invalid model path (Symlink violation)") | |
| except Exception as e: | |
| logging.error(f"Error resolving model path: {str(e)}") | |
| raise FileNotFoundError(f"Model path error: {str(e)}") | |
| if not final_path.exists(): | |
| logging.error(f"Model file not found: {final_path}") | |
| raise FileNotFoundError(f"Model not found: {final_path}") | |
| return YOLOv12FaceDetector(final_path.name) | |
| def allowed_file(filename): | |
| """Check if file extension is allowed""" | |
| return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS | |
| def is_image(filename): | |
| """Check if file is image""" | |
| ext = filename.rsplit(".", 1)[1].lower() | |
| return ext in {"jpg", "jpeg", "png", "gif"} | |
| def is_video(filename): | |
| """Check if file is video""" | |
| ext = filename.rsplit(".", 1)[1].lower() | |
| return ext in {"mp4", "avi", "mov", "mkv"} | |
| def index(): | |
| """Main page""" | |
| return render_template("index.html") | |
| def detect_image(): | |
| """Detect faces in uploaded image or live webcam frame""" | |
| try: | |
| if "file" not in request.files: | |
| return jsonify({"error": "No file provided"}), 400 | |
| file = request.files["file"] | |
| if file.filename == "": | |
| return jsonify({"error": "No file selected"}), 400 | |
| if not allowed_file(file.filename) or not is_image(file.filename): | |
| return jsonify({"error": "Only image files allowed"}), 400 | |
| # Get parameters | |
| model = request.form.get("model", "yolov12s-face-int8.onnx") | |
| blur = request.form.get("blur") == "true" | |
| # Nhận diện xem request có phải đến từ luồng Webcam hay không | |
| is_webcam = request.form.get("mode") == "webcam" or "webcam" in file.filename.lower() | |
| if not model.endswith(".onnx"): | |
| model = model.replace(".pt", ".onnx") | |
| if model not in ALLOWED_MODELS: | |
| app.logger.info(f"Invalid model '{model}' requested. Fallback to default.") | |
| model = "yolov12n-face-int8.onnx" if is_webcam else "yolov12s-face-int8.onnx" | |
| # Lấy instance detector (đã được cache ở RAM qua Singleton) | |
| detector = get_detector(model) | |
| # Giải mã ảnh trực tiếp từ bytes tốc độ cao | |
| image_data = file.read() | |
| image = decode_image_fast(image_data) | |
| if image is None: | |
| return jsonify({"error": "Invalid image data"}), 400 | |
| img_h, img_w = image.shape[:2] | |
| detections = detector.detect_faces(image, conf_threshold=0.32) | |
| # --- NHÁNH 1: WEBCAM (CHỈ TRẢ VỀ TỌA ĐỘ - ZERO SERVER-SIDE RENDERING) --- | |
| if is_webcam: | |
| return jsonify( | |
| { | |
| "success": True, | |
| "mode": "coordinates_only", | |
| "frame_width": img_w, | |
| "frame_height": img_h, | |
| "detections": { | |
| "count": len(detections), | |
| "boxes": [ | |
| { | |
| "x1": det["x1"], | |
| "y1": det["y1"], | |
| "x2": det["x2"], | |
| "y2": det["y2"], | |
| "w": det["w"], | |
| "h": det["h"], | |
| "conf": round(det["confidence"], 3), | |
| } | |
| for det in detections | |
| ], | |
| }, | |
| } | |
| ) | |
| # --- NHÁNH 2: ẢNH TĨNH UPLOAD THÔNG THƯỜNG --- | |
| if blur: | |
| result_image = detector.draw_faces(image, detections, show_conf=False, blur_faces=blur) | |
| else: | |
| result_image = detector.draw_faces(image, detections, show_conf=True, blur_faces=blur) | |
| if result_image is None: | |
| return jsonify({"error": "Failed to process image"}), 500 | |
| # Extract crops for gallery | |
| crops_base64 = [] | |
| if len(detections) > 0: | |
| crops = detector.get_face_crops(image, detections) | |
| for crop in crops: | |
| crop_bytes = encode_jpeg_fast(crop, quality=80) | |
| crops_base64.append(base64.b64encode(crop_bytes).decode()) | |
| # Convert result to base64 for display | |
| result_bytes = encode_jpeg_fast(result_image, quality=85) | |
| img_base64 = base64.b64encode(result_bytes).decode() | |
| # Prepare response | |
| response = { | |
| "success": True, | |
| "image": f"data:image/jpeg;base64,{img_base64}", | |
| "detections": { | |
| "count": len(detections), | |
| "faces": [ | |
| { | |
| "id": i + 1, | |
| "confidence": f"{det['confidence']:.2%}", | |
| "width": det["w"], | |
| "height": det["h"], | |
| "position": f"({det['x1']}, {det['y1']})", | |
| } | |
| for i, det in enumerate(detections) | |
| ], | |
| "crops": crops_base64, | |
| }, | |
| } | |
| return jsonify(response) | |
| except Exception: | |
| logging.exception("Error during image detection") | |
| return jsonify({"error": "Internal server error during image detection"}), 500 | |
| def detect_video(): | |
| """Detect faces in uploaded video""" | |
| try: | |
| if "file" not in request.files: | |
| return jsonify({"error": "No file provided"}), 400 | |
| file = request.files["file"] | |
| if file.filename == "": | |
| return jsonify({"error": "No file selected"}), 400 | |
| if not allowed_file(file.filename) or not is_video(file.filename): | |
| return jsonify({"error": "Only video files allowed"}), 400 | |
| # Get model selection | |
| raw_model = request.form.get("model", "yolov12n-face-int8.onnx") | |
| model = Path(raw_model).name | |
| if model.endswith(".pt"): | |
| model = model.replace(".pt", ".onnx") | |
| if not model.endswith(".onnx"): | |
| model = f"{model}.onnx" | |
| if model not in ALLOWED_MODELS: | |
| app.logger.info(f"Invalid model '{model}' requested. Fallback to default (Nano INT8).") | |
| model = "yolov12n-face-int8.onnx" | |
| # Save uploaded file | |
| filename = secure_filename(file.filename) | |
| input_path = UPLOAD_FOLDER / f"input_{filename}" | |
| output_path = UPLOAD_FOLDER / f"output_{filename}" | |
| file.save(input_path) | |
| # Detect faces in video | |
| detect_from_video( | |
| video_path=str(input_path), | |
| model_path=model, | |
| output_path=str(output_path), | |
| conf_threshold=0.32, | |
| ) | |
| # Return file info | |
| response = { | |
| "success": True, | |
| "message": "Video processing complete", | |
| "output_file": output_path.name, | |
| "download_url": f"/api/download/{output_path.name}", | |
| } | |
| return jsonify(response) | |
| except Exception: | |
| app.logger.exception("Error while processing video detection request") | |
| return jsonify({"error": "Internal server error"}), 500 | |
| def download_file(filename): | |
| try: | |
| filepath = UPLOAD_FOLDER / secure_filename(filename) | |
| if not filepath.exists(): | |
| return jsonify({"error": "File not found"}), 404 | |
| as_attachment = request.args.get("download") == "1" | |
| return send_file( | |
| filepath, | |
| mimetype="video/mp4", | |
| as_attachment=as_attachment, | |
| conditional=True, | |
| ) | |
| except Exception: | |
| app.logger.exception("Error while processing download request for %s", filename) | |
| return jsonify({"error": "Internal server error"}), 500 | |
| def get_models(): | |
| """Get ALL available ONNX models for dropdown selection""" | |
| models = { | |
| "nano": { | |
| "name": "yolov12n-face-int8.onnx", | |
| "label": "Nano (n) - Fastest (ONNX INT8)", | |
| "description": "Real-time speed, best for CPU/Webcam", | |
| "size": "Smallest", | |
| }, | |
| "small": { | |
| "name": "yolov12s-face-int8.onnx", | |
| "label": "Small (s) - Balanced (ONNX INT8)", | |
| "description": "Good balance of speed and accuracy", | |
| "size": "Small", | |
| }, | |
| "medium": { | |
| "name": "yolov12m-face-int8.onnx", | |
| "label": "Medium (m) - High Precision (ONNX INT8)", | |
| "description": "High accuracy, optimized CPU execution", | |
| "size": "Medium", | |
| }, | |
| "large": { | |
| "name": "yolov12l-face-int8.onnx", | |
| "label": "Large (l) - Max Accuracy (ONNX INT8)", | |
| "description": "Best detection quality", | |
| "size": "Large", | |
| }, | |
| } | |
| available: dict[str, Any] = {} | |
| for key, info in models.items(): | |
| model_path = MODELS_DIR / info["name"] | |
| if model_path.exists(): | |
| available[key] = info | |
| order = ["nano", "small", "medium", "large"] | |
| sorted_available = {k: available[k] for k in order if k in available} | |
| return jsonify(sorted_available) | |
| def submit_feedback(): | |
| try: | |
| data = request.json | |
| new_fb = Feedback( | |
| filename=data.get("filename"), | |
| model_name=data.get("model"), | |
| rating=data.get("rating"), | |
| comment=data.get("comment", ""), | |
| ) | |
| db.session.add(new_fb) | |
| db.session.commit() | |
| return jsonify({"success": True, "message": "Rating saved!"}) | |
| except Exception as e: | |
| app.logger.error(f"DB Error: {e}") | |
| return jsonify({"error": "Database error"}), 500 | |
| def health_check(): | |
| """Health check endpoint""" | |
| return jsonify({"status": "healthy", "service": "Face Detection ONNX API"}) | |
| def request_entity_too_large(error): | |
| """Handle file size exceeded""" | |
| return jsonify({"error": "File too large. Maximum 500MB allowed"}), 413 | |
| def internal_error(error): | |
| """Handle internal server error""" | |
| return jsonify({"error": "Internal server error"}), 500 | |
| def handle_rate_limit_error(e): | |
| """Handle rate limit exceeded errors""" | |
| app.logger.warning(f"Rate limit exceeded: {e.description}") | |
| return ( | |
| jsonify( | |
| { | |
| "error": "Too many requests", | |
| "message": f"Too fast! Please wait a moment. ({e.description})", | |
| } | |
| ), | |
| 429, | |
| ) | |
| if __name__ == "__main__": | |
| print("\n" + "=" * 70) | |
| print("🌐 Starting YOLOv12 Face Detection Web Server (ONNX Optimized)") | |
| print("=" * 70) | |
| print("\n📍 Server: http://localhost:7860") | |
| print("📁 Upload folder: ", UPLOAD_FOLDER) | |
| print("🔧 Models folder: ", MODELS_DIR) | |
| print("\n🎯 Available endpoints:") | |
| print(" GET / - Web interface") | |
| print(" POST /api/detect-image - Detect faces in image") | |
| print(" POST /api/detect-video - Detect faces in video") | |
| print(" GET /api/models - Get available ONNX models") | |
| print(" GET /api/health - Health check") | |
| print("\n" + "=" * 70 + "\n") | |
| # Determine debug mode from environment (default: disabled) | |
| debug_mode = os.getenv("FLASK_ENV") == "development" | |
| app.run(host="0.0.0.0", port=7860, debug=debug_mode, use_reloader=False) |