imgsearcher/app.py

#!/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('/')
def index():
    """首页"""
    return render_template('index.html')

@app.route('/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('/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('/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('/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('/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('/chat-with-image', methods=['GET'])
def chat_with_image_page():
    """图片对话页面"""
    return render_template('chat.html')

@app.route('/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('/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)