Mask humans, video to frames, missing code

README-pandora.md

@@ -25,35 +25,48 @@ pip install skylibs equilib
     > **camera 3**: right, under-exposed, serial: YN14111000
 
 2. Take one video capture with the two cameras well-exposed.
+
+    Use flash at the beginning (and at the end?) to use for synchronization.
+
 3. Take one video capture with Camera 1 (left) well-exposed and Camera 2 (right) fast-exposed.
-4. Run the script for automatic flash detection
+
+    Use flash at the beginning (and at the end?) to use for synchronization.
+
+4. Run the script for automatic flash detection / use Shotcut to manually align and trim the videos manually.
     
     It will temporally align the frames from the left and right camera and remove the initial frames with the flash.
 
+5. Extract frames using ffmpeg
+
+    ```
+    python mp4_videos_to_png_frames.py
+    ```
+
     It should output the following file structure
     ```
     left_colmap_baseline/
-        00000.png
-        00001.png
-        ...
-        00XXX.png
+        ....png
+        ....png
+        ....png
     right_colmap_baseline/
-        00000.png
-        00001.png
-        ...
-        00XXX.png
-    left_well_exposed/
-        00000.png
-        00001.png
-        ...
-        00XXX.png
-    right_fast_exposed/
-        00000.png
-        00001.png
-        ...
-        00XXX.png
+        ....png
+        ....png
+        ....png
+    left_e1/
+        ....png
+        ....png
+        ....png
+    right_right_e1/
+        ....png
+        ....png
+        ....png
+    ```
+    
+6. Run the code to mask out people and the stick. You may need to specify `--input-dir` and `--output-dir`
+
+    ```
+    python lantern_scripts/mask_humans2.py
     ```
-5. 
 
 
 ### OPTION 2: Synthetic data
@@ -65,11 +78,17 @@ python lantern_scripts/split_synthetic_colmap.py --has_exposure --left_dir ./Dua
 
 ## 2. Process data using COLMAP
 
+Example with synthetic data:
+
 ```
 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 
 ```
 
-> 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.
+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 
+```
 
 ## 2. Pre-training NeRF
 
@@ -81,6 +100,9 @@ After convergence, go to the next step.
 ns-train lantern-nerfacto --help
 ```
 
+> 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.
+
+
 ## 3. Running MomoNet
 
 (Using which frames?) A pretrained UNet image-to-image network called MomoNet will infill the missing dynamic range and store these predictions in new images.

lantern_scripts/mask_humans2.py

@@ -0,0 +1,114 @@
+from detectron2.utils.logger import setup_logger
+
+setup_logger()
+
+# import some common libraries
+import argparse
+import os
+
+import cv2
+import numpy as np
+import torch
+from lang_sam import LangSAM
+from PIL import Image
+
+# import some common detectron2 utilities
+from tqdm import tqdm
+
+FRAME_SUBFOLDERS_CONFIG = [
+    {
+        'name': 'left_colmap_baseline',
+        'flip_k': -1,
+        'mask_person': True,
+        'custom_mask': './data/lab_ground_floor/custom_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'
+    },
+    {
+        'name': 'left_e1',
+        'flip_k': -1,
+        'mask_person': True,
+        'custom_mask': './data/lab_ground_floor/custom_masks/left_mask.png'
+    },
+    {
+        'name': 'right_e2',
+        'flip_k': 1,
+        'mask_person': False,
+        'custom_mask': './data/lab_ground_floor/custom_masks/right_mask.png'
+    }
+]
+
+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')
+    
+    args = parser.parse_args()
+    out = args.output_dir
+    # out = f'{args.output_dir}/overlay'
+    if(not os.path.isdir(out)):
+        os.makedirs(out)
+    # stick = cv2.imread(args.stick)
+    for folder in FRAME_SUBFOLDERS_CONFIG:
+        assert os.path.isdir(os.path.join(args.input_dir, folder['name'])), f"Folder {folder['name']} does not exist"
+    
+    # imgs = os.listdir(args.input_dir)
+    
+    model = LangSAM()
+
+    for folder in FRAME_SUBFOLDERS_CONFIG:
+        print(f"Processing folder {folder['name']}")
+        # imgs = os.listdir(os.path.join(args.input_dir, folder['name']))
+        # filter png only
+        imgs = [f for f in os.listdir(os.path.join(args.input_dir, folder['name'])) if f.endswith(".png")]
+        custom_mask = cv2.imread(folder['custom_mask'])
+        custom_mask = cv2.cvtColor(custom_mask, cv2.COLOR_BGR2GRAY)
+        custom_mask = custom_mask.astype(np.bool_)
+        for i, file in enumerate(tqdm(imgs)):
+            original_image_pil = Image.open(os.path.join(args.input_dir, folder['name'], file)).convert("RGB")
+            original_image_np = np.array(original_image_pil)
+
+            if folder['mask_person']:
+                text_prompt = "person"
+                # 90 degrees rotat
+                image_np_rot = np.rot90(original_image_np, k=folder['flip_k'])
+                masks, boxes, phrases, logits = model.predict(Image.fromarray(image_np_rot), text_prompt)
+                masks = masks.detach().cpu().numpy()
+                masks = np.rot90(masks, k=-folder['flip_k'], axes=(1, 2))
+                
+                # aggregate masks
+                masks = np.max(masks, axis=0)
+
+            else:
+                masks = np.zeros((original_image_pil.height, original_image_pil.width))
+                
+            # custom mask (logical or)
+            masks = masks.astype(np.bool_)
+            masks = masks | custom_mask
+
+            # dilate mask
+            masks = masks.astype(np.uint8)
+            masks = cv2.dilate(masks, PERSON_DILATE_KERNEL)
+
+            # save masked image
+            out_image_np = original_image_np * (1 - masks)[:, :, None] + masks[:, :, None] * 255
+            out_image_np = out_image_np.astype(np.uint8)
+            out_masked_file_path = os.path.join(args.output_dir, folder['name'], file) + ".masked.png"
+            os.makedirs(os.path.dirname(out_masked_file_path), exist_ok=True)
+            out_masked_img = Image.fromarray(out_image_np)
+            out_masked_img.save(out_masked_file_path)
+
+            # masks to image
+            out_masked_file_path = os.path.join(args.output_dir, folder['name'], file)
+            os.makedirs(os.path.dirname(out_masked_file_path), exist_ok=True)
+            # invert mask
+            masks = 1 - masks
+            out_mask = Image.fromarray(masks * 255)
+            out_mask.save(out_masked_file_path)
+

lantern_scripts/mp4_videos_to_png_frames.py

@@ -0,0 +1,43 @@
+import subprocess
+import os
+import argparse
+
+
+# TODO: limit number of frames... maybe 3 per second?
+
+if __name__ == '__main__':
+    # List of input video files
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("--input_dir", type=str, default="./data/lab_ground_floor/trimmed_videos")
+
+    # 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")]
+
+    for video in videos:
+        # Get the video name without extension
+        video_name = os.path.splitext(os.path.basename(video))[0]
+
+        # Create a subfolder for the frames
+        output_folder = os.path.join(args.input_dir, f'{video_name}')
+        os.makedirs(output_folder, exist_ok=True)
+
+        # Run FFmpeg command to convert video to PNG frames
+        command = [
+            "ffmpeg",
+            "-i", video,
+            "-vf", "fps=29.97",
+            f"{output_folder}/output_%04d.png"
+        ]
+        subprocess.run(command, check=True)
+
+        # only keep frames where mod 10 = 0
+        for file in os.listdir(output_folder):
+            if not file.endswith(".png"):
+                continue
+            frame_number = int(os.path.splitext(file)[0].split("_")[1])
+            if frame_number % 10 != 0:
+                os.remove(os.path.join(output_folder, file))
+
+    print("Frames extraction completed.")
+

nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py

@@ -30,9 +30,9 @@ class BaseConverterToNerfstudioDataset(ABC):
     """Path the data, either a video file or a directory of images."""
     output_dir: Path
     """Path to the output directory."""
-    mask_dir: Path = None
+    mask_dir: Optional[Path] = None
     """Path to the mask directory."""
-    hdr_dir: Path = None
+    hdr_dir: Optional[Path] = None
     """Path to the HDR directory."""
     verbose: bool = False
     """If True, print extra logging."""

nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py

@@ -19,7 +19,9 @@ from pathlib import Path
 from typing import Dict, List, Literal, Optional, Tuple
 
 from nerfstudio.process_data import colmap_utils, hloc_utils, process_data_utils
-from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import BaseConverterToNerfstudioDataset
+from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import (
+    BaseConverterToNerfstudioDataset,
+)
 from nerfstudio.process_data.process_data_utils import CAMERA_MODELS
 from nerfstudio.utils import install_checks
 from nerfstudio.utils.rich_utils import CONSOLE
@@ -151,7 +153,7 @@ class ColmapConverterToNerfstudioDataset(BaseConverterToNerfstudioDataset):
     
     def _save_transforms(
         self,
-        num_frames: int,
+        num_frames: int = 0,
         image_id_to_depth_path: Optional[Dict[int, Path]] = None,
         camera_mask_path: Optional[Path] = None,
         camera_hdr_path: Optional[Path] = None,
@@ -177,7 +179,7 @@ class ColmapConverterToNerfstudioDataset(BaseConverterToNerfstudioDataset):
                     is_HDR=is_HDR,
                 )
                 summary_log.append(f"Colmap matched {num_matched_frames} images")
-            summary_log.append(colmap_utils.get_matching_summary(num_frames, num_matched_frames))
+            # summary_log.append(colmap_utils.get_matching_summary(num_frames, num_matched_frames))
 
         else:
             CONSOLE.log(

nerfstudio/process_data/colmap_utils.py

@@ -90,6 +90,7 @@ def run_colmap(
     colmap_dir: Path,
     camera_model: CameraModel,
     camera_mask_path: Optional[Path] = None,
+    mask_path: Optional[Path] = None,
     gpu: bool = True,
     verbose: bool = False,
     matching_method: Literal["vocab_tree", "exhaustive", "sequential"] = "vocab_tree",
@@ -101,7 +102,8 @@ def run_colmap(
         image_dir: Path to the directory containing the images.
         colmap_dir: Path to the output directory.
         camera_model: Camera model to use.
-        camera_mask_path: Path to the camera mask.
+        camera_mask_path: Path to the camera mask (a single image).
+        mask_path: Path for a directory containing masks for each image.
         gpu: If True, use GPU.
         verbose: If True, logs the output of the command.
         matching_method: Matching method to use.
@@ -124,6 +126,8 @@ def run_colmap(
     ]
     if camera_mask_path is not None:
         feature_extractor_cmd.append(f"--ImageReader.camera_mask_path {camera_mask_path}")
+    if mask_path is not None:
+        feature_extractor_cmd.append(f"--ImageReader.mask_path {mask_path}")
     feature_extractor_cmd = " ".join(feature_extractor_cmd)
     with status(msg="[bold yellow]Running COLMAP feature extractor...", spinner="moon", verbose=verbose):
         run_command(feature_extractor_cmd, verbose=verbose)
@@ -548,7 +552,7 @@ def colmap_to_json(
             "colmap_im_id": im_id,
         }
         if camera_mask_path is not None:
-            frame["mask_path"] = Path(camera_mask_path.absolute() / im_data.name).as_posix() # camera_mask_path.relative_to(camera_mask_path.parent.parent).as_posix()
+            frame["mask_path"] = (camera_mask_path/ im_data.name).as_posix() # camera_mask_path.relative_to(camera_mask_path.parent.parent).as_posix()
         if image_id_to_depth_path is not None:
             depth_path = image_id_to_depth_path[im_id]
             frame["depth_file_path"] = str(depth_path)

nerfstudio/process_data/equirect_utils.py

@@ -20,17 +20,22 @@ os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
 
 import json
 import math
+import shutil
 import sys
 from pathlib import Path
-from typing import List, Tuple
-import shutil
+from typing import List, Tuple, Optional
 
 import cv2
 import numpy as np
 import torch
 from numpy.linalg import inv
-from rich.progress import (BarColumn, Progress, TaskProgressColumn, TextColumn,
-                           TimeRemainingColumn)
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
 from scipy.spatial.transform import Rotation
 
 from equilib import Equi2Pers
@@ -115,12 +120,12 @@ def generate_planar_projections_from_equirectangular(
     image_dir: Path,
     planar_image_size: Tuple[int, int],
     samples_per_im: int,
-    mask_dir: Path,
-    hdr_dir: Path,
+    mask_dir: Optional[Path],
+    hdr_dir: Optional[Path] = None,
     crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
     is_HDR: bool = False,
     HDR_planar_image_size: Tuple[int, int] = (None, None)
-) -> Path:
+) -> Tuple[Path, Path, Path]:
     """Generate planar projections from an equirectangular image.
 
     Args:
@@ -148,7 +153,14 @@ def generate_planar_projections_from_equirectangular(
     if crop_factor[2] > 0:
         right_bound = 180 - 360 * crop_factor[2]
 
-    if samples_per_im == 8:
+    if samples_per_im == 1:
+        fov = 120
+        bound_arr = [0]
+        bound_arr = _crop_bound_arr_vertical(bound_arr, fov, crop_factor)
+        if bound_arr[0] is not None:
+            for i in np.arange(left_bound, right_bound, 360):
+                yaw_pitch_pairs.append((i, bound_arr[0]))
+    elif samples_per_im == 8:
         fov = 120
         bound_arr = [-45, 0, 45]
         bound_arr = _crop_bound_arr_vertical(bound_arr, fov, crop_factor)
@@ -285,7 +297,7 @@ def generate_planar_projections_from_equirectangular_with_two_exposures(
     e1_factor: float,
     e2_factor: float,
     crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
-) -> Path:
+) -> Tuple[Path, Path, Path, Path, Path]:
     """Generate planar projections from an equirectangular image.
 
     Args:
@@ -375,7 +387,7 @@ def generate_planar_projections_from_equirectangular_with_two_exposures(
                 im /=255.0
 
                 mask_e1_addr = os.path.join(mask_e1_dir, i)
-                mask_e1_addr = mask_e1_addr.replace('.png', '_e1.png')
+                # mask_e1_addr = mask_e1_addr.replace('.png', '_e1.png')
                 if not os.path.isfile(mask_e1_addr):
                     raise Exception("The corresponding mask E1 for " + i + " does not exist!!!")
                 mask_e1 = np.array(cv2.imread(mask_e1_addr))
@@ -383,7 +395,7 @@ def generate_planar_projections_from_equirectangular_with_two_exposures(
                 mask_e1 = torch.permute(mask_e1, (2, 0, 1))
 
                 mask_e2_addr = os.path.join(mask_e2_dir, i)
-                mask_e2_addr = mask_e2_addr.replace('.png', '_e2.png')
+                # mask_e2_addr = mask_e2_addr.replace('.png', '_e2.png')
                 if not os.path.isfile(mask_e2_addr):
                     raise Exception("The corresponding mask E2 for " + i + " does not exist!!!")
                 mask_e2 = np.array(cv2.imread(mask_e2_addr))