Пример #1
0
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()
Пример #2
0
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)
Пример #3
0
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()
Пример #4
0
    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)