diff --git a/nerfactory/datamanagers/dataparsers/friends_parser.py b/nerfactory/datamanagers/dataparsers/friends_parser.py index 728c2c743..3dd10844f 100644 --- a/nerfactory/datamanagers/dataparsers/friends_parser.py +++ b/nerfactory/datamanagers/dataparsers/friends_parser.py @@ -20,9 +20,11 @@ from pathlib import Path import numpy as np import torch +from PIL import Image from nerfactory.cameras.cameras import Cameras, CameraType from nerfactory.datamanagers.dataparsers.base import DataParser +from nerfactory.datamanagers.datasets import InputDataset from nerfactory.datamanagers.structs import ( DatasetInputs, PointCloud, @@ -37,6 +39,21 @@ from nerfactory.utils.colmap_utils import ( from nerfactory.utils.io import get_absolute_path, load_from_json +def get_semantics_and_masks(self: InputDataset, image_idx: int): + """function to process additional semantics and mask information""" + # handle mask + person_index = self.inputs.additional_inputs["semantics"].thing_classes.index("person") + thing_image_filename = self.inputs.additional_inputs["semantics"].thing_filenames[image_idx] + pil_image = Image.open(thing_image_filename) + thing_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] + mask = (thing_semantics != person_index).to(torch.float32) # 1 where valid + # handle semantics + stuff_image_filename = self.inputs.additional_inputs["semantics"].stuff_filenames[image_idx] + pil_image = Image.open(stuff_image_filename) + stuff_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] + return {"mask": mask, "semantics": stuff_semantics} + + @dataclass class Friends(DataParser): """Friends Dataset @@ -176,8 +193,8 @@ class Friends(DataParser): dataset_inputs = DatasetInputs( image_filenames=image_filenames, cameras=cameras, - semantics=semantics, point_cloud=point_cloud, scene_bounds=scene_bounds, + additional_inputs={"semantics": {"data": semantics, "func": get_semantics_and_masks}}, ) return dataset_inputs diff --git a/nerfactory/datamanagers/datasets.py b/nerfactory/datamanagers/datasets.py index 6f752b078..e988aeddf 100644 --- a/nerfactory/datamanagers/datasets.py +++ b/nerfactory/datamanagers/datasets.py @@ -17,7 +17,6 @@ Some dataset code. """ from __future__ import annotations -from abc import abstractmethod from typing import Dict, Union import numpy as np @@ -75,42 +74,17 @@ class InputDataset(Dataset): image = image[:, :, :3] return image - @abstractmethod - def get_mask(self, image_idx: int) -> Union[TensorType["image_height", "image_width", 1], None]: - """Returns a mask, which indicates which pixels are valid to use with nerf.""" - if self.inputs.semantics: - person_index = self.inputs.semantics.thing_classes.index("person") - thing_image_filename = self.inputs.semantics.thing_filenames[image_idx] - pil_image = Image.open(thing_image_filename) - thing_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] - mask = (thing_semantics != person_index).to(torch.float32) # 1 where valid - return mask - return None - - @abstractmethod - def get_semantics(self, image_idx: int) -> Union[TensorType["image_height", "image_width", "num_classes"], None]: - """Returns an image with semantic class values.""" - if self.inputs.semantics: - stuff_image_filename = self.inputs.semantics.stuff_filenames[image_idx] - pil_image = Image.open(stuff_image_filename) - stuff_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] - return stuff_semantics - return None - def get_data(self, image_idx) -> Dict: """Returns the ImageDataset data as a dictionary.""" image = self.get_image(image_idx) - mask = self.get_mask(image_idx) - semantics = self.get_semantics(image_idx) data = {"image_idx": image_idx} assert is_not_none(image) data["image"] = image - if mask is not None: - assert mask.shape[:2] == image.shape[:2] - data["mask"] = mask - if semantics is not None: - assert semantics.shape[:2] == image.shape[:2] - data["semantics"] = semantics + for _, data_func_dict in self.inputs.additional_inputs.items(): + assert "func" in data_func_dict, "Missing function to process data: specify `func` in `additional_inputs`" + func = data_func_dict["func"] + assert "data" in data_func_dict, "No data to process: specify `data` in `additional_inputs`" + data.update(func(data_func_dict["data"])) return data def __getitem__(self, image_idx): diff --git a/nerfactory/datamanagers/structs.py b/nerfactory/datamanagers/structs.py index 2d5fac18a..f30222fd4 100644 --- a/nerfactory/datamanagers/structs.py +++ b/nerfactory/datamanagers/structs.py @@ -18,12 +18,13 @@ Dataset input structures. from dataclasses import dataclass from pathlib import Path -from typing import Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Union import torch from torchtyping import TensorType from nerfactory.cameras.cameras import Cameras +from nerfactory.configs.utils import to_immutable_dict @dataclass @@ -121,12 +122,11 @@ class DatasetInputs: image_filenames: List[Path] cameras: Cameras - mask_filenames: Optional[List[Path]] = None - depth_filenames: Optional[List[Path]] = None - scene_bounds: SceneBounds = SceneBounds() - semantics: Optional[Semantics] = None - point_cloud: PointCloud = PointCloud() alpha_color: Optional[TensorType[3]] = None + scene_bounds: SceneBounds = SceneBounds() + point_cloud: PointCloud = PointCloud() + # we support additional input information/formats including mask/depth/semantics. + additional_inputs: Dict[str, Any] = to_immutable_dict({}) def as_dict(self) -> dict: """Returns the dataclass as a dictionary."""