303 lines
8.6 KiB
Python
303 lines
8.6 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
图片理解核心模块
|
||
基于 OpenAI GPT-4 Vision API 实现图片理解功能
|
||
"""
|
||
|
||
import base64
|
||
import json
|
||
import os
|
||
from typing import Optional
|
||
|
||
import requests
|
||
|
||
|
||
class ImageUnderstander:
|
||
"""图片理解器"""
|
||
|
||
def __init__(self, api_key: str, verbose: bool = False):
|
||
"""
|
||
初始化图片理解器
|
||
|
||
Args:
|
||
api_key: OpenAI API Key
|
||
verbose: 是否显示详细日志
|
||
"""
|
||
self.api_key = api_key
|
||
self.verbose = verbose
|
||
self.api_url = "https://api.openai.com/v1/chat/completions"
|
||
|
||
def _encode_image(self, image_path: str) -> str:
|
||
"""
|
||
将图片编码为 base64 格式
|
||
|
||
Args:
|
||
image_path: 图片路径
|
||
|
||
Returns:
|
||
base64 编码的字符串
|
||
"""
|
||
with open(image_path, "rb") as f:
|
||
return base64.b64encode(f.read()).decode("utf-8")
|
||
|
||
def _call_api(self, messages: list, max_tokens: int = 2000) -> str:
|
||
"""
|
||
调用 OpenAI API
|
||
|
||
Args:
|
||
messages: 消息列表
|
||
max_tokens: 最大返回token数
|
||
|
||
Returns:
|
||
API 返回的文本内容
|
||
"""
|
||
headers = {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
data = {
|
||
"model": "gpt-4o",
|
||
"messages": messages,
|
||
"max_tokens": max_tokens,
|
||
"temperature": 0.2
|
||
}
|
||
|
||
if self.verbose:
|
||
print(f"📡 正在调用 OpenAI API...")
|
||
|
||
response = requests.post(self.api_url, headers=headers, json=data)
|
||
|
||
if response.status_code != 200:
|
||
raise Exception(f"API 调用失败: {response.status_code} - {response.text}")
|
||
|
||
result = response.json()
|
||
return result["choices"][0]["message"]["content"]
|
||
|
||
def describe(self, image_path: str) -> str:
|
||
"""
|
||
描述图片内容
|
||
|
||
Args:
|
||
image_path: 图片路径
|
||
|
||
Returns:
|
||
图片描述文本
|
||
"""
|
||
if self.verbose:
|
||
print(f"🖼️ 正在描述图片...")
|
||
|
||
base64_image = self._encode_image(image_path)
|
||
|
||
messages = [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": "请详细描述这张图片，包括：场景、人物、物体、颜色、氛围等所有可见元素。"
|
||
},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/jpeg;base64,{base64_image}"
|
||
}
|
||
}
|
||
]
|
||
}
|
||
]
|
||
|
||
return self._call_api(messages)
|
||
|
||
def extract_text(self, image_path: str) -> str:
|
||
"""
|
||
提取图片中的文字 (OCR)
|
||
|
||
Args:
|
||
image_path: 图片路径
|
||
|
||
Returns:
|
||
提取的文字内容
|
||
"""
|
||
if self.verbose:
|
||
print(f"🔤 正在提取文字...")
|
||
|
||
base64_image = self._encode_image(image_path)
|
||
|
||
messages = [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": "请提取图片中的所有文字，包括标题、正文、标签、logo文字等。如果有手写体，请尽量识别。如果图片中没有文字，请明确说明。"
|
||
},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/jpeg;base64,{base64_image}"
|
||
}
|
||
}
|
||
]
|
||
}
|
||
]
|
||
|
||
return self._call_api(messages, max_tokens=3000)
|
||
|
||
def identify_objects(self, image_path: str) -> str:
|
||
"""
|
||
识别图片中的物体
|
||
|
||
Args:
|
||
image_path: 图片路径
|
||
|
||
Returns:
|
||
识别出的物体列表
|
||
"""
|
||
if self.verbose:
|
||
print(f"🎯 正在识别物体...")
|
||
|
||
base64_image = self._encode_image(image_path)
|
||
|
||
messages = [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": "请识别图片中的所有物体和元素，并列出它们。建议以列表形式输出，每个物体加上简短描述。"
|
||
},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/jpeg;base64,{base64_image}"
|
||
}
|
||
}
|
||
]
|
||
}
|
||
]
|
||
|
||
return self._call_api(messages)
|
||
|
||
def answer_question(self, image_path: str, question: str) -> str:
|
||
"""
|
||
针对图片回答问题
|
||
|
||
Args:
|
||
image_path: 图片路径
|
||
question: 问题
|
||
|
||
Returns:
|
||
问题的答案
|
||
"""
|
||
if self.verbose:
|
||
print(f"💬 正在回答问题: {question}")
|
||
|
||
base64_image = self._encode_image(image_path)
|
||
|
||
messages = [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": f"请根据图片内容回答这个问题: {question}"
|
||
},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/jpeg;base64,{base64_image}"
|
||
}
|
||
}
|
||
]
|
||
}
|
||
]
|
||
|
||
return self._call_api(messages)
|
||
|
||
def analyze(self, image_path: str, analysis_type: str = "comprehensive") -> dict:
|
||
"""
|
||
综合分析图片
|
||
|
||
Args:
|
||
image_path: 图片路径
|
||
analysis_type: 分析类型 (comprehensive/quick/deep)
|
||
|
||
Returns:
|
||
分析结果字典
|
||
"""
|
||
if self.verbose:
|
||
print(f"🔬 正在进行{analysis_type}分析...")
|
||
|
||
base64_image = self._encode_image(image_path)
|
||
|
||
if analysis_type == "comprehensive":
|
||
prompt = """
|
||
请对这张图片进行全面分析，包括：
|
||
1. 【场景描述】整体场景是什么？
|
||
2. 【主要物体】列出主要物体和它们的位置
|
||
3. 【文字内容】提取所有可见文字
|
||
4. 【颜色氛围】主要色调和氛围
|
||
5. 【推断信息】根据内容推断可能的主题/用途
|
||
|
||
请用JSON格式返回：
|
||
{
|
||
"description": "场景描述",
|
||
"objects": ["物体1", "物体2", ...],
|
||
"text": "提取的文字",
|
||
"colors": "主要颜色",
|
||
"mood": "氛围描述",
|
||
"推断": "可能的主题/用途"
|
||
}
|
||
"""
|
||
elif analysis_type == "quick":
|
||
prompt = "请用一句话简单描述这张图片是什么。"
|
||
else: # deep
|
||
prompt = """
|
||
请对这张图片进行深度分析，包括：
|
||
- 详细场景描述
|
||
- 每个物体的精确位置和特征
|
||
- 图片的整体构图和视觉层次
|
||
- 可能的隐含信息或象征意义
|
||
- 图片的质量和技术特点
|
||
|
||
请详细回答。
|
||
"""
|
||
|
||
messages = [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": prompt.strip()
|
||
},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/jpeg;base64,{base64_image}"
|
||
}
|
||
}
|
||
]
|
||
}
|
||
]
|
||
|
||
result_text = self._call_api(messages, max_tokens=3000)
|
||
|
||
# 尝试解析JSON
|
||
try:
|
||
# 尝试从返回中找到JSON
|
||
if "{" in result_text and "}" in result_text:
|
||
start = result_text.find("{")
|
||
end = result_text.rfind("}") + 1
|
||
json_str = result_text[start:end]
|
||
return json.loads(json_str)
|
||
except:
|
||
pass
|
||
|
||
# 如果不是JSON，返回原始文本
|
||
return {
|
||
"result": result_text,
|
||
"raw_format": True
|
||
}
|