Upload dataset/scaleup_dataset.py with huggingface_hub

dataset/scaleup_dataset.py

@@ -0,0 +1,637 @@
+import argparse
+import json
+import os
+import random
+import logging
+import shutil
+from typing import Dict, List, Optional, Tuple
+from PIL import Image
+from tqdm import tqdm
+from multiprocessing import Pool, cpu_count
+from functools import partial
+
+from scaleup_utils import (
+    load_jsonl,
+    save_jsonl,
+    load_blended_sample,
+    get_blended_sample_dirs,
+    compute_non_overlapping_box_xyxy,
+    compute_total_overlap,
+    create_layer_on_canvas,
+    build_spatial_aware_caption,
+    get_position_description,
+    get_box_size,
+    get_content_bbox,
+    load_caption_list,
+    get_laion_images_with_captions,
+    get_caption_images_with_text,
+    select_random_layers_from_samples,
+)
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Default canvas size
+CANVAS_SIZE = 1024
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Scale up PrismLayersPro-blended dataset')
+    parser.add_argument('--blended_dir', type=str, required=True,
+                        help='Path to PrismLayersPro-blended directory')
+    parser.add_argument('--laion_dir', type=str, required=True,
+                        help='Path to LAION aesthetic images directory')
+    parser.add_argument('--caption_dir', type=str, required=True,
+                        help='Path to caption images directory')
+    parser.add_argument('--caption_meta', type=str, required=True,
+                        help='Path to captions.jsonl with caption text')
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Output directory for scaled-up dataset')
+    parser.add_argument('--num_samples', type=int, default=100000,
+                        help='Number of new samples to generate')
+    parser.add_argument('--start_index', type=int, default=0,
+                        help='Starting sample index for output naming')
+    parser.add_argument('--seed', type=int, default=42,
+                        help='Random seed')
+    # Layer selection parameters
+    parser.add_argument('--min_donor_samples', type=int, default=2,
+                        help='Minimum number of samples to pick layers from')
+    parser.add_argument('--max_donor_samples', type=int, default=3,
+                        help='Maximum number of samples to pick layers from')
+    parser.add_argument('--min_layers_per_donor', type=int, default=1,
+                        help='Minimum layers to pick from each donor sample')
+    parser.add_argument('--max_layers_per_donor', type=int, default=2,
+                        help='Maximum layers to pick from each donor sample')
+    parser.add_argument('--added_layer_min_size', type=float, default=0.8,
+                        help='Minimum size ratio for added layers')
+    parser.add_argument('--added_layer_max_size', type=float, default=1.2,
+                        help='Maximum size ratio for added layers')
+    parser.add_argument('--laion_prob', type=float, default=0.1,
+                        help='Probability of including LAION image layer')
+    parser.add_argument('--caption_prob', type=float, default=0.2,
+                        help='Probability of including caption text image layer')
+    parser.add_argument('--laion_min_size', type=float, default=0.2,
+                        help='Minimum size ratio for LAION layer')
+    parser.add_argument('--laion_max_size', type=float, default=0.4,
+                        help='Maximum size ratio for LAION layer')
+    parser.add_argument('--caption_min_size', type=float, default=1.0,
+                        help='Minimum size ratio for caption layer')
+    parser.add_argument('--caption_max_size', type=float, default=1.2,
+                        help='Maximum size ratio for caption layer')
+    # Base layer removal parameters
+    parser.add_argument('--min_layers_to_remove', type=int, default=2,
+                        help='Minimum number of layers to remove from base sample')
+    parser.add_argument('--max_layers_to_remove', type=int, default=4,
+                        help='Maximum number of layers to remove from base sample')
+    # AlphaVAE layer parameters
+    parser.add_argument('--alphavae_dir', type=str, default=None,
+                        help='Path to AlphaVAE_frontview images directory')
+    parser.add_argument('--alphavae_prompts', type=str, default=None,
+                        help='Path to prompts.txt for AlphaVAE captions')
+    parser.add_argument('--alphavae_min_layers', type=int, default=0,
+                        help='Minimum number of AlphaVAE layers to add per sample')
+    parser.add_argument('--alphavae_max_layers', type=int, default=0,
+                        help='Maximum number of AlphaVAE layers to add per sample')
+    parser.add_argument('--alphavae_min_size', type=float, default=0.15,
+                        help='Minimum size ratio for AlphaVAE layers')
+    parser.add_argument('--alphavae_max_size', type=float, default=0.35,
+                        help='Maximum size ratio for AlphaVAE layers')
+    # Multiprocessing
+    parser.add_argument('--max_base_samples', type=int, default=None,
+                        help='Max number of base samples to use (sorted by name). '
+                             'E.g. 18000 to use sample_000000..sample_017999')
+    parser.add_argument('--skip_existing', action='store_true',
+                        help='Skip samples whose output directory already exists (for resuming)')
+    parser.add_argument('--num_workers', type=int, default=64,
+                        help='Number of parallel workers (0 = single-process)')
+    return parser.parse_args()
+
+
+def create_scaled_up_sample(
+    base_sample_dir: str,
+    all_sample_dirs: List[str],
+    laion_images: List[Tuple[str, str]],
+    caption_images: List[Tuple[str, str]],
+    alphavae_images: List[Tuple[str, str]],
+    output_dir: str,
+    sample_idx: int,
+    args: argparse.Namespace,
+) -> Optional[Dict]:
+    """
+    Create a new scaled-up sample by combining a base sample with layers from other samples.
+    
+    Returns metadata dict or None if failed.
+    """
+    # Load base sample
+    base_meta = load_blended_sample(base_sample_dir)
+    if base_meta is None:
+        logger.warning(f"Failed to load base sample: {base_sample_dir}")
+        return None
+    
+    canvas_size = base_meta.get('width', CANVAS_SIZE)
+    
+    # Create output sample directory
+    sample_name = f"sample_{sample_idx:06d}"
+    sample_output_dir = os.path.join(output_dir, sample_name)
+    os.makedirs(sample_output_dir, exist_ok=True)
+    
+    # Copy base_image
+    base_image = base_meta.get('base_image')
+    if base_image:
+        base_image.save(os.path.join(sample_output_dir, 'base_image.png'))
+    else:
+        base_image = Image.new('RGBA', (canvas_size, canvas_size), (0, 0, 0, 0))
+        base_image.save(os.path.join(sample_output_dir, 'base_image.png'))
+    
+    # Start with base composite
+    composite = base_image.copy()
+    
+    # Collect occupied boxes
+    occupied_boxes = []
+    
+    # New layers list
+    new_layers = []
+    current_layer_idx = 0
+    
+    # === Step 1: Copy layers from base sample (excluding laion_foreground and caption types) ===
+    # Also randomly remove 2-3 layers to keep total layer count reasonable
+    base_layers = base_meta.get('layers', [])
+    base_prism_layers = [l for l in base_layers if l.get('type') is None]
+    
+    # Randomly remove some layers from base
+    num_to_remove = random.randint(args.min_layers_to_remove, args.max_layers_to_remove)
+    num_to_remove = min(num_to_remove, max(0, len(base_prism_layers) - 1))  # Keep at least 1 layer
+    
+    removed_layer_indices = set()
+    if num_to_remove > 0 and len(base_prism_layers) > 1:
+        layers_to_remove = random.sample(base_prism_layers, num_to_remove)
+        removed_layer_indices = {l['layer_idx'] for l in layers_to_remove}
+    
+    # Filter out removed layers
+    base_prism_layers_filtered = [l for l in base_prism_layers if l['layer_idx'] not in removed_layer_indices]
+    
+    for layer in base_prism_layers_filtered:
+        orig_layer_idx = layer['layer_idx']
+        layer_img = base_meta.get('layer_images', {}).get(orig_layer_idx)
+        
+        if layer_img is None:
+            continue
+        
+        orig_box = layer.get('box', [0, 0, canvas_size, canvas_size])
+        caption = layer.get('caption', '')
+        orig_w = orig_box[2] - orig_box[0]
+        orig_h = orig_box[3] - orig_box[1]
+        
+        # Randomly reposition (layout-agnostic): find a non-overlapping spot
+        best_box = None
+        best_overlap_ratio = float('inf')
+        new_box = None
+        for _ in range(300):
+            x0 = random.randint(0, max(0, canvas_size - orig_w))
+            y0 = random.randint(0, max(0, canvas_size - orig_h))
+            candidate = [x0, y0, x0 + orig_w, y0 + orig_h]
+            box_area = orig_w * orig_h
+            if box_area <= 0:
+                continue
+            overlap = compute_total_overlap(candidate, occupied_boxes)
+            overlap_ratio = overlap / box_area
+            if overlap == 0:
+                new_box = candidate
+                break
+            if overlap_ratio < best_overlap_ratio:
+                best_overlap_ratio = overlap_ratio
+                best_box = candidate
+        if new_box is None:
+            new_box = best_box if best_box else [0, 0, orig_w, orig_h]
+        
+        # Place cropped layer onto full canvas at the new random position
+        layer_canvas = create_layer_on_canvas(layer_img, new_box, canvas_size)
+        
+        # Save layer with new index
+        layer_filename = f'layer_{current_layer_idx:02d}.png'
+        layer_canvas.save(os.path.join(sample_output_dir, layer_filename))
+        
+        # Composite
+        composite = Image.alpha_composite(composite, layer_canvas)
+        
+        # Record
+        w, h = get_box_size(new_box)
+        new_layers.append({
+            'layer_idx': current_layer_idx,
+            'caption': caption,
+            'box': new_box,
+            'width_dst': w,
+            'height_dst': h,
+            'image_path': layer_filename,
+            'source': 'base',
+            'source_sample': base_sample_dir,
+        })
+        
+        occupied_boxes.append(new_box)
+        current_layer_idx += 1
+    
+    # === Step 2: Add layers from other samples ===
+    num_donors = random.randint(args.min_donor_samples, args.max_donor_samples)
+    donor_layers = select_random_layers_from_samples(
+        all_sample_dirs,
+        exclude_sample=base_sample_dir,
+        num_samples_to_pick=num_donors,
+        num_layers_per_sample=(args.min_layers_per_donor, args.max_layers_per_donor)
+    )
+    
+    for layer_img, layer_info, source_sample in donor_layers:
+        caption = layer_info.get('caption', '')
+        
+        # Use the donor layer's original bounding box dimensions
+        orig_box = layer_info.get('box', [0, 0, canvas_size, canvas_size])
+        orig_w = orig_box[2] - orig_box[0]
+        orig_h = orig_box[3] - orig_box[1]
+        
+        # Find a non-overlapping position for a box of this exact size
+        best_box = None
+        best_overlap_ratio = float('inf')
+        new_box = None
+        for _ in range(300):
+            x0 = random.randint(0, max(0, canvas_size - orig_w))
+            y0 = random.randint(0, max(0, canvas_size - orig_h))
+            candidate = [x0, y0, x0 + orig_w, y0 + orig_h]
+            box_area = orig_w * orig_h
+            if box_area <= 0:
+                continue
+            overlap = compute_total_overlap(candidate, occupied_boxes)
+            overlap_ratio = overlap / box_area
+            if overlap == 0:
+                new_box = candidate
+                break
+            if overlap_ratio < best_overlap_ratio:
+                best_overlap_ratio = overlap_ratio
+                best_box = candidate
+        if new_box is None:
+            new_box = best_box if best_box else [0, 0, orig_w, orig_h]
+        
+        # Create layer on canvas at new position
+        layer_canvas = create_layer_on_canvas(layer_img, new_box, canvas_size)
+        
+        # Save layer
+        layer_filename = f'layer_{current_layer_idx:02d}.png'
+        layer_canvas.save(os.path.join(sample_output_dir, layer_filename))
+        
+        # Composite
+        composite = Image.alpha_composite(composite, layer_canvas)
+        
+        # Record
+        w, h = get_box_size(new_box)
+        new_layers.append({
+            'layer_idx': current_layer_idx,
+            'caption': caption,
+            'box': new_box,
+            'width_dst': w,
+            'height_dst': h,
+            'image_path': layer_filename,
+            'source': 'donor',
+            'source_sample': source_sample,
+            'original_layer_idx': layer_info.get('layer_idx'),
+        })
+        
+        occupied_boxes.append(new_box)
+        current_layer_idx += 1
+    
+    # === Step 2.5: Optionally add AlphaVAE layers (0 to max) ===
+    if alphavae_images and args.alphavae_max_layers > 0:
+        num_alpha = random.randint(args.alphavae_min_layers, args.alphavae_max_layers)
+        if num_alpha > 0:
+            selected_alpha = random.sample(alphavae_images, min(num_alpha, len(alphavae_images)))
+            for alpha_path, alpha_caption in selected_alpha:
+                try:
+                    alpha_img = Image.open(alpha_path).convert('RGBA')
+                except Exception as e:
+                    logger.warning(f"Failed to load AlphaVAE image: {alpha_path}, {e}")
+                    continue
+
+                alpha_box = compute_non_overlapping_box_xyxy(
+                    canvas_size, occupied_boxes,
+                    min_size_ratio=args.alphavae_min_size,
+                    max_size_ratio=args.alphavae_max_size,
+                    max_attempts=300,
+                    max_overlap_ratio=0.10,
+                    center_margin=32
+                )
+
+                alpha_layer = create_layer_on_canvas(alpha_img, alpha_box, canvas_size)
+
+                layer_filename = f'layer_{current_layer_idx:02d}.png'
+                alpha_layer.save(os.path.join(sample_output_dir, layer_filename))
+
+                composite = Image.alpha_composite(composite, alpha_layer)
+
+                w, h = get_box_size(alpha_box)
+                new_layers.append({
+                    'layer_idx': current_layer_idx,
+                    'caption': alpha_caption,
+                    'box': alpha_box,
+                    'width_dst': w,
+                    'height_dst': h,
+                    'image_path': layer_filename,
+                    'type': 'alphavae',
+                    'source_path': alpha_path,
+                })
+
+                occupied_boxes.append(alpha_box)
+                current_layer_idx += 1
+
+    # === Step 3: Optionally add LAION image ===
+    laion_caption = None
+    laion_path = None
+    if random.random() < args.laion_prob and laion_images:
+        laion_path, laion_caption = random.choice(laion_images)
+        try:
+            laion_img = Image.open(laion_path).convert('RGBA')
+            laion_orig_size = laion_img.size
+        except Exception as e:
+            logger.warning(f"Failed to load LAION image: {laion_path}, {e}")
+            laion_img = None
+        
+        if laion_img is not None:
+            laion_box = compute_non_overlapping_box_xyxy(
+                canvas_size, occupied_boxes,
+                min_size_ratio=args.laion_min_size,
+                max_size_ratio=args.laion_max_size,
+                max_attempts=300,
+                max_overlap_ratio=0.10,
+                center_margin=32
+            )
+            
+            # Create layer
+            laion_layer = create_layer_on_canvas(laion_img, laion_box, canvas_size)
+            
+            # Save
+            layer_filename = f'layer_{current_layer_idx:02d}.png'
+            laion_layer.save(os.path.join(sample_output_dir, layer_filename))
+            
+            # Composite
+            composite = Image.alpha_composite(composite, laion_layer)
+            
+            # Record
+            w, h = get_box_size(laion_box)
+            new_layers.append({
+                'layer_idx': current_layer_idx,
+                'caption': laion_caption,
+                'box': laion_box,
+                'width_dst': w,
+                'height_dst': h,
+                'image_path': layer_filename,
+                'type': 'laion_foreground',
+                'original_size': list(laion_orig_size),
+            })
+            
+            occupied_boxes.append(laion_box)
+            current_layer_idx += 1
+    
+    # === Step 4: Optionally add caption image ===
+    caption_text = None
+    caption_path = None
+    if random.random() < args.caption_prob and caption_images:
+        caption_path, caption_text = random.choice(caption_images)
+        try:
+            caption_img = Image.open(caption_path).convert('RGBA')
+            caption_orig_size = caption_img.size
+        except Exception as e:
+            logger.warning(f"Failed to load caption image: {caption_path}, {e}")
+            caption_img = None
+        
+        if caption_img is not None:
+            caption_box = compute_non_overlapping_box_xyxy(
+                canvas_size, occupied_boxes,
+                min_size_ratio=args.caption_min_size,
+                max_size_ratio=args.caption_max_size,
+                max_attempts=300,
+                max_overlap_ratio=0.10,
+                center_margin=32
+            )
+            
+            # Create layer
+            caption_layer = create_layer_on_canvas(caption_img, caption_box, canvas_size)
+            
+            # Compute tight bbox from actual non-transparent content
+            # Caption images have rectangular colored areas surrounded by transparency,
+            # so we use the actual content bounds instead of the placement box.
+            tight_box = get_content_bbox(caption_layer)
+            if tight_box is not None:
+                caption_box = tight_box
+            
+            # Save
+            layer_filename = f'layer_{current_layer_idx:02d}.png'
+            caption_layer.save(os.path.join(sample_output_dir, layer_filename))
+            
+            # Composite
+            composite = Image.alpha_composite(composite, caption_layer)
+            
+            # Record
+            w, h = get_box_size(caption_box)
+            new_layers.append({
+                'layer_idx': current_layer_idx,
+                'caption': f"Text: {caption_text}" if caption_text else "Text",
+                'box': caption_box,
+                'width_dst': w,
+                'height_dst': h,
+                'image_path': layer_filename,
+                'type': 'caption',
+                'original_size': list(caption_orig_size),
+            })
+            
+            occupied_boxes.append(caption_box)
+            current_layer_idx += 1
+    
+    # === Step 5: Save whole_image ===
+    composite.save(os.path.join(sample_output_dir, 'whole_image.png'))
+    
+    # === Step 6: Build spatial-aware whole caption ===
+    base_caption = base_meta.get('base_caption', '')
+    whole_caption = build_spatial_aware_caption(new_layers, canvas_size, base_caption)
+    
+    # === Step 7: Create metadata ===
+    metadata = {
+        'id': f'{sample_idx:09d}',
+        'style_category': base_meta.get('style_category', ''),
+        'whole_caption': whole_caption,
+        'base_caption': base_caption,
+        'layer_count': len(new_layers),
+        'layers': new_layers,
+        # Extra fields
+        'sample_dir': sample_name,
+        'width': canvas_size,
+        'height': canvas_size,
+        'base_sample': base_sample_dir,
+        'num_base_layers_removed': len(removed_layer_indices),
+        'num_donor_samples': num_donors,
+        'num_donated_layers': len(donor_layers),
+    }
+    
+    if laion_path:
+        metadata['laion_path'] = laion_path
+        metadata['laion_caption'] = laion_caption
+    if caption_path:
+        metadata['caption_path'] = caption_path
+        metadata['caption_text'] = caption_text
+    
+    # Save metadata
+    with open(os.path.join(sample_output_dir, 'metadata.json'), 'w', encoding='utf-8') as f:
+        json.dump(metadata, f, indent=2, ensure_ascii=False)
+    
+    return metadata
+
+
+def _worker_fn(task):
+    """Worker function for multiprocessing. Each task is a dict with all needed info."""
+    sample_idx = task['sample_idx']
+    base_sample_dir = task['base_sample_dir']
+    all_sample_dirs = task['all_sample_dirs']
+    laion_images = task['laion_images']
+    caption_images = task['caption_images']
+    alphavae_images = task['alphavae_images']
+    output_dir = task['output_dir']
+    args = task['args']
+    seed = task['seed']
+
+    if args.skip_existing:
+        meta_path = os.path.join(output_dir, f"sample_{sample_idx:06d}", 'metadata.json')
+        if os.path.exists(meta_path):
+            try:
+                with open(meta_path, 'r') as f:
+                    return json.load(f)
+            except Exception:
+                pass
+
+    random.seed(seed)
+
+    try:
+        metadata = create_scaled_up_sample(
+            base_sample_dir=base_sample_dir,
+            all_sample_dirs=all_sample_dirs,
+            laion_images=laion_images,
+            caption_images=caption_images,
+            alphavae_images=alphavae_images,
+            output_dir=output_dir,
+            sample_idx=sample_idx,
+            args=args,
+        )
+        return metadata
+    except Exception as e:
+        logger.error(f"Error processing sample {sample_idx}: {e}")
+        return None
+
+
+def main():
+    args = parse_args()
+    random.seed(args.seed)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Load existing blended samples
+    logger.info("Loading existing blended samples...")
+    all_sample_dirs = get_blended_sample_dirs(args.blended_dir, max_samples=args.max_base_samples)
+    logger.info(f"Found {len(all_sample_dirs)} existing samples"
+                + (f" (limited to first {args.max_base_samples})" if args.max_base_samples else ""))
+
+    if len(all_sample_dirs) < 10:
+        logger.error("Not enough existing samples to create scaled-up dataset!")
+        return
+
+    # Load caption list
+    logger.info("Loading caption list from captions.jsonl...")
+    caption_list = load_caption_list(args.caption_meta)
+    logger.info(f"Loaded {len(caption_list)} caption entries")
+
+    # Load LAION images (cap at 20000 for balanced diversity)
+    logger.info("Loading LAION images with captions...")
+    laion_images = get_laion_images_with_captions(args.laion_dir)
+    if len(laion_images) > 20000:
+        random.shuffle(laion_images)
+        laion_images = laion_images[:20000]
+    logger.info(f"Using {len(laion_images)} LAION images")
+
+    # Load caption images
+    logger.info("Loading caption images...")
+    caption_images = get_caption_images_with_text(args.caption_dir, caption_list)
+    logger.info(f"Found {len(caption_images)} caption images")
+
+    # Load AlphaVAE images (optional)
+    alphavae_images = []  # list of (image_path, caption)
+    if args.alphavae_dir and args.alphavae_prompts:
+        logger.info("Loading AlphaVAE images with prompts...")
+        with open(args.alphavae_prompts, 'r') as f:
+            alphavae_prompts = [l.strip() for l in f.readlines() if l.strip()]
+        alpha_files = sorted([
+            f for f in os.listdir(args.alphavae_dir)
+            if f.endswith('.png')
+        ])
+        for fname in alpha_files:
+            idx = int(fname.replace('.png', ''))
+            prompt_idx = idx // 5
+            if prompt_idx < len(alphavae_prompts):
+                caption = alphavae_prompts[prompt_idx]
+            else:
+                caption = ""
+            alphavae_images.append((os.path.join(args.alphavae_dir, fname), caption))
+        logger.info(f"Found {len(alphavae_images)} AlphaVAE images")
+
+    # Pre-generate tasks with deterministic per-sample seeds
+    rng = random.Random(args.seed)
+    tasks = []
+    for i in range(args.num_samples):
+        sample_idx = args.start_index + i
+        tasks.append({
+            'sample_idx': sample_idx,
+            'base_sample_dir': rng.choice(all_sample_dirs),
+            'all_sample_dirs': all_sample_dirs,
+            'laion_images': laion_images,
+            'caption_images': caption_images,
+            'alphavae_images': alphavae_images,
+            'output_dir': args.output_dir,
+            'args': args,
+            'seed': rng.randint(0, 2**31),
+        })
+
+    # Generate samples
+    all_metadata = []
+    failed_count = 0
+    num_workers = args.num_workers
+
+    if num_workers > 0:
+        logger.info(f"Using multiprocessing with {num_workers} workers")
+        with Pool(processes=num_workers) as pool:
+            for metadata in tqdm(
+                pool.imap_unordered(_worker_fn, tasks),
+                total=len(tasks),
+                desc="Generating samples"
+            ):
+                if metadata:
+                    all_metadata.append(metadata)
+                else:
+                    failed_count += 1
+    else:
+        logger.info("Using single-process mode")
+        for task in tqdm(tasks, desc="Generating samples"):
+            metadata = _worker_fn(task)
+            if metadata:
+                all_metadata.append(metadata)
+            else:
+                failed_count += 1
+
+    # Sort by sample index for deterministic output order
+    all_metadata.sort(key=lambda m: int(m['id']))
+
+    # Save index (scaleup_meta.jsonl)
+    index_path = os.path.join(args.output_dir, 'scaleup_meta.jsonl')
+    with open(index_path, 'w', encoding='utf-8') as f:
+        for meta in all_metadata:
+            f.write(json.dumps(meta, ensure_ascii=False) + '\n')
+
+    logger.info(f"Generated {len(all_metadata)} samples ({failed_count} failed)")
+    logger.info(f"Output saved to {args.output_dir}")
+    logger.info(f"Index saved to {index_path}")
+
+
+if __name__ == '__main__':
+    main()
+