sensitive-face-recognizer/app.py

import os
import tempfile

import face_recognition
import requests

from flask import Flask, request, jsonify, json


def create_app():
    app = Flask(__name__)

    def init_app():
        known_faces_dir = "./known_faces"
        load_known_faces(known_faces_dir)

    init_app()

    @app.route('/')
    def index():  # put application's code here
        return 'OK'

    @app.route('/recognize/url', methods=['POST'])
    def recognize_url():
        if 'url' not in request.json:
            return jsonify({'error': 'No URL provided'}), 400
        url = request.json.get('url')
        if not url:
            return jsonify({'error': 'Invalid URL'}), 400
        try:
            # 从 URL 下载文件
            response = requests.get(url)
            if response.status_code != 200:
                return jsonify({'error': 'Failed to download image from URL'}), 400
            # 创建一个临时文件保存下载的文件
            temp_file = tempfile.NamedTemporaryFile(delete=False)
            temp_file.write(response.content)
            temp_file.close()
            find = find_faces_in_image(temp_file.name)
            os.remove(temp_file.name)  # 删除临时文件
            results = return_json(find)
            return results
        except Exception as e:
            return jsonify({'error': str(e)}), 500

    @app.route('/recognize/file', methods=['POST'])
    def recognize_file():
        if 'file' not in request.files:
            return jsonify({'error': 'No file part'}), 400
        file = request.files['file']
        if file.filename == '':
            return jsonify({'error': 'No selected file'}), 400
        if file:
            # 创建一个临时文件保存上传的文件
            temp_file = tempfile.NamedTemporaryFile(delete=False)
            file.save(temp_file.name)
            find = find_faces_in_image(temp_file.name)
            os.remove(temp_file.name)# 删除临时文件
            results = return_json(find)
            return results
        return jsonify({'error': 'An error occurred'}), 500





    return app



def return_json(find):
    result = {
        "code": 0,
        "msg": "success",
        "data": None,  # 将 data 初始化为 None
        "success": True
    }
    if find:
        result["data"] = {"find_face_status": 1}
    else:
        result["data"] = {"find_face_status": 2}
    json_result = json.dumps(result)
    return json_result



# 全局变量存储已知人脸编码和名称
known_face_encodings = []
known_face_names = []


def load_known_faces(known_faces_dir):
    global known_face_encodings, known_face_names
    known_face_encodings = []
    known_face_names = []

    for filename in os.listdir(known_faces_dir):
        if filename.endswith(".jpg") or filename.endswith(".png"):
            image_path = os.path.join(known_faces_dir, filename)
            image = face_recognition.load_image_file(image_path)
            face_encoding = face_recognition.face_encodings(image)[0]
            known_face_encodings.append(face_encoding)
            known_face_names.append(os.path.splitext(filename)[0])


def find_faces_in_image(image_path):
    image = face_recognition.load_image_file(image_path)
    face_locations = face_recognition.face_locations(image)
    face_encodings = face_recognition.face_encodings(image, face_locations)

    found_faces = []
    for face_encoding in face_encodings:
        matches = face_recognition.compare_faces(known_face_encodings, face_encoding, tolerance=0.35)

        if True in matches:
            first_match_index = matches.index(True)
            name = known_face_names[first_match_index]
            found_faces.append(name)

    if found_faces:
        return True
    else:
        return False





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