diff --git a/README-pandora.md b/README-pandora.md index ca021530d..9d1d3037d 100644 --- a/README-pandora.md +++ b/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. diff --git a/lantern_scripts/mask_humans2.py b/lantern_scripts/mask_humans2.py new file mode 100644 index 000000000..89f210863 --- /dev/null +++ b/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) + diff --git a/lantern_scripts/mp4_videos_to_png_frames.py b/lantern_scripts/mp4_videos_to_png_frames.py new file mode 100644 index 000000000..ff01f06e4 --- /dev/null +++ b/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.") + diff --git a/nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py b/nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py index 4b45eb51f..fa443e8fe 100644 --- a/nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py +++ b/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.""" diff --git a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py index fefd54479..1f8ea8695 100644 --- a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py +++ b/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( diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index 2231d7e72..00d5aa178 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/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) diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py index d3ac3f605..eb9985711 100644 --- a/nerfstudio/process_data/equirect_utils.py +++ b/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))