72 lines
2.6 KiB
Python
72 lines
2.6 KiB
Python
"""图片文件解析器，使用 Vision API 识别图片内容"""
|
||
|
||
import base64
|
||
import os
|
||
from typing import List, Optional
|
||
|
||
from api_client import ApiClient, EXTENSION_MIME_MAP
|
||
from exceptions import ApiError, ParseError
|
||
from parsers.base import BaseParser
|
||
|
||
DEFAULT_VISION_PROMPT = """\
|
||
请识别并提取图片中的所有文字和关键信息。请按以下结构输出：
|
||
|
||
1. **产品/主题名称**：图片展示的主要产品或主题
|
||
2. **文字内容**：图片中所有可见的文字，保持原始排版
|
||
3. **关键信息**：成分、功效、用法用量、规格、价格等结构化信息
|
||
4. **图片描述**：简要描述图片的视觉内容（产品外观、包装等）
|
||
|
||
如果某项信息不存在，可以省略该项。"""
|
||
|
||
|
||
class ImageParser(BaseParser):
|
||
"""图片解析器，通过 Vision API 将图片转换为文本描述"""
|
||
|
||
def __init__(self, api_client: ApiClient, vision_prompt: Optional[str] = None):
|
||
self._api_client = api_client
|
||
self._vision_prompt = vision_prompt or DEFAULT_VISION_PROMPT
|
||
|
||
def supported_extensions(self) -> List[str]:
|
||
return [".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp"]
|
||
|
||
def parse(self, file_path: str) -> str:
|
||
"""
|
||
解析图片文件：读取二进制 → base64 编码 → 调用 Vision API → 返回文本描述。
|
||
|
||
会将文件名作为上下文提示传入 prompt，提高识别准确度。
|
||
"""
|
||
file_name = os.path.basename(file_path)
|
||
product_name = os.path.splitext(file_name)[0]
|
||
|
||
# 1. 读取图片文件
|
||
try:
|
||
with open(file_path, "rb") as f:
|
||
image_bytes = f.read()
|
||
except Exception as e:
|
||
raise ParseError(file_name, f"文件读取失败: {e}")
|
||
|
||
# 2. Base64 编码
|
||
image_base64 = base64.b64encode(image_bytes).decode("utf-8")
|
||
|
||
# 3. 根据扩展名确定 MIME 类型
|
||
ext = os.path.splitext(file_path)[1].lower()
|
||
mime_type = EXTENSION_MIME_MAP.get(ext, "image/png")
|
||
|
||
# 4. 构建带文件名上下文的 prompt
|
||
context_prompt = (
|
||
f"{self._vision_prompt}\n\n"
|
||
f"参考信息：该图片的文件名为「{product_name}」，可能与图片内容相关。"
|
||
)
|
||
|
||
# 5. 调用 Vision API
|
||
try:
|
||
result = self._api_client.vision(
|
||
system_prompt=context_prompt,
|
||
image_base64=image_base64,
|
||
mime_type=mime_type,
|
||
)
|
||
except ApiError as e:
|
||
raise ParseError(file_name, f"Vision API 调用失败: {e}")
|
||
|
||
return result
|