pytorch3d.implicitron.dataset specific datasets

specific datasets

class pytorch3d.implicitron.dataset.blender_dataset_map_provider.BlenderDatasetMapProvider(*args, **kwargs)[source]

Bases: SingleSceneDatasetMapProviderBase

Provides data for one scene from Blender synthetic dataset. Uses the code in load_blender.py

Members:

base_dir: directory holding the data for the scene. object_name: The name of the scene (e.g. “lego”). This is just used as a label.

It will typically be equal to the name of the directory self.base_dir.

path_manager_factory: Creates path manager which may be used for: interpreting paths.
n_known_frames_for_test: If set, training frames are included in the val: and test datasets, and this many random training frames are added to each test batch. If not set, test batches each contain just a single testing frame.

class pytorch3d.implicitron.dataset.json_index_dataset_map_provider.JsonIndexDatasetMapProvider(*args, **kwargs)[source]

Bases: DatasetMapProviderBase

Generates the training / validation and testing dataset objects for a dataset laid out on disk like Co3D, with annotations in json files.

Parameters:

category – The object category of the dataset.
task_str – “multisequence” or “singlesequence”.
dataset_root – The root folder of the dataset.
n_frames_per_sequence – Randomly sample #n_frames_per_sequence frames in each sequence.
test_on_train – Construct validation and test datasets from the training subset.
restrict_sequence_name – Restrict the dataset sequences to the ones present in the given list of names.
test_restrict_sequence_id – The ID of the loaded sequence. Active for task_str=’singlesequence’.
assert_single_seq – Assert that only frames from a single sequence are present in all generated datasets.
only_test_set – Load only the test set.
dataset_class_type – name of class (JsonIndexDataset or a subclass) to use for the dataset.
dataset_X_args (e.g. dataset_JsonIndexDataset_args) – arguments passed to all the dataset constructors.
path_manager_factory – (Optional) An object that generates an instance of PathManager that can translate provided file paths.
path_manager_factory_class_type – The class type of path_manager_factory.

category: str

task_str: str = 'singlesequence'

dataset_root: str = ''

n_frames_per_sequence: int = -1

test_on_train: bool = False

restrict_sequence_name: Tuple[str, ...] = ()

test_restrict_sequence_id: int = -1

assert_single_seq: bool = False

only_test_set: bool = False

dataset: JsonIndexDataset

dataset_class_type: str = 'JsonIndexDataset'

path_manager_factory: PathManagerFactory

path_manager_factory_class_type: str = 'PathManagerFactory'

classmethod dataset_tweak_args(type, args: DictConfig) → None[source]: Called by get_default_args(JsonIndexDatasetMapProvider) to not expose certain fields of each dataset class.

create_dataset()[source]: Prevent the member named dataset from being created.

get_dataset_map() → DatasetMap[source]

get_all_train_cameras() → CamerasBase | None[source]

class pytorch3d.implicitron.dataset.json_index_dataset_map_provider_v2.JsonIndexDatasetMapProviderV2(*args, **kwargs)[source]

Bases: DatasetMapProviderBase

Generates the training, validation, and testing dataset objects for a dataset laid out on disk like CO3Dv2, with annotations in gzipped json files.

The dataset is organized in the filesystem as follows:

self.dataset_root
    ├── <category_0>
    │   ├── <sequence_name_0>
    │   │   ├── depth_masks
    │   │   ├── depths
    │   │   ├── images
    │   │   ├── masks
    │   │   └── pointcloud.ply
    │   ├── <sequence_name_1>
    │   │   ├── depth_masks
    │   │   ├── depths
    │   │   ├── images
    │   │   ├── masks
    │   │   └── pointcloud.ply
    │   ├── ...
    │   ├── <sequence_name_N>
    │   ├── set_lists
    │       ├── set_lists_<subset_name_0>.json
    │       ├── set_lists_<subset_name_1>.json
    │       ├── ...
    │       ├── set_lists_<subset_name_M>.json
    │   ├── eval_batches
    │   │   ├── eval_batches_<subset_name_0>.json
    │   │   ├── eval_batches_<subset_name_1>.json
    │   │   ├── ...
    │   │   ├── eval_batches_<subset_name_M>.json
    │   ├── frame_annotations.jgz
    │   ├── sequence_annotations.jgz
    ├── <category_1>
    ├── ...
    ├── <category_K>

The dataset contains sequences named <sequence_name_i> from K categories with names <category_j>. Each category comprises sequence folders <category_k>/<sequence_name_i> containing the list of sequence images, depth maps, foreground masks, and valid-depth masks images, depths, masks, and depth_masks respectively. Furthermore, <category_k>/<sequence_name_i>/set_lists/ stores M json files set_lists_<subset_name_l>.json, each describing a certain sequence subset.

Users specify the loaded dataset subset by setting self.subset_name to one of the available subset names <subset_name_l>.

frame_annotations.jgz and sequence_annotations.jgz are gzipped json files containing the list of all frames and sequences of the given category stored as lists of FrameAnnotation and SequenceAnnotation objects respectivelly.

Each set_lists_<subset_name_l>.json file contains the following dictionary:

{
    "train": [
        (sequence_name: str, frame_number: int, image_path: str),
        ...
    ],
    "val": [
        (sequence_name: str, frame_number: int, image_path: str),
        ...
    ],
    "test": [
        (sequence_name: str, frame_number: int, image_path: str),
        ...
    ],
]

defining the list of frames (identified with their sequence_name and frame_number) in the “train”, “val”, and “test” subsets of the dataset. Note that frame_number can be obtained only from frame_annotations.jgz and does not necesarrily correspond to the numeric suffix of the corresponding image file name (e.g. a file <category_0>/<sequence_name_0>/images/frame00005.jpg can have its frame number set to 20, not 5).

Each eval_batches_<subset_name_l>.json file contains a list of evaluation examples in the following form:

[
    [  # batch 1
        (sequence_name: str, frame_number: int, image_path: str),
        ...
    ],
    [  # batch 1
        (sequence_name: str, frame_number: int, image_path: str),
        ...
    ],
]

Note that the evaluation examples always come from the “test” subset of the dataset. (test frames can repeat across batches).

Parameters:

category – Dataset categories to load expressed as a string of comma-separated category names (e.g. “apple,car,orange”).
subset_name – The name of the dataset subset. For CO3Dv2, these include e.g. “manyview_dev_0”, “fewview_test”, …
dataset_root – The root folder of the dataset.
test_on_train – Construct validation and test datasets from the training subset.
only_test_set – Load only the test set. Incompatible with test_on_train.
load_eval_batches – Load the file containing eval batches pointing to the test dataset.
n_known_frames_for_test – Add a certain number of known frames to each eval batch. Useful for evaluating models that require source views as input (e.g. NeRF-WCE / PixelNeRF).
dataset_args – Specifies additional arguments to the JsonIndexDataset constructor call.
path_manager_factory – (Optional) An object that generates an instance of PathManager that can translate provided file paths.
path_manager_factory_class_type – The class type of path_manager_factory.

category: str

subset_name: str

dataset_root: str = ''

test_on_train: bool = False

only_test_set: bool = False

load_eval_batches: bool = True

num_load_workers: int = 4

n_known_frames_for_test: int = 0

dataset_class_type: str = 'JsonIndexDataset'

dataset: JsonIndexDataset

path_manager_factory: PathManagerFactory

path_manager_factory_class_type: str = 'PathManagerFactory'

classmethod dataset_tweak_args(type, args: DictConfig) → None[source]: Called by get_default_args(JsonIndexDatasetMapProviderV2) to not expose certain fields of each dataset class.

create_dataset()[source]

get_dataset_map() → DatasetMap[source]

get_category_to_subset_name_list() → Dict[str, List[str]][source]

Returns a global dataset index containing the available subset names per category as a dictionary.

Returns:

category_to_subset_name_list –

A dictionary containing subset names available

per category of the following form:

{
    category_0: [category_0_subset_name_0, category_0_subset_name_1, ...],
    category_1: [category_1_subset_name_0, category_1_subset_name_1, ...],
    ...
}

get_all_train_cameras() → CamerasBase | None[source]

pytorch3d.implicitron.dataset.json_index_dataset_map_provider_v2.get_available_subset_names(dataset_root: str, category: str, path_manager: PathManager | None = None) → List[str][source]: Get the available subset names for a given category folder inside a root dataset folder dataset_root.

class pytorch3d.implicitron.dataset.llff_dataset_map_provider.LlffDatasetMapProvider(*args, **kwargs)[source]

Bases: SingleSceneDatasetMapProviderBase

Provides data for one scene from the LLFF dataset.

Members:

base_dir: directory holding the data for the scene. object_name: The name of the scene (e.g. “fern”). This is just used as a label.

It will typically be equal to the name of the directory self.base_dir.

path_manager_factory: Creates path manager which may be used for: interpreting paths.
n_known_frames_for_test: If set, training frames are included in the val: and test datasets, and this many random training frames are added to each test batch. If not set, test batches each contain just a single testing frame.

downscale_factor: determines image sizes.

downscale_factor: int = 4

class pytorch3d.implicitron.dataset.rendered_mesh_dataset_map_provider.RenderedMeshDatasetMapProvider(*args, **kwargs)[source]

Bases: DatasetMapProviderBase

A simple single-scene dataset based on PyTorch3D renders of a mesh. Provides num_views renders of the mesh as train, with no val and test. The renders are generated from viewpoints sampled at uniformly distributed azimuth intervals. The elevation is kept constant so that the camera’s vertical position coincides with the equator.

By default, uses Keenan Crane’s cow model, and the camera locations are set to make sense for that.

Although the rendering used to generate this dataset will use a GPU if one is available, the data it produces is on the CPU just like the data returned by implicitron’s other dataset map providers. This is because both datasets and models can be large, so implicitron’s training loop expects data on the CPU and only moves what it needs to the device.

For a more detailed explanation of this code, please refer to the docs/tutorials/fit_textured_mesh.ipynb notebook.

Members:

num_views: The number of generated renders. data_file: The folder that contains the mesh file. By default, finds

the cow mesh in the same repo as this code.

azimuth_range: number of degrees on each side of the start position to: take samples

distance: distance from camera centres to the origin. resolution: the common height and width of the output images. use_point_light: whether to use a particular point light as opposed

to ambient white.

gpu_idx: which gpu to use for rendering the mesh. path_manager_factory: (Optional) An object that generates an instance of

PathManager that can translate provided file paths.

path_manager_factory_class_type: The class type of path_manager_factory.

num_views: int = 40

data_file: str | None = None

azimuth_range: float = 180

distance: float = 2.7

resolution: int = 128

use_point_light: bool = True

gpu_idx: int | None = 0

path_manager_factory: PathManagerFactory

path_manager_factory_class_type: str = 'PathManagerFactory'

get_dataset_map() → DatasetMap[source]

get_all_train_cameras() → CamerasBase[source]