RagFlow_Plugs/test/test_image_compression_real.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试图片压缩修复 - 真实场景
使用更真实的大图片验证压缩方法
"""
import sys
import os
from io import BytesIO
from PIL import Image, ImageDraw, ImageFont
import random

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.minio.image_util import ImageUtil

def create_complex_test_image(width=3000, height=3000) -> BytesIO:
    """
    创建一个复杂的大尺寸测试图片，包含多种元素以增加文件大小
    
    Args:
        width: 图片宽度
        height: 图片高度
        
    Returns:
        BytesIO: 复杂大尺寸图片流
    """
    print(f"创建 {width}x{height} 的复杂测试图片...")
    
    # 创建一个白色背景图片
    img = Image.new('RGB', (width, height), color=(255, 255, 255))
    draw = ImageDraw.Draw(img)
    
    # 添加大量随机形状和颜色，增加图片复杂度
    for _ in range(10000):
        # 随机位置
        x1 = random.randint(0, width)
        y1 = random.randint(0, height)
        x2 = x1 + random.randint(10, 100)
        y2 = y1 + random.randint(10, 100)
        
        # 随机颜色
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
        
        # 随机形状
        shape_type = random.choice(['rectangle', 'ellipse', 'line'])
        if shape_type == 'rectangle':
            draw.rectangle([x1, y1, x2, y2], fill=color)
        elif shape_type == 'ellipse':
            draw.ellipse([x1, y1, x2, y2], fill=color)
        else:
            draw.line([x1, y1, x2, y2], fill=color, width=random.randint(1, 5))
    
    # 添加一些随机文本
    try:
        # 尝试使用默认字体
        font = ImageFont.load_default()
        for _ in range(1000):
            text = f"Test {random.randint(1, 1000)}"
            x = random.randint(0, width - 100)
            y = random.randint(0, height - 20)
            draw.text((x, y), text, fill=(0, 0, 0), font=font)
    except Exception as e:
        print(f"添加文本失败: {e}")
    
    # 将图片保存到BytesIO流，使用JPEG格式以获得更大的文件大小
    img_stream = BytesIO()
    img.save(img_stream, format='JPEG', quality=100)  # 使用最高质量生成大文件
    img_stream.seek(0)
    
    # 检查图片大小
    img_stream.seek(0, 2)
    size_kb = img_stream.tell() / 1024
    img_stream.seek(0)
    
    print(f"测试图片创建完成，大小为 {size_kb:.2f}KB")
    return img_stream

def test_image_compression():
    """
    测试图片压缩方法
    """
    print("开始测试图片压缩方法...")
    
    # 创建ImageUtil实例
    image_util = ImageUtil()
    
    # 测试不同尺寸的复杂图片压缩
    test_sizes = [
        (3000, 3000),   # 约 25MB
        (4000, 4000),   # 约 45MB
    ]
    
    for width, height in test_sizes:
        print(f"\n=== 测试 {width}x{height} 复杂图片压缩 ===")
        
        # 创建大尺寸测试图片
        img_stream = create_complex_test_image(width, height)
        
        # 调用压缩方法
        compressed_stream = image_util._compress_image(img_stream, "test_complex_image.jpg")
        
        # 检查压缩后的大小
        compressed_stream.seek(0, 2)
        compressed_size_kb = compressed_stream.tell() / 1024
        compressed_stream.seek(0)
        
        print(f"压缩后大小: {compressed_size_kb:.2f}KB")
        
        # 验证压缩结果
        if compressed_size_kb <= 5000:
            print("✅ 压缩成功！压缩后大小小于等于5000KB")
        else:
            print("❌ 压缩失败！压缩后大小仍大于5000KB")
    
    # 测试_compress_image_to_bytes方法
    print(f"\n=== 测试 _compress_image_to_bytes 方法 ===")
    img_stream = create_complex_test_image(4000, 4000)
    compressed_bytes = image_util._compress_image_to_bytes(img_stream)
    compressed_size_kb = len(compressed_bytes) / 1024
    print(f"压缩后字节大小: {compressed_size_kb:.2f}KB")
    
    if compressed_size_kb <= 5000:
        print("✅ _compress_image_to_bytes 压缩成功！")
    else:
        print("❌ _compress_image_to_bytes 压缩失败！")
    
    print("\n=== 所有测试完成 ===")

if __name__ == "__main__":
    test_image_compression()