Add files using upload-large-folder tool

c2ba671 verified 3 months ago

6.16 kB

#!/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.

# pyre-unsafe

"""
Script to visualize a previously trained model. Example call:
    pytorch3d_implicitron_visualizer \
    exp_dir='./exps/checkpoint_dir' visdom_show_preds=True visdom_port=8097 \
    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
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

from .experiment import Experiment


def visualize_reconstruction(
    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,
    split: str = "train",
    n_source_views: int = 9,
    n_eval_cameras: int = 40,
    visdom_show_preds: bool = False,
    visdom_server: str = "http://127.0.0.1",
    visdom_port: int = 8097,
    visdom_env: Optional[str] = None,
    **render_flyaround_kwargs,
) -> None:
    """
    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.
        render_flyaround_kwargs: Keyword arguments passed to the invoked `render_flyaround`
            function (see `pytorch3d.implicitron.models.visualization.render_flyaround`).
    """

    # 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")
    os.makedirs(output_directory, exist_ok=True)

    # 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
    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

    # Set the rendering image size
    model_factory_args = config.model_factory_ImplicitronModelFactory_args
    model_factory_args.force_resume = True
    model_args = model_factory_args.model_GenericModel_args
    model_args.render_image_width = render_size[0]
    model_args.render_image_height = render_size[1]

    # Load the previously trained model
    experiment = Experiment(**config)
    model = experiment.model_factory(exp_dir=exp_dir)
    device = torch.device("cuda")
    model.to(device)
    model.eval()

    # Setup the dataset
    data_source = experiment.data_source
    dataset_map, _ = data_source.get_datasets_and_dataloaders()
    dataset = dataset_map[split]
    if dataset is None:
        raise ValueError(f"{split} dataset not provided")

    if visdom_env is None:
        visdom_env = (
            "visualizer_" + config.training_loop_ImplicitronTrainingLoop_args.visdom_env
        )

    # iterate over the sequences in the dataset
    for sequence_name in dataset.sequence_names():
        with torch.no_grad():
            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)


enable_get_default_args(visualize_reconstruction)


def _get_config_from_experiment_directory(experiment_directory) -> DictConfig:
    cfg_file = os.path.join(experiment_directory, "expconfig.yaml")
    config = OmegaConf.load(cfg_file)
    # pyre-ignore[7]
    return OmegaConf.merge(get_default_args(Experiment), config)


def main(argv=sys.argv) -> None:
    # automatically parses arguments of visualize_reconstruction
    cfg = OmegaConf.create(get_default_args(visualize_reconstruction))
    cfg.update(OmegaConf.from_cli(argv))
    with torch.no_grad():
        visualize_reconstruction(**cfg)


if __name__ == "__main__":
    main()

	#!/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.

	# pyre-unsafe

	"""
	Script to visualize a previously trained model. Example call:

	pytorch3d_implicitron_visualizer \
	exp_dir='./exps/checkpoint_dir' visdom_show_preds=True visdom_port=8097 \
	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
	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

	from .experiment import Experiment


	def visualize_reconstruction(
	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,
	split: str = "train",
	n_source_views: int = 9,
	n_eval_cameras: int = 40,
	visdom_show_preds: bool = False,
	visdom_server: str = "http://127.0.0.1",
	visdom_port: int = 8097,
	visdom_env: Optional[str] = None,
	**render_flyaround_kwargs,
	) -> None:
	"""
	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.
	render_flyaround_kwargs: Keyword arguments passed to the invoked `render_flyaround`
	function (see `pytorch3d.implicitron.models.visualization.render_flyaround`).
	"""

	# 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")
	os.makedirs(output_directory, exist_ok=True)

	# 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
	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

	# Set the rendering image size
	model_factory_args = config.model_factory_ImplicitronModelFactory_args
	model_factory_args.force_resume = True
	model_args = model_factory_args.model_GenericModel_args
	model_args.render_image_width = render_size[0]
	model_args.render_image_height = render_size[1]

	# Load the previously trained model
	experiment = Experiment(**config)
	model = experiment.model_factory(exp_dir=exp_dir)
	device = torch.device("cuda")
	model.to(device)
	model.eval()

	# Setup the dataset
	data_source = experiment.data_source
	dataset_map, _ = data_source.get_datasets_and_dataloaders()
	dataset = dataset_map[split]
	if dataset is None:
	raise ValueError(f"{split} dataset not provided")

	if visdom_env is None:
	visdom_env = (
	"visualizer_" + config.training_loop_ImplicitronTrainingLoop_args.visdom_env
	)

	# iterate over the sequences in the dataset
	for sequence_name in dataset.sequence_names():
	with torch.no_grad():
	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)


	enable_get_default_args(visualize_reconstruction)


	def _get_config_from_experiment_directory(experiment_directory) -> DictConfig:
	cfg_file = os.path.join(experiment_directory, "expconfig.yaml")
	config = OmegaConf.load(cfg_file)
	# pyre-ignore[7]
	return OmegaConf.merge(get_default_args(Experiment), config)


	def main(argv=sys.argv) -> None:
	# automatically parses arguments of visualize_reconstruction
	cfg = OmegaConf.create(get_default_args(visualize_reconstruction))
	cfg.update(OmegaConf.from_cli(argv))
	with torch.no_grad():
	visualize_reconstruction(**cfg)


	if __name__ == "__main__":
	main()