knowledge/content_indentify/image_identifier.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图文识别脚本
主要功能：使用 Coze API 分析图片内容
"""

import os
import json
import time
import sys
from typing import Dict, Any, List, Optional
from dotenv import load_dotenv

# 导入自定义模块
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from coze.coze_hook import CozeHook


class ImageIdentifier:
    def __init__(self):
        # 加载环境变量
        load_dotenv()
        
        # 初始化Coze客户端
        self.coze = CozeHook()
    
    def extract_image_urls(self, formatted_content: Dict[str, Any]) -> List[str]:
        """提取图片URL列表"""
        image_urls = []
        image_url_list = formatted_content.get('image_url_list', [])
        
        for img_data in image_url_list:
            if isinstance(img_data, dict) and 'image_url' in img_data:
                image_urls.append(img_data['image_url'])
        
        return image_urls
    
    def analyze_images_with_coze(self, image_urls: List[str]) -> Dict[str, Any]:
        """使用Coze API分析图片内容"""
        try:
            if not image_urls:
                return {"images_comprehension": [], "error": "没有图片需要分析"}
            
            print(f"正在使用Coze API分析 {len(image_urls)} 张图片...")
            response = self.coze.run(image_urls)
            
            # 解析Coze响应
            if response and 'data' in response:
                try:
                    if isinstance(response['data'], str):
                        data = json.loads(response['data'])
                    else:
                        data = response['data']
                    
                    return {
                        "images_comprehension": data.get('images_comprehension', []),
                    }
                except json.JSONDecodeError:
                    return {"images_comprehension": [], "error": "Coze响应解析失败", "raw_response": response}
            else:
                return {"images_comprehension": [], "error": "Coze API响应异常", "raw_response": response}
                
        except Exception as e:
            print(f"Coze API调用失败: {e}")
            return {"images_comprehension": [], "error": f"Coze API调用失败: {str(e)}"}
    
    def process_images(self, formatted_content: Dict[str, Any]) -> Dict[str, Any]:
        """处理图片识别的主函数"""
        print("开始图片识别处理...")
        
        # 提取图片URL
        image_urls = self.extract_image_urls(formatted_content)
        print(f"提取到 {len(image_urls)} 张图片")
        
        if not image_urls:
            print("没有图片需要分析")
            return {"images_comprehension": [], "error": "没有图片需要分析"}
        
        # 分析图片
        result = self.analyze_images_with_coze(image_urls)
        
        if result.get("images_comprehension"):
            print(f"图片识别完成，共分析 {len(result['images_comprehension'])} 张图片")
        else:
            print("图片识别失败")
        
        return result


def main():
    """测试函数"""
    # 模拟数据
    test_content = {
        "image_url_list": [
            {
                "image_type": 2,
                "image_url": "http://example.com/image1.jpg"
            }
        ]
    }
    
    identifier = ImageIdentifier()
    result = identifier.process_images(
        test_content
    )
    
    print(f"识别结果: {json.dumps(result, ensure_ascii=False, indent=2)}")


if __name__ == '__main__':
    main()