generate_extended_dataset.py

"""
扩展版 UI 元素识别训练数据集生成器
生成更多样化的样本（500+）
"""

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

class ExtendedUIDatasetGenerator:
    def __init__(self, output_dir: str = "ui_dataset_extended"):
        self.output_dir = output_dir
        self.images_dir = os.path.join(output_dir, "images")
        self.annotations_dir = os.path.join(output_dir, "annotations")
        self.dataset = []
        
        # 创建输出目录
        os.makedirs(self.images_dir, exist_ok=True)
        os.makedirs(self.annotations_dir, exist_ok=True)
        
        # 颜色配置
        self.colors = {
            'primary': [(59, 130, 246), (99, 102, 241), (139, 92, 246)],  # 蓝色系
            'success': [(34, 197, 94), (22, 163, 74)],  # 绿色系
            'danger': [(239, 68, 68), (220, 38, 38)],  # 红色系
            'warning': [(245, 158, 11), (234, 88, 12)],  # 橙色系
            'neutral': [(107, 114, 128), (75, 85, 99)],  # 灰色系
            'text': [(17, 24, 39), (55, 65, 81)],  # 深色文本
            'border': [(209, 213, 219), (229, 231, 235)],  # 边框色
            'background': [(255, 255, 255), (249, 250, 251)]  # 背景色
        }
        
        # 文本内容
        self.texts = {
            'input': [
                'Enter your username', 'Password', 'Email address', 'Search...',
                'Type here...', 'Enter text', 'Username', 'Password', 'Email',
                'Phone number', 'Address', 'City', 'State', 'Zip code',
                'First name', 'Last name', 'Full name', 'Company name',
                'Website URL', 'Description', 'Comments', 'Message'
            ],
            'button': [
                'Submit', 'Save', 'Cancel', 'Delete', 'Edit', 'Update',
                'Create', 'Add', 'Remove', 'Login', 'Sign Up', 'Register',
                'Sign In', 'Log Out', 'Continue', 'Next', 'Previous',
                'Back', 'Forward', 'Search', 'Clear', 'Reset', 'Apply',
                'OK', 'Yes', 'No', 'Confirm', 'Download', 'Upload',
                'Send', 'Share', 'Copy', 'Paste', 'Cut'
            ],
            'link': [
                'Home', 'About', 'Contact', 'Services', 'Products',
                'Blog', 'News', 'FAQ', 'Help', 'Support', 'Documentation',
                'Terms', 'Privacy', 'Settings', 'Profile', 'Dashboard',
                'Login', 'Sign Up', 'Register', 'Forgot password?', 'Sign out'
            ],
            'select': [
                'Select option', 'Choose...', 'Please select', 'Dropdown',
                'Country', 'State', 'City', 'Language', 'Currency', 'Timezone'
            ],
            'checkbox': [
                'Remember me', 'I agree to terms', 'Subscribe to newsletter',
                'Keep me signed in', 'Save my preferences', 'Enable notifications'
            ],
            'radio': [
                'Male', 'Female', 'Other', 'Yes', 'No', 'Maybe',
                'Option 1', 'Option 2', 'Option 3'
            ]
        }
    
    def create_base_image(self, width: int = 800, height: int = 600) -> Image.Image:
        """创建基础图像"""
        bg_color = random.choice(self.colors['background'])
        img = Image.new('RGB', (width, height), tuple(bg_color))
        draw = ImageDraw.Draw(img)
        return img, draw
    
    def add_title(self, draw: ImageDraw.Draw, text: str, y: int = 20) -> None:
        """添加标题"""
        try:
            font = ImageFont.truetype("arial.ttf", 24)
        except:
            font = ImageFont.load_default()
        
        text_color = random.choice(self.colors['text'])
        draw.text((20, y), text, fill=tuple(text_color), font=font)
    
    def add_input_field(self, draw: ImageDraw.Draw, x: int, y: int, 
                       width: int = 300, height: int = 40, 
                       label: str = None) -> Dict:
        """添加输入框"""
        border_color = random.choice(self.colors['border'])
        bg_color = (255, 255, 255)
        
        # 绘制输入框
        draw.rectangle([x, y, x + width, y + height], 
                      outline=tuple(border_color), fill=bg_color, width=1)
        
        # 添加标签
        if label:
            try:
                font = ImageFont.truetype("arial.ttf", 14)
            except:
                font = ImageFont.load_default()
            
            text_color = random.choice(self.colors['text'])
            draw.text((x, y - 20), label, fill=tuple(text_color), font=font)
        
        # 添加占位符文本
        placeholder = random.choice(self.texts['input'])
        try:
            font = ImageFont.truetype("arial.ttf", 14)
        except:
            font = ImageFont.load_default()
        
        text_color = (156, 163, 175)  # 灰色占位符
        draw.text((x + 10, y + 10), placeholder, fill=text_color, font=font)
        
        return {
            'type': 'input',
            'x': x,
            'y': y,
            'width': width,
            'height': height,
            'label': placeholder,
            'description': f'{label}输入框' if label else '输入框'
        }
    
    def add_button(self, draw: ImageDraw.Draw, x: int, y: int, 
                  width: int = 120, height: int = 35, 
                  text: str = None, button_type: str = 'primary') -> Dict:
        """添加按钮"""
        if button_type == 'primary':
            bg_color = random.choice(self.colors['primary'])
            text_color = (255, 255, 255)
        elif button_type == 'success':
            bg_color = random.choice(self.colors['success'])
            text_color = (255, 255, 255)
        elif button_type == 'danger':
            bg_color = random.choice(self.colors['danger'])
            text_color = (255, 255, 255)
        elif button_type == 'warning':
            bg_color = random.choice(self.colors['warning'])
            text_color = (255, 255, 255)
        else:
            bg_color = random.choice(self.colors['neutral'])
            text_color = (255, 255, 255)
        
        # 绘制按钮
        draw.rectangle([x, y, x + width, y + height], 
                      fill=tuple(bg_color), outline=tuple(bg_color), width=0)
        
        # 添加文本
        if not text:
            text = random.choice(self.texts['button'])
        
        try:
            font = ImageFont.truetype("arial.ttf", 14)
        except:
            font = ImageFont.load_default()
        
        # 居中文本
        text_bbox = draw.textbbox((0, 0), text, font=font)
        text_width = text_bbox[2] - text_bbox[0]
        text_height = text_bbox[3] - text_bbox[1]
        text_x = x + (width - text_width) // 2
        text_y = y + (height - text_height) // 2
        
        draw.text((text_x, text_y), text, fill=text_color, font=font)
        
        return {
            'type': 'button',
            'x': x,
            'y': y,
            'width': width,
            'height': height,
            'label': text,
            'description': f'{button_type}按钮'
        }
    
    def add_link(self, draw: ImageDraw.Draw, x: int, y: int, 
                text: str = None) -> Dict:
        """添加链接"""
        if not text:
            text = random.choice(self.texts['link'])
        
        try:
            font = ImageFont.truetype("arial.ttf", 14)
        except:
            font = ImageFont.load_default()
        
        text_color = random.choice(self.colors['primary'])
        draw.text((x, y), text, fill=tuple(text_color), font=font)
        
        # 添加下划线
        text_bbox = draw.textbbox((x, y), text, font=font)
        text_width = text_bbox[2] - text_bbox[0]
        text_height = text_bbox[3] - text_bbox[1]
        draw.line([(x, text_bbox[3]), (x + text_width, text_bbox[3])], 
                 fill=tuple(text_color), width=1)
        
        return {
            'type': 'link',
            'x': x,
            'y': y,
            'width': text_width,
            'height': text_height,
            'label': text,
            'description': '链接'
        }
    
    def add_checkbox(self, draw: ImageDraw.Draw, x: int, y: int, 
                    label: str = None) -> Dict:
        """添加复选框"""
        size = 20
        border_color = random.choice(self.colors['border'])
        
        # 绘制复选框
        draw.rectangle([x, y, x + size, y + size], 
                      outline=tuple(border_color), fill=(255, 255, 255), width=1)
        
        # 添加标签
        if not label:
            label = random.choice(self.texts['checkbox'])
        
        try:
            font = ImageFont.truetype("arial.ttf", 14)
        except:
            font = ImageFont.load_default()
        
        text_color = random.choice(self.colors['text'])
        draw.text((x + size + 10, y), label, fill=tuple(text_color), font=font)
        
        return {
            'type': 'checkbox',
            'x': x,
            'y': y,
            'width': size,
            'height': size,
            'label': label,
            'description': '复选框'
        }
    
    def add_radio(self, draw: ImageDraw.Draw, x: int, y: int, 
                 label: str = None) -> Dict:
        """添加单选框"""
        size = 20
        border_color = random.choice(self.colors['border'])
        
        # 绘制单选框
        draw.ellipse([x, y, x + size, y + size], 
                    outline=tuple(border_color), fill=(255, 255, 255), width=1)
        
        # 添加标签
        if not label:
            label = random.choice(self.texts['radio'])
        
        try:
            font = ImageFont.truetype("arial.ttf", 14)
        except:
            font = ImageFont.load_default()
        
        text_color = random.choice(self.colors['text'])
        draw.text((x + size + 10, y), label, fill=tuple(text_color), font=font)
        
        return {
            'type': 'radio',
            'x': x,
            'y': y,
            'width': size,
            'height': size,
            'label': label,
            'description': '单选框'
        }
    
    def generate_form_page(self) -> Tuple[Image.Image, List[Dict]]:
        """生成表单页面"""
        img, draw = self.create_base_image(800, 600)
        self.add_title(draw, 'Contact Form')
        
        elements = []
        y = 80
        
        # 添加多个输入框
        num_inputs = random.randint(3, 6)
        labels = ['Name', 'Email', 'Phone', 'Subject', 'Message']
        
        for i in range(num_inputs):
            label = labels[i] if i < len(labels) else f'Field {i+1}'
            elem = self.add_input_field(draw, 50, y, 400, 40, label)
            elements.append(elem)
            y += 70
        
        # 添加按钮
        button_types = ['primary', 'success', 'danger']
        button_type = random.choice(button_types)
        elem = self.add_button(draw, 50, y, 120, 35, 'Submit', button_type)
        elements.append(elem)
        
        # 添加取消按钮
        elem = self.add_button(draw, 180, y, 120, 35, 'Cancel', 'neutral')
        elements.append(elem)
        
        return img, elements
    
    def generate_login_page(self) -> Tuple[Image.Image, List[Dict]]:
        """生成登录页面"""
        img, draw = self.create_base_image(800, 600)
        self.add_title(draw, 'Login')
        
        elements = []
        y = 100
        
        # 用户名输入框
        elem = self.add_input_field(draw, 200, y, 300, 40, 'Username')
        elements.append(elem)
        y += 70
        
        # 密码输入框
        elem = self.add_input_field(draw, 200, y, 300, 40, 'Password')
        elements.append(elem)
        y += 70
        
        # 记住我复选框
        elem = self.add_checkbox(draw, 200, y, 'Remember me')
        elements.append(elem)
        y += 50
        
        # 登录按钮
        elem = self.add_button(draw, 200, y, 300, 40, 'Login', 'primary')
        elements.append(elem)
        y += 60
        
        # 忘记密码链接
        elem = self.add_link(draw, 200, y, 'Forgot password?')
        elements.append(elem)
        
        return img, elements
    
    def generate_search_page(self) -> Tuple[Image.Image, List[Dict]]:
        """生成搜索页面"""
        img, draw = self.create_base_image(800, 600)
        self.add_title(draw, 'Search')
        
        elements = []
        y = 100
        
        # 搜索输入框
        elem = self.add_input_field(draw, 100, y, 500, 40, 'Search')
        elements.append(elem)
        y += 60
        
        # 搜索按钮
        elem = self.add_button(draw, 100, y, 120, 35, 'Search', 'primary')
        elements.append(elem)
        
        # 手气不错按钮
        elem = self.add_button(draw, 230, y, 150, 35, "I'm Feeling Lucky", 'neutral')
        elements.append(elem)
        
        return img, elements
    
    def generate_navigation_menu(self) -> Tuple[Image.Image, List[Dict]]:
        """生成导航菜单"""
        img, draw = self.create_base_image(800, 600)
        self.add_title(draw, 'Navigation Menu')
        
        elements = []
        y = 100
        
        # 添加多个链接
        links = ['Home', 'About', 'Services', 'Products', 'Blog', 'Contact']
        num_links = random.randint(4, 6)
        
        for i in range(num_links):
            link_text = links[i] if i < len(links) else f'Link {i+1}'
            elem = self.add_link(draw, 50, y, link_text)
            elements.append(elem)
            y += 50
        
        return img, elements
    
    def generate_dashboard(self) -> Tuple[Image.Image, List[Dict]]:
        """生成仪表板"""
        img, draw = self.create_base_image(800, 600)
        self.add_title(draw, 'Dashboard')
        
        elements = []
        
        # 顶部搜索框
        elem = self.add_input_field(draw, 500, 20, 250, 30, 'Search')
        elements.append(elem)
        
        # 侧边栏按钮
        y = 100
        sidebar_buttons = ['Dashboard', 'Users', 'Settings', 'Reports', 'Logout']
        for btn_text in sidebar_buttons:
            elem = self.add_button(draw, 20, y, 150, 35, btn_text, 'neutral')
            elements.append(elem)
            y += 50
        
        # 主内容区域
        y = 100
        elem = self.add_input_field(draw, 200, y, 400, 40, 'Title')
        elements.append(elem)
        y += 60
        
        elem = self.add_input_field(draw, 200, y, 400, 100, 'Content')
        elements.append(elem)
        y += 120
        
        # 操作按钮
        elem = self.add_button(draw, 200, y, 100, 35, 'Save', 'primary')
        elements.append(elem)
        
        elem = self.add_button(draw, 310, y, 100, 35, 'Cancel', 'neutral')
        elements.append(elem)
        
        return img, elements
    
    def generate_ecommerce_page(self) -> Tuple[Image.Image, List[Dict]]:
        """生成电商页面"""
        img, draw = self.create_base_image(800, 600)
        self.add_title(draw, 'Product Page')
        
        elements = []
        
        # 搜索框
        elem = self.add_input_field(draw, 500, 20, 250, 30, 'Search products')
        elements.append(elem)
        
        # 产品卡片
        y = 100
        for i in range(3):
            # 产品标题
            elem = self.add_input_field(draw, 50, y, 200, 30, f'Product {i+1}')
            elements.append(elem)
            y += 40
            
            # 价格（作为输入框模拟）
            elem = self.add_input_field(draw, 50, y, 100, 25, f'$99.99')
            elements.append(elem)
            y += 35
            
            # 添加到购物车按钮
            elem = self.add_button(draw, 50, y, 150, 30, 'Add to Cart', 'primary')
            elements.append(elem)
            y += 50
            
            # 购买按钮
            elem = self.add_button(draw, 210, y, 100, 30, 'Buy Now', 'success')
            elements.append(elem)
            y += 60
        
        return img, elements
    
    def generate_settings_page(self) -> Tuple[Image.Image, List[Dict]]:
        """生成设置页面"""
        img, draw = self.create_base_image(800, 600)
        self.add_title(draw, 'Settings')
        
        elements = []
        y = 100
        
        # 用户信息输入框
        elem = self.add_input_field(draw, 200, y, 300, 40, 'Username')
        elements.append(elem)
        y += 70
        
        elem = self.add_input_field(draw, 200, y, 300, 40, 'Email')
        elements.append(elem)
        y += 70
        
        # 复选框选项
        elem = self.add_checkbox(draw, 200, y, 'Receive email notifications')
        elements.append(elem)
        y += 40
        
        elem = self.add_checkbox(draw, 200, y, 'Enable two-factor authentication')
        elements.append(elem)
        y += 40
        
        # 单选框选项
        elem = self.add_radio(draw, 200, y, 'Light theme')
        elements.append(elem)
        y += 40
        
        elem = self.add_radio(draw, 200, y, 'Dark theme')
        elements.append(elem)
        y += 50
        
        # 保存按钮
        elem = self.add_button(draw, 200, y, 120, 35, 'Save Changes', 'primary')
        elements.append(elem)
        
        return img, elements
    
    def augment_image(self, img: Image.Image, elements: List[Dict]) -> Tuple[Image.Image, List[Dict]]:
        """图像增强"""
        # 随机选择增强方式
        augment_type = random.choice(['brightness', 'contrast', 'noise', 'blur'])
        
        if augment_type == 'brightness':
            factor = random.uniform(0.8, 1.2)
            img_array = np.array(img).astype(np.float32)
            img_array = np.clip(img_array * factor, 0, 255).astype(np.uint8)
            img = Image.fromarray(img_array)
        
        elif augment_type == 'contrast':
            factor = random.uniform(0.8, 1.2)
            img_array = np.array(img).astype(np.float32)
            img_array = np.clip((img_array - 128) * factor + 128, 0, 255).astype(np.uint8)
            img = Image.fromarray(img_array)
        
        elif augment_type == 'noise':
            img_array = np.array(img)
            noise = np.random.randint(-10, 11, img_array.shape)
            img_array = np.clip(img_array + noise, 0, 255).astype(np.uint8)
            img = Image.fromarray(img_array)
        
        elif augment_type == 'blur':
            img = img.filter(ImageFilter.GaussianBlur(radius=random.uniform(0.5, 1.5)))
        
        return img, elements
    
    def generate_dataset(self, num_samples: int = 500, 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),
            ('ecommerce', self.generate_ecommerce_page),
            ('settings', self.generate_settings_page)
        ]
        
        for i in range(num_samples):
            # 随机选择生成器
            gen_type, generator = random.choice(generators)
            
            # 生成图像
            img, elements = generator()
            
            # 图像增强
            if augment and random.random() < 0.3:
                from PIL import ImageFilter
                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) % 50 == 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}")
        
        # 统计信息
        self.print_statistics()
        
        return self.dataset
    
    def print_statistics(self):
        """打印统计信息"""
        print(f"\n📊 数据集统计:")
        print(f"   总样本数: {len(self.dataset)}")
        
        # 元素类型统计
        type_counts = {}
        for sample in self.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()):
            percentage = (count / sum(type_counts.values())) * 100
            print(f"      {elem_type}: {count} ({percentage:.1f}%)")

def main():
    """主函数"""
    print("=" * 60)
    print("🎨 扩展版 UI 元素识别训练数据集生成器")
    print("=" * 60)
    
    generator = ExtendedUIDatasetGenerator()
    generator.generate_dataset(num_samples=500, augment=True)
    
    print(f"\n✅ 数据集生成完成！")

if __name__ == "__main__":
    main()