#!/usr/bin/env python3 import json import os import sys import copy import argparse from pathlib import Path 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 ( generate_base_scene, find_blender, create_patched_utils, create_patched_render_script, create_render_from_json_script, save_scene, render_scene, IMAGE_COUNTERFACTUALS, NEGATIVE_COUNTERFACTUALS, COUNTERFACTUAL_TYPES ) def generate_all_counterfactual_examples(output_dir='output/counterfactual_examples', num_objects=5, render=False, use_gpu=0): print("="*70) print("GENERATING COUNTERFACTUAL EXAMPLES") print("="*70) script_dir = os.path.dirname(os.path.abspath(__file__)) project_root = os.path.dirname(script_dir) if not os.path.isabs(output_dir): output_dir = os.path.join(project_root, output_dir) os.makedirs(output_dir, exist_ok=True) scenes_dir = os.path.join(output_dir, 'scenes') images_dir = os.path.join(output_dir, 'images') if render else None if render: os.makedirs(images_dir, exist_ok=True) os.makedirs(scenes_dir, exist_ok=True) blender_path = find_blender() create_patched_utils() create_patched_render_script() if render: create_render_from_json_script() print(f"\n1. Generating base scene with {num_objects} objects...") base_scene = None for retry in range(3): base_scene = generate_base_scene(num_objects, blender_path, 0) if base_scene and len(base_scene.get('objects', [])) > 0: break print(f" Retry {retry + 1}/3...") if not base_scene or len(base_scene.get('objects', [])) == 0: print("ERROR: Failed to generate base scene") return original_path = os.path.join(scenes_dir, '00_original.json') save_scene(base_scene, original_path) print(f" [OK] Saved original scene: {original_path}") if render: original_image = os.path.join(images_dir, '00_original.png') print(f" Rendering original scene...") render_scene(blender_path, original_path, original_image, use_gpu=use_gpu) if os.path.exists(original_image): print(f" [OK] Rendered: 00_original.png") all_cf_types = {**IMAGE_COUNTERFACTUALS, **NEGATIVE_COUNTERFACTUALS} total_cfs = len(all_cf_types) print(f"\n2. Generating {total_cfs} counterfactual examples...") image_cf_index = 1 for cf_type, cf_func in sorted(IMAGE_COUNTERFACTUALS.items()): print(f"\n [Image CF {image_cf_index}/{len(IMAGE_COUNTERFACTUALS)}] Applying: {cf_type}") try: cf_scene, description = cf_func(copy.deepcopy(base_scene)) cf_scene['cf_metadata'] = { 'variant': f'image_cf_{image_cf_index}', 'is_counterfactual': True, 'cf_index': image_cf_index, 'cf_category': 'image_cf', 'cf_type': cf_type, 'cf_description': description, 'source_scene': '00_original' } filename = f"01_image_{image_cf_index:02d}_{cf_type}.json" cf_path = os.path.join(scenes_dir, filename) save_scene(cf_scene, cf_path) print(f" [OK] Saved: {filename}") print(f" Description: {description}") if render: cf_image = os.path.join(images_dir, filename.replace('.json', '.png')) render_scene(blender_path, cf_path, cf_image, use_gpu=use_gpu) if os.path.exists(cf_image): print(f" [OK] Rendered: {filename.replace('.json', '.png')}") image_cf_index += 1 except Exception as e: print(f" [ERROR] ERROR applying {cf_type}: {e}") import traceback traceback.print_exc() continue negative_cf_index = 1 for cf_type, cf_func in sorted(NEGATIVE_COUNTERFACTUALS.items()): print(f"\n [Negative CF {negative_cf_index}/{len(NEGATIVE_COUNTERFACTUALS)}] Applying: {cf_type}") try: if cf_type == 'add_noise': cf_scene, description = cf_func(copy.deepcopy(base_scene), min_noise_level='medium') else: cf_scene, description = cf_func(copy.deepcopy(base_scene)) cf_scene['cf_metadata'] = { 'variant': f'negative_cf_{negative_cf_index}', 'is_counterfactual': True, 'cf_index': negative_cf_index, 'cf_category': 'negative_cf', 'cf_type': cf_type, 'cf_description': description, 'source_scene': '00_original' } filename = f"02_negative_{negative_cf_index:02d}_{cf_type}.json" cf_path = os.path.join(scenes_dir, filename) save_scene(cf_scene, cf_path) print(f" [OK] Saved: {filename}") print(f" Description: {description}") if render: cf_image = os.path.join(images_dir, filename.replace('.json', '.png')) render_scene(blender_path, cf_path, cf_image, use_gpu=use_gpu) if os.path.exists(cf_image): print(f" [OK] Rendered: {filename.replace('.json', '.png')}") negative_cf_index += 1 except Exception as e: print(f" [ERROR] ERROR applying {cf_type}: {e}") import traceback traceback.print_exc() continue summary = { 'base_scene': '00_original.json', 'total_counterfactuals': len(all_cf_types), 'image_counterfactuals': len(IMAGE_COUNTERFACTUALS), 'negative_counterfactuals': len(NEGATIVE_COUNTERFACTUALS), 'counterfactuals': {} } image_cf_index = 1 for cf_type in sorted(IMAGE_COUNTERFACTUALS.keys()): filename = f"01_image_{image_cf_index:02d}_{cf_type}.json" summary['counterfactuals'][cf_type] = { 'filename': filename, 'category': 'image_cf', 'index': image_cf_index } image_cf_index += 1 negative_cf_index = 1 for cf_type in sorted(NEGATIVE_COUNTERFACTUALS.keys()): filename = f"02_negative_{negative_cf_index:02d}_{cf_type}.json" summary['counterfactuals'][cf_type] = { 'filename': filename, 'category': 'negative_cf', 'index': negative_cf_index } negative_cf_index += 1 summary_path = os.path.join(output_dir, 'summary.json') with open(summary_path, 'w') as f: json.dump(summary, f, indent=2) print("\n" + "="*70) print("SUMMARY") print("="*70) print(f"Base scene: {original_path}") print(f"Total counterfactuals generated: {len(all_cf_types)}") print(f" - Image CFs: {len(IMAGE_COUNTERFACTUALS)}") print(f" - Negative CFs: {len(NEGATIVE_COUNTERFACTUALS)}") print(f"\nAll files saved to: {output_dir}") print(f"Summary saved to: {summary_path}") if render: print(f"Images saved to: {images_dir}") print("="*70) if __name__ == '__main__': parser = argparse.ArgumentParser(description='Generate examples of each counterfactual type') parser.add_argument('--output_dir', type=str, default='output/counterfactual_examples', help='Output directory for examples (default: output/counterfactual_examples)') parser.add_argument('--num_objects', type=int, default=5, help='Number of objects in base scene (default: 5)') parser.add_argument('--render', action='store_true', help='Render scenes to images') parser.add_argument('--use_gpu', type=int, default=0, help='Use GPU rendering (0 = CPU, 1 = GPU, default: 0)') args = parser.parse_args() generate_all_counterfactual_examples(args.output_dir, args.num_objects, args.render, args.use_gpu)