diff --git a/README-pandora.md b/README-pandora.md index 9d1d3037d..462910199 100644 --- a/README-pandora.md +++ b/README-pandora.md @@ -5,7 +5,7 @@ ## 0. Setup ``` - +export PYTHONPATH=$pwd:$PYTHONPATH pip install skylibs equilib ``` @@ -39,7 +39,8 @@ pip install skylibs equilib 5. Extract frames using ffmpeg ``` - python mp4_videos_to_png_frames.py + python lantern_scripts/mp4_videos_to_png_frames.py --help + python lantern_scripts/mp4_videos_to_png_frames.py --input_dir data/lab_downstairs/trimmed_videos ``` It should output the following file structure @@ -56,16 +57,18 @@ pip install skylibs equilib ....png ....png ....png - right_right_e1/ + right_e2/ ....png ....png ....png ``` -6. Run the code to mask out people and the stick. You may need to specify `--input-dir` and `--output-dir` +6. Run the code to mask out people and the stick. + Make sure the masks in "./stick_masks" are valid ``` - python lantern_scripts/mask_humans2.py + python lantern_scripts/mask_humans2.py --help + python lantern_scripts/mask_humans2.py --input_dir data/lab_downstairs/trimmed_videos ``` @@ -81,13 +84,23 @@ python lantern_scripts/split_synthetic_colmap.py --has_exposure --left_dir ./Dua Example with synthetic data: ``` +export PYTHONPATH=$pwd:$PYTHONPATH ns-process-data lantern --data ./Dual_Cameras/split_colmap_subset --output-dir ./Dual_Cameras/split_colmap_subset/split_colmap_out5 --skip-image-processing --camera-type equirectangular --images-per-equirect 8 ``` Example with real data (note that `--mask-dir` was added): ``` -ns-process-data lantern --data ./data/lab_ground_floor/trimmed_videos --mask-dir ./data/lab_ground_floor/trimmed_videos_masks --output-dir ./data/lab_ground_floor/colmap_out --skip-image-processing --camera-type equirectangular --images-per-equirect 8 +export PYTHONPATH=$pwd:$PYTHONPATH +ns-process-data lantern --data ./data/lab_downstairs/captured_data --mask-dir ./data/lab_downstairs/masks --output-dir ./data/lab_downstairs/colmap_out --skip-image-processing --camera-type equirectangular --images-per-equirect 8 +``` + +After that, you'll need to copy manually (for now), the folder to `images`, for nerf training. + +Examples +``` +ns-process-data lantern --data ./data/lab_ground_floor/trimmed_videos_smart_subset --mask-dir ./data/lab_ground_floor/trimmed_videos_masks --output-dir ./data/lab_ground_floor/colmap_out_smart --skip-image-processing --camera-type equirectangular --images-per +-equirect 8 --skip-colmap --skip-perspective-transform ``` ## 2. Pre-training NeRF @@ -100,8 +113,15 @@ After convergence, go to the next step. ns-train lantern-nerfacto --help ``` +If you need to use less frames, add this argument: `--pipeline.datamanager.train-num-images-to-sample-from 200 --pipeline.datamanager.eval-num-images-to-sample-from 30` + > Thing to keep in mind: the way Nerfstudio splits train and test images may not be optimal: left might be in one dataset and right in another... that's a bit of a data leak. +Examples + +``` +ns-train lantern-nerfacto --data data/lab_ground_floor/colmap_out_smart --pipeline.datamanager.train-num-images-to-sample-from 200 --pipeline.datamanager.eval-num-images-to-sample-from 30 +``` ## 3. Running MomoNet diff --git a/lantern_scripts/extract_linearized.py b/lantern_scripts/extract_linearized.py index 38182036e..8f97819e5 100644 --- a/lantern_scripts/extract_linearized.py +++ b/lantern_scripts/extract_linearized.py @@ -8,7 +8,10 @@ import numpy as np # fonction that apply the correction on the colors on the input images -def apply_correction(coefs, img, CORRECTION_CURVE_TYPE) -> np.array: +def apply_correction(coefs, img, CORRECTION_CURVE_TYPE) -> np.ndarray: + # check that img is in range 0-255 + assert img.dtype == np.uint8 + B, G, R = coefs original_shape = img.shape img = img.astype(np.float32) diff --git a/lantern_scripts/mask_humans2.py b/lantern_scripts/mask_humans2.py index 89f210863..48592e774 100644 --- a/lantern_scripts/mask_humans2.py +++ b/lantern_scripts/mask_humans2.py @@ -20,25 +20,25 @@ FRAME_SUBFOLDERS_CONFIG = [ 'name': 'left_colmap_baseline', 'flip_k': -1, 'mask_person': True, - 'custom_mask': './data/lab_ground_floor/custom_masks/left_mask.png' + 'custom_mask': './stick_masks/left_mask.png' }, { 'name': 'right_colmap_baseline', 'flip_k': 1, 'mask_person': True, - 'custom_mask': './data/lab_ground_floor/custom_masks/right_mask.png' + 'custom_mask': './stick_masks/right_mask.png' }, { 'name': 'left_e1', 'flip_k': -1, 'mask_person': True, - 'custom_mask': './data/lab_ground_floor/custom_masks/left_mask.png' + 'custom_mask': './stick_masks/left_mask.png' }, { 'name': 'right_e2', 'flip_k': 1, 'mask_person': False, - 'custom_mask': './data/lab_ground_floor/custom_masks/right_mask.png' + 'custom_mask': './stick_masks/right_mask.png' } ] @@ -46,8 +46,8 @@ PERSON_DILATE_KERNEL = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (40, 40)) if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--input_dir', type=str, default='./data/lab_ground_floor/trimmed_videos', help='Directory containing all the frames') - parser.add_argument('--output_dir', type=str, default='./data/lab_ground_floor/trimmed_videos_masks') + parser.add_argument('--input_dir', type=str, default='./data/lab_ground_floor/trimmed_videos', help='Directory containing all the 4 directories containing frames') + parser.add_argument('--output_dir', type=str, default='./data/lab_ground_floor/trimmed_videos_masks', help='Directory to save masked frames') args = parser.parse_args() out = args.output_dir diff --git a/lantern_scripts/mp4_videos_to_png_frames.py b/lantern_scripts/mp4_videos_to_png_frames.py index ff01f06e4..6a5acfdea 100644 --- a/lantern_scripts/mp4_videos_to_png_frames.py +++ b/lantern_scripts/mp4_videos_to_png_frames.py @@ -1,7 +1,6 @@ -import subprocess -import os import argparse - +import os +import subprocess # TODO: limit number of frames... maybe 3 per second? @@ -13,10 +12,13 @@ if __name__ == '__main__': # get video files args = argparser.parse_args() videos = [os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir) if f.endswith(".mp4")] + if len(videos) == 0: + raise Exception("No video files found in input directory.") for video in videos: # Get the video name without extension video_name = os.path.splitext(os.path.basename(video))[0] + print(f"Extracting frames from {video}...") # Create a subfolder for the frames output_folder = os.path.join(args.input_dir, f'{video_name}') @@ -29,7 +31,7 @@ if __name__ == '__main__': "-vf", "fps=29.97", f"{output_folder}/output_%04d.png" ] - subprocess.run(command, check=True) + subprocess.run(command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) # only keep frames where mod 10 = 0 for file in os.listdir(output_folder): diff --git a/linearize-ricoh-z10-wb-3500.json b/linearize-ricoh-z10-wb-3500.json new file mode 100644 index 000000000..6ff6171b1 --- /dev/null +++ b/linearize-ricoh-z10-wb-3500.json @@ -0,0 +1 @@ +{"B": [54.09578323364258, 0.2747989296913147, 0.0019784104079008102], "G": [48.520992279052734, 0.05255260691046715, 0.0029351452831178904], "R": [47.87550735473633, 0.19775021076202393, 0.0023990024346858263], "type": "deg2"} \ No newline at end of file diff --git a/lantern_scripts/z10-color-correct-data.json b/linearize-ricoh-z10-wb-5000.json similarity index 100% rename from lantern_scripts/z10-color-correct-data.json rename to linearize-ricoh-z10-wb-5000.json diff --git a/nerfstudio/data/dataparsers/base_dataparser.py b/nerfstudio/data/dataparsers/base_dataparser.py index 56888e8ef..73d89346a 100644 --- a/nerfstudio/data/dataparsers/base_dataparser.py +++ b/nerfstudio/data/dataparsers/base_dataparser.py @@ -61,6 +61,8 @@ class DataparserOutputs: """Scene box of dataset. Used to bound the scene or provide the scene scale depending on model.""" mask_filenames: Optional[List[Path]] = None """Filenames for any masks that are required""" + saturation_mask_filenames: Optional[List[Path]] = None + """Filenames for any saturation masks that are required""" metadata: Dict[str, Any] = to_immutable_dict({}) """Dictionary of any metadata that be required for the given experiment. Will be processed by the InputDataset to create any additional tensors that may be required. diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index 582e0d8c3..e6b08b2d2 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -88,6 +88,7 @@ class Nerfstudio(DataParser): image_filenames = [] mask_filenames = [] + saturation_mask_filenames = [] depth_filenames = [] poses = [] exposures = [] @@ -155,6 +156,15 @@ class Nerfstudio(DataParser): downsample_folder_prefix="masks_", ) mask_filenames.append(mask_fname) + + if "saturation_mask_path" in frame: + saturation_mask_filepath = Path(frame["saturation_mask_path"]) + saturation_mask_fname = self._get_fname( + saturation_mask_filepath, + data_dir, + downsample_folder_prefix="saturation_masks_", + ) + saturation_mask_filenames.append(saturation_mask_fname) if "depth_file_path" in frame: depth_filepath = Path(frame["depth_file_path"]) @@ -171,6 +181,12 @@ class Nerfstudio(DataParser): Different number of image and mask filenames. You should check that mask_path is specified for every frame (or zero frames) in transforms.json. """ + assert len(saturation_mask_filenames) == 0 or ( + len(saturation_mask_filenames) == len(image_filenames) + ), """ + Different number of image and saturation mask filenames. + You should check that saturation_mask_path is specified for every frame (or zero frames) in transforms.json. + """ assert len(depth_filenames) == 0 or ( len(depth_filenames) == len(image_filenames) ), """ @@ -239,6 +255,7 @@ class Nerfstudio(DataParser): # Choose image_filenames and poses based on split, but after auto orient and scaling the poses. image_filenames = [image_filenames[i] for i in indices] mask_filenames = [mask_filenames[i] for i in indices] if len(mask_filenames) > 0 else [] + saturation_mask_filenames = [saturation_mask_filenames[i] for i in indices] if len(saturation_mask_filenames) > 0 else [] depth_filenames = [depth_filenames[i] for i in indices] if len(depth_filenames) > 0 else [] exposures = [exposures[i] for i in indices] if len(exposures) > 0 else [] @@ -306,6 +323,7 @@ class Nerfstudio(DataParser): cameras=cameras, scene_box=scene_box, mask_filenames=mask_filenames if len(mask_filenames) > 0 else None, + saturation_mask_filenames=saturation_mask_filenames if len(saturation_mask_filenames) > 0 else None, exposures=exposures if len(exposures) > 0 else None, dataparser_scale=scale_factor, dataparser_transform=transform_matrix, diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 1e817e9c8..27480dcc6 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -17,15 +17,8 @@ Code for sampling pixels. """ import random - -import torch -from jaxtyping import Int -from torch import Tensor -from math import ceil - from dataclasses import dataclass, field -from nerfstudio.data.utils.pixel_sampling_utils import erode_mask -import nerfstudio.utils.profiler as profiler +from math import ceil from typing import ( Dict, Optional, @@ -33,9 +26,15 @@ from typing import ( Union, ) +import torch +from jaxtyping import Int +from torch import Tensor + from nerfstudio.configs.base_config import ( InstantiateConfig, ) +from nerfstudio.data.utils.pixel_sampling_utils import erode_mask +from nerfstudio.utils import profiler @dataclass @@ -98,7 +97,7 @@ class PixelSampler: mask: mask of possible pixels in an image to sample from. """ # TODO Defining a flag for sampling even when there is a mask - if isinstance(mask, torch.Tensor) and False: + if isinstance(mask, torch.Tensor): # mask size: batch * h * w * 1 # Init state, cache masks indices. if self.states_of_mask == 0: diff --git a/nerfstudio/lantern/config.py b/nerfstudio/lantern/config.py index cf1a86d02..74650961f 100644 --- a/nerfstudio/lantern/config.py +++ b/nerfstudio/lantern/config.py @@ -6,18 +6,16 @@ Define custom method here that registers with Nerfstudio CLI. from __future__ import annotations -from nerfstudio.lantern.datamanager import HDRVanillaDataManagerConfig -from nerfstudio.lantern.model import LanternModelConfig from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.configs.base_config import ViewerConfig -from nerfstudio.data.dataparsers.nerfstudio_dataparser import \ - NerfstudioDataParserConfig -from nerfstudio.engine.optimizers import (AdamOptimizerConfig, - RAdamOptimizerConfig) +from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig +from nerfstudio.engine.optimizers import AdamOptimizerConfig from nerfstudio.engine.schedulers import ExponentialDecaySchedulerConfig from nerfstudio.engine.trainer import TrainerConfig +from nerfstudio.lantern.datamanager import HDRVanillaDataManagerConfig +from nerfstudio.lantern.model import LanternModelConfig from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig -from nerfstudio.plugins.types import MethodSpecification + def Get_lantern_config(): return TrainerConfig( diff --git a/nerfstudio/lantern/dataset.py b/nerfstudio/lantern/dataset.py index 27ab92476..07a46063d 100644 --- a/nerfstudio/lantern/dataset.py +++ b/nerfstudio/lantern/dataset.py @@ -111,6 +111,13 @@ class HDRInputDataset(InputDataset): 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 diff --git a/stick_masks/left_mask.png b/stick_masks/left_mask.png new file mode 100644 index 000000000..aca7abefd Binary files /dev/null and b/stick_masks/left_mask.png differ diff --git a/stick_masks/left_mask_bak.png b/stick_masks/left_mask_bak.png new file mode 100644 index 000000000..5c62b0c7c Binary files /dev/null and b/stick_masks/left_mask_bak.png differ diff --git a/stick_masks/right_mask.png b/stick_masks/right_mask.png new file mode 100644 index 000000000..b447835f5 Binary files /dev/null and b/stick_masks/right_mask.png differ