hangzhou_mat_manage/imgsearchimg/api/app.py

from flask import Flask, request, jsonify, render_template, send_from_directory
import os
from image_search import ImageSearchEngine
import magic
from urllib.request import urlretrieve
import time

app = Flask(__name__)

# 配置常量
UPLOAD_FOLDER = 'static/images'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp'}

# 搜索参数默认值
TOP_K = 5  # 默认返回的最大结果数
MIN_SCORE = 0.0  # 默认最小相似度分数
MAX_SCORE = 100.0  # 默认最大相似度分数

os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 初始化图像搜索引擎
search_engine = ImageSearchEngine()

def allowed_file(filename):
    """检查文件是否允许上传"""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def is_valid_image(file_path):
    """检查文件是否为有效的图片文件"""
    try:
        mime = magic.Magic(mime=True)
        file_type = mime.from_file(file_path)
        return file_type.startswith('image/')
    except Exception:
        return False

@app.route('/')
def index():
    """渲染主页"""
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    """处理图片上传请求"""
    try:
        # 获取并验证参数
        data = request.get_json()
        if not data:
            return jsonify({'error': '请求必须包含JSON数据'}), 400
        
        product_id = data.get('product_id')
        image_url = data.get('url')
        
        if not product_id:
            return jsonify({'error': '缺少product_id参数'}), 400
        if not image_url:
            return jsonify({'error': '缺少url参数'}), 400
        if not isinstance(image_url, str) or not (image_url.startswith('http://') or image_url.startswith('https://')):
            return jsonify({'error': '无效的图片URL'}), 400
        
        time_base = str(time.time() * 1000)
        image_path = os.path.join(UPLOAD_FOLDER, time_base + os.path.basename(image_url))
        urlretrieve(image_url, image_path)

        if allowed_file(image_path) and is_valid_image(image_path):
            # 添加图片到搜索引擎
            if search_engine.add_image_from_url(image_path, product_id):
                # 删除临时图片文件
                os.remove(image_path)
                return jsonify({'message': '上传成功', 'product_id': product_id})
            else:
                return jsonify({'error': '处理图片失败'}), 500
        else:
            return jsonify({'error': '图片错误'}), 500
            
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/search', methods=['POST'])
def search():
    """处理图片搜索请求"""
    try:
        # 获取并验证参数
        data = request.get_json()
        if not data:
            return jsonify({'error': '请求必须包含JSON数据'}), 400
        
        image_url = data.get('url')
        if not image_url:
            return jsonify({'error': '缺少url参数'}), 400
        if not isinstance(image_url, str) or not (image_url.startswith('http://') or image_url.startswith('https://')):
            return jsonify({'error': '无效的图片URL'}), 400
            
        # 获取可选参数
        limit = data.get('limit', TOP_K)
        min_score = data.get('min_score', MIN_SCORE)
        max_score = data.get('max_score', MAX_SCORE)
        
        # 验证参数类型和范围
        try:
            limit = int(limit)
            min_score = float(min_score)
            max_score = float(max_score)
            
            if limit <= 0:
                return jsonify({'error': 'limit必须大于0'}), 400
            if min_score < 0 or min_score > 100:
                return jsonify({'error': 'min_score必须在0到100之间'}), 400
            if max_score < 0 or max_score > 100:
                return jsonify({'error': 'max_score必须在0到100之间'}), 400
            if min_score > max_score:
                return jsonify({'error': 'min_score不能大于max_score'}), 400
        except ValueError:
            return jsonify({'error': '参数类型错误'}), 400
        
        start_download_time = time.time()
        time_base = str(time.time() * 1000)
        image_path = os.path.join(UPLOAD_FOLDER, time_base + os.path.basename(image_url))
        urlretrieve(image_url, image_path)
        end_download_time = time.time()
        print(f"下载耗时: { end_download_time - start_download_time } s",)

        if allowed_file(image_path) and is_valid_image(image_path):
            # 搜索相似图片，使用用户指定的 limit
            start_search_time = time.time()
            results = search_engine.search(image_path, top_k=limit)
            os.remove(image_path)
            # 格式化结果并过滤不在分数范围内的结果
            formatted_results = []
            for product_id, score in results:
                similarity = round(score * 100, 2)  # 转换为百分比
                if min_score <= similarity <= max_score:
                    formatted_results.append({
                        'product_id': product_id,
                        'similarity': similarity
                    })
            end_search_time = time.time()
            print(f"搜索耗时: { end_search_time - start_search_time } s",)
            return jsonify({'results': formatted_results})
        else:
            return jsonify({'error': '图片错误'})

    except Exception as e:
        return jsonify({'error': str(e)}), 500

# @app.route('/images')
# def list_images():
#     """列出所有已索引的图片"""
#     try:
#         images = []
#         for path in search_engine.image_paths:
#             filename = os.path.basename(path)
#             images.append({
#                 'filename': filename,
#                 'path': '/static/images/' + filename
#             })
#         return jsonify({'images': images})
#     except Exception as e:
#         return jsonify({'error': str(e)}), 500

@app.route('/remove/<product_id>', methods=['POST'])
def remove_image(product_id):
    """移除指定商品ID的图片特征"""
    try:
        # 从索引中移除
        if search_engine.remove_by_product_id(product_id):
            return jsonify({'message': '删除成功', 'product_id': product_id})
        else:
            return jsonify({'error': '商品ID不存在或删除失败'}), 404
            
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/clear/<password>', methods=['POST'])
def clear_index(password):
    """清空索引"""
    if password == "infish@2025":
        try:
            search_engine.clear()
        except Exception as e:
            return jsonify({'error': str(e)}), 500

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