icon_generation/test_split.py

#!/usr/bin/env python3
"""
测试split_icon.py的网格检测功能
"""
import cv2
from split_icon import detect_grid_cells, uniform_grid_split
import os

def test_image(image_path):
    """测试单个图像的网格检测"""
    if not os.path.exists(image_path):
        print(f"错误：文件不存在 {image_path}")
        return
    
    print(f"\n{'='*60}")
    print(f"测试图像: {os.path.basename(image_path)}")
    print(f"{'='*60}")
    
    # 读取图像
    img = cv2.imread(image_path)
    if img is None:
        print("错误：无法读取图像")
        return
    
    print(f"图像尺寸: {img.shape[1]}x{img.shape[0]}")
    
    # 测试网格检测
    boxes = detect_grid_cells(img, expected_cols=6, expected_rows=4)
    print(f"最终检测到的单元格数量: {len(boxes)}")
    
    # 可视化结果
    result_img = img.copy()
    for i, box in enumerate(boxes):
        x, y, w, h = box
        cv2.rectangle(result_img, (x, y), (x+w, y+h), (0, 255, 0), 2)
        # 在左上角添加序号
        cv2.putText(result_img, str(i+1), (x+5, y+20), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
    
    # 保存结果
    output_path = image_path.rsplit('.', 1)[0] + '_detected.png'
    cv2.imwrite(output_path, result_img)
    print(f"检测结果已保存到: {output_path}")

if __name__ == "__main__":
    # 测试batch_0001中的一张图像
    test_images = [
        "generated_icons/batch_0001/batch_0001_doodle.png",
        "generated_icons/batch_0001/batch_0001_hand_drawn_sketch.png",
        "generated_icons/batch_0001/batch_0001_isometric.png",
    ]
    
    for img_path in test_images:
        if os.path.exists(img_path):
            test_image(img_path)
            break