|
|
|
|
| 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()
|
|
|
|
|
