from flask import Flask, request, jsonify import os from image_search import ImageSearchEngine import time import config app = Flask(__name__) os.makedirs(config.UPLOAD_FOLDER, exist_ok=True) # 初始化图像搜索引擎 search_engine = ImageSearchEngine() def allowed_file(filename): """检查文件是否允许上传""" return '.' in filename and filename.rsplit('.', 1)[1].lower() in config.ALLOWED_EXTENSIONS @app.route('/add', methods=['POST']) def add_image_file(): """处理图片上传请求""" try: # 获取并验证参数 data = request.get_json() if not data: return jsonify({'error': '请求必须包含JSON数据'}), 400 product_id = data.get('imageId') 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 # 添加图片到搜索引擎 if search_engine.add_image_from_url(image_url, product_id): return jsonify({'message': '上传成功', 'product_id': product_id}) 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 print(f"search payload, type:{type(data)}, data: {data}") # 获取可选参数 limit = data.get('limit', config.TOP_K) min_score = data.get('min_score', config.MIN_SCORE) max_score = data.get('max_score', config.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 > 1.0: return jsonify({'error': 'min_score必须在0到1之间'}), 400 if max_score < 0 or max_score > 1.0: return jsonify({'error': 'max_score必须在0到1之间'}), 400 if min_score > max_score: return jsonify({'error': 'min_score不能大于max_score'}), 400 except ValueError: return jsonify({'error': '参数类型错误'}), 400 print(f"search apply limit:{limit} min_score:{min_score} max_score:{max_score}") start_search_time = time.time() results = search_engine.search(image_url, top_k=limit) # 格式化结果并过滤不在分数范围内的结果 formatted_results = [] for product_id, score in results: similarity = score # 转换为百分比 if min_score <= similarity <= max_score: formatted_results.append({ 'Id': product_id, 'Score': similarity }) end_search_time = time.time() print(f"搜索耗时: { end_search_time - start_search_time } s",) return jsonify(formatted_results), 200 except Exception as e: return jsonify({'error': str(e)}), 500 @app.route('/remove/', 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/', 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)