generate_ui_dataset.py

"""
训练数据集扩展工具
生成多样化的 UI 元素样本，用于改进 AI 模型的泛化能力
"""

import os
import json
import random
from typing import List, Dict, Tuple
from PIL import Image, ImageDraw, ImageFont
import numpy as np

class UIDatasetGenerator:
    def __init__(self, output_dir: str = "ui_dataset"):
        self.output_dir = output_dir
        self.images_dir = os.path.join(output_dir, "images")
        self.annotations_dir = os.path.join(output_dir, "annotations")
        
        os.makedirs(self.images_dir, exist_ok=True)
        os.makedirs(self.annotations_dir, exist_ok=True)
        
        self.dataset = []
    
    def get_font(self, size: int = 14) -> ImageFont.FreeTypeFont:
        """获取字体"""
        try:
            return ImageFont.truetype("arial.ttf", size)
        except:
            try:
                return ImageFont.truetype("DejaVuSans.ttf", size)
            except:
                return ImageFont.load_default()
    
    def generate_form_page(self, width: int = 800, height: int = 600) -> Tuple[Image.Image, List[Dict]]:
        """生成表单页面"""
        img = Image.new('RGB', (width, height), color='white')
        draw = ImageDraw.Draw(img)
        
        font = self.get_font(14)
        title_font = self.get_font(20)
        
        elements = []
        
        # 标题
        titles = ["Registration Form", "Contact Us", "Sign Up", "Create Account"]
        title = random.choice(titles)
        draw.text((50, 30), title, fill='black', font=title_font)
        
        # 输入框
        input_labels = [
            ("Username", "Enter your username"),
            ("Email", "Enter your email address"),
            ("Password", "Enter your password"),
            ("Confirm Password", "Confirm your password"),
            ("Phone", "Enter your phone number"),
            ("Address", "Enter your address")
        ]
        
        y_pos = 80
        num_inputs = random.randint(3, 6)
        
        for i in range(num_inputs):
            label_text, placeholder = input_labels[i]
            
            # 标签
            draw.text((50, y_pos), label_text, fill='black', font=font)
            
            # 输入框
            input_y = y_pos + 25
            input_width = random.randint(300, 500)
            input_height = random.randint(30, 40)
            
            draw.rectangle([50, input_y, 50 + input_width, input_y + input_height], 
                        outline='black', width=2)
            draw.text((60, input_y + 5), placeholder, fill='gray', font=font)
            
            elements.append({
                'type': 'input',
                'x': 50,
                'y': input_y,
                'width': input_width,
                'height': input_height,
                'label': placeholder,
                'description': f'{label_text} 输入框'
            })
            
            y_pos = input_y + input_height + 20
        
        # 按钮
        button_labels = ["Submit", "Register", "Sign Up", "Create Account", "Cancel"]
        num_buttons = random.randint(1, 3)
        
        button_y = y_pos + 20
        button_x = 50
        
        for i in range(num_buttons):
            button_label = button_labels[i]
            button_width = random.randint(120, 180)
            button_height = random.randint(35, 45)
            
            # 随机颜色
            colors = ['#4285f4', '#34a853', '#fbbc05', '#ea4335', '#999999']
            color = random.choice(colors)
            
            draw.rectangle([button_x, button_y, button_x + button_width, button_y + button_height], 
                        fill=color)
            draw.text((button_x + 20, button_y + 10), button_label, fill='white', font=font)
            
            elements.append({
                'type': 'button',
                'x': button_x,
                'y': button_y,
                'width': button_width,
                'height': button_height,
                'label': button_label,
                'description': f'{button_label} 按钮'
            })
            
            button_x += button_width + 20
        
        return img, elements
    
    def generate_login_page(self, width: int = 400, height: int = 500) -> Tuple[Image.Image, List[Dict]]:
        """生成登录页面"""
        img = Image.new('RGB', (width, height), color='white')
        draw = ImageDraw.Draw(img)
        
        font = self.get_font(14)
        title_font = self.get_font(20)
        
        elements = []
        
        # 标题
        draw.text((140, 50), "Login", fill='black', font=title_font)
        
        # 用户名输入框
        input_y = 120
        input_width = 300
        input_height = 40
        
        draw.rectangle([50, input_y, 50 + input_width, input_y + input_height], 
                    outline='black', width=2)
        draw.text((60, input_y + 10), "Username", fill='gray', font=font)
        
        elements.append({
            'type': 'input',
            'x': 50,
            'y': input_y,
            'width': input_width,
            'height': input_height,
            'label': 'Username',
            'description': '用户名输入框'
        })
        
        # 密码输入框
        input_y = 180
        draw.rectangle([50, input_y, 50 + input_width, input_y + input_height], 
                    outline='black', width=2)
        draw.text((60, input_y + 10), "Password", fill='gray', font=font)
        
        elements.append({
            'type': 'input',
            'x': 50,
            'y': input_y,
            'width': input_width,
            'height': input_height,
            'label': 'Password',
            'description': '密码输入框'
        })
        
        # 记住我复选框
        checkbox_y = 250
        checkbox_size = 20
        draw.rectangle([50, checkbox_y, 50 + checkbox_size, checkbox_y + checkbox_size], 
                    outline='black', width=2)
        draw.text((80, checkbox_y), "Remember me", fill='black', font=font)
        
        elements.append({
            'type': 'checkbox',
            'x': 50,
            'y': checkbox_y,
            'width': checkbox_size,
            'height': checkbox_size,
            'label': 'Remember me',
            'description': '记住我复选框'
        })
        
        # 登录按钮
        button_y = 300
        button_width = 300
        button_height = 50
        
        draw.rectangle([50, button_y, 50 + button_width, button_y + button_height], 
                    fill='#4285f4')
        draw.text((140, button_y + 15), "Login", fill='white', font=title_font)
        
        elements.append({
            'type': 'button',
            'x': 50,
            'y': button_y,
            'width': button_width,
            'height': button_height,
            'label': 'Login',
            'description': '登录按钮'
        })
        
        # 忘记密码链接
        link_y = 370
        draw.text((120, link_y), "Forgot password?", fill='#0066cc', font=font)
        
        elements.append({
            'type': 'link',
            'x': 120,
            'y': link_y,
            'width': 150,
            'height': 20,
            'label': 'Forgot password?',
            'description': '忘记密码链接'
        })
        
        return img, elements
    
    def generate_search_page(self, width: int = 800, height: int = 400) -> Tuple[Image.Image, List[Dict]]:
        """生成搜索页面"""
        img = Image.new('RGB', (width, height), color='white')
        draw = ImageDraw.Draw(img)
        
        font = self.get_font(16)
        title_font = self.get_font(24)
        
        elements = []
        
        # Logo/标题
        draw.text((300, 50), "Search Engine", fill='#4285f4', font=title_font)
        
        # 搜索输入框
        input_y = 120
        input_width = 600
        input_height = 50
        
        draw.rectangle([100, input_y, 100 + input_width, input_y + input_height], 
                    outline='#dfe1e5', width=2)
        draw.rectangle([100, input_y, 100 + input_width, input_y + input_height], 
                    outline='#4285f4', width=2)
        draw.text((115, input_y + 15), "Search...", fill='gray', font=font)
        
        elements.append({
            'type': 'input',
            'x': 100,
            'y': input_y,
            'width': input_width,
            'height': input_height,
            'label': 'Search...',
            'description': '搜索输入框'
        })
        
        # 搜索按钮
        button_y = 190
        button_width = 150
        button_height = 40
        
        draw.rectangle([100, button_y, 100 + button_width, button_y + button_height], 
                    fill='#f8f9fa')
        draw.text((120, button_y + 10), "Google Search", fill='#3c4043', font=font)
        
        elements.append({
            'type': 'button',
            'x': 100,
            'y': button_y,
            'width': button_width,
            'height': button_height,
            'label': 'Google Search',
            'description': '搜索按钮'
        })
        
        # 手气不错按钮
        button_x = 270
        draw.rectangle([button_x, button_y, button_x + button_width, button_y + button_height], 
                    fill='#f8f9fa')
        draw.text((290, button_y + 10), "I'm Feeling Lucky", fill='#3c4043', font=font)
        
        elements.append({
            'type': 'button',
            'x': button_x,
            'y': button_y,
            'width': button_width,
            'height': button_height,
            'label': "I'm Feeling Lucky",
            'description': '手气不错按钮'
        })
        
        return img, elements
    
    def generate_navigation_menu(self, width: int = 1000, height: int = 100) -> Tuple[Image.Image, List[Dict]]:
        """生成导航菜单"""
        img = Image.new('RGB', (width, height), color='white')
        draw = ImageDraw.Draw(img)
        
        font = self.get_font(14)
        
        elements = []
        
        # 导航链接
        link_sets = [
            ["Home", "About", "Services", "Products", "Contact"],
            ["Dashboard", "Analytics", "Reports", "Settings", "Help"],
            ["Products", "Solutions", "Pricing", "Support", "Blog"]
        ]
        
        links = random.choice(link_sets)
        x_pos = 50
        
        for link in links:
            # 绘制蓝色文本（模拟链接）
            draw.text((x_pos, 40), link, fill='#0066cc', font=font)
            # 添加下划线
            text_width = len(link) * 9
            draw.line([(x_pos, 60), (x_pos + text_width, 60)], fill='#0066cc', width=2)
            
            elements.append({
                'type': 'link',
                'x': x_pos,
                'y': 40,
                'width': text_width,
                'height': 20,
                'label': link,
                'description': f'{link} 导航链接'
            })
            
            x_pos += text_width + 50
        
        return img, elements
    
    def generate_dashboard(self, width: int = 1200, height: int = 800) -> Tuple[Image.Image, List[Dict]]:
        """生成仪表板"""
        img = Image.new('RGB', (width, height), color='#f5f5f5')
        draw = ImageDraw.Draw(img)
        
        font = self.get_font(12)
        title_font = self.get_font(18)
        
        elements = []
        
        # 顶部导航栏
        draw.rectangle([0, 0, width, 60], fill='#333333')
        draw.text((20, 20), "Dashboard", fill='white', font=title_font)
        
        # 搜索框
        search_y = 15
        search_width = 250
        search_height = 30
        
        draw.rectangle([900, search_y, 900 + search_width, search_y + search_height], 
                    outline='white', width=2)
        draw.text((910, search_y + 5), "Search...", fill='gray', font=font)
        
        elements.append({
            'type': 'input',
            'x': 900,
            'y': search_y,
            'width': search_width,
            'height': search_height,
            'label': 'Search',
            'description': '搜索框'
        })
        
        # 侧边栏按钮
        sidebar_buttons = ["Dashboard", "Analytics", "Reports", "Settings"]
        y_pos = 80
        
        for button in sidebar_buttons:
            button_width = 200
            button_height = 45
            
            draw.rectangle([0, y_pos, button_width, y_pos + button_height], fill='#444444')
            draw.text((20, y_pos + 15), button, fill='white', font=font)
            
            elements.append({
                'type': 'button',
                'x': 0,
                'y': y_pos,
                'width': button_width,
                'height': button_height,
                'label': button,
                'description': f'{button} 导航按钮'
            })
            
            y_pos += button_height + 5
        
        # 主内容区域 - 输入框
        input_y = 100
        input_width = 350
        input_height = 30
        
        draw.rectangle([250, input_y, 250 + input_width, input_y + input_height], 
                    outline='black', width=2)
        draw.text((260, input_y + 5), "Filter by date", fill='gray', font=font)
        
        elements.append({
            'type': 'input',
            'x': 250,
            'y': input_y,
            'width': input_width,
            'height': input_height,
            'label': 'Filter by date',
            'description': '日期过滤输入框'
        })
        
        # 导出按钮
        button_y = 100
        button_width = 150
        button_height = 30
        
        draw.rectangle([650, button_y, 650 + button_width, button_y + button_height], 
                    fill='#4CAF50')
        draw.text((690, button_y + 5), "Export", fill='white', font=font)
        
        elements.append({
            'type': 'button',
            'x': 650,
            'y': button_y,
            'width': button_width,
            'height': button_height,
            'label': 'Export',
            'description': '导出按钮'
        })
        
        # 刷新按钮
        button_x = 820
        draw.rectangle([button_x, button_y, button_x + button_width, button_y + button_height], 
                    fill='#2196F3')
        draw.text((860, button_y + 5), "Refresh", fill='white', font=font)
        
        elements.append({
            'type': 'button',
            'x': button_x,
            'y': button_y,
            'width': button_width,
            'height': button_height,
            'label': 'Refresh',
            'description': '刷新按钮'
        })
        
        return img, elements
    
    def augment_image(self, img: Image.Image, elements: List[Dict]) -> Tuple[Image.Image, List[Dict]]:
        """图像增强"""
        augmented_elements = []
        
        # 随机选择增强方式
        augment_type = random.choice(['none', 'brightness', 'contrast', 'noise'])
        
        if augment_type == 'brightness':
            # 亮度调整
            factor = random.uniform(0.8, 1.2)
            img = Image.eval(img, lambda x: int(x * factor))
            augmented_elements = elements.copy()
        
        elif augment_type == 'contrast':
            # 对比度调整
            factor = random.uniform(0.8, 1.2)
            img = Image.eval(img, lambda x: int((x - 128) * factor + 128))
            augmented_elements = elements.copy()
        
        elif augment_type == 'noise':
            # 添加噪声
            img_array = np.array(img)
            noise = np.random.randint(-10, 10, img_array.shape)
            img_array = np.clip(img_array + noise, 0, 255)
            img = Image.fromarray(img_array.astype(np.uint8))
            augmented_elements = elements.copy()
        
        else:
            augmented_elements = elements.copy()
        
        return img, augmented_elements
    
    def generate_dataset(self, num_samples: int = 100, augment: bool = True) -> List[Dict]:
        """生成完整数据集"""
        print(f"🚀 开始生成 {num_samples} 个样本...")
        
        generators = [
            ('form', self.generate_form_page),
            ('login', self.generate_login_page),
            ('search', self.generate_search_page),
            ('navigation', self.generate_navigation_menu),
            ('dashboard', self.generate_dashboard)
        ]
        
        for i in range(num_samples):
            # 随机选择生成器
            gen_type, generator = random.choice(generators)
            
            # 生成图像
            img, elements = generator()
            
            # 图像增强
            if augment and random.random() < 0.3:
                img, elements = self.augment_image(img, elements)
                gen_type = f"{gen_type}_augmented"
            
            # 保存图像
            image_filename = f"{gen_type}_{i:04d}.png"
            image_path = os.path.join(self.images_dir, image_filename)
            img.save(image_path)
            
            # 保存标注
            annotation = {
                'image': image_filename,
                'width': img.width,
                'height': img.height,
                'elements': elements
            }
            
            annotation_filename = f"{gen_type}_{i:04d}.json"
            annotation_path = os.path.join(self.annotations_dir, annotation_filename)
            
            with open(annotation_path, 'w', encoding='utf-8') as f:
                json.dump(annotation, f, indent=2, ensure_ascii=False)
            
            self.dataset.append(annotation)
            
            if (i + 1) % 10 == 0:
                print(f"✅ 已生成 {i + 1}/{num_samples} 个样本")
        
        # 保存数据集索引
        index_path = os.path.join(self.output_dir, "dataset_index.json")
        with open(index_path, 'w', encoding='utf-8') as f:
            json.dump({
                'total_samples': num_samples,
                'samples': self.dataset
            }, f, indent=2, ensure_ascii=False)
        
        print(f"✅ 数据集生成完成！")
        print(f"📁 图像目录: {self.images_dir}")
        print(f"📁 标注目录: {self.annotations_dir}")
        print(f"📋 索引文件: {index_path}")
        
        return self.dataset

def main():
    """主函数"""
    print("=" * 60)
    print("🎨 UI 元素识别训练数据集生成器")
    print("=" * 60)
    
    generator = UIDatasetGenerator()
    
    # 生成数据集
    dataset = generator.generate_dataset(num_samples=100, augment=True)
    
    print("\n📊 数据集统计:")
    print(f"   总样本数: {len(dataset)}")
    
    # 统计元素类型
    type_counts = {}
    for sample in dataset:
        for elem in sample['elements']:
            elem_type = elem['type']
            type_counts[elem_type] = type_counts.get(elem_type, 0) + 1
    
    print(f"\n   元素类型分布:")
    for elem_type, count in sorted(type_counts.items()):
        print(f"      {elem_type}: {count}")
    
    print("\n✅ 数据集生成完成！")

if __name__ == "__main__":
    main()