imgsearcher/app/api/baidu_image_search.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import base64
import json
import os
import requests
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

class BaiduImageSearch:
    """百度相似图片搜索API封装类"""
    
    def __init__(self):
        """初始化，获取API密钥"""
        self.api_key = os.getenv('BAIDU_API_KEY')
        self.secret_key = os.getenv('BAIDU_SECRET_KEY')
        self.access_token = None
        self.get_access_token()
    
    def get_access_token(self):
        """获取百度API的access_token"""
        url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {
            "grant_type": "client_credentials", 
            "client_id": self.api_key, 
            "client_secret": self.secret_key
        }
        response = requests.post(url, params=params)
        if response.status_code == 200:
            result = response.json()
            self.access_token = result.get('access_token')
            return self.access_token
        else:
            raise Exception(f"获取access_token失败: {response.text}")
    
    def add_image(self, image_path=None, image_base64=None, url=None, brief=None, tags=None):
        """
        添加图片到图库
        
        Args:
            image_path: 本地图片路径
            image_base64: 图片的base64编码
            url: 图片URL
            brief: 图片摘要信息，最长256B，如：{"name":"图片名称", "id":"123"}
            tags: 分类信息，1 - 65535范围内的整数	tag间以逗号分隔，最多2个tag，2个tag无层级关系，检索时支持逻辑运算。样例："100,11" ；检索时可圈定分类维度进行检索
            
        Returns:
            dict: API返回的结果
        """
        request_url = f"https://aip.baidubce.com/rest/2.0/image-classify/v1/realtime_search/similar/add?access_token={self.access_token}"
        
        params = {}
        
        # 设置brief信息（必选）
        if brief:
            if isinstance(brief, dict):
                brief = json.dumps(brief, ensure_ascii=False)
            params['brief'] = brief
        else:
            raise ValueError("brief参数是必须的")
        
        # 设置图片信息（三选一）
        if image_path:
            with open(image_path, 'rb') as f:
                image = base64.b64encode(f.read())
            params['image'] = image
        elif image_base64:
            params['image'] = image_base64
        elif url:
            params['url'] = url
        else:
            raise ValueError("必须提供image_path、image_base64或url其中之一")
        
        # 设置可选的tags
        if tags:
            params['tags'] = tags
        
        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
        response = requests.post(request_url, data=params, headers=headers)
        print(f"添加图片请求: {params}")
        print(f"添加图片响应: {response.text}")
        
        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"添加图片失败: {response.text}")
    
    def search_image(self, image_path=None, image_base64=None, url=None, tags=None, tag_logic=None, type_filter=None, pn=0, rn=300):
        """
        检索相似图片
        
        Args:
            image_path: 本地图片路径
            image_base64: 图片的base64编码
            url: 图片URL
            tags: 分类信息过滤，如："1,2"
            tag_logic: 标签逻辑，0表示逻辑and，1表示逻辑or
            type_filter: 类型过滤，用于按图片类型进行筛选
            pn: 分页起始位置，默认0
            rn: 返回结果数量，默认300，最大1000
            
        Returns:
            dict: API返回的结果
        """
        request_url = f"https://aip.baidubce.com/rest/2.0/image-classify/v1/realtime_search/similar/search?access_token={self.access_token}"
        
        params = {}
        
        # 设置图片信息（三选一）
        if image_path:
            with open(image_path, 'rb') as f:
                image = base64.b64encode(f.read())
            params['image'] = image
        elif image_base64:
            params['image'] = image_base64
        elif url:
            params['url'] = url
        else:
            raise ValueError("必须提供image_path、image_base64或url其中之一")
        
        # 设置可选参数
        if tags:
            params['tags'] = tags
        if tag_logic is not None:
            params['tag_logic'] = tag_logic
        if pn is not None:
            params['pn'] = pn
        if rn is not None:
            params['rn'] = rn
            
        # 处理搜索结果的类型过滤
        self.type_filter = type_filter
        
        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
        # print(f"搜索图片请求: {params}")
        response = requests.post(request_url, data=params, headers=headers)
        # print(f"搜索图片响应: {response.text}")
        
        if response.status_code == 200:
            result = response.json()
            
            # 如果有类型过滤，对结果进行过滤
            if self.type_filter and 'result' in result and result['result']:
                filtered_results = []
                for item in result['result']:
                    try:
                        brief = item.get('brief', '{"type":""}')
                        
                        # 确保brief是字符串
                        if not isinstance(brief, str):
                            brief = str(brief)
                            
                        try:
                            brief_info = json.loads(brief)
                        except json.JSONDecodeError as e:
                            print(f"JSON解析错误: {e}, brief: {brief}")
                            continue
                            
                        if brief_info.get('type', '') == self.type_filter:
                            filtered_results.append(item)
                    except Exception as e:
                        print(f"处理搜索结果项出错: {e}")
                        continue
                result['result'] = filtered_results
                result['result_num'] = len(filtered_results)
            
            # 分析结果，确定最可信的类型
            if 'result' in result and result['result']:
                result['most_reliable_types'] = self._determine_most_reliable_types(result['result'])
            
            return result
        else:
            raise Exception(f"检索图片失败: {response.text}")
    
    def delete_image(self, image_path=None, image_base64=None, url=None, cont_sign=None):
        """
        删除图库中的图片
        
        Args:
            image_path: 本地图片路径
            image_base64: 图片的base64编码
            url: 图片URL
            cont_sign: 图片签名，支持批量删除，格式如："932301884,1068006219;316336521,553141152"
            
        Returns:
            dict: API返回的结果
        """
        request_url = f"https://aip.baidubce.com/rest/2.0/image-classify/v1/realtime_search/similar/delete?access_token={self.access_token}"
        
        params = {}
        
        # 设置图片信息（三选一）
        if image_path:
            with open(image_path, 'rb') as f:
                image = base64.b64encode(f.read())
            params['image'] = image
        elif image_base64:
            params['image'] = image_base64
        elif url:
            params['url'] = url
        elif cont_sign:
            params['cont_sign'] = cont_sign
        else:
            raise ValueError("必须提供image_path、image_base64、url或cont_sign其中之一")
        
        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
        response = requests.post(request_url, data=params, headers=headers)
        
        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"删除图片失败: {response.text}")
    
    def _determine_most_reliable_types(self, results):
        """
        使用多种方法确定最可信的类型
        只使用相似度前十的结果进行计算
        
        Args:
            results: 搜索结果列表
            
        Returns:
            dict: 包含不同方法确定的最可信类型
        """
        # 防止空结果
        if not results:
            return {
                "method1": "",
                "method2": "",
                "method3": "",
                "method4": "",
                "method5": ""
            }
            
        # 只使用相似度前十的结果
        # 按相似度降序排序并取前10个
        sorted_results = sorted(results, key=lambda x: x.get('score', 0), reverse=True)
        top_results = sorted_results[:min(10, len(sorted_results))]
        
        print(f"使用前{len(top_results)}个结果进行最可信类型分析")
        # 提取前十个结果的类型和分数
        type_scores = []
        for item in top_results:
            try:
                score = item.get('score', 0)
                brief = item.get('brief', '{"type":""}')
                
                # 确保brief是字符串
                if not isinstance(brief, str):
                    brief = str(brief)
                
                try:
                    brief_info = json.loads(brief)
                except json.JSONDecodeError as e:
                    print(f"JSON解析错误: {e}, brief: {brief}")
                    continue
                    
                item_type = brief_info.get('type', '')
                if item_type:  # 只考虑有类型的结果
                    type_scores.append((item_type, score))
            except Exception as e:
                print(f"处理搜索结果项出错: {e}, item: {item}")
                continue
        
        if not type_scores:
            return {
                "method1": "",
                "method2": "",
                "method3": "",
                "method4": "",
                "method5": ""
            }
        
        # 方法一：基于最高相似度分数选择类型
        method1_type = ""
        max_score = -1
        for t, s in type_scores:
            if s > max_score:
                max_score = s
                method1_type = t
        
        # 方法二：基于加权投票选择类型
        type_weight = {}
        for t, s in type_scores:
            if t not in type_weight:
                type_weight[t] = 0
            type_weight[t] += s
        
        try:
            method2_type = max(type_weight.items(), key=lambda x: x[1])[0] if type_weight else ""
        except ValueError:
            method2_type = ""
        
        # 方法三：基于阈值过滤和加权投票选择类型
        # 只考虑相似度大于0.6的结果
        threshold = 0.6
        filtered_type_weight = {}
        for t, s in type_scores:
            if s >= threshold:
                if t not in filtered_type_weight:
                    filtered_type_weight[t] = 0
                filtered_type_weight[t] += s
        
        try:
            method3_type = max(filtered_type_weight.items(), key=lambda x: x[1])[0] if filtered_type_weight else ""
        except ValueError:
            method3_type = ""
        
        # 方法四：基于多数投票选择类型
        type_count = {}
        for t, _ in type_scores:
            if t not in type_count:
                type_count[t] = 0
            type_count[t] += 1
        
        try:
            method4_type = max(type_count.items(), key=lambda x: x[1])[0] if type_count else ""
        except ValueError:
            method4_type = ""
        
        # 方法五：基于加权多数投票选择类型
        # 权重 = 计数 * 平均分数
        type_weighted_count = {}
        type_total_score = {}
        
        # 初始化类型计数和总分数
        for t, s in type_scores:
            if t not in type_total_score:
                type_total_score[t] = 0
            type_total_score[t] += s
        
        # 计算加权分数
        for t in type_count.keys():
            if type_count[t] > 0:
                avg_score = type_total_score.get(t, 0) / type_count[t]
                type_weighted_count[t] = type_count[t] * avg_score
        
        try:
            method5_type = max(type_weighted_count.items(), key=lambda x: x[1])[0] if type_weighted_count else ""
        except ValueError:
            method5_type = ""
        
        return {
            "method1": method1_type,  # 最高分数法
            "method2": method2_type,  # 加权投票法
            "method3": method3_type,  # 阈值过滤加权投票法
            "method4": method4_type,  # 多数投票法
            "method5": method5_type   # 加权多数投票法
        }
    
    def update_image(self, image_path=None, image_base64=None, url=None, cont_sign=None, brief=None, tags=None):
        """
        更新图库中图片的摘要和分类信息
        
        Args:
            image_path: 本地图片路径
            image_base64: 图片的base64编码
            url: 图片URL
            cont_sign: 图片签名
            brief: 更新的摘要信息，最长256B
            tags: 更新的分类信息，最多2个tag，如："1,2"
            
        Returns:
            dict: API返回的结果
        """
        request_url = f"https://aip.baidubce.com/rest/2.0/image-classify/v1/realtime_search/similar/update?access_token={self.access_token}"
        
        params = {}
        
        # 设置图片信息（三选一）
        if image_path:
            with open(image_path, 'rb') as f:
                image = base64.b64encode(f.read())
            params['image'] = image
        elif image_base64:
            params['image'] = image_base64
        elif url:
            params['url'] = url
        elif cont_sign:
            params['cont_sign'] = cont_sign
        else:
            raise ValueError("必须提供image_path、image_base64、url或cont_sign其中之一")
        
        # 设置更新信息
        if brief:
            if isinstance(brief, dict):
                brief = json.dumps(brief, ensure_ascii=False)
            params['brief'] = brief
        
        if tags:
            params['tags'] = tags
        
        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
        response = requests.post(request_url, data=params, headers=headers)
        
        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"更新图片失败: {response.text}")