Train own dataset

[AndreaCastiella]

Hello! Nice work! I would like to test the method using my own dataset. What steps should I follow to prepare the data in order to train NeuSky? Thank you!

[malfonsoarquimea]

I would also be interested in that regard!

[JADGardner]

Hi! Thanks for checking out Neusky.

For your own data you'll need a cityscapes segmentation mask for each image, in the paper we use, ViT-Adapter, but any cityscapes segmentation method should work.

I've added two Python files you'll need to modify to process your data.

neusky.data.dataparser.custom_neusky_dataparser.py containing CustomNeuskyDataparser
neusky.data.datasets.custom_neusky_dataset.py containing CustomNeuskyDataset

CustomNeuskyDataparser sets up cameras in the Nerfstudio coordinate convention and locates the image and segmentation files.
CustomNeuskyDataset returns the required image and masks for a given index.

The mask returned by CustomNeuskyDataset should be a [H, W, 4] tensor made up of the following masks:

static_object_mask = get_static_from_cityscapes() # [H, W, 1]
foreground_mask = get_foreground_from_cityscapes() # [H, W, 1]
ground_mask = get_ground_from_cityscapes() # [H, W, 1]
sky_mask = get_sky_from_cityscapes() # [H, W, 1]
mask = torch.cat([static_object_mask, foreground_mask, ground_mask, sky_mask], dim=-1)

You'll then need to update this line in the config to use your custom dataparser:

neusky/neusky/configs/neusky_config.py

Line 47 in 96de93e

dataparser=NeRFOSRCityScapesDataParserConfig(

And update these lines to use your custom dataset:

neusky/neusky/data/datamanagers/neusky_datamanager.py

Lines 141 to 157 in 96de93e

	def create_train_dataset(self) -> NeuSkyDataset:
	return NeuSkyDataset(
	dataparser_outputs=self.train_dataparser_outputs,
	scale_factor=self.config.camera_res_scale_factor,
	split="train",
	)
	def create_eval_dataset(self) -> NeuSkyDataset:
	self.test_outputs = self.dataparser.get_dataparser_outputs("test")
	self.val_outputs = self.dataparser.get_dataparser_outputs("val")
	# self.num_test = len(test_outputs.image_filenames)
	# self.num_val = len(val_outputs.image_filenames)
	return NeuSkyDataset(
	dataparser_outputs=self.test_outputs if self.test_mode == "test" else self.val_outputs,
	scale_factor=self.config.camera_res_scale_factor,
	split=self.test_split,
	)

I'll probably update the Datamanger to make the dataset swappable via the config rather than changing lines in the code as shown above but this will work for now. I hope this helps. Let me know if you have any issues.

[AndreaCastiella]

Thank you very much! I will try it and let you know if I get it to work. If I find something useful for other users I will post it too :)

[AndreaCastiella]

Hello again! I've successfully trained Neusky using a nerfstudio-data type dataset. I created a new dataparser "nerfstudio_cityscapes_dataparser" and updated the config file to utilize this parser. There was no need to modify the Dataset and DataManager, which was quite straightforward. This is what I changed in neusky/neusky/configs/neusky_config.py:

from neusky.data.dataparsers.nerfstudio_cityscapes_dataparser import NerfstudioCityScapesDataParserConfig
...
dataparser=NerfstudioCityScapesDataParserConfig(
    auto_scale_poses=True,
    crop_to_equal_size=True,
    pad_to_equal_size=False,
    scene_scale=1.0,  # AABB
    mask_vegetation=True,
    mask_out_of_view_frustum_objects=True,
),

An important note for anyone following this approach: you'll need to add the mask paths into the transforms.json file. I've prepared a script that simplifies adding these paths. Ensure the masks are stored in a folder named segmented, and that each segmented image is named identically to its corresponding original image.

Here's how the data folder should be organized:

data/
├── images/
│ ├── frame_00001.png
│ ├── frame_00002.png
│ ├── ...
│ └── frame_NNNNN.png
├── segmented/
│ ├── frame_00001_mask.png
│ ├── frame_00002_mask.png
│ ├── ...
│ └── frame_NNNNN_mask.png
└── transforms.json

import json
 
def add_mask_path_to_json(json_file_path: str, json_file_path_out: str, image_folder: str = "images/", mask_folder: str = "segmented/"):
    # Read the JSON file
    with open(json_file_path, 'r') as file:
        data = json.load(file)
 
    # Iterate through each frame and add the mask_path
    for frame in data['frames']:
        image_path = frame['file_path']
        # Replace folder name in the same path
        mask_path = image_path.replace(image_folder, mask_folder)
        frame['mask_path'] = mask_path
 
    # Write the modified data back to the JSON file
    with open(json_file_path_out, 'w') as file:
        json.dump(data, file, indent=4)
 
 
if __name__ == "__main__":
    # Replace with the path to your transforms.json file
    json_file_path = 'path_to_transforms.json'
    json_file_path_out = 'path_to_transforms_out.json'
    add_mask_path_to_json(json_file_path, json_file_path_out)

I'm looking forward to sharing the nerfstudio_cityscapes_dataparser code with the community, but I want to make sure it aligns with your interests. If you find this contribution valuable, please guide me on the best way to share it :).