base_rag/test_image_ocr.py

107 lines
3.2 KiB
Python

#!/usr/bin/env python3
"""
测试优化后的图片OCR功能
"""
import sys
import os
import asyncio
from pathlib import Path
# 添加源码路径
sys.path.append(os.path.join(os.path.dirname(__file__), "src"))
from base_rag.image_processor import ImageProcessor
async def test_image_ocr():
"""测试不同模式下的图片OCR功能"""
# 测试配置
configs = [
{"type": "local", "engine": "easyocr"},
{"type": "basic"},
# {"type": "api", "api_url": "http://localhost:8000/image2text"} # 需要实际API
]
print("🧪 开始测试图片OCR功能")
print("=" * 50)
# 寻找测试图片
test_images = []
# 检查常见的图片位置
possible_paths = [
"./test_files/",
"./demo_documents/",
"./examples/",
"./"
]
image_extensions = ['.png', '.jpg', '.jpeg', '.gif', '.bmp']
for path in possible_paths:
if os.path.exists(path):
for file in os.listdir(path):
if any(file.lower().endswith(ext) for ext in image_extensions):
test_images.append(os.path.join(path, file))
if not test_images:
print("⚠️ 未找到测试图片,创建示例图片...")
# 创建一个简单的测试图片
try:
from PIL import Image, ImageDraw, ImageFont
# 创建包含文字的测试图片
img = Image.new('RGB', (400, 200), color='white')
draw = ImageDraw.Draw(img)
# 尝试使用默认字体
try:
font = ImageFont.truetype("/System/Library/Fonts/Arial.ttf", 24)
except:
font = ImageFont.load_default()
# 添加测试文字
test_text = "Hello World!\nPython OCR Test\n测试中文识别"
draw.text((50, 50), test_text, fill='black', font=font)
test_image_path = "./test_ocr_image.png"
img.save(test_image_path)
test_images = [test_image_path]
print(f"✅ 创建测试图片: {test_image_path}")
except Exception as e:
print(f"❌ 创建测试图片失败: {e}")
return
print(f"📸 找到 {len(test_images)} 个测试图片")
# 测试每种配置
for i, config in enumerate(configs, 1):
print(f"\n🔧 测试配置 {i}: {config}")
print("-" * 30)
try:
processor = ImageProcessor(config)
# 处理每个测试图片
for img_path in test_images[:2]: # 限制测试图片数量
print(f"\n📷 处理图片: {os.path.basename(img_path)}")
if os.path.exists(img_path):
result = processor.extract_image_description(img_path)
print(f"结果:\n{result}")
else:
print(f"❌ 图片不存在: {img_path}")
except Exception as e:
print(f"❌ 配置 {config} 测试失败: {e}")
print("\n" + "=" * 50)
print("🏁 测试完成")
if __name__ == "__main__":
asyncio.run(test_image_ocr())