icon_generation/batch_icon_generator.py

import os
import json
from google import genai
from google.genai import types
from pathlib import Path
import time
from multiprocessing import Pool
from functools import partial

# API配置
API_KEY = os.getenv("GEMINI_API_KEY") or os.getenv("OPENAI_API_KEY", "")
client = genai.Client(
    api_key=API_KEY,
    http_options={"base_url": os.getenv("GEMINI_BASE_URL", "https://aihubmix.com/gemini")},
)

# 配置参数
ASPECT_RATIO = "3:2"
BATCH_SIZE = 24
MIN_BATCH_SIZE = 6
OUTPUT_DIR = "generated_icons"
NUM_PROCESSES = 3

def load_domain_attributes(file_path):
    """读取domain_attributes.txt文件，返回扁平化的(domain, attribute)对列表"""
    domain_attribute_pairs = []
    current_domain = None
    
    with open(file_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    
    for line in lines:
        line = line.strip()
        if not line:
            continue
        
        # 检查是否是domain行（没有逗号的行）
        if ',' not in line:
            current_domain = line
        else:
            # 这是attributes行
            if current_domain:
                attributes = [attr.strip() for attr in line.split(',')]
                for attr in attributes:
                    domain_attribute_pairs.append((current_domain, attr))
    
    return domain_attribute_pairs

def load_templates(file_path):
    """读取template_batch.json文件"""
    with open(file_path, 'r', encoding='utf-8') as f:
        templates = json.load(f)
    return templates

def batch_domain_attribute_pairs(pairs, batch_size=BATCH_SIZE):
    """将domain-attribute pairs分批，每批batch_size个"""
    batches = []
    for i in range(0, len(pairs), batch_size):
        batch = pairs[i:i + batch_size]
        batches.append(batch)
    return batches

def generate_prompt(template, domain_attribute_pairs):
    """生成prompt，将模板中的占位符替换为实际内容"""
    # 格式化为: "domain1: attribute1, domain2: attribute2, ..."
    pairs_text = ", ".join([f"{domain}: {attr}" for domain, attr in domain_attribute_pairs])
    prompt = template.replace("{DOMAIN_ATTRIBUTE_PAIRS}", pairs_text)
    
    # 添加固定的布局要求
    actual_count = len(domain_attribute_pairs)
    layout_requirement = (
        f" The output must be a single image with an exact 6:4 aspect ratio (landscape orientation, width greater than height). "
        f"The image must contain exactly {actual_count} icons arranged in a strict grid of 6 columns (horizontal, left to right) and 4 rows (vertical, top to bottom). "
        f"Do not rotate, transpose, or alter the grid orientation. "
        f"No text, letters, numbers, labels, captions, or icon titles. Each icon must not include any titles or written elements. "
        f"Use a pure white background only."
    )
    
    prompt = prompt + layout_requirement
    return prompt

def generate_icons_for_batch(batch_idx, domain_attribute_pairs, templates, output_dir):
    """为单个batch生成所有风格的icons"""
    batch_dir = os.path.join(output_dir, f"batch_{batch_idx:04d}")
    os.makedirs(batch_dir, exist_ok=True)
    
    print(f"\n{'='*80}")
    print(f"批次 {batch_idx} (含 {len(domain_attribute_pairs)} 个 domain-attribute pairs)")
    
    success_count = 0
    failed_count = 0
    skipped_count = 0
    
    # 为每个style生成icons
    for style_name, template in templates.items():
        print(f"\n  风格: {style_name}")
        
        # 生成文件名
        image_filename = f"batch_{batch_idx:04d}_{style_name}.png"
        annotation_filename = f"batch_{batch_idx:04d}_{style_name}.txt"
        
        image_path = os.path.join(batch_dir, image_filename)
        annotation_path = os.path.join(batch_dir, annotation_filename)
        
        # 断点续传：检查文件是否已存在
        if os.path.exists(image_path) and os.path.exists(annotation_path):
            print(f"    ⏭️  跳过（文件已存在）: {image_filename}")
            skipped_count += 1
            success_count += 1  # 已存在的文件计入成功数
            continue
        
        try:
            # 生成prompt
            prompt = generate_prompt(template, domain_attribute_pairs)
            
            # 每个进程需要创建自己的API客户端
            client = genai.Client(
                api_key=API_KEY,
                http_options={"base_url": "https://aihubmix.com/gemini"},
            )
            
            # 调用API生成图像
            response = client.models.generate_content(
                model="gemini-3-pro-image-preview",
                contents=prompt,
                config=types.GenerateContentConfig(
                    response_modalities=['TEXT', 'IMAGE'],
                    image_config=types.ImageConfig(
                        aspect_ratio=ASPECT_RATIO
                    ),
                ),
            )
            
            # 保存图像和文本
            for part in response.parts:
                if part.text:
                    print(f"    生成说明: {part.text[:100]}...")
                elif image := part.as_image():
                    image.save(image_path)
                    print(f"    ✅ 图像已保存: {image_filename}")
                    
                    # 保存标注
                    save_annotation(annotation_path, domain_attribute_pairs, style_name)
                    print(f"    ✅ 标注已保存: {annotation_filename}")
            
            success_count += 1
            
        except Exception as e:
            print(f"    ❌ 生成失败: {str(e)}")
            failed_count += 1
            continue
    
    if skipped_count > 0:
        print(f"\n  批次 {batch_idx} 完成: ✅ {success_count} 成功 (含 {skipped_count} 个跳过), ❌ {failed_count} 失败")
    else:
        print(f"\n  批次 {batch_idx} 完成: ✅ {success_count} 成功, ❌ {failed_count} 失败")
    return success_count, failed_count

def save_annotation(file_path, domain_attribute_pairs, style):
    """保存txt标注文件，格式：第一行style，后续每行domain, attribute"""
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(f"Style: {style}\n")
        for domain, attr in domain_attribute_pairs:
            f.write(f"{domain}, {attr}\n")

def process_single_batch(args):
    """处理单个batch的所有风格（用于并发处理）"""
    batch_idx, batch, templates, output_dir = args
    
    # 跳过少于MIN_BATCH_SIZE的最后一批（如果不是第一批）
    if len(batch) < MIN_BATCH_SIZE and batch_idx > 1:
        print(f"\n⚠️ 跳过批次 {batch_idx}: 只有 {len(batch)} 个pairs (少于最小值 {MIN_BATCH_SIZE})")
        return 0, 0, True  # success_count, failed_count, skipped
    
    success, failed = generate_icons_for_batch(batch_idx, batch, templates, output_dir)
    return success, failed, False

def main():
    """主函数"""
    print("="*80)
    print("批量图标生成器 (扁平化模式)")
    print("="*80)
    
    # 检查文件是否存在
    domain_file = "domain_attributes.txt"
    template_file = "template_batch.json"
    
    if not os.path.exists(domain_file):
        print(f"❌ 错误: 找不到文件 {domain_file}")
        return
    
    if not os.path.exists(template_file):
        print(f"❌ 错误: 找不到文件 {template_file}")
        return
    
    # 创建输出目录
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    
    # 加载数据
    print(f"\n📖 加载domain-attribute pairs...")
    domain_attribute_pairs = load_domain_attributes(domain_file)
    print(f"   共加载 {len(domain_attribute_pairs)} 个 domain-attribute pairs")
    
    print(f"\n📖 加载模板...")
    templates = load_templates(template_file)
    print(f"   共加载 {len(templates)} 个风格模板: {', '.join(templates.keys())}")
    
    # 配置信息
    print(f"\n⚙️  配置:")
    print(f"   长宽比: {ASPECT_RATIO}")
    print(f"   单批最大数量: {BATCH_SIZE}")
    print(f"   单批最小数量: {MIN_BATCH_SIZE}")
    print(f"   并发进程数: {NUM_PROCESSES}")
    print(f"   输出目录: {OUTPUT_DIR}")
    
    # 分批
    batches = batch_domain_attribute_pairs(domain_attribute_pairs, BATCH_SIZE)
    print(f"\n📦 分成 {len(batches)} 个批次")
    
    # 扫描已存在的文件（断点续传预检）
    print(f"\n🔍 扫描已存在的文件...")
    existing_files = 0
    total_expected_files = 0
    for batch_idx, batch in enumerate(batches, 1):
        if len(batch) < MIN_BATCH_SIZE and batch_idx > 1:
            continue
        batch_dir = os.path.join(OUTPUT_DIR, f"batch_{batch_idx:04d}")
        for style_name in templates.keys():
            total_expected_files += 1
            image_filename = f"batch_{batch_idx:04d}_{style_name}.png"
            annotation_filename = f"batch_{batch_idx:04d}_{style_name}.txt"
            image_path = os.path.join(batch_dir, image_filename)
            annotation_path = os.path.join(batch_dir, annotation_filename)
            if os.path.exists(image_path) and os.path.exists(annotation_path):
                existing_files += 1
    
    print(f"   已存在: {existing_files}/{total_expected_files} 个文件")
    print(f"   需要生成: {total_expected_files - existing_files} 个文件")
    
    # 准备并发任务
    tasks = []
    for batch_idx, batch in enumerate(batches, 1):
        tasks.append((batch_idx, batch, templates, OUTPUT_DIR))
    
    # 使用进程池并发处理多个batch
    print(f"\n🚀 使用 {NUM_PROCESSES} 个进程并发处理批次...")
    print(f"💡 提示: 已存在的文件将自动跳过（断点续传）")
    total_success = 0
    total_failed = 0
    
    with Pool(processes=NUM_PROCESSES) as pool:
        results = pool.map(process_single_batch, tasks)
    
    # 统计结果
    for success, failed, skipped in results:
        if not skipped:
            total_success += success
            total_failed += failed
    
    print(f"\n{'='*80}")
    print("✅ 全部生成完成！")
    print(f"总计: ✅ {total_success} 成功, ❌ {total_failed} 失败")
    print(f"输出目录: {OUTPUT_DIR}")
    print("="*80)

if __name__ == "__main__":
    main()