simplify dataset inputs; move processing to dataset class (#355)

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

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):

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."""