diff --git a/nerfstudio/lantern/dataset.py b/nerfstudio/lantern/dataset.py index f5f0b362d..f0e1fd49c 100644 --- a/nerfstudio/lantern/dataset.py +++ b/nerfstudio/lantern/dataset.py @@ -151,6 +151,123 @@ class HDRNerfactoInputDataset(InputDataset): return self._dataparser_outputs.image_filenames +class HDRNerfactoWoCrfInputDataset(InputDataset): + """Dataset that returns HDR images, OpenEXR. + + Args: + dataparser_outputs: description of where and how to read input images. + scale_factor: The scaling factor for the dataparser outputs + """ + + # exclude_batch_keys_from_device: List[str] = ["image", "mask"] + # cameras: Cameras + + def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0): + super().__init__(dataparser_outputs, scale_factor) + + def __len__(self): + return len(self._dataparser_outputs.image_filenames) + + def get_numpy_image(self, image_idx: int) -> npt.NDArray[np.uint8]: + """Returns the image of shape (H, W, 3 or 4), with np.float32 values. + + Args: + image_idx: The image index in the dataset. + """ + image_filename = self._dataparser_outputs.image_filenames[image_idx] + #print('image_filename', image_filename) + hdr_image = cv2.imread(str(image_filename), cv2.IMREAD_UNCHANGED) + # print('hdr_image', hdr_image.shape, hdr_image.dtype) + # print('hdr_image', hdr_image.max(), hdr_image.min()) + # print('hdr_image', hdr_image[0,0,:]) + + if self.scale_factor != 1.0: + # Here cv2 img shape is (height, width, BGR) + height, width = hdr_image.shape[:2] + newsize = (int(width * self.scale_factor), int(height * self.scale_factor)) + hdr_image = cv2.resize(hdr_image, newsize, interpolation = cv2.INTER_LINEAR) + hdr_image = cv2.cvtColor(hdr_image, cv2.COLOR_BGR2RGB) + hdr_image = hdr_image.astype("float32") # shape is (h, w) or (h, w, 3 or 4) + if len(hdr_image.shape) == 2: + hdr_image = hdr_image[:, :, None].repeat(3, axis=2) + assert len(hdr_image.shape) == 3 + assert hdr_image.dtype == np.float32 + assert hdr_image.shape[2] in [3, 4], f"Image shape of {hdr_image.shape} is in correct." + return hdr_image + + def get_image(self, image_idx: int) -> Float[Tensor, "image_height image_width num_channels"]: + """Returns a 3 channel image. + + Args: + image_idx: The image index in the dataset. + """ + image = torch.from_numpy(self.get_numpy_image(image_idx).astype("float32")) + #print('a', image.max()) + # image = image * self._dataparser_outputs.exposures[image_idx] + # print('b', image.max()) + + # if self._dataparser_outputs.alpha_color is not None and image.shape[-1] == 4: + # image = image[:, :, :3] * image[:, :, -1:] + self._dataparser_outputs.alpha_color * (1.0 - image[:, :, -1:]) + return image + + def get_data(self, image_idx: int) -> Dict: + """Returns the ImageDataset data as a dictionary. + + Args: + image_idx: The image index in the dataset. + """ + image = self.get_image(image_idx) + data = {"image_idx": image_idx, "image": image} + + # TODO: this doesnt work for training + image_filename = self._dataparser_outputs.image_filenames[image_idx] + data["image_filename"] = str(image_filename.name) + + if self._dataparser_outputs.mask_filenames is not None: + mask_filepath = self._dataparser_outputs.mask_filenames[image_idx] + data["mask"] = get_image_mask_tensor_from_path(filepath=mask_filepath, scale_factor=self.scale_factor) + assert ( + data["mask"].shape[:2] == data["image"].shape[:2] + ), f"Mask and image have different shapes. Got {data['mask'].shape[:2]} and {data['image'].shape[:2]}" + + if self._dataparser_outputs.saturation_mask_filenames is not None: + saturation_mask_filepath = self._dataparser_outputs.saturation_mask_filenames[image_idx] + data["saturation_mask"] = get_image_mask_tensor_from_path(filepath=saturation_mask_filepath, scale_factor=self.scale_factor) + assert ( + data["saturation_mask"].shape[:2] == data["image"].shape[:2] + ), f"Saturation mask and image have different shapes. Got {data['saturation_mask'].shape[:2]} and {data['image'].shape[:2]}" + + if self._dataparser_outputs.exposures is not None: + exposure_value = self._dataparser_outputs.exposures[image_idx] + data["exposure"] = exposure_value + metadata = self.get_metadata(data) + data.update(metadata) + return data + + def get_metadata(self, data: Dict) -> Dict: + """Method that can be used to process any additional metadata that may be part of the model inputs. + + Args: + image_idx: The image index in the dataset. + """ + del data + return {} + + def __getitem__(self, image_idx: int) -> Dict: + data = self.get_data(image_idx) + return data + + @property + def image_filenames(self) -> List[Path]: + """ + Returns image filenames for this dataset. + The order of filenames is the same as in the Cameras object for easy mapping. + """ + + return self._dataparser_outputs.image_filenames + + + class HDRInputDataset(InputDataset): """Dataset that returns HDR images, OpenEXR.