AnonymousECCV15285/MMIB_Counterfactual_image_generation_tool
0
1from __future__ import print_function
2import math, sys, random, argparse, json, os, tempfile
3from datetime import datetime as dt
4from collections import Counter
5try:
6 from PIL import Image, ImageFilter
7except ImportError:
8 Image = None
9 ImageFilter = None
10
11INSIDE_BLENDER = True
12try:
13 import bpy, bpy_extras
14 from mathutils import Vector
15except ImportError as e:
16 INSIDE_BLENDER = False
17
18def extract_args(input_argv=None):
19 if input_argv is None:
20 input_argv = sys.argv
21 output_argv = []
22 if '--' in input_argv:
23 idx = input_argv.index('--')
24 output_argv = input_argv[(idx + 1):]
25 return output_argv
26
27def parse_args(parser, argv=None):
28 return parser.parse_args(extract_args(argv))
29
30def delete_object(obj):
31 if not INSIDE_BLENDER:
32 return
33 bpy.ops.object.select_all(action='DESELECT')
34 obj.select_set(True)
35 bpy.context.view_layer.objects.active = obj
36 bpy.ops.object.delete()
37
38def get_camera_coords(cam, pos):
39 if not INSIDE_BLENDER:
40 return (0, 0, 0)
41 scene = bpy.context.scene
42 x, y, z = bpy_extras.object_utils.world_to_camera_view(scene, cam, pos)
43 scale = scene.render.resolution_percentage / 100.0
44 w = int(scale * scene.render.resolution_x)
45 h = int(scale * scene.render.resolution_y)
46 px = int(round(x * w))
47 py = int(round(h - y * h))
48 return (px, py, z)
49
50def set_layer(obj, layer_idx):
51 if not INSIDE_BLENDER:
52 return
53 obj.layers[layer_idx] = True
54 for i in range(len(obj.layers)):
55 obj.layers[i] = (i == layer_idx)
56
57def add_object(object_dir, name, scale, loc, theta=0):
58 if not INSIDE_BLENDER:
59 return
60 object_dir = os.path.abspath(os.path.normpath(object_dir))
61 if not os.path.exists(object_dir):
62 print(f"ERROR: Object directory does not exist: {object_dir}")
63 return
64 count = 0
65 for obj in bpy.data.objects:
66 if obj.name.startswith(name):
67 count += 1
68
69 blend_file = os.path.join(object_dir, '%s.blend' % name)
70 blend_file = os.path.abspath(blend_file).replace('\\', '/')
71 if not os.path.exists(blend_file):
72 print(f"ERROR: Blend file does not exist: {blend_file}")
73 return
74 directory = blend_file + '/Object/'
75 try:
76 bpy.ops.wm.append(
77 directory=directory,
78 filename=name,
79 filter_blender=True
80 )
81 except Exception as e:
82 error_msg = str(e)
83 print(f"ERROR: Failed to load object {name} from {directory}: {error_msg}")
84 print(f" Error type: {type(e).__name__}")
85 raise
86
87 new_name = '%s_%d' % (name, count)
88 bpy.data.objects[name].name = new_name
89
90 x, y = loc
91 bpy.context.view_layer.objects.active = bpy.data.objects[new_name]
92 bpy.context.object.rotation_euler[2] = theta
93 bpy.ops.transform.resize(value=(scale, scale, scale))
94 bpy.ops.transform.translate(value=(x, y, scale))
95
96def load_materials(material_dir):
97 if not INSIDE_BLENDER:
98 return
99 material_dir = os.path.abspath(os.path.normpath(material_dir))
100 if not os.path.exists(material_dir):
101 print(f"ERROR: Material directory does not exist: {material_dir}")
102 return
103 for fn in os.listdir(material_dir):
104 if not fn.endswith('.blend'): continue
105 name = os.path.splitext(fn)[0]
106 blend_file = os.path.join(material_dir, fn)
107 blend_file = os.path.abspath(blend_file).replace('\\', '/')
108 if not os.path.exists(blend_file):
109 print(f"ERROR: Blend file does not exist: {blend_file}")
110 continue
111 directory = blend_file + '/NodeTree/'
112 try:
113 bpy.ops.wm.append(
114 directory=directory,
115 filename=name,
116 filter_blender=True
117 )
118 except Exception as e:
119 error_msg = str(e)
120 print(f"ERROR: Failed to load material {name} from {directory}: {error_msg}")
121 print(f" Error type: {type(e).__name__}")
122 raise
123
124def apply_filter_to_image(image_path, filter_type, filter_strength):
125 image_path = os.path.abspath(image_path)
126 if Image is None:
127 print(f"ERROR: PIL/Image not available, cannot apply filter {filter_type}")
128 return
129 if not os.path.exists(image_path):
130 print(f"ERROR: Image file does not exist: {image_path}")
131 return
132
133 try:
134 img = Image.open(image_path)
135
136 if filter_type == 'blur':
137 radius = max(1, int(filter_strength))
138 img = img.filter(ImageFilter.GaussianBlur(radius=radius))
139
140 elif filter_type == 'vignette':
141 width, height = img.size
142 center_x, center_y = width // 2, height // 2
143 max_dist = math.sqrt(center_x**2 + center_y**2)
144
145 img = img.convert('RGB')
146 pixels = img.load()
147 for y in range(height):
148 for x in range(width):
149 dist = math.sqrt((x - center_x)**2 + (y - center_y)**2)
150 factor = 1.0 - (dist / max_dist) * (filter_strength / 5.0)
151 factor = max(0.0, min(1.0, factor))
152
153 r, g, b = pixels[x, y]
154 pixels[x, y] = (int(r * factor), int(g * factor), int(b * factor))
155
156 elif filter_type == 'fisheye':
157 width, height = img.size
158 center_x, center_y = width / 2.0, height / 2.0
159 max_radius = min(center_x, center_y)
160
161 img = img.convert('RGB')
162 output = Image.new('RGB', (width, height))
163 out_pixels = output.load()
164 in_pixels = img.load()
165
166 for y in range(height):
167 for x in range(width):
168 dx = (x - center_x) / max_radius
169 dy = (y - center_y) / max_radius
170 distance = math.sqrt(dx*dx + dy*dy)
171
172 if distance > 1.0:
173 out_pixels[x, y] = (0, 0, 0)
174 else:
175 theta = math.atan2(dy, dx)
176 r_normalized = distance
177 r_distorted = r_normalized * (1.0 + filter_strength * (1.0 - r_normalized))
178 r_distorted = min(1.0, r_distorted)
179
180 src_x = int(center_x + r_distorted * max_radius * math.cos(theta))
181 src_y = int(center_y + r_distorted * max_radius * math.sin(theta))
182
183 if 0 <= src_x < width and 0 <= src_y < height:
184 out_pixels[x, y] = in_pixels[src_x, src_y]
185 else:
186 out_pixels[x, y] = (0, 0, 0)
187
188 img = output
189
190 img.save(image_path)
191 print(f"[OK] Applied {filter_type} filter (strength: {filter_strength:.2f})")
192 except Exception as e:
193 import traceback
194 print(f"ERROR applying filter {filter_type}: {e}")
195 traceback.print_exc()
196 raise
197
198def add_material(name, **properties):
199 if not INSIDE_BLENDER:
200 return
201 mat_count = len(bpy.data.materials)
202 bpy.ops.material.new()
203 mat = bpy.data.materials['Material']
204 mat.name = 'Material_%d' % mat_count
205 obj = bpy.context.active_object
206 assert len(obj.data.materials) == 0
207 obj.data.materials.append(mat)
208
209 output_node = None
210 for n in mat.node_tree.nodes:
211 if n.name == 'Material Output':
212 output_node = n
213 break
214
215 group_node = mat.node_tree.nodes.new('ShaderNodeGroup')
216 group_node.node_tree = bpy.data.node_groups[name]
217
218 for inp in group_node.inputs:
219 if inp.name in properties:
220 inp.default_value = properties[inp.name]
221
222 mat.node_tree.links.new(
223 group_node.outputs['Shader'],
224 output_node.inputs['Surface'],
225 )
226
227parser = argparse.ArgumentParser()
228parser.add_argument('--scene_file', default=None,
229 help="Optional JSON file to load scene from. If provided, renders from JSON instead of generating random scenes.")
230parser.add_argument('--base_scene_blendfile', default='data/base_scene.blend',
231 help="Base blender file on which all scenes are based; includes " +
232 "ground plane, lights, and camera.")
233parser.add_argument('--properties_json', default='data/properties.json',
234 help="JSON file defining objects, materials, sizes, and colors. " +
235 "The \"colors\" field maps from CLEVR color names to RGB values; " +
236 "The \"sizes\" field maps from CLEVR size names to scalars used to " +
237 "rescale object models; the \"materials\" and \"shapes\" fields map " +
238 "from CLEVR material and shape names to .blend files in the " +
239 "--object_material_dir and --shape_dir directories respectively.")
240parser.add_argument('--shape_dir', default='data/shapes',
241 help="Directory where .blend files for object models are stored")
242parser.add_argument('--material_dir', default='data/materials',
243 help="Directory where .blend files for materials are stored")
244parser.add_argument('--shape_color_combos_json', default=None,
245 help="Optional path to a JSON file mapping shape names to a list of " +
246 "allowed color names for that shape. This allows rendering images " +
247 "for CLEVR-CoGenT.")
248
249parser.add_argument('--min_objects', default=3, type=int,
250 help="The minimum number of objects to place in each scene")
251parser.add_argument('--max_objects', default=10, type=int,
252 help="The maximum number of objects to place in each scene")
253parser.add_argument('--min_dist', default=0.15, type=float,
254 help="The minimum allowed distance between object centers")
255parser.add_argument('--margin', default=0.2, type=float,
256 help="Along all cardinal directions (left, right, front, back), all " +
257 "objects will be at least this distance apart. This makes resolving " +
258 "spatial relationships slightly less ambiguous.")
259parser.add_argument('--min_pixels_per_object', default=50, type=int,
260 help="All objects will have at least this many visible pixels in the " +
261 "final rendered images; this ensures that no objects are fully " +
262 "occluded by other objects.")
263parser.add_argument('--max_retries', default=100, type=int,
264 help="The number of times to try placing an object before giving up and " +
265 "re-placing all objects in the scene.")
266
267parser.add_argument('--start_idx', default=0, type=int,
268 help="The index at which to start for numbering rendered images. Setting " +
269 "this to non-zero values allows you to distribute rendering across " +
270 "multiple machines and recombine the results later.")
271parser.add_argument('--num_images', default=5, type=int,
272 help="The number of images to render")
273parser.add_argument('--filename_prefix', default='CLEVR',
274 help="This prefix will be prepended to the rendered images and JSON scenes")
275parser.add_argument('--split', default='new',
276 help="Name of the split for which we are rendering. This will be added to " +
277 "the names of rendered images, and will also be stored in the JSON " +
278 "scene structure for each image.")
279parser.add_argument('--output_image_dir', default='../output/images/',
280 help="The directory where output images will be stored. It will be " +
281 "created if it does not exist.")
282parser.add_argument('--output_scene_dir', default='../output/scenes/',
283 help="The directory where output JSON scene structures will be stored. " +
284 "It will be created if it does not exist.")
285parser.add_argument('--output_scene_file', default='../output/CLEVR_scenes.json',
286 help="Path to write a single JSON file containing all scene information")
287parser.add_argument('--output_blend_dir', default='output/blendfiles',
288 help="The directory where blender scene files will be stored, if the " +
289 "user requested that these files be saved using the " +
290 "--save_blendfiles flag; in this case it will be created if it does " +
291 "not already exist.")
292parser.add_argument('--save_blendfiles', type=int, default=0,
293 help="Setting --save_blendfiles 1 will cause the blender scene file for " +
294 "each generated image to be stored in the directory specified by " +
295 "the --output_blend_dir flag. These files are not saved by default " +
296 "because they take up ~5-10MB each.")
297parser.add_argument('--version', default='1.0',
298 help="String to store in the \"version\" field of the generated JSON file")
299parser.add_argument('--license',
300 default="Creative Commons Attribution (CC-BY 4.0)",
301 help="String to store in the \"license\" field of the generated JSON file")
302parser.add_argument('--date', default=dt.today().strftime("%m/%d/%Y"),
303 help="String to store in the \"date\" field of the generated JSON file; " +
304 "defaults to today's date")
305
306parser.add_argument('--use_gpu', default=0, type=int,
307 help="Setting --use_gpu 1 enables GPU-accelerated rendering using CUDA. " +
308 "You must have an NVIDIA GPU with the CUDA toolkit installed for " +
309 "to work.")
310parser.add_argument('--width', default=320, type=int,
311 help="The width (in pixels) for the rendered images")
312parser.add_argument('--height', default=240, type=int,
313 help="The height (in pixels) for the rendered images")
314parser.add_argument('--key_light_jitter', default=1.0, type=float,
315 help="The magnitude of random jitter to add to the key light position.")
316parser.add_argument('--fill_light_jitter', default=1.0, type=float,
317 help="The magnitude of random jitter to add to the fill light position.")
318parser.add_argument('--back_light_jitter', default=1.0, type=float,
319 help="The magnitude of random jitter to add to the back light position.")
320parser.add_argument('--camera_jitter', default=0.5, type=float,
321 help="The magnitude of random jitter to add to the camera position")
322parser.add_argument('--render_num_samples', default=512, type=int,
323 help="The number of samples to use when rendering. Larger values will " +
324 "result in nicer images but will cause rendering to take longer.")
325parser.add_argument('--render_min_bounces', default=8, type=int,
326 help="The minimum number of bounces to use for rendering.")
327parser.add_argument('--render_max_bounces', default=8, type=int,
328 help="The maximum number of bounces to use for rendering.")
329parser.add_argument('--render_tile_size', default=256, type=int,
330 help="The tile size to use for rendering. This should not affect the " +
331 "quality of the rendered image but may affect the speed; CPU-based " +
332 "rendering may achieve better performance using smaller tile sizes " +
333 "while larger tile sizes may be optimal for GPU-based rendering.")
334parser.add_argument('--output_image', default=None,
335 help="Output image path (used when rendering from JSON)")
336
337MIN_VISIBLE_FRACTION = 0.001
338MIN_VISIBLE_FRACTION_PARTIAL_OCCLUSION = 0.0005
339MIN_PIXELS_FLOOR = 50
340
341BASE_MIN_VISIBILITY_FRACTION = 0.9
342CF_OCCLUSION_MIN_VISIBILITY_FRACTION = 0.3
343CF_OCCLUSION_MAX_VISIBILITY_FRACTION = 0.5
344CF_OCCLUSION_HARD_MIN_FRACTION = 0.2
345
346
347def min_visible_pixels(width, height, fraction=MIN_VISIBLE_FRACTION, floor=MIN_PIXELS_FLOOR):
348 return max(floor, int(width * height * fraction))
349
350
351BACKGROUND_COLORS = {
352 'default': None,
353 'gray': (0.5, 0.5, 0.5),
354 'blue': (0.2, 0.4, 0.8),
355 'green': (0.2, 0.6, 0.3),
356 'brown': (0.4, 0.3, 0.2),
357 'purple': (0.5, 0.3, 0.6),
358 'orange': (0.8, 0.5, 0.2),
359 'white': (0.9, 0.9, 0.9),
360 'dark_gray': (0.2, 0.2, 0.2),
361 'red': (0.7, 0.2, 0.2),
362 'yellow': (0.8, 0.8, 0.3),
363 'cyan': (0.3, 0.7, 0.8),
364}
365
366LIGHTING_PRESETS = {
367 'default': {'key': 1.0, 'fill': 0.5, 'back': 0.3},
368 'bright': {'key': 12.0, 'fill': 6.0, 'back': 4.0},
369 'dim': {'key': 0.008, 'fill': 0.004, 'back': 0.002},
370 'warm': {'key': 5.0, 'fill': 0.8, 'back': 0.3, 'color': (1.0, 0.5, 0.2)},
371 'cool': {'key': 4.0, 'fill': 2.0, 'back': 1.5, 'color': (0.2, 0.5, 1.0)},
372 'dramatic': {'key': 15.0, 'fill': 0.005, 'back': 0.002},
373}
374
375def set_background_color(color_name):
376 """Set the world background color"""
377 if not INSIDE_BLENDER:
378 return
379 if color_name not in BACKGROUND_COLORS or BACKGROUND_COLORS[color_name] is None:
380 return
381
382 rgb = BACKGROUND_COLORS[color_name]
383
384 world = bpy.context.scene.world
385 if world is None:
386 world = bpy.data.worlds.new("World")
387 bpy.context.scene.world = world
388
389 world.use_nodes = True
390 nodes = world.node_tree.nodes
391
392 bg_node = None
393 for node in nodes:
394 if node.type == 'BACKGROUND':
395 bg_node = node
396 break
397
398 if bg_node is None:
399 bg_node = nodes.new(type='ShaderNodeBackground')
400
401 bg_node.inputs['Color'].default_value = (rgb[0], rgb[1], rgb[2], 1.0)
402 bg_node.inputs['Strength'].default_value = 1.0
403
404 print(f"Set background color to {color_name}: RGB{rgb}")
405
406def set_ground_color(color_name):
407 if not INSIDE_BLENDER:
408 return
409 if color_name not in BACKGROUND_COLORS or BACKGROUND_COLORS[color_name] is None:
410 return
411
412 rgb = BACKGROUND_COLORS[color_name]
413
414 ground = None
415 for obj in bpy.data.objects:
416 if 'ground' in obj.name.lower() or 'plane' in obj.name.lower():
417 ground = obj
418 break
419
420 if ground is None:
421 return
422
423 if len(ground.data.materials) == 0:
424 mat = bpy.data.materials.new(name="Ground_Material")
425 ground.data.materials.append(mat)
426 else:
427 mat = ground.data.materials[0]
428
429 mat.use_nodes = True
430 nodes = mat.node_tree.nodes
431
432 bsdf = None
433 for node in nodes:
434 if node.type == 'BSDF_PRINCIPLED':
435 bsdf = node
436 break
437
438 if bsdf:
439 bsdf.inputs['Base Color'].default_value = (rgb[0], rgb[1], rgb[2], 1.0)
440 print(f"Set ground color to {color_name}")
441
442def set_lighting(lighting_name):
443 """Set lighting conditions"""
444 if not INSIDE_BLENDER:
445 return
446 if lighting_name not in LIGHTING_PRESETS:
447 return
448
449 preset = LIGHTING_PRESETS[lighting_name]
450
451 lamp_names = ['Lamp_Key', 'Lamp_Fill', 'Lamp_Back']
452 intensity_keys = ['key', 'fill', 'back']
453
454 for lamp_name, int_key in zip(lamp_names, intensity_keys):
455 if lamp_name in bpy.data.objects:
456 lamp_obj = bpy.data.objects[lamp_name]
457 if lamp_obj.data and hasattr(lamp_obj.data, 'energy'):
458 base_energy = lamp_obj.data.energy
459 lamp_obj.data.energy = base_energy * preset.get(int_key, 1.0)
460
461 if 'color' in preset and hasattr(lamp_obj.data, 'color'):
462 lamp_obj.data.color = preset['color']
463
464 print(f"Set lighting to {lighting_name}")
465
466def render_from_json(args):
467 if not INSIDE_BLENDER:
468 print("ERROR: render_from_json must be run inside Blender")
469 return
470
471 output_dir = os.path.dirname(args.output_image) if args.output_image else '.'
472 if output_dir and not os.path.exists(output_dir):
473 os.makedirs(output_dir)
474
475 with open(args.scene_file, 'r') as f:
476 scene_struct = json.load(f)
477
478 num_objects = len(scene_struct.get('objects', []))
479 print(f"Scene has {num_objects} objects")
480
481 base_scene_path = os.path.abspath(args.base_scene_blendfile)
482 bpy.ops.wm.open_mainfile(filepath=base_scene_path)
483
484 try:
485 load_materials(args.material_dir)
486 except Exception as e:
487 print(f"Warning: Could not load materials: {e}")
488
489 background_color = scene_struct.get('background_color', None)
490 if background_color:
491 set_background_color(background_color)
492 set_ground_color(background_color)
493
494 lighting = scene_struct.get('lighting', None)
495 if lighting:
496 set_lighting(lighting)
497
498 render_args = bpy.context.scene.render
499 render_args.engine = "CYCLES"
500 render_args.filepath = args.output_image
501 render_args.resolution_x = args.width
502 render_args.resolution_y = args.height
503 render_args.resolution_percentage = 100
504
505 if args.use_gpu == 1:
506 try:
507 bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
508 bpy.context.scene.cycles.device = 'GPU'
509 print("[OK] GPU rendering enabled")
510 except Exception as e:
511 print(f"Warning: Could not enable GPU: {e}")
512
513 bpy.context.scene.cycles.samples = args.render_num_samples
514
515 filter_type = scene_struct.get('filter_type')
516 filter_strength = scene_struct.get('filter_strength', 1.0)
517
518 if filter_type == 'fisheye':
519 camera = bpy.data.objects.get('Camera')
520 if camera and camera.data:
521 cam_data = camera.data
522 if cam_data.type == 'PERSP':
523 cam_data.lens = cam_data.lens * 0.7
524 print(f"[OK] Zoomed out camera for fisheye: lens={cam_data.lens:.1f}mm")
525
526 with open(args.properties_json, 'r') as f:
527 properties = json.load(f)
528 color_name_to_rgba = {}
529 for name, rgb in properties['colors'].items():
530 rgba = [float(c) / 255.0 for c in rgb] + [1.0]
531 color_name_to_rgba[name] = rgba
532 size_mapping = properties['sizes']
533
534 shape_semantic_to_file = properties['shapes']
535 material_semantic_to_file = properties['materials']
536
537 blender_objects = []
538 print("Adding objects to scene...")
539 for i, obj_info in enumerate(scene_struct.get('objects', [])):
540 x, y, z = obj_info['3d_coords']
541 r = size_mapping[obj_info['size']]
542 semantic_shape = obj_info['shape']
543
544 if semantic_shape == 'cube':
545 r /= math.sqrt(2)
546
547 if semantic_shape not in shape_semantic_to_file:
548 print(f"ERROR: Shape '{semantic_shape}' not found")
549 continue
550
551 shape_file_name = shape_semantic_to_file[semantic_shape]
552
553 try:
554 add_object(args.shape_dir, shape_file_name, r, (x, y), theta=obj_info['rotation'])
555 except Exception as e:
556 print(f"Error adding object {i}: {e}")
557 continue
558 if INSIDE_BLENDER and bpy.context.object:
559 blender_objects.append(bpy.context.object)
560
561 rgba = color_name_to_rgba[obj_info['color']]
562 semantic_material = obj_info['material']
563
564 if semantic_material not in material_semantic_to_file:
565 print(f"ERROR: Material '{semantic_material}' not found")
566 continue
567
568 mat_file_name = material_semantic_to_file[semantic_material]
569
570 try:
571 add_material(mat_file_name, Color=rgba)
572 except Exception as e:
573 print(f"Warning: Could not add material: {e}")
574
575 if blender_objects:
576 cf_meta = scene_struct.get('cf_metadata') or {}
577 cf_type = cf_meta.get('cf_type', '')
578
579 visibility_info = None
580 if INSIDE_BLENDER and Image is not None:
581 try:
582 visibility_info = compute_visibility_fractions(blender_objects)
583 except Exception as e:
584 print(f"Warning: compute_visibility_fractions failed during render: {e}")
585 visibility_info = None
586
587 all_visible = True
588 fail_reason = 'unknown visibility failure'
589
590 if visibility_info is not None:
591 ratios, scene_counts, full_counts = visibility_info
592
593 if cf_type == 'occlusion_change':
594 too_hidden = [
595 (i, r) for i, r in enumerate(ratios)
596 if full_counts[i] > 0 and r < CF_OCCLUSION_HARD_MIN_FRACTION
597 ]
598 band_objects = [
599 (i, r) for i, r in enumerate(ratios)
600 if full_counts[i] > 0 and CF_OCCLUSION_MIN_VISIBILITY_FRACTION <= r <= CF_OCCLUSION_MAX_VISIBILITY_FRACTION
601 ]
602
603 if too_hidden:
604 all_visible = False
605 min_r = min(r for (_, r) in too_hidden)
606 fail_reason = (f'at least one object is too occluded in occlusion_change CF; '
607 f'min visibility fraction={min_r:.3f} '
608 f'(required >= {CF_OCCLUSION_HARD_MIN_FRACTION})')
609 elif not band_objects:
610 all_visible = False
611 fail_reason = (f'no object falls into required occlusion band '
612 f'[{CF_OCCLUSION_MIN_VISIBILITY_FRACTION}, '
613 f'{CF_OCCLUSION_MAX_VISIBILITY_FRACTION}]')
614 else:
615 all_visible = True
616
617 else:
618 too_occluded = [
619 (i, r) for i, r in enumerate(ratios)
620 if full_counts[i] > 0 and r < BASE_MIN_VISIBILITY_FRACTION
621 ]
622 if too_occluded:
623 all_visible = False
624 min_r = min(r for (_, r) in too_occluded)
625 fail_reason = (f'at least one object is too occluded in base scene; '
626 f'min visibility fraction={min_r:.3f} '
627 f'(required >= {BASE_MIN_VISIBILITY_FRACTION})')
628 else:
629 all_visible = True
630
631 else:
632 # Fallback to legacy absolute pixel-based visibility when we cannot
633 # compute per-object relative visibility (e.g., PIL not available).
634 w = getattr(args, 'width', 320)
635 h = getattr(args, 'height', 240)
636 if cf_type == 'occlusion_change':
637 min_pixels = min_visible_pixels(w, h, MIN_VISIBLE_FRACTION_PARTIAL_OCCLUSION, MIN_PIXELS_FLOOR)
638 else:
639 base = min_visible_pixels(w, h, MIN_VISIBLE_FRACTION, MIN_PIXELS_FLOOR)
640 min_pixels = max(getattr(args, 'min_pixels_per_object', MIN_PIXELS_FLOOR), base)
641 all_visible = check_visibility(blender_objects, min_pixels)
642 if not all_visible:
643 fail_reason = 'at least one object has too few visible pixels'
644
645 if not all_visible:
646 print(f'Visibility check failed: {fail_reason}')
647 for obj in blender_objects:
648 try:
649 delete_object(obj)
650 except Exception:
651 pass
652 sys.exit(1)
653
654 filter_type = scene_struct.get('filter_type')
655 filter_strength = scene_struct.get('filter_strength', 1.0)
656
657 print(f"Rendering to {args.output_image}...")
658
659 try:
660 bpy.ops.render.render(write_still=True)
661 print("[OK] Rendering complete!")
662 except Exception as e:
663 print(f"Error during rendering: {e}")
664 sys.exit(1)
665
666 post_filter_type = scene_struct.get('filter_type')
667 if post_filter_type and post_filter_type != 'fisheye':
668 if Image is None:
669 print(f"Warning: PIL not available, cannot apply post-filter {post_filter_type}")
670 elif not os.path.exists(args.output_image):
671 print(f"Warning: Output image does not exist: {args.output_image}")
672 else:
673 try:
674 post_filter_strength = scene_struct.get('filter_strength', 1.0)
675 apply_filter_to_image(args.output_image, post_filter_type, post_filter_strength)
676 except Exception as e:
677 import traceback
678 print(f"Warning: Failed to apply post-filter {post_filter_type}: {e}")
679 traceback.print_exc()
680
681def main(args):
682 if args.scene_file:
683 render_from_json(args)
684 return
685
686 num_digits = 6
687 prefix = '%s_%s_' % (args.filename_prefix, args.split)
688 img_template = '%s%%0%dd.png' % (prefix, num_digits)
689 scene_template = '%s%%0%dd.json' % (prefix, num_digits)
690 blend_template = '%s%%0%dd.blend' % (prefix, num_digits)
691 img_template = os.path.join(args.output_image_dir, img_template)
692 scene_template = os.path.join(args.output_scene_dir, scene_template)
693 blend_template = os.path.join(args.output_blend_dir, blend_template)
694
695 if not os.path.isdir(args.output_image_dir):
696 os.makedirs(args.output_image_dir)
697 if not os.path.isdir(args.output_scene_dir):
698 os.makedirs(args.output_scene_dir)
699 if args.save_blendfiles == 1 and not os.path.isdir(args.output_blend_dir):
700 os.makedirs(args.output_blend_dir)
701
702 all_scene_paths = []
703 for i in range(args.num_images):
704 img_path = img_template % (i + args.start_idx)
705 scene_path = scene_template % (i + args.start_idx)
706 all_scene_paths.append(scene_path)
707 blend_path = None
708 if args.save_blendfiles == 1:
709 blend_path = blend_template % (i + args.start_idx)
710 num_objects = random.randint(args.min_objects, args.max_objects)
711 render_scene(args,
712 num_objects=num_objects,
713 output_index=(i + args.start_idx),
714 output_split=args.split,
715 output_image=img_path,
716 output_scene=scene_path,
717 output_blendfile=blend_path,
718 )
719
720 all_scenes = []
721 for scene_path in all_scene_paths:
722 with open(scene_path, 'r') as f:
723 all_scenes.append(json.load(f))
724 output = {
725 'info': {
726 'date': args.date,
727 'version': args.version,
728 'split': args.split,
729 'license': args.license,
730 },
731 'scenes': all_scenes
732 }
733 if args.output_scene_file:
734 output_dir = os.path.dirname(os.path.abspath(args.output_scene_file))
735 if output_dir:
736 os.makedirs(output_dir, exist_ok=True)
737 with open(args.output_scene_file, 'w') as f:
738 json.dump(output, f)
739
740
741
742def render_scene(args,
743 num_objects=5,
744 output_index=0,
745 output_split='none',
746 output_image='render.png',
747 output_scene='render_json',
748 output_blendfile=None,
749 ):
750
751 base_scene_path = os.path.abspath(args.base_scene_blendfile)
752 bpy.ops.wm.open_mainfile(filepath=base_scene_path)
753 load_materials(args.material_dir)
754
755 render_args = bpy.context.scene.render
756 render_args.engine = "CYCLES"
757 render_args.filepath = output_image
758 render_args.resolution_x = args.width
759 render_args.resolution_y = args.height
760 render_args.resolution_percentage = 100
761 if args.use_gpu == 1:
762 bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
763 bpy.context.preferences.addons['cycles'].preferences.get_devices()
764 for device in bpy.context.preferences.addons['cycles'].preferences.devices:
765 device.use = True
766
767 bpy.data.worlds['World'].cycles.sample_as_light = True
768 bpy.context.scene.cycles.blur_glossy = 2.0
769 bpy.context.scene.cycles.samples = args.render_num_samples
770 bpy.context.scene.cycles.transparent_min_bounces = args.render_min_bounces
771 bpy.context.scene.cycles.transparent_max_bounces = args.render_max_bounces
772 if args.use_gpu == 1:
773 bpy.context.scene.cycles.device = 'GPU'
774
775 scene_struct = {
776 'split': output_split,
777 'image_index': output_index,
778 'image_filename': os.path.basename(output_image),
779 'objects': [],
780 'directions': {},
781 }
782
783 bpy.ops.mesh.primitive_plane_add(size=10, location=(0, 0, 0))
784 plane = bpy.context.object
785
786 def rand(L):
787 return 2.0 * L * (random.random() - 0.5)
788
789 if args.camera_jitter > 0:
790 for i in range(3):
791 bpy.data.objects['Camera'].location[i] += rand(args.camera_jitter)
792
793 camera = bpy.data.objects['Camera']
794 plane_normal = plane.data.vertices[0].normal
795 cam_behind = camera.matrix_world.to_quaternion() @ Vector((0, 0, -1))
796 cam_left = camera.matrix_world.to_quaternion() @ Vector((-1, 0, 0))
797 cam_up = camera.matrix_world.to_quaternion() @ Vector((0, 1, 0))
798 plane_behind = (cam_behind - cam_behind.project(plane_normal)).normalized()
799 plane_left = (cam_left - cam_left.project(plane_normal)).normalized()
800 plane_up = cam_up.project(plane_normal).normalized()
801
802 delete_object(plane)
803
804 scene_struct['directions']['behind'] = tuple(plane_behind)
805 scene_struct['directions']['front'] = tuple(-plane_behind)
806 scene_struct['directions']['left'] = tuple(plane_left)
807 scene_struct['directions']['right'] = tuple(-plane_left)
808 scene_struct['directions']['above'] = tuple(plane_up)
809 scene_struct['directions']['below'] = tuple(-plane_up)
810
811 if args.key_light_jitter > 0:
812 for i in range(3):
813 bpy.data.objects['Lamp_Key'].location[i] += rand(args.key_light_jitter)
814 if args.back_light_jitter > 0:
815 for i in range(3):
816 bpy.data.objects['Lamp_Back'].location[i] += rand(args.back_light_jitter)
817 if args.fill_light_jitter > 0:
818 for i in range(3):
819 bpy.data.objects['Lamp_Fill'].location[i] += rand(args.fill_light_jitter)
820
821 objects, blender_objects = add_random_objects(scene_struct, num_objects, args, camera)
822 scene_struct['objects'] = objects
823 scene_struct['relationships'] = compute_all_relationships(scene_struct)
824 while True:
825 try:
826 bpy.ops.render.render(write_still=True)
827 break
828 except Exception as e:
829 print(e)
830
831 with open(output_scene, 'w') as f:
832 json.dump(scene_struct, f, indent=2)
833
834 if output_blendfile is not None:
835 bpy.ops.wm.save_as_mainfile(filepath=output_blendfile)
836
837
838def add_random_objects(scene_struct, num_objects, args, camera, max_scene_attempts=10):
839 scene_attempt = 0
840 while scene_attempt < max_scene_attempts:
841 scene_attempt += 1
842
843 with open(args.properties_json, 'r') as f:
844 properties = json.load(f)
845 color_name_to_rgba = {}
846 for name, rgb in properties['colors'].items():
847 rgba = [float(c) / 255.0 for c in rgb] + [1.0]
848 color_name_to_rgba[name] = rgba
849 material_mapping = [(v, k) for k, v in properties['materials'].items()]
850 object_mapping = [(v, k) for k, v in properties['shapes'].items()]
851 size_mapping = list(properties['sizes'].items())
852
853 shape_color_combos = None
854 if args.shape_color_combos_json is not None:
855 with open(args.shape_color_combos_json, 'r') as f:
856 shape_color_combos = list(json.load(f).items())
857
858 positions = []
859 objects = []
860 blender_objects = []
861 for i in range(num_objects):
862 size_name, r = random.choice(size_mapping)
863
864 num_tries = 0
865 while True:
866 num_tries += 1
867 if num_tries > args.max_retries:
868 for obj in blender_objects:
869 delete_object(obj)
870 break
871 x = random.uniform(-3, 3)
872 y = random.uniform(-3, 3)
873 dists_good = True
874 margins_good = True
875 for (xx, yy, rr) in positions:
876 dx, dy = x - xx, y - yy
877 dist = math.sqrt(dx * dx + dy * dy)
878 if dist - r - rr < args.min_dist:
879 dists_good = False
880 break
881 for direction_name in ['left', 'right', 'front', 'behind']:
882 direction_vec = scene_struct['directions'][direction_name]
883 assert direction_vec[2] == 0
884 margin = dx * direction_vec[0] + dy * direction_vec[1]
885 if 0 < margin < args.margin:
886 print(margin, args.margin, direction_name)
887 print('BROKEN MARGIN!')
888 margins_good = False
889 break
890 if not margins_good:
891 break
892
893 if dists_good and margins_good:
894 break
895
896 if num_tries > args.max_retries:
897 break
898
899 if shape_color_combos is None:
900 obj_name, obj_name_out = random.choice(object_mapping)
901 color_name, rgba = random.choice(list(color_name_to_rgba.items()))
902 else:
903 obj_name_out, color_choices = random.choice(shape_color_combos)
904 color_name = random.choice(color_choices)
905 obj_name = [k for k, v in object_mapping if v == obj_name_out][0]
906 rgba = color_name_to_rgba[color_name]
907
908 if obj_name == 'Cube':
909 r /= math.sqrt(2)
910
911 theta = 360.0 * random.random()
912 add_object(args.shape_dir, obj_name, r, (x, y), theta=theta)
913 obj = bpy.context.object
914 blender_objects.append(obj)
915 positions.append((x, y, r))
916
917 mat_name, mat_name_out = random.choice(material_mapping)
918 add_material(mat_name, Color=rgba)
919
920 pixel_coords = get_camera_coords(camera, obj.location)
921 objects.append({
922 'shape': obj_name_out,
923 'size': size_name,
924 'material': mat_name_out,
925 '3d_coords': tuple(obj.location),
926 'rotation': theta,
927 'pixel_coords': pixel_coords,
928 'color': color_name,
929 })
930
931 if len(objects) < num_objects:
932 continue
933
934 visibility_info = None
935 if INSIDE_BLENDER and Image is not None:
936 try:
937 visibility_info = compute_visibility_fractions(blender_objects)
938 except Exception as e:
939 print(f"Warning: compute_visibility_fractions failed during scene generation: {e}")
940 visibility_info = None
941
942 all_visible = True
943 if visibility_info is not None:
944 ratios, scene_counts, full_counts = visibility_info
945 min_ratio = min((r for r in ratios if full_counts[ratios.index(r)] > 0), default=1.0)
946 all_visible = all(
947 (full_counts[i] == 0) or (ratios[i] >= BASE_MIN_VISIBILITY_FRACTION)
948 for i in range(len(ratios))
949 )
950 if not all_visible:
951 print(f'Some objects are too occluded in generated scene; '
952 f'min visibility fraction={min_ratio:.3f} (required >= {BASE_MIN_VISIBILITY_FRACTION})')
953 else:
954 # Fallback to legacy absolute pixel-based visibility when PIL or Blender context is unavailable.
955 min_pixels = max(args.min_pixels_per_object, min_visible_pixels(args.width, args.height))
956 all_visible = check_visibility(blender_objects, min_pixels)
957
958 if not all_visible:
959 print('Some objects are occluded; replacing objects')
960 for obj in blender_objects:
961 delete_object(obj)
962 continue
963
964 return objects, blender_objects
965
966 raise RuntimeError(f"Failed to generate a valid scene after {max_scene_attempts} attempts")
967
968
969def compute_all_relationships(scene_struct, eps=0.2):
970 """
971 Computes relationships between all pairs of objects in the scene.
972
973 Returns a dictionary mapping string relationship names to lists of lists of
974 integers, where output[rel][i] gives a list of object indices that have the
975 relationship rel with object i. For example if j is in output['left'][i] then
976 object j is left of object j.
977 """
978 all_relationships = {}
979 for name, direction_vec in scene_struct['directions'].items():
980 if name == 'above' or name == 'below': continue
981 all_relationships[name] = []
982 for i, obj1 in enumerate(scene_struct['objects']):
983 coords1 = obj1['3d_coords']
984 related = set()
985 for j, obj2 in enumerate(scene_struct['objects']):
986 if obj1 == obj2: continue
987 coords2 = obj2['3d_coords']
988 diff = [coords2[k] - coords1[k] for k in [0, 1, 2]]
989 dot = sum(diff[k] * direction_vec[k] for k in [0, 1, 2])
990 if dot > eps:
991 related.add(j)
992 all_relationships[name].append(sorted(list(related)))
993 return all_relationships
994
995
996def compute_visibility_fractions(blender_objects):
997 if not INSIDE_BLENDER or not blender_objects:
998 return None
999 if Image is None:
1000 return None
1001
1002 # First pass: all objects together (occluded counts).
1003 fd, path = tempfile.mkstemp(suffix='.png')
1004 os.close(fd)
1005 try:
1006 colors_list = render_shadeless(blender_objects, path, use_distinct_colors=True)
1007 img = Image.open(path).convert('RGB')
1008 w, h = img.size
1009 pix = img.load()
1010 color_to_idx = {}
1011 for i, (r, g, b) in enumerate(colors_list):
1012 key = (round(r * 255), round(g * 255), round(b * 255))
1013 color_to_idx[key] = i
1014 scene_counts = [0] * len(blender_objects)
1015 for y in range(h):
1016 for x in range(w):
1017 key = (pix[x, y][0], pix[x, y][1], pix[x, y][2])
1018 if key in color_to_idx:
1019 scene_counts[color_to_idx[key]] += 1
1020 finally:
1021 try:
1022 os.remove(path)
1023 except Exception:
1024 pass
1025
1026 # Second pass: per-object "full area" with other objects hidden.
1027 full_counts = []
1028 original_hide_render = [obj.hide_render for obj in blender_objects]
1029 try:
1030 for idx, obj in enumerate(blender_objects):
1031 # Hide all other objects, ensure this one is visible.
1032 for j, other in enumerate(blender_objects):
1033 if j == idx:
1034 other.hide_render = False
1035 else:
1036 other.hide_render = True
1037
1038 fd_i, path_i = tempfile.mkstemp(suffix='.png')
1039 os.close(fd_i)
1040 try:
1041 colors_list = render_shadeless([obj], path_i, use_distinct_colors=True)
1042 img = Image.open(path_i).convert('RGB')
1043 w, h = img.size
1044 pix = img.load()
1045 color_to_idx = {}
1046 for i, (r, g, b) in enumerate(colors_list):
1047 key = (round(r * 255), round(g * 255), round(b * 255))
1048 color_to_idx[key] = i
1049 count = 0
1050 for y in range(h):
1051 for x in range(w):
1052 key = (pix[x, y][0], pix[x, y][1], pix[x, y][2])
1053 if key in color_to_idx:
1054 count += 1
1055 full_counts.append(count)
1056 finally:
1057 try:
1058 os.remove(path_i)
1059 except Exception:
1060 pass
1061 finally:
1062 # Restore previous hide_render flags.
1063 for obj, prev in zip(blender_objects, original_hide_render):
1064 obj.hide_render = prev
1065
1066 visibility = []
1067 for scene_c, full_c in zip(scene_counts, full_counts):
1068 if full_c <= 0:
1069 visibility.append(0.0)
1070 else:
1071 visibility.append(float(scene_c) / float(full_c))
1072
1073 return visibility, scene_counts, full_counts
1074
1075
1076def check_visibility(blender_objects, min_pixels_per_object):
1077 """
1078 Legacy absolute pixel-count visibility check, kept as a fallback when
1079 relative per-object visibility cannot be computed.
1080 """
1081 if not INSIDE_BLENDER or not blender_objects:
1082 return True
1083 if Image is None:
1084 return True
1085 fd, path = tempfile.mkstemp(suffix='.png')
1086 os.close(fd)
1087 try:
1088 colors_list = render_shadeless(blender_objects, path, use_distinct_colors=True)
1089 img = Image.open(path).convert('RGB')
1090 w, h = img.size
1091 pix = img.load()
1092 color_to_idx = {}
1093 for i, (r, g, b) in enumerate(colors_list):
1094 key = (round(r * 255), round(g * 255), round(b * 255))
1095 color_to_idx[key] = i
1096 counts = [0] * len(blender_objects)
1097 for y in range(h):
1098 for x in range(w):
1099 key = (pix[x, y][0], pix[x, y][1], pix[x, y][2])
1100 if key in color_to_idx:
1101 counts[color_to_idx[key]] += 1
1102 all_visible = all(c >= min_pixels_per_object for c in counts)
1103 return all_visible
1104 finally:
1105 try:
1106 os.remove(path)
1107 except Exception:
1108 pass
1109
1110
1111def render_shadeless(blender_objects, path='flat.png', use_distinct_colors=False):
1112 """
1113 Render a version of the scene with shading disabled and unique materials
1114 assigned to all objects. The image itself is written to path. This is used to ensure
1115 that all objects will be visible in the final rendered scene (when check_visibility is enabled).
1116 Returns a list of (r,g,b) colors in object order (for visibility counting when use_distinct_colors=True).
1117 """
1118 render_args = bpy.context.scene.render
1119
1120 old_filepath = render_args.filepath
1121 old_engine = render_args.engine
1122
1123 render_args.filepath = path
1124 render_args.engine = 'BLENDER_EEVEE_NEXT'
1125
1126 view_layer = bpy.context.scene.view_layers[0]
1127 old_use_pass_combined = view_layer.use_pass_combined
1128
1129 for obj_name in ['Lamp_Key', 'Lamp_Fill', 'Lamp_Back', 'Ground']:
1130 if obj_name in bpy.data.objects:
1131 obj = bpy.data.objects[obj_name]
1132 obj.hide_render = True
1133
1134 n = len(blender_objects)
1135 object_colors = [] if use_distinct_colors else set()
1136 old_materials = []
1137 for i, obj in enumerate(blender_objects):
1138 if len(obj.data.materials) > 0:
1139 old_materials.append(obj.data.materials[0])
1140 else:
1141 old_materials.append(None)
1142
1143 mat = bpy.data.materials.new(name='Material_%d' % i)
1144 mat.use_nodes = True
1145 nodes = mat.node_tree.nodes
1146 nodes.clear()
1147
1148 node_emission = nodes.new(type='ShaderNodeEmission')
1149 node_output = nodes.new(type='ShaderNodeOutputMaterial')
1150
1151 if use_distinct_colors:
1152 r = (i + 1) / (n + 1)
1153 g, b = 0.5, 0.5
1154 object_colors.append((r, g, b))
1155 else:
1156 while True:
1157 r, g, b = [random.random() for _ in range(3)]
1158 if (r, g, b) not in object_colors:
1159 break
1160 object_colors.add((r, g, b))
1161
1162 node_emission.inputs['Color'].default_value = (r, g, b, 1.0)
1163 mat.node_tree.links.new(node_emission.outputs['Emission'], node_output.inputs['Surface'])
1164
1165 if len(obj.data.materials) > 0:
1166 obj.data.materials[0] = mat
1167 else:
1168 obj.data.materials.append(mat)
1169
1170 bpy.ops.render.render(write_still=True)
1171
1172 for mat, obj in zip(old_materials, blender_objects):
1173 if mat is not None:
1174 obj.data.materials[0] = mat
1175 elif len(obj.data.materials) > 0:
1176 obj.data.materials.clear()
1177
1178 for obj_name in ['Lamp_Key', 'Lamp_Fill', 'Lamp_Back', 'Ground']:
1179 if obj_name in bpy.data.objects:
1180 obj = bpy.data.objects[obj_name]
1181 obj.hide_render = False
1182
1183 render_args.filepath = old_filepath
1184 render_args.engine = old_engine
1185
1186 return object_colors
1187
1188
1189if __name__ == '__main__':
1190 if INSIDE_BLENDER:
1191 argv = extract_args()
1192 args = parser.parse_args(argv)
1193 main(args)
1194 elif '--help' in sys.argv or '-h' in sys.argv:
1195 parser.print_help()
1196 else:
1197 print('This script is intended to be called from blender like this:')
1198 print()
1199 print('blender --background --python render_images.py -- [args]')
1200 print()