sensitive-face-recognizer/app.py

import os
import tempfile

import face_recognition
import requests
import logging
import jieba.analyse

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, url)
            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, file.filename)
            os.remove(temp_file.name)  # 删除临时文件
            results = return_json(find)
            return results
        return jsonify({'error': 'An error occurred'}), 500

    @app.route('/title/split', methods=['POST'])
    def get_title_split():
        if 'title' not in request.json:
            return jsonify({'error': 'No title provided'}), 400
        title = request.json.get('title')
        if not title:
            return jsonify({'error': 'not title'}), 400
        return evaluate(title, 3)

    return app


# 调用 setup_logging 函数创建并配置日志记录器
def setup_logging():
    # 创建一个 Logger 对象
    logger = logging.getLogger('sensitive-face-recognizer_logger')
    # 设置日志级别，可根据需要调整为 DEBUG、INFO、WARNING、ERROR 或 CRITICAL
    logger.setLevel(logging.INFO)

    # 创建一个 Handler 对象，这里使用 StreamHandler 输出到控制台
    console_handler = logging.StreamHandler()
    # 创建一个文件处理器，输出到文件
    file_handler = logging.FileHandler('app.log')

    # 创建一个 Formatter 对象，设置日志格式
    formatter = logging.Formatter('%(levelname)s - %(message)s')
    console_handler.setFormatter(formatter)
    file_handler.setFormatter(formatter)

    # 将 Handler 添加到 Logger 中
    logger.addHandler(console_handler)
    logger.addHandler(file_handler)

    return logger


logger = setup_logging()


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 get_face_similarity(known_face_encoding, face_encoding):
    # 计算两张脸的面部编码之间的距离
    face_distances = face_recognition.face_distance([known_face_encoding], face_encoding)
    # 计算相似度，使用 1 - 距离作为相似度
    similarity = 1 - face_distances
    return similarity


def find_faces_in_image(image_path, url):
    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)
            known_face_encoding = known_face_encodings[first_match_index]
            similarity = get_face_similarity(known_face_encoding, face_encoding)
            logger.info(f'image_name={name} similarity={similarity} url={url}')
    if found_faces:
        return True
    else:
        return False


def load_top_tags():
    kv = {}
    try:
        with open('top_tags.txt', 'r', encoding='utf-8') as cache_file:
            for index, line in enumerate(cache_file):
                if index == 10000:
                    break
                line = line.strip()
                if not line:
                    continue
                k, _ = line.split('\t')
                kv[k] = index
    except FileNotFoundError:
        print("未找到 top_tags.txt 文件。")
    return kv


kv = load_top_tags()


def evaluate(title, top_size):
    if title is None or len(title) == 0:
        return ""
    keys = jieba.analyse.extract_tags(title, topK=top_size, withWeight=False, allowPOS=('n', 'v'))
    if keys is None or len(keys) == 0:
        return ""
    keys_filter = []
    for k in keys:
        if k in kv:
            keys_filter.append(k)
    if len(keys_filter) == 0:
        return ""
    return ",".join(keys_filter)


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