scripts/generate_scenes.py

#!/usr/bin/env python3

import sys
import os

script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(script_dir)
sys.path.insert(0, project_root)

from pipeline import (
    find_blender,
    create_patched_render_script,
    create_run_directory,
    generate_base_scene,
    generate_counterfactuals,
    save_scene,
    save_checkpoint,
    load_checkpoint,
    get_completed_scenes_from_folder,
    COUNTERFACTUAL_TYPES,
    list_counterfactual_types
)
import argparse
import random
import json
import shutil
from datetime import datetime

def main():
    parser = argparse.ArgumentParser(
        description='Generate scene JSON files with counterfactuals (no rendering)'
    )
    
    parser.add_argument('--num_scenes', type=int, default=5,
                       help='Number of scene sets to generate')
    parser.add_argument('--num_objects', type=int, default=None,
                       help='Fixed number of objects per scene (overrides min/max)')
    parser.add_argument('--min_objects', type=int, default=3,
                       help='Minimum object count (if num_objects not given)')
    parser.add_argument('--max_objects', type=int, default=7,
                       help='Maximum object count (if num_objects not given)')
    parser.add_argument('--num_counterfactuals', type=int, default=2,
                       help='Number of counterfactual variants per scene')
    parser.add_argument('--blender_path', type=str, default=None,
                       help='Path to Blender executable (auto-detected if not provided)')
    parser.add_argument('--output_dir', type=str, default='output',
                       help='Base directory for all runs')
    parser.add_argument('--run_name', type=str, default=None,
                       help='Optional custom name for this run')
    parser.add_argument('--resume', action='store_true',
                       help='Resume from last checkpoint (requires --run_name)')
    parser.add_argument('--cf_types', nargs='+', 
                   choices=[
                       'change_color', 'change_shape', 'change_size', 
                       'change_material', 'change_position', 
                       'add_object', 'remove_object', 'swap_attribute', 'occlusion_change', 'relational_flip',
                       'replace_object',
                       'change_background', 
                       'change_lighting', 'add_noise',
                       'apply_fisheye', 'apply_blur', 'apply_vignette', 'apply_chromatic_aberration'
                   ],
                   help='Specific counterfactual types to use')
    parser.add_argument('--semantic_only', action='store_true',
                   help='Generate only Semantic/Image counterfactuals (no Negative CFs)')
    parser.add_argument('--negative_only', action='store_true',
                   help='Generate only Negative counterfactuals (no Semantic CFs)')
    parser.add_argument('--same_cf_type', action='store_true',
                   help='Use the same counterfactual type for all variants')
    parser.add_argument('--min_cf_change_score', type=float, default=1.0,
                   help='Minimum heuristic change score for counterfactuals')
    parser.add_argument('--max_cf_attempts', type=int, default=10,
                   help='Max retries per counterfactual to meet --min_cf_change_score')
    parser.add_argument('--min_noise_level', type=str, default='light',
                   choices=['light', 'medium', 'heavy'],
                   help='Minimum noise level when using add_noise counterfactual')
    parser.add_argument('--list_cf_types', action='store_true',
                   help='List all available counterfactual types and exit')
    
    args = parser.parse_args()
    
    if args.list_cf_types:
        list_counterfactual_types()
        return
    
    if args.resume and not args.run_name:
        print("ERROR: --run_name is required when using --resume")
        return
    
    blender_path = args.blender_path or find_blender()
    print(f"Using Blender: {blender_path}")
    
    print("\nPreparing scripts...")
    create_patched_render_script()
    
    run_dir = create_run_directory(args.output_dir, args.run_name)
    temp_run_id = os.path.basename(run_dir)
    print(f"\n{'='*70}")
    print(f"RUN DIRECTORY: {run_dir}")
    print(f"{'='*70}")

    scenes_dir = os.path.join(run_dir, 'scenes')
    os.makedirs(scenes_dir, exist_ok=True)
    
    checkpoint_file = os.path.join(run_dir, 'checkpoint.json')
    
    completed_scenes = set()
    if args.resume:
        completed_scenes = load_checkpoint(checkpoint_file)
        existing_scenes = get_completed_scenes_from_folder(scenes_dir)
        completed_scenes.update(existing_scenes)
        
        if completed_scenes:
            print(f"\n[RESUME] Found {len(completed_scenes)} already completed scenes")
        else:
            print("\n[WARNING] Resume flag set but no checkpoint found, starting fresh")
    
    print("\n" + "="*70)
    print(f"GENERATING {args.num_scenes} SCENE SETS (JSON ONLY)")
    print(f"Each with {args.num_counterfactuals} counterfactual variants")
    print("="*70)
    
    successful_scenes = 0
    
    for i in range(args.num_scenes):
        if i in completed_scenes:
            print(f"\n[SKIP] Skipping scene {i} (already completed)")
            successful_scenes += 1
            continue
        
        print(f"\n{'='*70}")
        print(f"SCENE SET {i+1}/{args.num_scenes} (Scene #{i})")
        print(f"{'='*70}")
        
        if args.num_objects is not None:
            num_objects = args.num_objects
        else:
            num_objects = random.randint(args.min_objects, args.max_objects)
        
        base_scene = None
        for retry in range(3):
            base_scene = generate_base_scene(num_objects, blender_path, i, temp_run_dir=temp_run_id)
            if base_scene and len(base_scene['objects']) > 0:
                break
            print(f"  Retry {retry + 1}/3...")
        
        if not base_scene or len(base_scene['objects']) == 0:
            print(f"  [FAILED] Failed to generate scene {i+1}")
            continue
        
        successful_scenes += 1
        
        print(f"  Creating {args.num_counterfactuals} counterfactuals...")
        counterfactuals = generate_counterfactuals(
            base_scene,
            args.num_counterfactuals,
            cf_types=args.cf_types,
            same_cf_type=args.same_cf_type,
            min_change_score=args.min_cf_change_score,
            max_cf_attempts=args.max_cf_attempts,
            min_noise_level=args.min_noise_level,
            semantic_only=args.semantic_only,
            negative_only=args.negative_only
        )
        
        for idx, cf in enumerate(counterfactuals):
            cf_cat = cf.get('cf_category', 'unknown')
            print(f"    CF{idx+1} [{cf_cat}] ({cf['type']}): {cf['description']}")
        
        scene_num = i + 1
        scene_prefix = f"scene_{scene_num:04d}"
        scene_paths = {'original': os.path.join(scenes_dir, f"{scene_prefix}_original.json")}
        
        base_scene['cf_metadata'] = {
            'variant': 'original',
            'is_counterfactual': False,
            'cf_index': None,
            'cf_category': 'original',
            'cf_type': None,
            'cf_description': None,
            'source_scene': scene_prefix,
        }
        save_scene(base_scene, scene_paths['original'])
        
        for idx, cf in enumerate(counterfactuals):
            cf_name = f"cf{idx+1}"
            scene_paths[cf_name] = os.path.join(scenes_dir, f"{scene_prefix}_{cf_name}.json")
            
            cf_scene = cf['scene']
            cf_scene['cf_metadata'] = {
                'variant': cf_name,
                'is_counterfactual': True,
                'cf_index': idx + 1,
                'cf_category': cf.get('cf_category', 'unknown'),
                'cf_type': cf.get('type', None),
                'cf_description': cf.get('description', None),
                'change_score': cf.get('change_score', None),
                'change_attempts': cf.get('change_attempts', None),
                'source_scene': scene_prefix,
            }
            save_scene(cf_scene, scene_paths[cf_name])
        
        print(f"  [OK] Saved {len(counterfactuals) + 1} scene files")
        
        completed_scenes.add(i)
        save_checkpoint(checkpoint_file, list(completed_scenes))
    
    metadata = {
        'timestamp': datetime.now().isoformat(),
        'num_scenes': args.num_scenes,
        'num_counterfactuals': args.num_counterfactuals,
        'successful_scenes': successful_scenes,
        'successful_renders': 0,
        'cf_types': args.cf_types if args.cf_types else 'default',
        'semantic_only': args.semantic_only,
        'negative_only': args.negative_only,
    }
    
    metadata_path = os.path.join(run_dir, 'run_metadata.json')
    with open(metadata_path, 'w') as f:
        json.dump(metadata, f, indent=2)
    
    temp_run_path = os.path.join(os.getcwd(), 'temp_output', temp_run_id)
    if os.path.exists(temp_run_path):
        shutil.rmtree(temp_run_path)
    if os.path.exists('render_images_patched.py'):
        os.remove('render_images_patched.py')
    
    print("\n" + "="*70)
    print("SCENE COMPLETE")
    print("="*70)
    print(f"Run directory: {run_dir}")
    print(f"Successfully generated: {successful_scenes}/{args.num_scenes} scene sets")
    print(f"\nOutput:")
    print(f"  Scene files: {scenes_dir}/")
    print(f"  Metadata: {metadata_path}")
    print(f"  Checkpoint: {checkpoint_file}")
    print(f"\nNext step: Run 'python pipeline.py --render_only --run_name {args.run_name}' to render these scenes")
    print("="*70)

if __name__ == '__main__':
    main()