From f7133687360660e00a36510839bf44976f02e612 Mon Sep 17 00:00:00 2001 From: jacknugent Date: Thu, 3 Oct 2024 09:59:45 -0400 Subject: [PATCH] material segmentation render stage --- infinigen/core/rendering/render.py | 87 ++++++++++++++++++---- infinigen_examples/configs_nature/base.gin | 7 +- 2 files changed, 75 insertions(+), 19 deletions(-) diff --git a/infinigen/core/rendering/render.py b/infinigen/core/rendering/render.py index 2adce25a6..8626fa764 100644 --- a/infinigen/core/rendering/render.py +++ b/infinigen/core/rendering/render.py @@ -90,6 +90,17 @@ def set_pass_indices(): return tree_output +def set_material_pass_indices(): + output_material_properties = {} + mat_index = 1 + for mat in bpy.data.materials: + if mat.pass_index == 0: + mat.pass_index = mat_index + mat_index += 1 + output_material_properties[mat.name] = {"pass_index": mat.pass_index} + return output_material_properties + + # Can be pasted directly into the blender console def make_clay(): clay_material = bpy.data.materials.new(name="clay") @@ -145,14 +156,25 @@ def configure_compositor_output( passes_to_save, saving_ground_truth, ): - file_output_node = nw.new_node( + file_output_node_png = nw.new_node( Nodes.OutputFile, attrs={ "base_path": str(frames_folder), - "format.file_format": "OPEN_EXR" if saving_ground_truth else "PNG", + "format.file_format": "PNG", "format.color_mode": "RGB", }, ) + file_output_node_exr = nw.new_node( + Nodes.OutputFile, + attrs={ + "base_path": str(frames_folder), + "format.file_format": "OPEN_EXR", + "format.color_mode": "RGB", + }, + ) + default_file_output_node = ( + file_output_node_exr if saving_ground_truth else file_output_node_png + ) file_slot_list = [] viewlayer = bpy.context.scene.view_layers["ViewLayer"] render_layers = nw.new_node(Nodes.RenderLayers) @@ -161,6 +183,13 @@ def configure_compositor_output( setattr(viewlayer, f"use_pass_{viewlayer_pass}", True) else: setattr(viewlayer.cycles, f"use_pass_{viewlayer_pass}", True) + # must save the material pass index as EXR + file_output_node = ( + default_file_output_node + if viewlayer_pass != "material_index" + else file_output_node_exr + ) + slot_input = file_output_node.file_slots.new(socket_name) render_socket = render_layers.outputs[socket_name] if viewlayer_pass == "vector": @@ -173,24 +202,15 @@ def configure_compositor_output( nw.links.new(render_socket, slot_input) file_slot_list.append(file_output_node.file_slots[slot_input.name]) - slot_input = file_output_node.file_slots["Image"] + slot_input = default_file_output_node.file_slots["Image"] image = image_denoised if image_denoised is not None else image_noisy - nw.links.new(image, file_output_node.inputs["Image"]) + nw.links.new(image, default_file_output_node.inputs["Image"]) if saving_ground_truth: slot_input.path = "UniqueInstances" else: - image_exr_output_node = nw.new_node( - Nodes.OutputFile, - attrs={ - "base_path": str(frames_folder), - "format.file_format": "OPEN_EXR", - "format.color_mode": "RGB", - }, - ) - rgb_exr_slot_input = file_output_node.file_slots["Image"] - nw.links.new(image, image_exr_output_node.inputs["Image"]) - file_slot_list.append(image_exr_output_node.file_slots[rgb_exr_slot_input.path]) - file_slot_list.append(file_output_node.file_slots[slot_input.path]) + nw.links.new(image, file_output_node_exr.inputs["Image"]) + file_slot_list.append(file_output_node_exr.file_slots[slot_input.path]) + file_slot_list.append(default_file_output_node.file_slots[slot_input.path]) return file_slot_list @@ -308,6 +328,22 @@ def postprocess_blendergt_outputs(frames_folder, output_stem): uniq_inst_path.unlink() +def postprocess_materialgt_output(frames_folder, output_stem): + # Save material segmentation visualization if present + ma_seg_dst_path = frames_folder / f"IndexMA{output_stem}.exr" + if ma_seg_dst_path.is_file(): + ma_seg_mask_array = load_seg_mask(ma_seg_dst_path) + np.save( + ma_seg_dst_path.with_name(f"MaterialSegmentation{output_stem}.npy"), + ma_seg_mask_array, + ) + imwrite( + ma_seg_dst_path.with_name(f"MaterialSegmentation{output_stem}.png"), + colorize_int_array(ma_seg_mask_array), + ) + ma_seg_dst_path.unlink() + + def configure_compositor( frames_folder: Path, passes_to_save: list, @@ -380,6 +416,22 @@ def render_image( with Timer("Flat Shading"): global_flat_shading() + else: + segment_materials = "material_index" in (x[0] for x in passes_to_save) + if segment_materials: + with Timer("Set material indices"): + material_data = set_material_pass_indices() + json_object = json.dumps(material_data, indent=4) + first_frame = bpy.context.scene.frame_start + suffix = get_suffix( + dict( + cam_rig=camera_rig_id, + resample=0, + frame=first_frame, + subcam=subcam_id, + ) + ) + (frames_folder / f"Materials{suffix}.json").write_text(json_object) if not bpy.context.scene.use_nodes: bpy.context.scene.use_nodes = True @@ -422,6 +474,9 @@ def render_image( output_folder=frames_folder, frame=frame, ) + bpy.context.scene.frame_set(frame) + suffix = get_suffix(dict(frame=frame, **indices)) + postprocess_materialgt_output(frames_folder, suffix) for file in tmp_dir.glob("*.png"): file.unlink() diff --git a/infinigen_examples/configs_nature/base.gin b/infinigen_examples/configs_nature/base.gin index 9f4a58f51..a85931921 100644 --- a/infinigen_examples/configs_nature/base.gin +++ b/infinigen_examples/configs_nature/base.gin @@ -50,7 +50,8 @@ full/render_image.passes_to_save = [ ['volume_direct', 'VolumeDir'], ['emit', 'Emit'], ['environment', 'Env'], - ['ambient_occlusion', 'AO'] + ['ambient_occlusion', 'AO'], + ['material_index', 'IndexMA'], ] flat/render_image.passes_to_save = [ ['z', 'Depth'], @@ -78,7 +79,7 @@ camera.camera_pose_proposal.altitude = ("weighted_choice", camera.camera_pose_proposal.pitch = ("clip_gaussian", 90, 30, 20, 160) -# WARNING: Large camera rig translations or rotations require special handling. +# WARNING: Large camera rig translations or rotations require special handling. # if your cameras are not all approximately forward facing within a few centimeters, you must either: # - configure the pipeline to generate assets / terrain for each camera separately, rather than sharing it between the whole rig # - or, treat your camera rig as multiple camera rigs each with one camera, and implement code to positon them correctly @@ -86,4 +87,4 @@ camera.spawn_camera_rigs.n_camera_rigs = 1 camera.spawn_camera_rigs.camera_rig_config = [ {'loc': (0, 0, 0), 'rot_euler': (0, 0, 0)}, {'loc': (0.075, 0, 0), 'rot_euler': (0, 0, 0)} -] \ No newline at end of file +]