def _get_frame_gt(dataset_objects: List[bpy.types.Mesh],
unit_scaling: float,
ignore_dist_thres: float,
destination_frame: List[str] = ["X", "-Y", "-Z"]):
""" Returns GT pose annotations between active camera and objects.
:param dataset_objects: Save annotations for these objects.
:param unit_scaling: 1000. for outputting poses in mm
:param ignore_dist_thres: Distance between camera and object after which object is ignored. Mostly due to failed physics.
:param destination_frame: Transform poses from Blender internal coordinates to OpenCV coordinates
:return: A list of GT camera-object pose annotations for scene_gt.json
"""
H_c2w_opencv = Matrix(
WriterUtility.get_cam_attribute(
bpy.context.scene.camera,
'cam2world_matrix',
local_frame_change=destination_frame))
frame_gt = []
for obj in dataset_objects:
H_m2w = Matrix(
WriterUtility.get_common_attribute(obj, 'matrix_world'))
cam_H_m2c = H_c2w_opencv.inverted() @ H_m2w
cam_R_m2c = cam_H_m2c.to_quaternion().to_matrix()
cam_t_m2c = cam_H_m2c.to_translation()
# ignore examples that fell through the plane
if not np.linalg.norm(list(cam_t_m2c)) > ignore_dist_thres:
cam_t_m2c = list(cam_t_m2c * unit_scaling)
frame_gt.append({
'cam_R_m2c':
list(cam_R_m2c[0]) + list(cam_R_m2c[1]) +
list(cam_R_m2c[2]),
'cam_t_m2c':
cam_t_m2c,
'obj_id':
obj["category_id"]
})
else:
print('ignored obj, ', obj["category_id"], 'because either ')
print(
'(1) it is further away than parameter "ignore_dist_thres: ",',
ignore_dist_thres)
print(
'(e.g. because it fell through a plane during physics sim)'
)
print('or')
print('(2) the object pose has not been given in meters')
return frame_gt
def _get_frame_camera(save_world2cam,
depth_scale=1.0,
unit_scaling=1000.,
destination_frame=["X", "-Y", "-Z"]):
""" Returns camera parameters for the active camera.
:param save_world2cam: If true, camera to world transformations "cam_R_w2c", "cam_t_w2c" are saved in scene_camera.json
:param depth_scale: Multiply the uint16 output depth image with this factor to get depth in mm.
:param unit_scaling: 1000. for outputting poses in mm
:param destination_frame: Transform poses from Blender internal coordinates to OpenCV coordinates
:return: dict containing info for scene_camera.json
"""
cam_K = WriterUtility.get_cam_attribute(bpy.context.scene.camera,
'cam_K')
frame_camera_dict = {
'cam_K': cam_K[0] + cam_K[1] + cam_K[2],
'depth_scale': depth_scale
}
if save_world2cam:
H_c2w_opencv = Matrix(
WriterUtility.get_cam_attribute(
bpy.context.scene.camera,
'cam2world_matrix',
local_frame_change=destination_frame))
H_w2c_opencv = H_c2w_opencv.inverted()
R_w2c_opencv = H_w2c_opencv.to_quaternion().to_matrix()
t_w2c_opencv = H_w2c_opencv.to_translation() * unit_scaling
frame_camera_dict['cam_R_w2c'] = list(R_w2c_opencv[0]) + list(
R_w2c_opencv[1]) + list(R_w2c_opencv[2])
frame_camera_dict['cam_t_w2c'] = list(t_w2c_opencv)
return frame_camera_dict
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 _write_camera(camera_path: str, depth_scale: float = 1.0):
""" Writes camera.json into dataset_dir.
:param camera_path: Path to camera.json
:param depth_scale: Multiply the uint16 output depth image with this factor to get depth in mm.
"""
cam_K = WriterUtility.get_cam_attribute(bpy.context.scene.camera,
'cam_K')
camera = {
'cx': cam_K[0][2],
'cy': cam_K[1][2],
'depth_scale': depth_scale,
'fx': cam_K[0][0],
'fy': cam_K[1][1],
'height': bpy.context.scene.render.resolution_y,
'width': bpy.context.scene.render.resolution_x
}
BopWriterUtility._save_json(camera_path, camera)
def _load_and_postprocess(self, file_path, key, version="1.0.0"):
"""
Loads an image and post process it.
:param file_path: Image path. Type: string.
:param key: The image's key with regards to the hdf5 file. Type: string.
:param version: The version number original data. Type: String. Default: 1.0.0.
:return: The post-processed image that was loaded using the file path.
"""
data = WriterUtility.load_output_file(resolve_path(file_path),
self.write_alpha_channel,
remove=False)
data, new_key, new_version = self._apply_postprocessing(
key, data, version)
if isinstance(data, np.ndarray):
print("Key: " + key + " - shape: " + str(data.shape) +
" - dtype: " + str(data.dtype) + " - path: " + file_path)
else:
print("Key: " + key + " - path: " + file_path)
return data, new_key, new_version
def run(self):
if self._avoid_output:
print("Avoid output is on, no output produced!")
return
if self._append_to_existing_output:
frame_offset = 0
# Look for hdf5 file with highest index
for path in os.listdir(self._output_dir):
if path.endswith(".hdf5"):
index = path[:-len(".hdf5")]
if index.isdigit():
frame_offset = max(frame_offset, int(index) + 1)
else:
frame_offset = 0
# Go through all frames
for frame in range(bpy.context.scene.frame_start,
bpy.context.scene.frame_end):
# Create output hdf5 file
hdf5_path = os.path.join(self._output_dir,
str(frame + frame_offset) + ".hdf5")
with h5py.File(hdf5_path, "w") as f:
if not GlobalStorage.is_in_storage("output"):
print("No output was designed in prior models!")
return
# Go through all the output types
print("Merging data for frame " + str(frame) + " into " +
hdf5_path)
for output_type in GlobalStorage.get("output"):
# Build path (path attribute is format string)
file_path = output_type["path"]
if '%' in file_path:
file_path = file_path % frame
# Check if file exists
if not os.path.exists(file_path):
# If not try stereo suffixes
path_l, path_r = WriterUtility._get_stereo_path_pair(
file_path)
if not os.path.exists(path_l) or not os.path.exists(
path_r):
raise Exception("File not found: " + file_path)
else:
use_stereo = True
else:
use_stereo = False
if use_stereo:
path_l, path_r = WriterUtility._get_stereo_path_pair(
file_path)
img_l, new_key, new_version = self._load_and_postprocess(
path_l, output_type["key"], output_type["version"])
img_r, new_key, new_version = self._load_and_postprocess(
path_r, output_type["key"], output_type["version"])
if self.config.get_bool("stereo_separate_keys", False):
WriterUtility._write_to_hdf_file(
f, new_key + "_0", img_l)
WriterUtility._write_to_hdf_file(
f, new_key + "_1", img_r)
else:
data = np.array([img_l, img_r])
WriterUtility._write_to_hdf_file(f, new_key, data)
else:
data, new_key, new_version = self._load_and_postprocess(
file_path, output_type["key"],
output_type["version"])
WriterUtility._write_to_hdf_file(f, new_key, data)
WriterUtility._write_to_hdf_file(f, new_key + "_version",
np.string_([new_version]))
blender_proc_version = Utility.get_current_version()
if blender_proc_version:
WriterUtility._write_to_hdf_file(
f, "blender_proc_version",
np.string_(blender_proc_version))
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 _write_frames(chunks_dir: str,
dataset_objects: list,
depths: List[np.ndarray] = [],
colors: List[np.ndarray] = [],
color_file_format: str = "PNG",
depth_scale: float = 1.0,
frames_per_chunk: int = 1000,
m2mm: bool = True,
ignore_dist_thres: float = 100.,
save_world2cam: bool = True,
jpg_quality: int = 95):
"""Write each frame's ground truth into chunk directory in BOP format
:param chunks_dir: Path to the output directory of the current chunk.
:param dataset_objects: Save annotations for these objects.
:param depths: List of depth images in m to save
:param colors: List of color images to save
:param color_file_format: File type to save color images. Available: "PNG", "JPEG"
:param jpg_quality: If color_file_format is "JPEG", save with the given quality.
:param depth_scale: Multiply the uint16 output depth image with this factor to get depth in mm. Used to trade-off between depth accuracy
and maximum depth value. Default corresponds to 65.54m maximum depth and 1mm accuracy.
:param ignore_dist_thres: Distance between camera and object after which object is ignored. Mostly due to failed physics.
:param m2mm: Original bop annotations and models are in mm. If true, we convert the gt annotations to mm here. This
is needed if BopLoader option mm2m is used.
:param frames_per_chunk: Number of frames saved in each chunk (called scene in BOP)
"""
# Format of the depth images.
depth_ext = '.png'
rgb_tpath = os.path.join(chunks_dir, '{chunk_id:06d}', 'rgb',
'{im_id:06d}' + '{im_type}')
depth_tpath = os.path.join(chunks_dir, '{chunk_id:06d}', 'depth',
'{im_id:06d}' + depth_ext)
chunk_camera_tpath = os.path.join(chunks_dir, '{chunk_id:06d}',
'scene_camera.json')
chunk_gt_tpath = os.path.join(chunks_dir, '{chunk_id:06d}',
'scene_gt.json')
# Paths to the already existing chunk folders (such folders may exist
# when appending to an existing dataset).
chunk_dirs = sorted(glob.glob(os.path.join(chunks_dir, '*')))
chunk_dirs = [d for d in chunk_dirs if os.path.isdir(d)]
# Get ID's of the last already existing chunk and frame.
curr_chunk_id = 0
curr_frame_id = 0
if len(chunk_dirs):
last_chunk_dir = sorted(chunk_dirs)[-1]
last_chunk_gt_fpath = os.path.join(last_chunk_dir, 'scene_gt.json')
chunk_gt = BopWriterUtility._load_json(last_chunk_gt_fpath,
keys_to_int=True)
# Last chunk and frame ID's.
last_chunk_id = int(os.path.basename(last_chunk_dir))
last_frame_id = int(sorted(chunk_gt.keys())[-1])
# Current chunk and frame ID's.
curr_chunk_id = last_chunk_id
curr_frame_id = last_frame_id + 1
if curr_frame_id % frames_per_chunk == 0:
curr_chunk_id += 1
curr_frame_id = 0
# Initialize structures for the GT annotations and camera info.
chunk_gt = {}
chunk_camera = {}
if curr_frame_id != 0:
# Load GT and camera info of the chunk we are appending to.
chunk_gt = BopWriterUtility._load_json(
chunk_gt_tpath.format(chunk_id=curr_chunk_id),
keys_to_int=True)
chunk_camera = BopWriterUtility._load_json(
chunk_camera_tpath.format(chunk_id=curr_chunk_id),
keys_to_int=True)
# Go through all frames.
num_new_frames = bpy.context.scene.frame_end - bpy.context.scene.frame_start
if len(depths) != len(colors) != num_new_frames:
raise Exception(
"The amount of images stored in the depths/colors does not correspond to the amount"
"of images specified by frame_start to frame_end.")
for frame_id in range(bpy.context.scene.frame_start,
bpy.context.scene.frame_end):
# Activate frame.
bpy.context.scene.frame_set(frame_id)
# Reset data structures and prepare folders for a new chunk.
if curr_frame_id == 0:
chunk_gt = {}
chunk_camera = {}
os.makedirs(
os.path.dirname(
rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=0,
im_type='PNG')))
os.makedirs(
os.path.dirname(
depth_tpath.format(chunk_id=curr_chunk_id, im_id=0)))
# Get GT annotations and camera info for the current frame.
# Output translation gt in m or mm
unit_scaling = 1000. if m2mm else 1.
chunk_gt[curr_frame_id] = BopWriterUtility._get_frame_gt(
dataset_objects, unit_scaling, ignore_dist_thres)
chunk_camera[curr_frame_id] = BopWriterUtility._get_frame_camera(
save_world2cam, depth_scale, unit_scaling)
if colors:
color_rgb = colors[frame_id]
color_bgr = color_rgb[..., ::-1].copy()
if color_file_format == 'PNG':
rgb_fpath = rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id,
im_type='.png')
cv2.imwrite(rgb_fpath, color_bgr)
elif color_file_format == 'JPEG':
rgb_fpath = rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id,
im_type='.jpg')
cv2.imwrite(rgb_fpath, color_bgr,
[int(cv2.IMWRITE_JPEG_QUALITY), jpg_quality])
else:
rgb_output = Utility.find_registered_output_by_key("colors")
if rgb_output is None:
raise Exception("RGB image has not been rendered.")
color_ext = '.png' if rgb_output['path'].endswith(
'png') else '.jpg'
# Copy the resulting RGB image.
rgb_fpath = rgb_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id,
im_type=color_ext)
shutil.copyfile(rgb_output['path'] % frame_id, rgb_fpath)
if depths:
depth = depths[frame_id]
else:
# Load the resulting dist image.
dist_output = Utility.find_registered_output_by_key("distance")
if dist_output is None:
raise Exception("Distance image has not been rendered.")
distance = WriterUtility.load_output_file(resolve_path(
dist_output['path'] % frame_id),
remove=False)
depth = dist2depth(distance)
# Scale the depth to retain a higher precision (the depth is saved
# as a 16-bit PNG image with range 0-65535).
depth_mm = 1000.0 * depth # [m] -> [mm]
depth_mm_scaled = depth_mm / float(depth_scale)
# Save the scaled depth image.
depth_fpath = depth_tpath.format(chunk_id=curr_chunk_id,
im_id=curr_frame_id)
BopWriterUtility._save_depth(depth_fpath, depth_mm_scaled)
# Save the chunk info if we are at the end of a chunk or at the last new frame.
if ((curr_frame_id + 1) % frames_per_chunk == 0) or \
(frame_id == num_new_frames - 1):
# Save GT annotations.
BopWriterUtility._save_json(
chunk_gt_tpath.format(chunk_id=curr_chunk_id), chunk_gt)
# Save camera info.
BopWriterUtility._save_json(
chunk_camera_tpath.format(chunk_id=curr_chunk_id),
chunk_camera)
# Update ID's.
curr_chunk_id += 1
curr_frame_id = 0
else:
curr_frame_id += 1