cutThenThink/examples/processor_example.py

#!/usr/bin/env python3
"""
处理流程整合示例

演示如何使用 ImageProcessor 进行完整的图片处理流程
"""

import sys
import logging
from pathlib import Path

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.core.processor import ImageProcessor, ProcessCallback, ProcessResult, create_markdown_result
from src.config.settings import get_settings
from src.utils.logger import init_logger, get_logger


# 初始化日志
log_dir = project_root / "logs"
init_logger(log_dir=log_dir, level="INFO", colored_console=True)
logger = get_logger(__name__)


class DemoProcessCallback(ProcessCallback):
    """
    演示用回调类
    """

    def on_start(self, message: str = "开始处理"):
        print(f"\n{'=' * 60}")
        print(f"🚀 {message}")
        print(f"{'=' * 60}\n")

    def on_ocr_start(self, message: str = "开始 OCR 识别"):
        print(f"📸 {message}...")

    def on_ocr_complete(self, result):
        print(f"   ✅ 识别完成: {len(result.results)} 行文本")
        print(f"   📊 置信度: {result.total_confidence:.2%}")
        print(f"   📝 预览: {result.full_text[:50]}...")

    def on_ai_start(self, message: str = "开始 AI 分类"):
        print(f"\n🤖 {message}...")

    def on_ai_complete(self, result):
        emoji_map = {
            "TODO": "✅",
            "NOTE": "📝",
            "IDEA": "💡",
            "REF": "📚",
            "FUNNY": "😄",
            "TEXT": "📄"
        }
        emoji = emoji_map.get(result.category.value, "📄")
        print(f"   {emoji} 分类: {result.category.value}")
        print(f"   📊 置信度: {result.confidence:.2%}")
        print(f"   📌 标题: {result.title}")
        print(f"   🏷️  标签: {', '.join(result.tags)}")

    def on_save_start(self, message: str = "开始保存到数据库"):
        print(f"\n💾 {message}...")

    def on_save_complete(self, record_id: int):
        print(f"   ✅ 保存成功: 记录 ID = {record_id}")

    def on_error(self, message: str, exception=None):
        print(f"\n❌ 错误: {message}")
        if exception:
            print(f"   异常类型: {type(exception).__name__}")

    def on_complete(self, result: ProcessResult):
        print(f"\n{'=' * 60}")
        if result.success:
            print(f"✨ 处理成功!")
            print(f"   耗时: {result.process_time:.2f} 秒")
            print(f"   步骤: {' -> '.join(result.steps_completed)}")
        else:
            print(f"❌ 处理失败")
            if result.error_message:
                print(f"   错误: {result.error_message}")
        if result.warnings:
            print(f"\n⚠️  警告:")
            for warning in result.warnings:
                print(f"   - {warning}")
        print(f"{'=' * 60}\n")


def process_single_image_demo(image_path: str):
    """
    处理单张图片的演示

    Args:
        image_path: 图片路径
    """
    print(f"\n处理图片: {image_path}")

    # 加载配置
    settings = get_settings()

    # 创建处理器
    callback = DemoProcessCallback()
    processor = ImageProcessor(
        ocr_config={
            'mode': settings.ocr.mode.value,
            'lang': settings.ocr.lang,
            'use_gpu': settings.ocr.use_gpu
        },
        ai_config=settings.ai,
        db_path=str(project_root / "data" / "cutnthink.db"),
        callback=callback
    )

    # 处理图片
    result = processor.process_image(image_path)

    # 显示 Markdown 结果
    if result.ai_result:
        markdown = create_markdown_result(result.ai_result, result.ocr_result.full_text if result.ocr_result else "")
        print("\n" + "=" * 60)
        print("Markdown 格式结果:")
        print("=" * 60)
        print(markdown)
        print("=" * 60 + "\n")

    return result


def main():
    """主函数"""
    print("""
╔══════════════════════════════════════════════════════════╗
║         CutThenThink - 处理流程整合示例                  ║
║         OCR -> AI -> 存储 完整流程演示                   ║
╚══════════════════════════════════════════════════════════╝
    """)

    # 检查命令行参数
    if len(sys.argv) < 2:
        print("用法: python processor_example.py <图片路径>")
        print("\n示例:")
        print("  python processor_example.py /path/to/image.png")
        print("  python processor_example.py /path/to/image.jpg")
        return

    image_path = sys.argv[1]

    # 检查文件是否存在
    if not Path(image_path).exists():
        print(f"❌ 错误: 文件不存在: {image_path}")
        return

    # 处理图片
    try:
        result = process_single_image_demo(image_path)

        if result.success:
            print("\n🎉 处理完成!")
        else:
            print("\n⚠️  处理未完全成功，请检查日志")

    except Exception as e:
        logger.error(f"处理失败: {e}", exc_info=True)
        print(f"\n❌ 处理失败: {e}")


if __name__ == "__main__":
    main()