icon_generation/backup/image_batch_generator_backup.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

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 line_scan_remove_background(self, image: Image.Image, bg_color: tuple, tolerance: int = 30, min_consecutive: int = 5) -> Image.Image:
        """逐行扫描去除连续的背景色区域"""
        # 转换为RGBA模式
        if image.mode != 'RGBA':
            image = image.convert('RGBA')
        
        # 转换为numpy数组
        data = np.array(image)
        width, height = image.size
        
        def color_distance(c1, c2):
            """计算颜色距离"""
            return np.sqrt(sum((a - b) ** 2 for a, b in zip(c1[:3], c2[:3])))
        
        def is_background_color(pixel_color):
            """判断是否为背景色"""
            return color_distance(pixel_color[:3], bg_color) <= tolerance
        
        def process_line(line_data, is_horizontal=True):
            """处理一行或一列的数据，去除连续的背景色区域"""
            line_length = len(line_data)
            i = 0
            
            while i < line_length:
                # 检查当前像素是否为背景色
                if is_background_color(line_data[i]):
                    # 找到连续背景色区域的结束位置
                    consecutive_start = i
                    while i < line_length and is_background_color(line_data[i]):
                        i += 1
                    consecutive_end = i
                    consecutive_length = consecutive_end - consecutive_start
                    
                    # 如果连续背景色区域超过阈值，设置为透明
                    if consecutive_length > min_consecutive:
                        for j in range(consecutive_start, consecutive_end):
                            line_data[j] = (0, 0, 0, 0)  # 设置为透明
                else:
                    i += 1
            
            return line_data
        
        # 逐行扫描（水平方向）
        print(f"    🔍 逐行扫描（水平方向）...")
        for y in range(height):
            row_data = data[y, :].copy()
            processed_row = process_line(row_data, is_horizontal=True)
            data[y, :] = processed_row
        
        # 逐列扫描（垂直方向）
        print(f"    🔍 逐列扫描（垂直方向）...")
        for x in range(width):
            col_data = data[:, x].copy()
            processed_col = process_line(col_data, is_horizontal=False)
            data[:, x] = processed_col
        
        # 转换回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}")
        
        # 使用逐行扫描去除背景
        processed_image = self.line_scan_remove_background(image, bg_color, tolerance=30, min_consecutive=5)
        
        # 裁剪透明边界
        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
        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:
            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(10):
        result = generator.run_pipeline()
    print(f"📊 结果: {result['generated_images']} 张图像已保存")


if __name__ == "__main__":
    main()