def resize_frames(src_dir, dst_dir, rig, first, last, threshold=None):
"""Resizes a frame to the appropriate pyramid level sizes. Files are saved in
level_0/[camera], ..., level_9/[camera] in the destination directory.
Args:
src_dir (str): Path to the source directory.
dst_dir (str): Path to the destination directory.
rig (dict[str, _]): Rig descriptor object.
first (str): Name of the first frame to render.
last (str): Name of the last frame to render.
threshold (int): Threshold to be used for binary thresholding. No thresholding
is performed if None is passed in.
"""
num_workers = mp.cpu_count()
try:
mp.set_start_method(
"spawn",
force=True) # see: https://github.com/opencv/opencv/issues/5150
except RuntimeError:
pass
pool = mp.Pool(num_workers)
for frame in range(int(first), int(last) + 1):
for camera in rig["cameras"]:
verify_frame(src_dir, camera["id"], get_frame_name(frame))
pool.apply_async(
resize_camera,
args=(
src_dir,
dst_dir,
camera["id"],
camera["resolution"],
get_frame_name(frame),
threshold,
),
)
pool.close()
pool.join()
def _run_bin(msg):
"""Runs the binary associated with the message. The execution assumes the worker is
running in a configured Docker container.
Args:
msg (dict[str, str]): Message received from RabbitMQ publisher.
"""
msg_cp = copy(msg)
# The binary flag convention includes the "last" frame
msg_cp["last"] = get_frame_name(int(msg["last"]))
app_name = msg_cp["app"].split(":")[0]
relevant_flags = [flag["name"] for flag in bin_to_flags[app_name]]
cmd_flags = " ".join([
f"--{flag}={msg_cp[flag]}" for flag in relevant_flags
if flag in msg_cp and msg_cp[flag] != ""
])
# Order is determined to prevent substrings from being accidentally replaced
input_root = msg_cp["input_root"].rstrip("/")
output_root = msg_cp["output_root"].rstrip("/")
root_order = ([output_root, input_root]
if input_root in output_root else [input_root, output_root])
root_to_docker = {
input_root: config.DOCKER_INPUT_ROOT,
output_root: config.DOCKER_OUTPUT_ROOT,
}
for root in root_order:
if not os.path.exists(root):
cmd_flags = cmd_flags.replace(root, root_to_docker[root])
bin_path = os.path.join(config.DOCKER_BUILD_ROOT, "bin", app_name + ".exe")
cmd = f"set GLOG_alsologtostderr=1 && set GLOG_stderrthreshold=0 && {bin_path} {cmd_flags}"
run_command(cmd)
def main():
"""Runs the main render pipeline with the parameters passed in through command line args."""
base_params = {
flag: value
for flag, value in FLAGS.flag_values_dict().items()
if flag in setup.flag_names
}
if not FLAGS.skip_setup and FLAGS.cloud == "":
setup.setup_workers(base_params)
# If the address is a Samba or local endpoint, we have the addresses mapped in Docker
# Access to the externally visible endpoint is only necessary in remote cases
input_protocol = Address(FLAGS.input_root).protocol
base_params["num_levels"] = len(config.WIDTHS)
if input_protocol is None or input_protocol == "smb":
base_params["input_root"] = config.DOCKER_INPUT_ROOT
base_params["output_root"] = config.DOCKER_OUTPUT_ROOT
base_params["rig"] = FLAGS.rig.replace(FLAGS.input_root,
base_params["input_root"])
input_image_types = {
"background_color",
"background_disp",
"color",
"foreground_masks",
}
for image_type in input_image_types:
set_input_param(base_params, image_type)
# frame_chunks use the Python range standard where first is included but last excluded
frame_chunks = [{
"first":
get_frame_name(frame),
"last":
get_frame_name(min(int(FLAGS.last), frame + FLAGS.chunk_size - 1)),
} for frame in range(int(FLAGS.first),
int(FLAGS.last) + 1, FLAGS.chunk_size)]
if FLAGS.background_frame == "":
background_frame = None
else:
background_frame = [{
"first":
get_frame_name(int(FLAGS.background_frame)),
"last":
get_frame_name(int(FLAGS.background_frame)),
}]
pipeline = Pipeline(FLAGS.master, base_params, frame_chunks,
background_frame, FLAGS.force_recompute)
# We need resized colors to compute foreground masks
pipeline_stages = [
(pipeline.precompute_resizes, FLAGS.run_precompute_resizes),
(pipeline.generate_foreground_masks,
FLAGS.run_generate_foreground_masks),
]
if FLAGS.use_foreground_masks:
# Resize foreground masks
pipeline_stages.append((
pipeline.precompute_resizes_foreground,
FLAGS.run_precompute_resizes_foreground,
))
pipeline_stages.append(
(pipeline.depth_estimation, FLAGS.run_depth_estimation))
if FLAGS.format == "6dof":
pipeline_stages += [
(pipeline.convert_to_binary, FLAGS.run_convert_to_binary),
(pipeline.fusion, FLAGS.run_fusion),
]
else:
pipeline_stages.append(
(pipeline.simple_mesh_renderer, FLAGS.run_simple_mesh_renderer))
pipeline.run(pipeline_stages)
setup.cleanup_workers()
def depth_estimation(self):
"""Runs distributed depth estimation with temporal filtering."""
post_resize_params = copy(self.base_params)
if self.base_params["disparity_type"] == "disparity":
post_resize_params["color"] = os.path.join(
self.base_params["input_root"],
image_type_paths[config.type_to_levels_type["color"]],
)
post_resize_params["foreground_masks"] = os.path.join(
self.base_params["input_root"],
image_type_paths[
config.type_to_levels_type["foreground_masks"]],
)
post_resize_params["background_disp"] = os.path.join(
self.base_params["input_root"],
image_type_paths[
config.type_to_levels_type["background_disp"]],
)
elif self.base_params["disparity_type"] == "background_disp":
post_resize_params["color"] = os.path.join(
self.base_params["input_root"],
image_type_paths[config.type_to_levels_type["color"]],
)
start_level = (post_resize_params["level_start"]
if post_resize_params["level_start"] != -1 else
len(config.WIDTHS) - 1)
if post_resize_params["level_end"] != -1:
end_level = post_resize_params["level_end"]
else:
for level, width in enumerate(config.WIDTHS):
if post_resize_params["resolution"] >= width:
end_level = level
break
# Ranges for temporal filtering (only used if performing temporal filtering)
filter_ranges = [{
"first":
frame_chunk["first"],
"last":
frame_chunk["last"],
"filter_first":
get_frame_name(
max(
int(post_resize_params["first"]),
int(frame_chunk["first"]) -
post_resize_params["time_radius"],
)),
"filter_last":
get_frame_name(
min(
int(post_resize_params["last"]),
int(frame_chunk["last"]) +
post_resize_params["time_radius"],
)),
} for frame_chunk in self.frame_chunks]
for level in range(start_level, end_level - 1, -1):
depth_params = copy(post_resize_params)
if level != end_level:
depth_params[
"output_formats"] = "pfm" # Force only PFM at non-finest levels
depth_params.update({
"app":
f"DerpCLI: Level {level}",
"level_start":
level,
"level_end":
level,
"image_type":
depth_params["disparity_type"],
"dst_level":
level,
"dst_image_type":
depth_params["disparity_type"],
})
self.run_halted_queue(depth_params, self.frame_chunks)
if post_resize_params["do_temporal_filter"]:
filter_params = copy(post_resize_params)
filter_params.update({
"app":
"TemporalBilateralFilter",
"level":
level,
"use_foreground_masks":
post_resize_params["do_temporal_masking"],
"dst_level":
level,
"dst_image_type":
"disparity_time_filtered",
})
self.run_halted_queue(filter_params, filter_ranges)
transfer_params = copy(post_resize_params)
transfer_params.update({
"app": "Transfer",
"src_level": level,
"src_image_type": "disparity_time_filtered",
"dst_level": level,
"dst_image_type": "disparity",
"force_recompute": True,
})
self.run_halted_queue(transfer_params, self.frame_chunks)
if post_resize_params["resolution"] > config.WIDTHS[end_level]:
# The upsampling color level is the smallest one larger than our last level
dst_level = end_level - 1 if end_level > 0 else None
upsample_params = copy(self.base_params)
upsample_params.update({
"app":
"UpsampleDisparity",
"level":
end_level,
"image_type":
post_resize_params["disparity_type"],
"dst_level":
dst_level,
"dst_image_type":
config.type_to_upsample_type["disparity"],
})
if post_resize_params["disparity_type"] == "background_disp":
frame_chunks = self.background_frame
elif post_resize_params["disparity_type"] == "disparity":
frame_chunks = self.frame_chunks
else:
raise Exception(
f"Invalid disparity type {post_resize_params['disparity_type']}"
)
self.run_halted_queue(upsample_params, frame_chunks)
transfer_params = copy(post_resize_params)
transfer_params.update({
"app":
"Transfer",
"src_level":
None,
"src_image_type":
config.type_to_upsample_type["disparity"],
"dst_level":
None,
"dst_image_type":
post_resize_params["disparity_type"],
})
self.run_halted_queue(transfer_params, frame_chunks)
else:
transfer_params = copy(post_resize_params)
transfer_params.update({
"app":
"Transfer",
"src_level":
end_level,
"src_image_type":
post_resize_params["disparity_type"],
"dst_level":
None,
"dst_image_type":
post_resize_params["disparity_type"],
})
self.run_halted_queue(transfer_params, self.frame_chunks)