Spaces:
Sleeping
Sleeping
| """ | |
| create_synthetic_dataset.py — Generate synthetic training data locally. | |
| Creates simple image editing pairs for quick model training and testing. | |
| """ | |
| import os | |
| import numpy as np | |
| from PIL import Image, ImageDraw, ImageFilter, ImageEnhance | |
| import pathlib | |
| def create_synthetic_pair(idx: int, size: int = 256): | |
| """Create a synthetic image editing pair (input, edited output).""" | |
| # Create a base image with random shapes and colors | |
| img = Image.new('RGB', (size, size), color=(255, 255, 255)) | |
| draw = ImageDraw.Draw(img) | |
| # Draw random shapes | |
| np.random.seed(idx) # Deterministic for reproducibility | |
| for _ in range(5): | |
| x0, y0 = np.random.randint(0, size, 2) | |
| x1, y1 = x0 + np.random.randint(20, 100), y0 + np.random.randint(20, 100) | |
| color = tuple(np.random.randint(0, 256, 3)) | |
| draw.rectangle([x0, y0, x1, y1], fill=color, outline=(0, 0, 0), width=2) | |
| # Add text | |
| draw.text((10, 10), f"Sample {idx}", fill=(0, 0, 0)) | |
| # Create edited version by applying a transformation | |
| edited = img.copy() | |
| edited = ImageEnhance.Color(edited).enhance(1.3) # Increase saturation | |
| edited = ImageEnhance.Brightness(edited).enhance(0.9) # Slight darkening | |
| edited = edited.filter(ImageFilter.SMOOTH) | |
| return img, edited | |
| def create_dataset(num_samples: int = 100, output_dir: str = "./data/synthetic"): | |
| """Create and save synthetic training pairs in LocalEditDataset format.""" | |
| os.makedirs(output_dir, exist_ok=True) | |
| # Create directory structure expected by LocalEditDataset | |
| input_dir = os.path.join(output_dir, 'input') | |
| edited_dir = os.path.join(output_dir, 'edited') | |
| os.makedirs(input_dir, exist_ok=True) | |
| os.makedirs(edited_dir, exist_ok=True) | |
| prompts = [] | |
| prompt_templates = [ | |
| 'enhance image quality', | |
| 'increase saturation', | |
| 'brighten the image', | |
| 'apply artistic filter', | |
| 'improve contrast', | |
| 'warm color tone', | |
| 'cool color tone', | |
| 'add detail', | |
| 'smooth texture', | |
| 'vibrant colors', | |
| ] | |
| for i in range(num_samples): | |
| img, edited = create_synthetic_pair(i) | |
| img_path = os.path.join(input_dir, f'{i:06d}.png') | |
| edited_path = os.path.join(edited_dir, f'{i:06d}.png') | |
| img.save(img_path) | |
| edited.save(edited_path) | |
| prompt = prompt_templates[i % len(prompt_templates)] | |
| prompts.append(prompt) | |
| if (i + 1) % 20 == 0: | |
| print(f' Created {i+1}/{num_samples} pairs') | |
| # Save prompts file | |
| prompts_path = os.path.join(output_dir, 'prompts.txt') | |
| with open(prompts_path, 'w') as f: | |
| f.write('\n'.join(prompts)) | |
| print(f'✓ Dataset created at {output_dir}') | |
| print(f' - {num_samples} input images') | |
| print(f' - {num_samples} edited images') | |
| print(f' - prompts.txt with instructions') | |
| if __name__ == "__main__": | |
| print("Creating synthetic training dataset...") | |
| create_dataset(num_samples=100, output_dir="./data/synthetic") | |