Added lots of stuff

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
 

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)

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

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

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

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.

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,

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:    

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(

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