YOLOv12_HFD / src /web_app.py
RevDra's picture
Update src/web_app.py
8670ed5
Raw History Blame Contribute Delete
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"}
@app.route("/")
def index():
"""Main page"""
return render_template("index.html")
@app.route("/api/detect-image", methods=["POST"])
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
@app.route("/api/detect-video", methods=["POST"])
@limiter.limit("5 per hour")
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
@app.route("/api/download/<filename>", methods=["GET"])
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
@app.route("/api/models", methods=["GET"])
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)
@app.route("/api/feedback", methods=["POST"])
@limiter.limit("5 per minute")
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
@app.route("/api/health", methods=["GET"])
def health_check():
"""Health check endpoint"""
return jsonify({"status": "healthy", "service": "Face Detection ONNX API"})
@app.errorhandler(413)
def request_entity_too_large(error):
"""Handle file size exceeded"""
return jsonify({"error": "File too large. Maximum 500MB allowed"}), 413
@app.errorhandler(500)
def internal_error(error):
"""Handle internal server error"""
return jsonify({"error": "Internal server error"}), 500
@app.errorhandler(RateLimitExceeded)
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)