icon_generation/backup/image_batch_generator.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图像批量生成Pipeline
根据topic_style.json配置，批量生成适合infographic装饰的图像
"""

import os
import json
import random
import sys
import time
from typing import Dict, List, Tuple
from openai import OpenAI
from concurrent.futures import ThreadPoolExecutor
from google import genai
from google.genai import types
from PIL import Image, ImageDraw
from io import BytesIO
import numpy as np
from collections import Counter

# 添加项目根目录到路径
# sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# from config import api_key, base_url

api_key = 'xxx'
base_url = "https://aihubmix.com/v1"

class ImageBatchGenerator:
    def __init__(self):
        """初始化生成器"""
        # OpenAI client for text generation
        self.openai_client = OpenAI(
            api_key=api_key,
            base_url=base_url,
        )
        
        # Gemini client for image generation
        self.genai_client = genai.Client(
            api_key=api_key,
            http_options={"base_url": "https://aihubmix.com/gemini"},
        )
        
        # 加载topic_style配置
        self.config_path = os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            'generator', 'topic_style.json'
        )
        self.load_config()
        
        # 输出目录
        self.output_dir = os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            'gen_output'
        )
        os.makedirs(self.output_dir, exist_ok=True)
        
        # 设计prompt模板
        self.design_prompt_template = """
[TASK START]
OBJECTIVE: Generate a text-to-image prompt for a single, isolated clipart icon based on the provided inputs.

INPUTS:
Topic: {topic}
Style Keyword: {style_keyword}
Concept: {concept}

PROCESS:
Write a text-to-image prompt describing this concept, rendered using the specified Style Keyword.

CONSTRAINTS:
- The output must be a single icon or a small, unified group of objects
- The icon MUST be isolated on a pure white background (#FFFFFF)
- No shadows, textures or patterns in the background
- The background must be completely clean and empty
- The final prompt must be concise and descriptive

REQUIRED OUTPUT:
[The final text-to-image prompt, make sure to specify "on pure white background" in the prompt]

[TASK END]
"""
        
        # 概念生成prompt
        self.concept_generation_prompt = """
Generate 10 different concrete concepts for the topic "{topic}".

Requirements:
1. Each concept must be a specific, tangible object or clear visual scene
2. Use detailed descriptions (e.g. "stethoscope on medical chart" vs "medical")
3. Focus on real-world items, tools, places or situations
4. Each concept should be immediately recognizable and relatable
5. Concepts should work well as simple icons or decorative elements
6. Keep descriptions concise but specific

Return in this format:
1. [concept1]
2. [concept2]
3. [concept3]
...
10. [concept10]
"""
        
        # 设计评判prompt
        self.design_evaluation_prompt = """
Evaluate the following design concepts and select the 5 best ones for infographic decoration.

Evaluation criteria:
1. Visual clarity: Easy to recognize and understand
2. Decorative value: Suitable as decorative elements without interfering with main information
3. Universality: Broad applicability

Design concept list:
{concepts}

Select the 5 best concepts and return in this format:
Selected concepts:
1. [concept name]
2. [concept name]
3. [concept name]
4. [concept name]
5. [concept name]
"""

    def load_config(self):
        """加载topic_style配置文件"""
        with open(self.config_path, 'r', encoding='utf-8') as f:
            self.config = json.load(f)
        print(f"✅ 加载配置: {len(self.config)} 个风格类别")

    def select_random_category_and_elements(self) -> Tuple[str, str, str]:
        """随机选择category、keyword和topic"""
        category = random.choice(list(self.config.keys()))
        category_data = self.config[category]
        keyword = random.choice(category_data['keywords'])
        topic = random.choice(category_data['topics'])
        
        print(f"🎯 选中: {category} | {keyword} | {topic}")
        return category, keyword, topic

    def generate_concepts(self, topic: str) -> List[str]:
        """使用ChatGPT生成10个概念"""
        print(f"🧠 生成概念...")
        
        response = self.openai_client.chat.completions.create(
            model="gpt-5-mini",
            messages=[
                {"role": "user", "content": self.concept_generation_prompt.format(topic=topic)}
            ],
            temperature=0.8
        )
        
        content = response.choices[0].message.content
        
        # 解析概念列表 - 修复方括号解析问题
        concepts = []
        lines = content.strip().split('\n')
        for line in lines:
            line = line.strip()
            if line and (line[0].isdigit() or line.startswith('-')):
                # 提取方括号内的内容
                if '[' in line and ']' in line:
                    start = line.find('[')
                    end = line.find(']')
                    if start != -1 and end != -1 and end > start:
                        concept = line[start+1:end].strip()
                        if concept:
                            concepts.append(concept)
                else:
                    # 如果没有方括号，提取序号后的内容
                    concept = line.split('.', 1)[-1].strip()
                    if concept:
                        concepts.append(concept)
        
        print(f"✅ 生成 {len(concepts)} 个概念")
        return concepts[:10]

    def evaluate_and_select_concepts(self, concepts: List[str]) -> List[str]:
        """评判并选择5个最佳概念"""
        print(f"🔍 评判概念...")
        
        concepts_text = ""
        for i, concept in enumerate(concepts, 1):
            concepts_text += f"{i}. {concept}\n"
        
        response = self.openai_client.chat.completions.create(
            model="gpt-5-mini",
            messages=[
                {"role": "user", "content": self.design_evaluation_prompt.format(concepts=concepts_text)}
            ],
            temperature=0.3
        )
        
        content = response.choices[0].message.content
        
        # 解析选中的概念
        selected_concepts = []
        lines = content.strip().split('\n')
        
        for line in lines:
            line = line.strip()
            if line and line[0].isdigit() and '.' in line:
                concept_name = line.split('.', 1)[1].strip()
                # 在原始概念中查找匹配
                for concept in concepts:
                    if concept_name.lower() in concept.lower() or concept.lower() in concept_name.lower():
                        if concept not in selected_concepts:
                            selected_concepts.append(concept)
                            break
        
        # 如果解析不足5个，随机补充
        if len(selected_concepts) < 5:
            remaining = [c for c in concepts if c not in selected_concepts]
            selected_concepts.extend(random.sample(remaining, min(5 - len(selected_concepts), len(remaining))))
        
        print(f"✅ 选中 {len(selected_concepts[:5])} 个概念")
        return selected_concepts[:5]

    def detect_background_color(self, image: Image.Image) -> tuple:
        """检测图像的背景颜色，返回(背景色, 是否为杂乱背景)"""
        # 获取图像尺寸
        width, height = image.size
        
        # 采样边界点
        sample_points = []
        
        # 四个角
        sample_points.extend([
            (0, 0), (width-1, 0), (0, height-1), (width-1, height-1)
        ])
        
        # 边界中点
        sample_points.extend([
            (width//2, 0), (width//2, height-1),  # 上下边中点
            (0, height//2), (width-1, height//2)  # 左右边中点
        ])
        
        # 边界线采样（每边采样10个点）
        for i in range(1, 10):
            ratio = i / 10.0
            # 上边
            sample_points.append((int(width * ratio), 0))
            # 下边
            sample_points.append((int(width * ratio), height-1))
            # 左边
            sample_points.append((0, int(height * ratio)))
            # 右边
            sample_points.append((width-1, int(height * ratio)))
        
        # 获取所有采样点的颜色
        colors = []
        for x, y in sample_points:
            if 0 <= x < width and 0 <= y < height:
                pixel = image.getpixel((x, y))
                if isinstance(pixel, int):  # 灰度图
                    colors.append((pixel, pixel, pixel))
                elif len(pixel) >= 3:  # RGB或RGBA
                    colors.append(pixel[:3])
        
        # 统计颜色众数
        color_counts = Counter(colors)
        if color_counts:
            most_common_color, most_common_count = color_counts.most_common(1)[0]
            total_samples = len(colors)
            
            # 计算众数颜色占比
            ratio = most_common_count / total_samples
            
            # 如果众数颜色占比小于50%，认为背景杂乱
            is_messy = ratio < 0.5
            
            return most_common_color, is_messy
        
        # 默认返回白色，非杂乱
        return (255, 255, 255), False

    def optimized_flood_fill_remove_background(self, image: Image.Image, bg_color: tuple, tolerance: int = 30) -> Image.Image:
        """使用优化的flood fill算法从边界去除背景色"""
        # 转换为RGBA模式
        if image.mode != 'RGBA':
            image = image.convert('RGBA')
        
        # 转换为numpy数组
        data = np.array(image, dtype=np.uint8)
        height, width = data.shape[:2]
        
        # 创建访问标记数组
        visited = np.zeros((height, width), dtype=bool)
        
        # 预计算颜色距离的平方（避免开方运算）
        def color_distance_squared(c1, c2):
            """计算颜色距离的平方，避免开方运算提高性能"""
            return sum((int(a) - int(b)) ** 2 for a, b in zip(c1[:3], c2[:3]))
        
        tolerance_squared = tolerance * tolerance
        
        def is_background_color(pixel_color):
            """判断是否为背景色，使用平方距离比较"""
            return color_distance_squared(pixel_color[:3], bg_color) <= tolerance_squared
        
        def optimized_flood_fill(start_x, start_y):
            """优化的flood fill算法，使用栈而非递归，批量处理"""
            if (start_y >= height or start_x >= width or 
                start_y < 0 or start_x < 0 or 
                visited[start_y, start_x]):
                return
            
            # 使用deque作为栈，性能更好
            from collections import deque
            stack = deque([(start_x, start_y)])
            pixels_to_clear = []
            
            while stack:
                x, y = stack.pop()
                
                # 边界检查
                if x < 0 or x >= width or y < 0 or y >= height or visited[y, x]:
                    continue
                
                current_color = data[y, x]
                
                # 检查颜色是否在容差范围内
                if not is_background_color(current_color):
                    continue
                
                # 标记为已访问
                visited[y, x] = True
                pixels_to_clear.append((x, y))
                
                # 添加相邻像素到栈中（4连通）
                stack.extend([
                    (x+1, y), (x-1, y), (x, y+1), (x, y-1)
                ])
            
            # 批量设置像素为透明
            for x, y in pixels_to_clear:
                data[y, x] = (0, 0, 0, 0)
        
        print(f"    🌊 优化Flood Fill处理...")
        
        # 从边界开始flood fill，优化边界遍历
        # 上边和下边
        for x in range(0, width, 2):  # 每隔一个像素采样，提高性能
            optimized_flood_fill(x, 0)
            optimized_flood_fill(x, height-1)
        
        # 左边和右边
        for y in range(0, height, 2):  # 每隔一个像素采样，提高性能
            optimized_flood_fill(0, y)
            optimized_flood_fill(width-1, y)
        
        # 补充处理边界的奇数位置
        for x in range(1, width, 2):
            if not visited[0, x]:
                optimized_flood_fill(x, 0)
            if not visited[height-1, x]:
                optimized_flood_fill(x, height-1)
        
        for y in range(1, height, 2):
            if not visited[y, 0]:
                optimized_flood_fill(0, y)
            if not visited[y, width-1]:
                optimized_flood_fill(width-1, y)
        
        # 转换回PIL图像
        return Image.fromarray(data, 'RGBA')

    def crop_transparent_borders(self, image: Image.Image) -> Image.Image:
        """裁剪透明边界，去除多余区域"""
        if image.mode != 'RGBA':
            return image
        
        # 转换为numpy数组
        data = np.array(image)
        
        # 获取alpha通道
        alpha = data[:, :, 3]
        
        # 找到非透明像素的边界
        non_transparent = np.where(alpha > 0)
        
        if len(non_transparent[0]) == 0:
            # 如果图像完全透明，返回最小尺寸
            return image.crop((0, 0, 1, 1))
        
        # 计算边界框
        min_y, max_y = non_transparent[0].min(), non_transparent[0].max()
        min_x, max_x = non_transparent[1].min(), non_transparent[1].max()
        
        # 添加小的边距（5像素）
        padding = 5
        width, height = image.size
        
        min_x = max(0, min_x - padding)
        min_y = max(0, min_y - padding)
        max_x = min(width - 1, max_x + padding)
        max_y = min(height - 1, max_y + padding)
        
        # 裁剪图像
        cropped = image.crop((min_x, min_y, max_x + 1, max_y + 1))
        
        return cropped

    def post_process_image(self, image: Image.Image) -> Image.Image:
        """后处理图像：去除背景并裁剪多余区域，如果背景杂乱则返回None"""
        print(f"    🔧 后处理图像...")
        
        # 检测背景颜色和杂乱程度
        bg_color, is_messy = self.detect_background_color(image)
        
        if is_messy:
            print(f"    ❌ 检测到杂乱背景，抛弃此图片")
            return None
        
        print(f"    📊 检测到背景色: {bg_color}")
        
        # 使用优化的flood fill去除背景
        processed_image = self.optimized_flood_fill_remove_background(image, bg_color, tolerance=30)
        
        # 裁剪透明边界
        cropped_image = self.crop_transparent_borders(processed_image)
        
        original_size = image.size
        final_size = cropped_image.size
        print(f"    ✂️ 尺寸调整: {original_size} → {final_size}")
        
        return cropped_image

    def generate_prompt_and_image(self, concept: str, topic: str, keyword: str, category: str) -> str:
        """为单个概念生成prompt并生成图像"""
        print(f"  🎨 处理: {concept[:50]}...")
        
        # 生成设计prompt
        prompt = self.design_prompt_template.format(
            topic=topic,
            style_keyword=keyword,
            concept=concept
        )
        
        response = self.openai_client.chat.completions.create(
            model="gpt-5-mini",
            messages=[
                {"role": "user", "content": prompt}
            ],
            temperature=0.7
        )
        
        image_prompt = response.choices[0].message.content.strip()
        
        # 生成图像使用imagen-4.0，带重试机制
        max_retries = 5
        retry_delay = 5  # 秒
        response = None
        
        for attempt in range(max_retries):
            try:
                print(f"    🖼️ 生成图像 (尝试 {attempt + 1}/{max_retries})...")
                response = self.genai_client.models.generate_images(
                    model='imagen-4.0-fast-generate-001',
                    prompt=image_prompt,
                    config=types.GenerateImagesConfig(
                        number_of_images=1,
                        aspect_ratio="1:1",
                    )
                )
                # 如果成功，跳出重试循环
                if response and hasattr(response, 'generated_images') and response.generated_images:
                    print(f"    ✅ 图像生成成功")
                    break
                else:
                    print(f"    ⚠️ 图像生成返回空结果")
                    if attempt < max_retries - 1:
                        print(f"    ⏳ 等待 {retry_delay} 秒后重试...")
                        time.sleep(retry_delay)
                    
            except Exception as e:
                print(f"    ❌ 图像生成失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}")
                if attempt < max_retries - 1:
                    print(f"    ⏳ 等待 {retry_delay} 秒后重试...")
                    time.sleep(retry_delay)
                else:
                    print(f"    💀 所有重试均失败，放弃生成此图像")
                    return None
        
        # 保存图像
        if response and hasattr(response, 'generated_images') and response.generated_images:
            generated_image = response.generated_images[0]
            image = Image.open(BytesIO(generated_image.image.image_bytes))
            
            # 后处理图像：去除背景
            processed_image = self.post_process_image(image)
            
            # 如果图像被抛弃（杂乱背景），返回None
            if processed_image is None:
                print(f"    🗑️ 图片已抛弃")
                return None
            
            # 构建文件名 - 使用连字符连接，下划线替换空格
            safe_topic = topic.replace(' ', '_')
            safe_category = category.replace(' ', '_')
            safe_concept = concept[:30].replace(' ', '_')
            
            # 移除非字母数字和允许的字符
            safe_topic = "".join(c for c in safe_topic if c.isalnum() or c in ('_', '-')).strip('_-')
            safe_category = "".join(c for c in safe_category if c.isalnum() or c in ('_', '-')).strip('_-')
            safe_concept = "".join(c for c in safe_concept if c.isalnum() or c in ('_', '-')).strip('_-')
            
            timestamp = int(time.time())
            filename = f"{safe_topic}-{safe_category}-{safe_concept}-{timestamp}.png"
            filepath = os.path.join(self.output_dir, filename)
            
            processed_image.save(filepath)
            print(f"    ✅ 保存: {os.path.basename(filepath)}")
            return filepath
        
        return None

    def run_pipeline(self) -> Dict:
        """运行完整的批量生成pipeline"""
        print("🚀 开始图像批量生成Pipeline")
        
        # 1. 随机选择category、keyword和topic
        category, keyword, topic = self.select_random_category_and_elements()
        
        # 2. 生成10个概念
        concepts = self.generate_concepts(topic)
        
        # 3. 评判并选择5个最佳概念
        selected_concepts = self.evaluate_and_select_concepts(concepts)
        
        # 4. 并行生成prompt和图像
        print(f"🖼️ 并行生成 {len(selected_concepts)} 张图像...")
        generated_files = []
        
        with ThreadPoolExecutor(max_workers=3) as executor:
            futures = []
            for concept in selected_concepts:
                future = executor.submit(
                    self.generate_prompt_and_image, 
                    concept, topic, keyword, category
                )
                futures.append(future)
            
            for future in futures:
                result = future.result()
                if result:  # 只有成功生成且未被抛弃的图片才会被添加
                    generated_files.append(result)
        
        print(f"✅ 完成! 生成 {len(generated_files)} 张图像")
        
        return {
            'category': category,
            'keyword': keyword,
            'topic': topic,
            'generated_images': len(generated_files),
            'output_files': generated_files
        }


def main():
    """主函数"""
    random.seed(int(time.time()))
    generator = ImageBatchGenerator()
    for i in range(1000):
        result = generator.run_pipeline()
    print(f"📊 结果: {result['generated_images']} 张图像已保存")


if __name__ == "__main__":
    main()