325 lines
11 KiB
Python
325 lines
11 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
import os
|
||
import json
|
||
import base64
|
||
from flask import Flask, render_template, request, jsonify, redirect, url_for, session
|
||
from flask_cors import CORS
|
||
from werkzeug.utils import secure_filename
|
||
from app.api.baidu_image_search import BaiduImageSearch
|
||
from app.api.image_utils import ImageUtils
|
||
from app.api.azure_openai import AzureOpenAI
|
||
from app.api.type_manager_mongo import TypeManagerMongo
|
||
|
||
app = Flask(__name__, template_folder='app/templates', static_folder='app/static')
|
||
CORS(app) # 启用CORS支持跨域请求
|
||
|
||
# 设置会话密钥
|
||
app.secret_key = os.urandom(24)
|
||
|
||
# 配置上传文件夹
|
||
UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
|
||
if not os.path.exists(UPLOAD_FOLDER):
|
||
os.makedirs(UPLOAD_FOLDER)
|
||
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
|
||
|
||
# 允许的文件扩展名
|
||
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
|
||
|
||
# 初始化百度图像搜索API
|
||
image_search_api = BaiduImageSearch()
|
||
|
||
# 初始化Azure OpenAI API
|
||
try:
|
||
azure_openai_api = AzureOpenAI()
|
||
except ValueError as e:
|
||
print(f"警告: {str(e)}")
|
||
azure_openai_api = None
|
||
|
||
# 初始化类型管理器
|
||
type_manager = TypeManagerMongo()
|
||
|
||
def allowed_file(filename):
|
||
"""检查文件扩展名是否允许"""
|
||
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
|
||
|
||
@app.route('/imgsearcherApi/')
|
||
def index():
|
||
"""首页"""
|
||
return render_template('index.html')
|
||
|
||
@app.route('/imgsearcherApi/upload', methods=['POST'])
|
||
def upload_image():
|
||
"""上传图片到图库"""
|
||
if 'file' not in request.files:
|
||
return jsonify({'error': '没有文件上传'}), 400
|
||
|
||
file = request.files['file']
|
||
if file.filename == '':
|
||
return jsonify({'error': '没有选择文件'}), 400
|
||
|
||
if not allowed_file(file.filename):
|
||
return jsonify({'error': '不支持的文件类型'}), 400
|
||
|
||
# 获取表单数据
|
||
image_type = request.form.get('type', '')
|
||
description = request.form.get('description', '')
|
||
name = request.form.get('name', '')
|
||
tags = request.form.get('tags', '1') # 默认标签为1
|
||
|
||
# 将类型和描述信息保存到本地
|
||
type_manager.add_type(image_type, description)
|
||
|
||
# 创建brief信息（不包含描述信息）
|
||
brief = {
|
||
'type': image_type,
|
||
'name': name
|
||
}
|
||
|
||
# 保存文件
|
||
filename = secure_filename(file.filename)
|
||
file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
|
||
file.save(file_path)
|
||
|
||
try:
|
||
# 调用API添加图片
|
||
result = image_search_api.add_image(
|
||
image_path=file_path,
|
||
brief=brief,
|
||
tags=tags
|
||
)
|
||
|
||
# 添加本地存储路径到结果
|
||
result['file_path'] = file_path
|
||
|
||
return jsonify(result)
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/search', methods=['POST'])
|
||
def search_image():
|
||
"""搜索相似图片"""
|
||
if 'file' not in request.files:
|
||
return jsonify({'error': '没有文件上传'}), 400
|
||
|
||
file = request.files['file']
|
||
if file.filename == '':
|
||
return jsonify({'error': '没有选择文件'}), 400
|
||
|
||
if not allowed_file(file.filename):
|
||
return jsonify({'error': '不支持的文件类型'}), 400
|
||
|
||
# 获取表单数据
|
||
image_type = request.form.get('type', '')
|
||
tags = request.form.get('tags', '1') # 默认标签为1
|
||
tag_logic = request.form.get('tag_logic', '0')
|
||
|
||
# 保存文件
|
||
filename = secure_filename(file.filename)
|
||
file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
|
||
file.save(file_path)
|
||
|
||
try:
|
||
# 调用API搜索图片
|
||
result = image_search_api.search_image(
|
||
image_path=file_path,
|
||
tags=tags,
|
||
tag_logic=tag_logic,
|
||
type_filter=image_type
|
||
)
|
||
|
||
# 从本地获取类型描述信息，并添加到结果中
|
||
if 'result' in result and result['result']:
|
||
for item in result['result']:
|
||
try:
|
||
brief = item.get('brief', '{}')
|
||
if isinstance(brief, str):
|
||
brief_dict = json.loads(brief)
|
||
item_type = brief_dict.get('type', '')
|
||
|
||
# 从本地获取描述信息
|
||
description = type_manager.get_description(item_type)
|
||
|
||
# 将描述信息添加到brief中
|
||
brief_dict['description'] = description
|
||
item['brief'] = json.dumps(brief_dict, ensure_ascii=False)
|
||
except Exception as e:
|
||
print(f"处理搜索结果项出错: {e}")
|
||
continue
|
||
|
||
return jsonify(result)
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/delete', methods=['POST'])
|
||
def delete_image():
|
||
"""删除图库中的图片"""
|
||
data = request.get_json()
|
||
|
||
if not data or 'cont_sign' not in data:
|
||
return jsonify({'error': '缺少必要参数'}), 400
|
||
|
||
cont_sign = data['cont_sign']
|
||
|
||
try:
|
||
# 调用API删除图片
|
||
result = image_search_api.delete_image(cont_sign=cont_sign)
|
||
return jsonify(result)
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/update', methods=['POST'])
|
||
def update_image():
|
||
"""更新图库中的图片信息"""
|
||
data = request.get_json()
|
||
|
||
if not data or 'cont_sign' not in data:
|
||
return jsonify({'error': '缺少必要参数'}), 400
|
||
|
||
cont_sign = data['cont_sign']
|
||
brief = data.get('brief')
|
||
tags = data.get('tags')
|
||
|
||
# 如果提供了brief信息，将类型和描述信息保存到本地
|
||
if brief and isinstance(brief, dict):
|
||
image_type = brief.get('type', '')
|
||
description = brief.get('description', '')
|
||
|
||
# 如果有描述信息，则更新到本地
|
||
if image_type and description:
|
||
type_manager.add_type(image_type, description)
|
||
|
||
# 从要上传的brief中移除描述信息
|
||
brief.pop('description', None)
|
||
|
||
try:
|
||
# 调用API更新图片
|
||
result = image_search_api.update_image(
|
||
cont_sign=cont_sign,
|
||
brief=brief,
|
||
tags=tags
|
||
)
|
||
return jsonify(result)
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/api/token')
|
||
def get_token():
|
||
"""获取API访问令牌"""
|
||
try:
|
||
token = image_search_api.get_access_token()
|
||
return jsonify({'access_token': token})
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/chat-with-image', methods=['GET'])
|
||
def chat_with_image_page():
|
||
"""图片对话页面"""
|
||
return render_template('chat.html')
|
||
|
||
@app.route('/imgsearcherApi/api/upload-chat-image', methods=['POST'])
|
||
def upload_chat_image():
|
||
"""上传图片用于对话"""
|
||
if 'file' not in request.files:
|
||
return jsonify({'error': '没有文件上传'}), 400
|
||
|
||
file = request.files['file']
|
||
if file.filename == '':
|
||
return jsonify({'error': '没有选择文件'}), 400
|
||
|
||
if not allowed_file(file.filename):
|
||
return jsonify({'error': '不支持的文件类型'}), 400
|
||
|
||
# 保存文件
|
||
filename = secure_filename(file.filename)
|
||
file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
|
||
file.save(file_path)
|
||
|
||
# 存储图片路径到会话
|
||
session['chat_image_path'] = file_path
|
||
|
||
# 清空对话历史
|
||
session['conversation_history'] = []
|
||
|
||
try:
|
||
# 使用百度图像搜索API搜索图片
|
||
search_result = image_search_api.search_image(image_path=file_path)
|
||
|
||
# 使用方法三确定最可信的类型
|
||
reliable_types = image_search_api._determine_most_reliable_types(search_result.get('result', []))
|
||
image_type = reliable_types.get('method3', '')
|
||
|
||
# 从本地获取描述信息
|
||
description = type_manager.get_description(image_type)
|
||
|
||
# 如果本地没有描述信息，尝试从搜索结果中获取
|
||
if not description:
|
||
for item in search_result.get('result', []):
|
||
brief = item.get('brief', '{}')
|
||
if isinstance(brief, str):
|
||
try:
|
||
brief_dict = json.loads(brief)
|
||
if brief_dict.get('type') == image_type:
|
||
description = brief_dict.get('description', '')
|
||
if description:
|
||
# 将获取到的描述信息保存到本地
|
||
type_manager.add_type(image_type, description)
|
||
break
|
||
except:
|
||
continue
|
||
|
||
return jsonify({
|
||
'success': True,
|
||
'image_path': file_path,
|
||
'image_type': image_type,
|
||
'description': description
|
||
})
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
@app.route('/imgsearcherApi/api/chat', methods=['POST'])
|
||
def chat():
|
||
"""与图片进行对话"""
|
||
if not azure_openai_api:
|
||
return jsonify({'error': 'Azure OpenAI API未正确配置'}), 500
|
||
|
||
data = request.get_json()
|
||
if not data or 'message' not in data:
|
||
return jsonify({'error': '缺少消息内容'}), 400
|
||
|
||
message = data['message']
|
||
image_path = session.get('chat_image_path')
|
||
|
||
if not image_path or not os.path.exists(image_path):
|
||
return jsonify({'error': '没有上传图片或图片已失效'}), 400
|
||
|
||
# 获取对话历史
|
||
conversation_history = session.get('conversation_history', [])
|
||
|
||
try:
|
||
# 调用Azure OpenAI API进行对话
|
||
response = azure_openai_api.chat_with_image(
|
||
image_path=image_path,
|
||
message=message,
|
||
conversation_history=conversation_history
|
||
)
|
||
|
||
# 提取回复内容
|
||
reply = response['choices'][0]['message']['content']
|
||
|
||
# 更新对话历史
|
||
conversation_history.append({"role": "user", "content": message})
|
||
conversation_history.append({"role": "assistant", "content": reply})
|
||
session['conversation_history'] = conversation_history
|
||
|
||
return jsonify({
|
||
'reply': reply,
|
||
'conversation_history': conversation_history
|
||
})
|
||
except Exception as e:
|
||
return jsonify({'error': str(e)}), 500
|
||
|
||
if __name__ == '__main__':
|
||
app.run(debug=True, host='0.0.0.0', port=5001)
|