多模态工程实战：构建图文理解流水线

import asyncio
import base64
from pathlib import Path
from openai import AsyncOpenAI
from pydantic import BaseModel, Field
from tenacity import retry, stop_after_attempt, wait_exponential

client = AsyncOpenAI()

class DocumentExtraction(BaseModel):
    """文档结构化提取结果"""
    title: str = Field(description="文档标题")
    content_type: str = Field(description="内容类型: invoice/receipt/contract/report")
    key_fields: dict[str, str] = Field(description="关键字段键值对")
    tables: list[dict] = Field(default_factory=list, description="提取的表格数据")
    confidence: float = Field(ge=0, le=1, description="提取置信度")

def encode_image(image_path: str) -> str:
    """将图像编码为 base64"""
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

def resize_for_token_efficiency(image_path: str, max_dim: int = 2048) -> str:
    """自适应裁剪，控制视觉 token 数量"""
    from PIL import Image
    import io

    img = Image.open(image_path)
    if max(img.size) > max_dim:
        ratio = max_dim / max(img.size)
        new_size = (int(img.width * ratio), int(img.height * ratio))
        img = img.resize(new_size, Image.LANCZOS)

    buffer = io.BytesIO()
    img.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode("utf-8")

@retry(stop=stop_after_attempt(3), wait=wait_exponential(min=1, max=10))
async def extract_document(image_path: str) -> DocumentExtraction:
    """从图像中提取结构化文档信息"""
    image_b64 = resize_for_token_efficiency(image_path)

    response = await client.beta.chat.completions.parse(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": (
                    "你是一个专业的文档解析引擎。从图像中提取所有结构化信息，"
                    "包括标题、关键字段和表格数据。保持原始数据精度。"
                ),
            },
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "请解析这份文档并提取所有结构化信息。"},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/png;base64,{image_b64}",
                            "detail": "high",
                        },
                    },
                ],
            },
        ],
        response_format=DocumentExtraction,
        temperature=0,
    )

    return response.choices[0].message.parsed

async def batch_process(image_dir: str, concurrency: int = 5) -> list[DocumentExtraction]:
    """批量并发处理，控制并发数避免限流"""
    semaphore = asyncio.Semaphore(concurrency)
    images = list(Path(image_dir).glob("*.png")) + list(Path(image_dir).glob("*.jpg"))

    async def process_one(path: str):
        async with semaphore:
            return await extract_document(str(path))

    results = await asyncio.gather(*[process_one(img) for img in images])
    return results

import torch
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
from PIL import Image
from fastapi import FastAPI, UploadFile
from pydantic import BaseModel
import uvicorn

app = FastAPI(title="VLM Pipeline Service")

# 模型加载（启动时执行一次）
model = Qwen2VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2-VL-7B-Instruct",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="flash_attention_2",
)
processor = AutoProcessor.from_pretrained(
    "Qwen/Qwen2-VL-7B-Instruct",
    min_pixels=256 * 28 * 28,
    max_pixels=1280 * 28 * 28,
)

class VLMRequest(BaseModel):
    prompt: str = "描述这张图片的内容并提取所有文字信息。"
    max_tokens: int = 2048
    temperature: float = 0.1

@app.post("/v1/understand")
async def understand_image(file: UploadFile, request: VLMRequest = VLMRequest()):
    """图像理解接口"""
    image = Image.open(file.file).convert("RGB")

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": request.prompt},
            ],
        }
    ]

    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=request.max_tokens,
            temperature=request.temperature,
            do_sample=request.temperature > 0,
        )

    generated_ids = output_ids[:, inputs.input_ids.shape[1]:]
    response_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

    return {
        "content": response_text,
        "usage": {
            "input_tokens": inputs.input_ids.shape[1],
            "output_tokens": generated_ids.shape[1],
        },
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

import Anthropic from "@anthropic-ai/sdk";
import sharp from "sharp";
import crypto from "crypto";

interface DocumentPage {
  pageNumber: number;
  imageBuffer: Buffer;
  width: number;
  height: number;
}

interface ExtractedDocument {
  pages: PageResult[];
  metadata: DocumentMetadata;
  processingTime: number;
}

interface PageResult {
  pageNumber: number;
  text: string;
  tables: TableData[];
  figures: FigureDescription[];
  layout: LayoutBlock[];
}

interface TableData {
  headers: string[];
  rows: string[][];
  caption?: string;
}

interface FigureDescription {
  description: string;
  boundingBox: { x: number; y: number; w: number; h: number };
}

interface LayoutBlock {
  type: "title" | "paragraph" | "list" | "table" | "figure" | "header" | "footer";
  content: string;
  order: number;
}

interface DocumentMetadata {
  totalPages: number;
  documentType: string;
  language: string;
  cacheKey: string;
}

class DocumentPipeline {
  private client: Anthropic;
  private cache: Map<string, ExtractedDocument> = new Map();

  constructor() {
    this.client = new Anthropic();
  }

  async processDocument(pdfBuffer: Buffer): Promise<ExtractedDocument> {
    const startTime = Date.now();

    // 1. 生成缓存键
    const cacheKey = crypto.createHash("sha256").update(pdfBuffer).digest("hex");
    if (this.cache.has(cacheKey)) {
      return this.cache.get(cacheKey)!;
    }

    // 2. PDF 转图像（每页独立处理）
    const pages = await this.pdfToImages(pdfBuffer);

    // 3. 并行处理所有页面
    const pageResults = await Promise.all(
      pages.map((page) => this.processPage(page))
    );

    // 4. 跨页关联与合并
    const mergedResults = this.mergeResults(pageResults);

    const result: ExtractedDocument = {
      pages: mergedResults,
      metadata: {
        totalPages: pages.length,
        documentType: await this.classifyDocument(pages[0]),
        language: "auto-detected",
        cacheKey,
      },
      processingTime: Date.now() - startTime,
    };

    this.cache.set(cacheKey, result);
    return result;
  }

  private async processPage(page: DocumentPage): Promise<PageResult> {
    // 自适应分辨率：短边不低于 768px，长边不超过 2048px
    const optimized = await this.optimizeResolution(page.imageBuffer);

    const response = await this.client.messages.create({
      model: "claude-sonnet-4-20250514",
      max_tokens: 4096,
      messages: [
        {
          role: "user",
          content: [
            {
              type: "image",
              source: {
                type: "base64",
                media_type: "image/png",
                data: optimized.toString("base64"),
              },
            },
            {
              type: "text",
              text: `分析这个文档页面，返回JSON格式结果：
{
  "text": "完整OCR文本",
  "tables": [{"headers": [...], "rows": [[...], ...], "caption": "..."}],
  "figures": [{"description": "...", "boundingBox": {"x":0,"y":0,"w":0,"h":0}}],
  "layout": [{"type": "title|paragraph|list|table|figure", "content": "...", "order": 1}]
}
只返回JSON，不要其他内容。`,
            },
          ],
        },
      ],
    });

    const content = response.content[0];
    if (content.type === "text") {
      const parsed = JSON.parse(content.text);
      return { pageNumber: page.pageNumber, ...parsed };
    }
    throw new Error("Unexpected response format");
  }

  private async optimizeResolution(buffer: Buffer): Promise<Buffer> {
    const metadata = await sharp(buffer).metadata();
    const { width = 0, height = 0 } = metadata;
    const maxDim = 2048;
    const minDim = 768;

    let targetWidth = width;
    let targetHeight = height;

    if (Math.max(width, height) > maxDim) {
      const ratio = maxDim / Math.max(width, height);
      targetWidth = Math.round(width * ratio);
      targetHeight = Math.round(height * ratio);
    } else if (Math.min(width, height) < minDim) {
      const ratio = minDim / Math.min(width, height);
      targetWidth = Math.round(width * ratio);
      targetHeight = Math.round(height * ratio);
    }

    return sharp(buffer).resize(targetWidth, targetHeight).png().toBuffer();
  }

  private async classifyDocument(firstPage: DocumentPage): Promise<string> {
    // 使用首页快速分类文档类型
    return "report"; // 简化示例
  }

  private async pdfToImages(buffer: Buffer): Promise<DocumentPage[]> {
    // PDF 渲染为图像（实际使用 pdf-to-img 或 poppler）
    return [];
  }

  private mergeResults(pages: PageResult[]): PageResult[] {
    // 跨页表格合并、段落延续处理
    return pages;
  }
}

export { DocumentPipeline, ExtractedDocument, PageResult };