knowledge-agent/indentify/image_identifier.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图文识别脚本
主要功能：使用 Gemini API 进行图片OCR识别
"""

import os
import json
import time
import sys
from typing import Dict, Any, List, Optional
from dotenv import load_dotenv
import google.generativeai as genai
from PIL import Image
import requests
from io import BytesIO

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


class ImageIdentifier:
    def __init__(self):
        # 加载环境变量
        load_dotenv()
        
        # 初始化Gemini API
        api_key = os.getenv('GEMINI_API_KEY')
        if not api_key:
            raise ValueError("请在环境变量中设置 GEMINI_API_KEY")
        
        genai.configure(api_key=api_key)
        self.model = genai.GenerativeModel('gemini-2.5-flash')
    
    def download_image(self, image_url: str) -> Optional[Image.Image]:
        """下载图片并转换为PIL Image对象"""
        try:
            response = requests.get(image_url, timeout=10)
            response.raise_for_status()
            image = Image.open(BytesIO(response.content))
            return image
        except Exception as e:
            print(f"下载图片失败 {image_url}: {e}")
            return None
    
    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_image_with_gemini(self, image: Image.Image) -> Dict[str, Any]:
        """使用Gemini API分析单张图片内容"""
        try:
            # 构建OCR提示词
            prompt = """
            #### 人设
            你是一名图像文字理解专家，请对输入的文章图片进行精准的文字提取和结构化整理。

            #### 任务要求如下：
            1. 仅提取图片中可见的文字内容，不需要改写、总结或推理隐藏信息。
            2. 如果图片包含结构（如表格、图表、标题、段落等），请按结构输出。
            3. 所有提取的内容需保持原始顺序和排版上下文的逻辑。
            4. 不需要进行OCR校正，只需要原样提取图中文字。
            5. 舍弃图片中和标题不相关的文字
            6. 对于结构不明确或自由排列的文字，按照从上到下、从左到右的顺序依次提取。
            """
            
            response = self.model.generate_content([prompt, image])
            
            return {
                "text_content": response.text,
                "success": True
            }
            
        except Exception as e:
            print(f"Gemini API调用失败: {e}")
            return {
                "text_content": "",
                "success": False,
                "error": str(e)
            }
    
    def analyze_images_with_gemini(self, image_urls: List[str]) -> Dict[str, Any]:
        """使用Gemini API分析多张图片内容"""
        try:
            if not image_urls:
                return {"images_comprehension": [], "error": "没有图片需要分析"}
            
            print(f"正在使用Gemini API分析 {len(image_urls)} 张图片...")
            results = []
            
            for i, image_url in enumerate(image_urls):
                print(f"正在处理第 {i+1} 张图片: {image_url}")
                
                # 下载图片
                image = self.download_image(image_url)
                if image is None:
                    results.append({
                        "image_url": image_url,
                        "text_content": "",
                        "success": False,
                        "error": "图片下载失败"
                    })
                    continue
                
                # 分析图片
                result = self.analyze_image_with_gemini(image)
                result["image_url"] = image_url
                results.append(result)
                
                # 添加延迟避免API限制
                time.sleep(1)
            
            return {
                "images_comprehension": results
            }
                
        except Exception as e:
            print(f"Gemini API批量调用失败: {e}")
            return {"images_comprehension": [], "error": f"Gemini API调用失败: {str(e)}"}
    
    def process_images(self, formatted_content: Dict[str, Any]) -> Dict[str, Any]:
        """处理图片识别的主函数"""
        print("开始图片OCR识别处理...")
        
        # 提取图片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_gemini(image_urls)
        
        if result.get("images_comprehension"):
            successful_count = sum(1 for img in result['images_comprehension'] if img.get('success', False))
            print(f"图片OCR识别完成，成功分析 {successful_count}/{len(result['images_comprehension'])} 张图片")
        else:
            print("图片OCR识别失败")
        
        return result


def main():
    """测试函数"""
    # 模拟数据
    test_content = {
        "image_url_list": [
            {
                "image_type": 2,
                "image_url": "http://rescdn.yishihui.com/pipeline/image/ea4f33e9-9e36-4124-aaec-138ea9bcadd9.jpg"
            }
        ]
    }
    
    try:
        identifier = ImageIdentifier()
        result = identifier.process_images(test_content)
        
        print(f"识别结果: {json.dumps(result, ensure_ascii=False, indent=2)}")
    except Exception as e:
        print(f"初始化失败: {e}")


if __name__ == '__main__':
    main()