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