visualize_reconstruction.py 6 KB
Newer Older
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
1
2
3
4
5
6
7
#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

8
9
"""
Script to visualize a previously trained model. Example call:
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
10

11
12
    pytorch3d_implicitron_visualizer \
    exp_dir='./exps/checkpoint_dir' visdom_show_preds=True visdom_port=8097 \
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
13
14
15
16
17
18
19
20
21
    n_eval_cameras=40 render_size="[64,64]" video_size="[256,256]"
"""

import os
import sys
from typing import Optional, Tuple

import numpy as np
import torch
22
23
24
from omegaconf import DictConfig, OmegaConf
from pytorch3d.implicitron.models.visualization.render_flyaround import render_flyaround
from pytorch3d.implicitron.tools.config import enable_get_default_args, get_default_args
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
25

26
from .experiment import Experiment
27

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
28

29
def visualize_reconstruction(
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
30
31
32
33
34
    exp_dir: str = "",
    restrict_sequence_name: Optional[str] = None,
    output_directory: Optional[str] = None,
    render_size: Tuple[int, int] = (512, 512),
    video_size: Optional[Tuple[int, int]] = None,
35
    split: str = "train",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
36
37
38
    n_source_views: int = 9,
    n_eval_cameras: int = 40,
    visdom_show_preds: bool = False,
39
40
    visdom_server: str = "http://127.0.0.1",
    visdom_port: int = 8097,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
41
    visdom_env: Optional[str] = None,
42
    **render_flyaround_kwargs,
43
) -> None:
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
    """
    Given an `exp_dir` containing a trained Implicitron model, generates videos consisting
    of renderes of sequences from the dataset used to train and evaluate the trained
    Implicitron model.

    Args:
        exp_dir: Implicitron experiment directory.
        restrict_sequence_name: If set, defines the list of sequences to visualize.
        output_directory: If set, defines a custom directory to output visualizations to.
        render_size: The size (HxW) of the generated renders.
        video_size: The size (HxW) of the output video.
        split: The dataset split to use for visualization.
            Can be "train" / "val" / "test".
        n_source_views: The number of source views added to each rendered batch. These
            views are required inputs for models such as NeRFormer / NeRF-WCE.
        n_eval_cameras: The number of cameras each fly-around trajectory.
        visdom_show_preds: If `True`, outputs visualizations to visdom.
        visdom_server: The address of the visdom server.
        visdom_port: The port of the visdom server.
        visdom_env: If set, defines a custom name for the visdom environment.
64
65
        render_flyaround_kwargs: Keyword arguments passed to the invoked `render_flyaround`
            function (see `pytorch3d.implicitron.models.visualization.render_flyaround`).
66
67
    """

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
68
69
70
71
    # In case an output directory is specified use it. If no output_directory
    # is specified create a vis folder inside the experiment directory
    if output_directory is None:
        output_directory = os.path.join(exp_dir, "vis")
72
    os.makedirs(output_directory, exist_ok=True)
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
73
74
75
76
77
78
79
80
81
82

    # Set the random seeds
    torch.manual_seed(0)
    np.random.seed(0)

    # Get the config from the experiment_directory,
    # and overwrite relevant fields
    config = _get_config_from_experiment_directory(exp_dir)
    config.exp_dir = exp_dir
    # important so that the CO3D dataset gets loaded in full
83
84
85
86
87
88
89
90
91
    data_source_args = config.data_source_ImplicitronDataSource_args
    if "dataset_map_provider_JsonIndexDatasetMapProvider_args" in data_source_args:
        dataset_args = (
            data_source_args.dataset_map_provider_JsonIndexDatasetMapProvider_args
        )
        dataset_args.test_on_train = False
        if restrict_sequence_name is not None:
            dataset_args.restrict_sequence_name = restrict_sequence_name

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
92
    # Set the rendering image size
93
    model_factory_args = config.model_factory_ImplicitronModelFactory_args
94
    model_factory_args.force_resume = True
95
96
97
    model_args = model_factory_args.model_GenericModel_args
    model_args.render_image_width = render_size[0]
    model_args.render_image_height = render_size[1]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
98
99

    # Load the previously trained model
100
101
102
103
    experiment = Experiment(**config)
    model = experiment.model_factory(exp_dir=exp_dir)
    device = torch.device("cuda")
    model.to(device)
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
104
105
106
    model.eval()

    # Setup the dataset
107
108
    data_source = experiment.data_source
    dataset_map, _ = data_source.get_datasets_and_dataloaders()
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
109
    dataset = dataset_map[split]
110
111
    if dataset is None:
        raise ValueError(f"{split} dataset not provided")
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
112

113
114
115
116
117
    if visdom_env is None:
        visdom_env = (
            "visualizer_" + config.training_loop_ImplicitronTrainingLoop_args.visdom_env
        )

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
118
    # iterate over the sequences in the dataset
119
    for sequence_name in dataset.sequence_names():
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
120
        with torch.no_grad():
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
            render_kwargs = {
                "dataset": dataset,
                "sequence_name": sequence_name,
                "model": model,
                "output_video_path": os.path.join(output_directory, "video"),
                "n_source_views": n_source_views,
                "visdom_show_preds": visdom_show_preds,
                "n_flyaround_poses": n_eval_cameras,
                "visdom_server": visdom_server,
                "visdom_port": visdom_port,
                "visdom_environment": visdom_env,
                "video_resize": video_size,
                "device": device,
                **render_flyaround_kwargs,
            }
            render_flyaround(**render_kwargs)
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
137
138


139
140
141
142
enable_get_default_args(visualize_reconstruction)


def _get_config_from_experiment_directory(experiment_directory) -> DictConfig:
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
143
144
    cfg_file = os.path.join(experiment_directory, "expconfig.yaml")
    config = OmegaConf.load(cfg_file)
145
    # pyre-ignore[7]
146
    return OmegaConf.merge(get_default_args(Experiment), config)
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
147
148


149
def main(argv=sys.argv) -> None:
150
151
    # automatically parses arguments of visualize_reconstruction
    cfg = OmegaConf.create(get_default_args(visualize_reconstruction))
152
    cfg.update(OmegaConf.from_cli(argv))
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
153
    with torch.no_grad():
154
        visualize_reconstruction(**cfg)
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
155
156
157


if __name__ == "__main__":
158
    main()