API proposal

#18
by pravdomil - opened

rerender exposes:

def generate_first_result(
        state: GlobalState,
        cfg: Config,
        input_image: numpy.ndarray,
) -> torch.Tensor

def generate_next_result(
        cfg: Config,
        state: GlobalState,
        first_image: numpy.ndarray,
        first_result: torch.Tensor,
        previous_image: numpy.ndarray,
        previous_result: torch.Tensor,
        i: int,
        image: numpy.ndarray,
) -> torch.Tensor

then it can be used as follows:

def main(cfg: Config):
    state = get_state(cfg)

    import decord
    reader = decord.VideoReader(cfg.input_path, width=cfg.image_resolution, height=cfg.image_resolution)

    first_image = reader.next().asnumpy()
    first_result = generate_first_result(state, cfg, first_image)

    previous_image = first_image
    previous_result = first_result

    writer = moviepy.video.io.ffmpeg_writer.FFMPEG_VideoWriter(
        cfg.output_path,
        (cfg.image_resolution, cfg.image_resolution),
        reader.get_avg_fps() / cfg.frame_skip,
    )

    frame_indexes = range(cfg.start_frame, cfg.end_frame, cfg.frame_skip)
    for i, index in enumerate(frame_indexes):
        print(str(round(i / len(frame_indexes) * 100)) + "%")

        reader.seek_accurate(index)
        image = reader.next().asnumpy()

        result = generate_next_result(cfg, state, first_image, first_result, previous_image, previous_result, i, image)

        writer.write_frame(torch_to_numpy(result)[0])

        previous_image = image
        previous_result = result

    writer.close()

Sign up or log in to comment