detect_face.py

from flask import Flask, request, jsonify, send_file
import cv2
import numpy as np
from PIL import Image
import io
import torch
import os

# Disable fork validation check
torch.hub._validate_not_a_forked_repo = lambda a, b, c: True

# Load YOLOv5-face model
model_path = os.path.join(os.path.dirname(__file__), 'yolov5n-face.pt')
try:
    model = torch.hub.load('deepcam-cn/yolov5-face', 'custom', model_path, torch.device('cpu')) 
    model.conf = 0.3  # Confidence threshold for faces
    model.eval()
except Exception as e:
    print(f"Error loading model: {str(e)}")
    raise

app = Flask(__name__)

@app.route('/detect', methods=['POST'])
def detect_faces():
    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
    
    try:
        # Read the image file
        image_bytes = file.read()
        processed_image = process_image(image_bytes)
        # Return the processed image
        return send_file(processed_image, mimetype='image/jpeg')
    except Exception as e:
        print(f"Error in /detect endpoint: {str(e)}")
        return jsonify({"error": str(e)}), 500

def process_image(image_bytes):
    try:
        # Convert bytes to PIL Image
        image = Image.open(io.BytesIO(image_bytes))
        
        # Run inference
        results = model(image)
        
        # Get bounding boxes
        boxes = results.xyxy[0].cpu().numpy()
        
        # Draw boxes on image
        img_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        
        for box in boxes:
            x1, y1, x2, y2, conf, cls = box
            # Draw rectangle around face
            cv2.rectangle(img_np, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
            # Add confidence score
            cv2.putText(img_np, f'Face {conf:.2f}', 
                        (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 
                        0.5, (0, 255, 0), 2)
        
        # Convert back to bytes
        is_success, buffer = cv2.imencode(".jpg", img_np)
        if is_success:
            return io.BytesIO(buffer.tobytes())
        else:
            raise Exception("Failed to encode image")
    
    except Exception as e:
        print(f"Error processing image: {str(e)}")
        raise

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