def enable_diffuse_color_output(output_dir: Optional[str] = None, file_prefix: str = "diffuse_",
output_key: str = "diffuse"):
""" Enables writing diffuse color (albedo) images.
Diffuse color images will be written in the form of .png files during the next rendering.
:param output_dir: The directory to write files to, if this is None the temporary directory is used.
:param file_prefix: The prefix to use for writing the files.
:param output_key: The key to use for registering the diffuse color output.
"""
if output_dir is None:
output_dir = Utility.get_temporary_directory()
bpy.context.scene.render.use_compositing = True
bpy.context.scene.use_nodes = True
tree = bpy.context.scene.node_tree
links = tree.links
bpy.context.view_layer.use_pass_diffuse_color = True
render_layer_node = Utility.get_the_one_node_with_type(tree.nodes, 'CompositorNodeRLayers')
final_output = render_layer_node.outputs["DiffCol"]
output_file = tree.nodes.new('CompositorNodeOutputFile')
output_file.base_path = output_dir
output_file.format.file_format = "PNG"
output_file.file_slots.values()[0].path = file_prefix
links.new(final_output, output_file.inputs['Image'])
Utility.add_output_entry({
"key": output_key,
"path": os.path.join(output_dir, file_prefix) + "%04d" + ".png",
"version": "2.0.0"
})
def build_convex_decomposition_collision_shape(
self,
vhacd_path: str,
temp_dir: str = None,
cache_dir: str = "blenderproc_resources/decomposition_cache"):
""" Builds a collision shape of the object by decomposing it into near convex parts using V-HACD
:param vhacd_path: The directory in which vhacd should be installed or is already installed.
:param temp_dir: The temp dir to use for storing the object files created by v-hacd.
:param cache_dir: If a directory is given, convex decompositions are stored there named after the meshes hash. If the same mesh is decomposed a second time, the result is loaded from the cache and the actual decomposition is skipped.
"""
if platform == "win32":
raise Exception("This is currently not supported under Windows")
if temp_dir is None:
temp_dir = Utility.get_temporary_directory()
# Decompose the object
parts = convex_decomposition(self,
temp_dir,
resolve_path(vhacd_path),
cache_dir=resolve_path(cache_dir))
parts = [MeshObject(p) for p in parts]
# Make the convex parts children of this object, enable their rigid body component and hide them
for part in parts:
part.set_parent(self)
part.enable_rigidbody(True, "CONVEX_HULL")
part.hide()
def enable_distance_output(output_dir: Optional[str] = None, file_prefix: str = "distance_",
output_key: str = "distance", distance_start: float = 0.1, distance_range: float = 25.0,
distance_falloff: str = "LINEAR"):
""" Enables writing distance images.
Distance images will be written in the form of .exr files during the next rendering.
:param output_dir: The directory to write files to, if this is None the temporary directory is used.
:param file_prefix: The prefix to use for writing the files.
:param output_key: The key to use for registering the distance output.
:param distance_start: Starting distance of the distance, measured from the camera.
:param distance_range: Total distance in which the distance is measured. \
distance_end = distance_start + distance_range.
:param distance_falloff: Type of transition used to fade distance. Available: [LINEAR, QUADRATIC, INVERSE_QUADRATIC]
"""
if output_dir is None:
output_dir = Utility.get_temporary_directory()
bpy.context.scene.render.use_compositing = True
bpy.context.scene.use_nodes = True
GlobalStorage.add("renderer_distance_end", distance_start + distance_range)
tree = bpy.context.scene.node_tree
links = tree.links
# Use existing render layer
render_layer_node = Utility.get_the_one_node_with_type(tree.nodes, 'CompositorNodeRLayers')
# Set mist pass limits
bpy.context.scene.world.mist_settings.start = distance_start
bpy.context.scene.world.mist_settings.depth = distance_range
bpy.context.scene.world.mist_settings.falloff = distance_falloff
bpy.context.view_layer.use_pass_mist = True # Enable distance pass
# Create a mapper node to map from 0-1 to SI units
mapper_node = tree.nodes.new("CompositorNodeMapRange")
links.new(render_layer_node.outputs["Mist"], mapper_node.inputs['Value'])
# map the values 0-1 to range distance_start to distance_range
mapper_node.inputs['To Min'].default_value = distance_start
mapper_node.inputs['To Max'].default_value = distance_start + distance_range
final_output = mapper_node.outputs['Value']
# Build output node
output_file = tree.nodes.new("CompositorNodeOutputFile")
output_file.base_path = output_dir
output_file.format.file_format = "OPEN_EXR"
output_file.file_slots.values()[0].path = file_prefix
# Feed the Z-Buffer or Mist output of the render layer to the input of the file IO layer
links.new(final_output, output_file.inputs['Image'])
Utility.add_output_entry({
"key": output_key,
"path": os.path.join(output_dir, file_prefix) + "%04d" + ".exr",
"version": "2.0.0",
"trim_redundant_channels": True
})
def _default_init(self):
"""
These operations are called during all modules inits
"""
self._output_dir = resolve_path(
self.config.get_string("output_dir", ""))
os.makedirs(self._output_dir, exist_ok=True)
self._temp_dir = Utility.get_temporary_directory()
self._avoid_output = self.config.get_bool("avoid_output", False)
def render(output_dir: Optional[str] = None, file_prefix: str = "rgb_", output_key: Optional[str] = "colors",
load_keys: Optional[Set[str]] = None, return_data: bool = True,
keys_with_alpha_channel: Optional[Set[str]] = None) -> Dict[str, Union[np.ndarray, List[np.ndarray]]]:
""" Render all frames.
This will go through all frames from scene.frame_start to scene.frame_end and render each of them.
:param output_dir: The directory to write files to, if this is None the temporary directory is used. \
The temporary directory is usually in the shared memory (only true for linux).
:param file_prefix: The prefix to use for writing the images.
:param output_key: The key to use for registering the output.
:param load_keys: Set of output keys to load when available
:param return_data: Whether to load and return generated data. Backwards compatibility to config-based pipeline.
:param keys_with_alpha_channel: A set containing all keys whose alpha channels should be loaded.
:return: dict of lists of raw renderer output. Keys can be 'distance', 'colors', 'normals'
"""
if output_dir is None:
output_dir = Utility.get_temporary_directory()
if load_keys is None:
load_keys = {'colors', 'distance', 'depth', 'normals', 'diffuse'}
keys_with_alpha_channel = {'colors'} if bpy.context.scene.render.film_transparent else None
if output_key is not None:
Utility.add_output_entry({
"key": output_key,
"path": os.path.join(output_dir, file_prefix) + "%04d" +
map_file_format_to_file_ending(bpy.context.scene.render.image_settings.file_format),
"version": "2.0.0"
})
load_keys.add(output_key)
bpy.context.scene.render.filepath = os.path.join(output_dir, file_prefix)
# Skip if there is nothing to render
if bpy.context.scene.frame_end != bpy.context.scene.frame_start:
if len(get_all_blender_mesh_objects()) == 0:
raise Exception("There are no mesh-objects to render, "
"please load an object before invoking the renderer.")
# As frame_end is pointing to the next free frame, decrease it by one, as
# blender will render all frames in [frame_start, frame_ned]
bpy.context.scene.frame_end -= 1
bpy.ops.render.render(animation=True, write_still=True)
# Revert changes
bpy.context.scene.frame_end += 1
return WriterUtility.load_registered_outputs(load_keys, keys_with_alpha_channel) if return_data else {}
def enable_depth_output(output_dir: Optional[str] = None, file_prefix: str = "depth_", output_key: str = "depth"):
""" Enables writing depth images.
Depth images will be written in the form of .exr files during the next rendering.
:param output_dir: The directory to write files to.
:param file_prefix: The prefix to use for writing the files.
:param output_key: The key to use for registering the depth output.
"""
if output_dir is None:
output_dir = Utility.get_temporary_directory()
bpy.context.scene.render.use_compositing = True
bpy.context.scene.use_nodes = True
tree = bpy.context.scene.node_tree
links = tree.links
# Use existing render layer
render_layer_node = Utility.get_the_one_node_with_type(tree.nodes, 'CompositorNodeRLayers')
# Enable z-buffer pass
bpy.context.view_layer.use_pass_z = True
# Build output node
output_file = tree.nodes.new("CompositorNodeOutputFile")
output_file.base_path = output_dir
output_file.format.file_format = "OPEN_EXR"
output_file.file_slots.values()[0].path = file_prefix
# Feed the Z-Buffer output of the render layer to the input of the file IO layer
links.new(render_layer_node.outputs["Depth"], output_file.inputs['Image'])
Utility.add_output_entry({
"key": output_key,
"path": os.path.join(output_dir, file_prefix) + "%04d" + ".exr",
"version": "2.0.0",
"trim_redundant_channels": True
})
def enable_normals_output(output_dir: Optional[str] = None, file_prefix: str = "normals_",
output_key: str = "normals"):
""" Enables writing normal images.
Normal images will be written in the form of .exr files during the next rendering.
:param output_dir: The directory to write files to, if this is None the temporary directory is used.
:param file_prefix: The prefix to use for writing the files.
:param output_key: The key to use for registering the normal output.
"""
if output_dir is None:
output_dir = Utility.get_temporary_directory()
bpy.context.scene.render.use_compositing = True
bpy.context.scene.use_nodes = True
tree = bpy.context.scene.node_tree
links = tree.links
# Use existing render layer
render_layer_node = Utility.get_the_one_node_with_type(tree.nodes, 'CompositorNodeRLayers')
separate_rgba = tree.nodes.new("CompositorNodeSepRGBA")
space_between_nodes_x = 200
space_between_nodes_y = -300
separate_rgba.location.x = space_between_nodes_x
separate_rgba.location.y = space_between_nodes_y
links.new(render_layer_node.outputs["Normal"], separate_rgba.inputs["Image"])
combine_rgba = tree.nodes.new("CompositorNodeCombRGBA")
combine_rgba.location.x = space_between_nodes_x * 14
c_channels = ["R", "G", "B"]
offset = space_between_nodes_x * 2
multiplication_values: List[List[bpy.types.Node]] = [[], [], []]
channel_results = {}
for row_index, channel in enumerate(c_channels):
# matrix multiplication
mulitpliers = []
for column in range(3):
multiply = tree.nodes.new("CompositorNodeMath")
multiply.operation = "MULTIPLY"
multiply.inputs[1].default_value = 0 # setting at the end for all frames
multiply.location.x = column * space_between_nodes_x + offset
multiply.location.y = row_index * space_between_nodes_y
links.new(separate_rgba.outputs[c_channels[column]], multiply.inputs[0])
mulitpliers.append(multiply)
multiplication_values[row_index].append(multiply)
first_add = tree.nodes.new("CompositorNodeMath")
first_add.operation = "ADD"
first_add.location.x = space_between_nodes_x * 5 + offset
first_add.location.y = row_index * space_between_nodes_y
links.new(mulitpliers[0].outputs["Value"], first_add.inputs[0])
links.new(mulitpliers[1].outputs["Value"], first_add.inputs[1])
second_add = tree.nodes.new("CompositorNodeMath")
second_add.operation = "ADD"
second_add.location.x = space_between_nodes_x * 6 + offset
second_add.location.y = row_index * space_between_nodes_y
links.new(first_add.outputs["Value"], second_add.inputs[0])
links.new(mulitpliers[2].outputs["Value"], second_add.inputs[1])
channel_results[channel] = second_add
# set the matrix accordingly
rot_around_x_axis = mathutils.Matrix.Rotation(math.radians(-90.0), 4, 'X')
for frame in range(bpy.context.scene.frame_start, bpy.context.scene.frame_end):
used_rotation_matrix = CameraUtility.get_camera_pose(frame) @ rot_around_x_axis
for row_index in range(3):
for column_index in range(3):
current_multiply = multiplication_values[row_index][column_index]
current_multiply.inputs[1].default_value = used_rotation_matrix[column_index][row_index]
current_multiply.inputs[1].keyframe_insert(data_path='default_value', frame=frame)
offset = 8 * space_between_nodes_x
for index, channel in enumerate(c_channels):
multiply = tree.nodes.new("CompositorNodeMath")
multiply.operation = "MULTIPLY"
multiply.location.x = space_between_nodes_x * 2 + offset
multiply.location.y = index * space_between_nodes_y
links.new(channel_results[channel].outputs["Value"], multiply.inputs[0])
if channel == "G":
multiply.inputs[1].default_value = -0.5
else:
multiply.inputs[1].default_value = 0.5
add = tree.nodes.new("CompositorNodeMath")
add.operation = "ADD"
add.location.x = space_between_nodes_x * 3 + offset
add.location.y = index * space_between_nodes_y
links.new(multiply.outputs["Value"], add.inputs[0])
add.inputs[1].default_value = 0.5
output_channel = channel
if channel == "G":
output_channel = "B"
elif channel == "B":
output_channel = "G"
links.new(add.outputs["Value"], combine_rgba.inputs[output_channel])
output_file = tree.nodes.new("CompositorNodeOutputFile")
output_file.base_path = output_dir
output_file.format.file_format = "OPEN_EXR"
output_file.file_slots.values()[0].path = file_prefix
output_file.location.x = space_between_nodes_x * 15
links.new(combine_rgba.outputs["Image"], output_file.inputs["Image"])
Utility.add_output_entry({
"key": output_key,
"path": os.path.join(output_dir, file_prefix) + "%04d" + ".exr",
"version": "2.0.0"
})
def load_AMASS(data_path: str,
sub_dataset_id: str,
temp_dir: str = None,
body_model_gender: str = None,
subject_id: str = "",
sequence_id: int = -1,
frame_id: int = -1,
num_betas: int = 10,
num_dmpls: int = 8) -> List[MeshObject]:
"""
use the pose parameters to generate the mesh and loads it to the scene.
:param data_path: The path to the AMASS Dataset folder in resources folder.
:param sub_dataset_id: Identifier for the sub dataset, the dataset which the human pose object should be extracted from.
Available: ['CMU', 'Transitions_mocap', 'MPI_Limits', 'SSM_synced', 'TotalCapture',
'Eyes_Japan_Dataset', 'MPI_mosh', 'MPI_HDM05', 'HumanEva', 'ACCAD', 'EKUT', 'SFU', 'KIT', 'H36M', 'TCD_handMocap', 'BML']
:param temp_dir: A temp directory which is used for writing the temporary .obj file.
:param body_model_gender: The model gender, pose will represented using male, female or neutral body shape.
Available:[male, female, neutral]. If None is selected a random one is choosen.
:param subject_id: Type of motion from which the pose should be extracted, this is dataset dependent parameter.
If left empty a random subject id is picked.
:param sequence_id: Sequence id in the dataset, sequences are the motion recorded to represent certain action.
If set to -1 a random sequence id is selected.
:param frame_id: Frame id in a selected motion sequence. If none is selected a random one is picked
:param num_betas: Number of body parameters
:param num_dmpls: Number of DMPL parameters
:return: The list of loaded mesh objects.
"""
if body_model_gender is None:
body_model_gender = random.choice(["male", "female", "neutral"])
if temp_dir is None:
temp_dir = Utility.get_temporary_directory()
# Install required additonal packages
SetupUtility.setup_pip([
"git+https://github.com/abahnasy/smplx",
"git+https://github.com/abahnasy/human_body_prior"
])
# Get the currently supported mocap datasets by this loader
taxonomy_file_path = os.path.join(data_path, "taxonomy.json")
supported_mocap_datasets = AMASSLoader._get_supported_mocap_datasets(
taxonomy_file_path, data_path)
# selected_obj = self._files_with_fitting_ids
pose_body, betas = AMASSLoader._get_pose_parameters(
supported_mocap_datasets, num_betas, sub_dataset_id, subject_id,
sequence_id, frame_id)
# load parametric Model
body_model, faces = AMASSLoader._load_parametric_body_model(
data_path, body_model_gender, num_betas, num_dmpls)
# Generate Body representations using SMPL model
body_repr = body_model(pose_body=pose_body, betas=betas)
# Generate .obj file represents the selected pose
generated_obj = AMASSLoader._write_body_mesh_to_obj_file(
body_repr, faces, temp_dir)
loaded_obj = load_obj(generated_obj)
AMASSLoader._correct_materials(loaded_obj)
# set the shading mode explicitly to smooth
for obj in loaded_obj:
obj.set_shading_mode("SMOOTH")
# removes the x axis rotation found in all ShapeNet objects, this is caused by importing .obj files
# the object has the same pose as before, just that the rotation_euler is now [0, 0, 0]
for obj in loaded_obj:
obj.persist_transformation_into_mesh(location=False,
rotation=True,
scale=False)
# move the origin of the object to the world origin and on top of the X-Y plane
# makes it easier to place them later on, this does not change the `.location`
for obj in loaded_obj:
obj.move_origin_to_bottom_mean_point()
bpy.ops.object.select_all(action='DESELECT')
return loaded_obj
def render_optical_flow(output_dir: str = None, temp_dir: str = None, get_forward_flow: bool = True,
get_backward_flow: bool = True, blender_image_coordinate_style: bool = False,
forward_flow_output_file_prefix: str = "forward_flow_",
forward_flow_output_key: str = "forward_flow",
backward_flow_output_file_prefix: str = "backward_flow_",
backward_flow_output_key: str = "backward_flow", return_data: bool = True) -> \
Dict[str, Union[np.ndarray, List[np.ndarray]]]:
""" Renders the optical flow (forward and backward) for all frames.
:param output_dir: The directory to write images to.
:param temp_dir: The directory to write intermediate data to.
:param get_forward_flow: Whether to render forward optical flow.
:param get_backward_flow: Whether to render backward optical flow.
:param blender_image_coordinate_style: Whether to specify the image coordinate system at the bottom left (blender default; True) or top left (standard convention; False).
:param forward_flow_output_file_prefix: The file prefix that should be used when writing forward flow to a file.
:param forward_flow_output_key: The key which should be used for storing forward optical flow values.
:param backward_flow_output_file_prefix: The file prefix that should be used when writing backward flow to a file.
:param backward_flow_output_key: The key which should be used for storing backward optical flow values.
:param return_data: Whether to load and return generated data. Backwards compatibility to config-based pipeline.
:return: dict of lists of raw renderer outputs. Keys can be 'forward_flow', 'backward_flow'
"""
if get_forward_flow is False and get_backward_flow is False:
raise Exception(
"Take the FlowRenderer Module out of the config if both forward and backward flow are set to False!"
)
if output_dir is None:
output_dir = Utility.get_temporary_directory()
if temp_dir is None:
temp_dir = Utility.get_temporary_directory()
with Utility.UndoAfterExecution():
RendererUtility._render_init()
RendererUtility.set_samples(1)
RendererUtility.set_adaptive_sampling(0)
RendererUtility.set_denoiser(None)
RendererUtility.set_light_bounces(1, 0, 0, 1, 0, 8, 0)
FlowRendererUtility._output_vector_field(get_forward_flow,
get_backward_flow, output_dir)
# only need to render once; both fwd and bwd flow will be saved
temporary_fwd_flow_file_path = os.path.join(temp_dir, 'fwd_flow_')
temporary_bwd_flow_file_path = os.path.join(temp_dir, 'bwd_flow_')
RendererUtility.render(temp_dir, "bwd_flow_", None, load_keys=set())
# After rendering: convert to optical flow or calculate hsv visualization, if desired
for frame in range(bpy.context.scene.frame_start,
bpy.context.scene.frame_end):
# temporarily save respective vector fields
if get_forward_flow:
file_path = temporary_fwd_flow_file_path + "%04d" % frame + ".exr"
fwd_flow_field = load_image(file_path,
num_channels=4).astype(np.float32)
if not blender_image_coordinate_style:
fwd_flow_field[:, :, 1] = fwd_flow_field[:, :, 1] * -1
fname = os.path.join(
output_dir,
forward_flow_output_file_prefix) + '%04d' % frame
forward_flow = fwd_flow_field * -1 # invert forward flow to point at next frame
np.save(fname + '.npy', forward_flow[:, :, :2])
if get_backward_flow:
file_path = temporary_bwd_flow_file_path + "%04d" % frame + ".exr"
bwd_flow_field = load_image(file_path,
num_channels=4).astype(np.float32)
if not blender_image_coordinate_style:
bwd_flow_field[:, :, 1] = bwd_flow_field[:, :, 1] * -1
fname = os.path.join(
output_dir,
backward_flow_output_file_prefix) + '%04d' % frame
np.save(fname + '.npy', bwd_flow_field[:, :, :2])
load_keys = set()
# register desired outputs
if get_forward_flow:
Utility.register_output(output_dir, forward_flow_output_file_prefix,
forward_flow_output_key, '.npy', '2.0.0')
load_keys.add(forward_flow_output_key)
if get_backward_flow:
Utility.register_output(output_dir, backward_flow_output_file_prefix,
backward_flow_output_key, '.npy', '2.0.0')
load_keys.add(backward_flow_output_key)
return WriterUtility.load_registered_outputs(
load_keys) if return_data else {}
def load_bop(bop_dataset_path: str, sys_paths: Union[List[str], str], temp_dir: str = None, model_type: str = "", cam_type: str = "", split: str = "test", scene_id: int = -1, obj_ids: list = [], sample_objects: bool = False, num_of_objs_to_sample: int = None, obj_instances_limit: int = -1, move_origin_to_x_y_plane: bool = False, source_frame: list = ["X", "-Y", "-Z"], mm2m: bool = False) -> List[MeshObject]:
""" Loads the 3D models of any BOP dataset and allows replicating BOP scenes
- Interfaces with the bob_toolkit, allows loading of train, val and test splits
- Relative camera poses are loaded/computed with respect to a reference model
- Sets real camera intrinsics
:param bop_dataset_path: Full path to a specific bop dataset e.g. /home/user/bop/tless.
:param sys_paths: System paths to append. Can be a string or a list of strings.
:param temp_dir: A temp directory which is used for writing the temporary .ply file.
:param model_type: Optionally, specify type of BOP model. Available: [reconst, cad or eval].
:param cam_type: Camera type. If not defined, dataset-specific default camera type is used.
:param split: Optionally, test or val split depending on BOP dataset.
:param scene_id: Optionally, specify BOP dataset scene to synthetically replicate. Default: -1 (no scene is replicated,
only BOP Objects are loaded).
:param obj_ids: List of object ids to load. Default: [] (all objects from the given BOP dataset if scene_id is not
specified).
:param sample_objects: Toggles object sampling from the specified dataset.
:param num_of_objs_to_sample: Amount of objects to sample from the specified dataset. If this amount is bigger than the dataset
actually contains, then all objects will be loaded.
:param obj_instances_limit: Limits the amount of object copies when sampling. Default: -1 (no limit).
:param move_origin_to_x_y_plane: Move center of the object to the lower side of the object, this will not work when used in combination with
pose estimation tasks! This is designed for the use-case where BOP objects are used as filler objects in
the background.
:param source_frame: Can be used if the given positions and rotations are specified in frames different from the blender
frame. Has to be a list of three strings. Example: ['X', '-Z', 'Y']: Point (1,2,3) will be transformed
to (1, -3, 2). Available: ['X', 'Y', 'Z', '-X', '-Y', '-Z'].
:param mm2m: Specify whether to convert poses and models to meters.
:return: The list of loaded mesh objects.
"""
# Make sure sys_paths is a list
if not isinstance(sys_paths, list):
sys_paths = [sys_paths]
for sys_path in sys_paths:
if 'bop_toolkit' in sys_path:
sys.path.append(sys_path)
if temp_dir is None:
temp_dir = Utility.get_temporary_directory()
scale = 0.001 if mm2m else 1
bop_dataset_name = os.path.basename(bop_dataset_path)
has_external_texture = bop_dataset_name in ["ycbv", "ruapc"]
if obj_ids or sample_objects:
allow_duplication = True
else:
allow_duplication = False
datasets_path = os.path.dirname(bop_dataset_path)
dataset = os.path.basename(bop_dataset_path)
print("bob: {}, dataset_path: {}".format(bop_dataset_path, datasets_path))
print("dataset: {}".format(dataset))
try:
from bop_toolkit_lib import dataset_params, inout
except ImportError as error:
print('ERROR: Please download the bop_toolkit package and add it to sys_paths in config!')
print('https://github.com/thodan/bop_toolkit')
raise error
model_p = dataset_params.get_model_params(datasets_path, dataset, model_type=model_type if model_type else None)
cam_p = dataset_params.get_camera_params(datasets_path, dataset, cam_type=cam_type if cam_type else None)
try:
split_p = dataset_params.get_split_params(datasets_path, dataset, split=split)
except ValueError:
raise Exception("Wrong path or {} split does not exist in {}.".format(split, dataset))
bpy.context.scene.render.resolution_x = cam_p['im_size'][0]
bpy.context.scene.render.resolution_y = cam_p['im_size'][1]
loaded_objects = []
# only load all/selected objects here, use other modules for setting poses
# e.g. camera.CameraSampler / object.ObjectPoseSampler
if scene_id == -1:
# TLESS exception because images are cropped
if bop_dataset_name in ['tless']:
cam_p['K'][0, 2] = split_p['im_size'][0] / 2
cam_p['K'][1, 2] = split_p['im_size'][1] / 2
# set camera intrinsics
CameraUtility.set_intrinsics_from_K_matrix(cam_p['K'], split_p['im_size'][0], split_p['im_size'][1])
obj_ids = obj_ids if obj_ids else model_p['obj_ids']
# if sampling is enabled
if sample_objects:
loaded_ids = {}
loaded_amount = 0
if obj_instances_limit != -1 and len(obj_ids) * obj_instances_limit < num_of_objs_to_sample:
raise RuntimeError("{}'s {} split contains {} objects, {} object where requested to sample with "
"an instances limit of {}. Raise the limit amount or decrease the requested "
"amount of objects.".format(bop_dataset_path, split, len(obj_ids),
num_of_objs_to_sample,
obj_instances_limit))
while loaded_amount != num_of_objs_to_sample:
random_id = choice(obj_ids)
if random_id not in loaded_ids.keys():
loaded_ids.update({random_id: 0})
# if there is no limit or if there is one, but it is not reached for this particular object
if obj_instances_limit == -1 or loaded_ids[random_id] < obj_instances_limit:
cur_obj = BopLoader._load_mesh(random_id, model_p, bop_dataset_name, has_external_texture, temp_dir, allow_duplication, scale)
loaded_ids[random_id] += 1
loaded_amount += 1
loaded_objects.append(cur_obj)
else:
print("ID {} was loaded {} times with limit of {}. Total loaded amount {} while {} are "
"being requested".format(random_id, loaded_ids[random_id], obj_instances_limit,
loaded_amount, num_of_objs_to_sample))
else:
for obj_id in obj_ids:
cur_obj = BopLoader._load_mesh(obj_id, model_p, bop_dataset_name, has_external_texture, temp_dir, allow_duplication, scale)
loaded_objects.append(cur_obj)
# replicate scene: load scene objects, object poses, camera intrinsics and camera poses
else:
sc_gt = inout.load_scene_gt(split_p['scene_gt_tpath'].format(**{'scene_id': scene_id}))
sc_camera = inout.load_json(split_p['scene_camera_tpath'].format(**{'scene_id': scene_id}))
for i, (cam_id, insts) in enumerate(sc_gt.items()):
cam_K, cam_H_m2c_ref = BopLoader._get_ref_cam_extrinsics_intrinsics(sc_camera, cam_id, insts, scale)
if i == 0:
# define world = first camera
cam_H_m2w_ref = cam_H_m2c_ref.copy()
cur_objs = []
# load scene objects and set their poses
for inst in insts:
cur_objs.append(BopLoader._load_mesh(inst['obj_id'], model_p, bop_dataset_name, has_external_texture, temp_dir, allow_duplication, scale))
BopLoader.set_object_pose(cur_objs[-1], inst, scale)
cam_H_c2w = BopLoader._compute_camera_to_world_trafo(cam_H_m2w_ref, cam_H_m2c_ref, source_frame)
# set camera intrinsics
CameraUtility.set_intrinsics_from_K_matrix(cam_K, split_p['im_size'][0], split_p['im_size'][1])
# set camera extrinsics as next frame
frame_id = CameraUtility.add_camera_pose(cam_H_c2w)
# Add key frame for camera shift, as it changes from frame to frame in the tless replication
cam = bpy.context.scene.camera.data
cam.keyframe_insert(data_path='shift_x', frame=frame_id)
cam.keyframe_insert(data_path='shift_y', frame=frame_id)
# Copy object poses to key frame (to be sure)
for cur_obj in cur_objs:
BopLoader._insert_key_frames(cur_obj, frame_id)
# move the origin of the object to the world origin and on top of the X-Y plane
# makes it easier to place them later on, this does not change the `.location`
# This is only useful if the BOP objects are not used in a pose estimation scenario.
if move_origin_to_x_y_plane:
for obj in loaded_objects:
obj.move_origin_to_bottom_mean_point()
return loaded_objects
def render_segmap(output_dir: Optional[str] = None, temp_dir: Optional[str] = None,
map_by: Union[str, List[str]] = "class",
default_values: Optional[Dict[str, int]] = None, file_prefix: str = "segmap_",
output_key: str = "segmap", segcolormap_output_file_prefix: str = "instance_attribute_map_",
segcolormap_output_key: str = "segcolormap", use_alpha_channel: bool = False,
render_colorspace_size_per_dimension: int = 2048) -> \
Dict[str, Union[np.ndarray, List[np.ndarray]]]:
""" Renders segmentation maps for all frames
:param output_dir: The directory to write images to.
:param temp_dir: The directory to write intermediate data to.
:param map_by: The attributes to be used for color mapping.
:param default_values: The default values used for the keys used in attributes, if None is {"class": 0}.
:param file_prefix: The prefix to use for writing the images.
:param output_key: The key to use for registering the output.
:param segcolormap_output_file_prefix: The prefix to use for writing the segmentation-color map csv.
:param segcolormap_output_key: The key to use for registering the segmentation-color map output.
:param use_alpha_channel: If true, the alpha channel stored in .png textures is used.
:param render_colorspace_size_per_dimension: As we use float16 for storing the rendering, the interval of \
integers which can be precisely stored is [-2048, 2048]. As \
blender does not allow negative values for colors, we use \
[0, 2048] ** 3 as our color space which allows ~8 billion \
different colors/objects. This should be enough.
:return: dict of lists of segmaps and (for instance segmentation) segcolormaps
"""
if output_dir is None:
output_dir = Utility.get_temporary_directory()
if temp_dir is None:
temp_dir = Utility.get_temporary_directory()
if default_values is None:
default_values = {"class": 0}
with Utility.UndoAfterExecution():
RendererUtility._render_init()
RendererUtility.set_samples(1)
RendererUtility.set_adaptive_sampling(0)
RendererUtility.set_denoiser(None)
RendererUtility.set_light_bounces(1, 0, 0, 1, 0, 8, 0)
attributes = map_by
if 'class' in default_values:
default_values['cp_category_id'] = default_values['class']
# Get objects with meshes (i.e. not lights or cameras)
objs_with_mats = get_all_blender_mesh_objects()
colors, num_splits_per_dimension, objects = _colorize_objects_for_instance_segmentation(
objs_with_mats, use_alpha_channel,
render_colorspace_size_per_dimension)
bpy.context.scene.cycles.filter_width = 0.0
if use_alpha_channel:
MaterialLoaderUtility.add_alpha_channel_to_textures(
blurry_edges=False)
# Determine path for temporary and for final output
temporary_segmentation_file_path = os.path.join(temp_dir, "seg_")
final_segmentation_file_path = os.path.join(output_dir, file_prefix)
RendererUtility.set_output_format("OPEN_EXR", 16)
RendererUtility.render(temp_dir, "seg_", None, return_data=False)
# Find optimal dtype of output based on max index
for dtype in [np.uint8, np.uint16, np.uint32]:
optimal_dtype = dtype
if np.iinfo(optimal_dtype).max >= len(colors) - 1:
break
if isinstance(attributes, str):
# only one result is requested
result_channels = 1
attributes = [attributes]
elif isinstance(attributes, list):
result_channels = len(attributes)
else:
raise Exception(
"The type of this is not supported here: {}".format(
attributes))
# define them for the avoid rendering case
there_was_an_instance_rendering = False
list_of_attributes: List[str] = []
# Check if stereo is enabled
if bpy.context.scene.render.use_multiview:
suffixes = ["_L", "_R"]
else:
suffixes = [""]
return_dict: Dict[str, Union[np.ndarray, List[np.ndarray]]] = {}
save_in_csv_attributes: Dict[int, Dict[str, Any]] = {}
# After rendering
for frame in range(
bpy.context.scene.frame_start,
bpy.context.scene.frame_end): # for each rendered frame
save_in_csv_attributes: Dict[int, Dict[str, Any]] = {}
for suffix in suffixes:
file_path = temporary_segmentation_file_path + (
"%04d" % frame) + suffix + ".exr"
segmentation = load_image(file_path)
print(file_path, segmentation.shape)
segmap = Utility.map_back_from_equally_spaced_equidistant_values(
segmentation, num_splits_per_dimension,
render_colorspace_size_per_dimension)
segmap = segmap.astype(optimal_dtype)
object_ids = np.unique(segmap)
max_id = np.max(object_ids)
if max_id >= len(objects):
raise Exception(
"There are more object colors than there are objects")
combined_result_map = []
there_was_an_instance_rendering = False
list_of_attributes = []
channels = []
for channel_id in range(result_channels):
num_default_values = 0
resulting_map = np.zeros(
(segmap.shape[0], segmap.shape[1]),
dtype=optimal_dtype)
was_used = False
current_attribute = attributes[channel_id]
org_attribute = current_attribute
# if the class is used the category_id attribute is evaluated
if current_attribute == "class":
current_attribute = "cp_category_id"
# in the instance case the resulting ids are directly used
if current_attribute == "instance":
there_was_an_instance_rendering = True
resulting_map = segmap
was_used = True
else:
if current_attribute != "cp_category_id":
list_of_attributes.append(current_attribute)
# for the current attribute remove cp_ and _csv, if present
attribute = current_attribute
if attribute.startswith("cp_"):
attribute = attribute[len("cp_"):]
# check if a default value was specified
default_value_set = False
if current_attribute in default_values or attribute in default_values:
default_value_set = True
if current_attribute in default_values:
default_value = default_values[
current_attribute]
elif attribute in default_values:
default_value = default_values[attribute]
# iterate over all object ids
for object_id in object_ids:
# Convert np.uint8 to int, such that the save_in_csv_attributes dict can later be serialized
object_id = int(object_id)
# get the corresponding object via the id
current_obj = objects[object_id]
# if the current obj has a attribute with that name -> get it
if hasattr(current_obj, attribute):
value = getattr(current_obj, attribute)
# if the current object has a custom property with that name -> get it
elif current_attribute.startswith(
"cp_") and attribute in current_obj:
value = current_obj[attribute]
elif current_attribute.startswith("cf_"):
if current_attribute == "cf_basename":
value = current_obj.name
if "." in value:
value = value[:value.rfind(".")]
elif default_value_set:
# if none of the above applies use the default value
value = default_value
num_default_values += 1
else:
# if the requested current_attribute is not a custom property or a attribute
# or there is a default value stored
# it throws an exception
raise Exception(
"The obj: {} does not have the "
"attribute: {}, striped: {}. Maybe try a default "
"value.".format(current_obj.name,
current_attribute,
attribute))
# save everything which is not instance also in the .csv
if isinstance(
value,
(int, float, np.integer, np.floating)):
was_used = True
resulting_map[segmap == object_id] = value
if object_id in save_in_csv_attributes:
save_in_csv_attributes[object_id][
attribute] = value
else:
save_in_csv_attributes[object_id] = {
attribute: value
}
if was_used and num_default_values < len(object_ids):
channels.append(org_attribute)
combined_result_map.append(resulting_map)
return_dict.setdefault(
"{}_segmaps{}".format(org_attribute, suffix),
[]).append(resulting_map)
fname = final_segmentation_file_path + ("%04d" %
frame) + suffix
# combine all resulting images to one image
resulting_map = np.stack(combined_result_map, axis=2)
# remove the unneeded third dimension
if resulting_map.shape[2] == 1:
resulting_map = resulting_map[:, :, 0]
# TODO: Remove unnecessary save when we give up backwards compatibility
np.save(fname, resulting_map)
if there_was_an_instance_rendering:
mappings = []
for object_id, attribute_dict in save_in_csv_attributes.items(
):
mappings.append({"idx": object_id, **attribute_dict})
return_dict.setdefault("instance_attribute_maps",
[]).append(mappings)
# write color mappings to file
# TODO: Remove unnecessary csv file when we give up backwards compatibility
csv_file_path = os.path.join(
output_dir,
segcolormap_output_file_prefix + ("%04d.csv" % frame))
with open(csv_file_path, 'w', newline='') as csvfile:
# get from the first element the used field names
fieldnames = ["idx"]
# get all used object element keys
for object_element in save_in_csv_attributes.values():
fieldnames.extend(list(object_element.keys()))
break
for channel_name in channels:
fieldnames.append("channel_{}".format(channel_name))
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
# save for each object all values in one row
for obj_idx, object_element in save_in_csv_attributes.items(
):
object_element["idx"] = obj_idx
for i, channel_name in enumerate(channels):
object_element["channel_{}".format(
channel_name)] = i
writer.writerow(object_element)
else:
if len(list_of_attributes) > 0:
raise Exception(
"There were attributes specified in the may_by, which could not be saved as "
"there was no \"instance\" may_by key used. This is true for this/these "
"keys: {}".format(", ".join(list_of_attributes)))
# if there was no instance rendering no .csv file is generated!
# delete all saved infos about .csv
save_in_csv_attributes = {}
Utility.register_output(output_dir, file_prefix, output_key, ".npy",
"2.0.0")
if save_in_csv_attributes:
Utility.register_output(output_dir, segcolormap_output_file_prefix,
segcolormap_output_key, ".csv", "2.0.0")
return return_dict
