From e504c1f2f83321e8779fc8552763cc4de0b50f18 Mon Sep 17 00:00:00 2001 From: mokad6 Date: Wed, 17 Apr 2024 12:16:01 -0400 Subject: [PATCH] Working on LVSN lab sequences - filtering pose outliers - renedr rgb_fast --- README-pandora.md | 2 +- lantern_scripts/mask_humans2.py | 8 +- nerfstudio/data/pixel_samplers.py | 97 ++++++++++++------- nerfstudio/fields/nerfacto_field.py | 6 +- nerfstudio/lantern/config.py | 15 +-- nerfstudio/lantern/dataset.py | 4 + nerfstudio/lantern/model.py | 17 +++- nerfstudio/process_data/equirect_utils.py | 4 +- .../process_data/openSFM_to_nerfstudio.py | 77 +++++++++++---- nerfstudio/scripts/render.py | 52 +++++++--- nerfstudio/utils/colormaps.py | 18 ++-- nerfstudio/viewer/server/control_panel.py | 2 +- 12 files changed, 202 insertions(+), 100 deletions(-) diff --git a/README-pandora.md b/README-pandora.md index 462910199..89cb052cc 100644 --- a/README-pandora.md +++ b/README-pandora.md @@ -68,7 +68,7 @@ pip install skylibs equilib Make sure the masks in "./stick_masks" are valid ``` python lantern_scripts/mask_humans2.py --help - python lantern_scripts/mask_humans2.py --input_dir data/lab_downstairs/trimmed_videos + python lantern_scripts/mask_humans2.py --input_dir data/lab_downstairs/trimmed_videos --output_dir ``` diff --git a/lantern_scripts/mask_humans2.py b/lantern_scripts/mask_humans2.py index 48592e774..dc1d4db1a 100644 --- a/lantern_scripts/mask_humans2.py +++ b/lantern_scripts/mask_humans2.py @@ -17,13 +17,13 @@ from tqdm import tqdm FRAME_SUBFOLDERS_CONFIG = [ { - 'name': 'left_colmap_baseline', + 'name': 'left_sfm', 'flip_k': -1, 'mask_person': True, 'custom_mask': './stick_masks/left_mask.png' }, { - 'name': 'right_colmap_baseline', + 'name': 'right_sfm', 'flip_k': 1, 'mask_person': True, '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 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') + parser.add_argument('--input_dir', type=str, help='Directory containing all the 4 directories containing frames') + parser.add_argument('--output_dir', type=str, help='Directory to save masked frames') args = parser.parse_args() out = args.output_dir diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 7bd526510..78be90da0 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -61,8 +61,10 @@ class PixelSampler: config: PixelSamplerConfig is_well_exposed_computed: bool = False is_fast_exposed_computed: bool = False + is_light_sources_computed: bool = False well_exposed_indices: torch.Tensor fast_exposed_indices: torch.Tensor + light_sources_indices: torch.Tensor fast_image_batch = {} well_image_batch = {} def __init__(self, config: PixelSamplerConfig, **kwargs) -> None: @@ -85,32 +87,34 @@ class PixelSampler: def sample_considering_mask(self, mask, device, batch_size): num_masks = mask.shape[0] - size_of_disk = 50 - num_disks = ceil(num_masks / size_of_disk) - # slice mask to 100 disk - random_idx_disk = int(torch.randint(0, num_disks, (1,))) indices = torch.tensor([]).cuda().to(device) nonzero_indices_whole = torch.tensor([]).to(device) - - for current_disk, start in enumerate(range(0, num_masks, size_of_disk)): - is_last_disk = ((start//size_of_disk) == num_disks-1) - end = start + size_of_disk if not is_last_disk else num_masks - current_masks = mask[start:end,...].to(device) - nonzero_indices = torch.nonzero(current_masks, as_tuple=False) - # Add offset to indices:local indices --> global - nonzero_indices[:,0] += start + + if num_masks > 0: + size_of_disk = 50 + num_disks = ceil(num_masks / size_of_disk) + # slice mask to 100 disk + random_idx_disk = int(torch.randint(0, num_disks, (1,))) - num_samples = batch_size // (num_disks-1) if current_disk != random_idx_disk \ - else batch_size % (num_disks-1) - chosen_indices = torch.randint(0, nonzero_indices.shape[0], (num_samples,), device=device) - - indice_this_disk = nonzero_indices[chosen_indices] - indices = torch.cat((indices, indice_this_disk), dim=0).long() - nonzero_indices_whole = torch.cat((nonzero_indices_whole, nonzero_indices.cpu()), dim=0) - - del current_masks - del nonzero_indices - del indice_this_disk + for current_disk, start in enumerate(range(0, num_masks, size_of_disk)): + is_last_disk = ((start//size_of_disk) == num_disks-1) + end = start + size_of_disk if not is_last_disk else num_masks + current_masks = mask[start:end,...].to(device) + nonzero_indices = torch.nonzero(current_masks, as_tuple=False) + # Add offset to indices:local indices --> global + nonzero_indices[:,0] += start + + num_samples = batch_size // (num_disks-1) if current_disk != random_idx_disk \ + else batch_size % (num_disks-1) + chosen_indices = torch.randint(0, nonzero_indices.shape[0], (num_samples,), device=device) + + indice_this_disk = nonzero_indices[chosen_indices] + indices = torch.cat((indices, indice_this_disk), dim=0).long() + nonzero_indices_whole = torch.cat((nonzero_indices_whole, nonzero_indices.cpu()), dim=0) + + del current_masks + del nonzero_indices + del indice_this_disk return indices, nonzero_indices_whole @@ -126,6 +130,7 @@ class PixelSampler: device: Union[torch.device, str] = "cpu", is_fast_expo: bool = False, only_light_sources: bool = False, + validity: Optional[Tensor] = None, ) -> Int[Tensor, "batch_size 3"]: """ Naive pixel sampler, uniformly samples across all possible pixels of all possible images. @@ -145,14 +150,27 @@ class PixelSampler: mask = mask.squeeze(-1) if is_fast_expo: - if self.is_fast_exposed_computed: - chosen_indices = torch.randint(0, self.fast_exposed_indices.shape[0], (batch_size,)) - return self.fast_exposed_indices[chosen_indices].to(device) + if only_light_sources: + if self.is_light_sources_computed: + chosen_indices = torch.randint(0, self.light_sources_indices.shape[0], (batch_size,)) + return self.light_sources_indices[chosen_indices].to(device) + else: + assert validity is not None, "Validy needs to be defined for the lught sources samples" + validity = validity.squeeze(-1).cpu() + mask_combine = torch.logical_and(mask, validity) + indices, nonzero_indices_whole = self.sample_considering_mask(mask_combine, device, batch_size) + self.light_sources_indices = nonzero_indices_whole.long().cpu() + self.is_light_sources_computed = True + return indices else: - indices, nonzero_indices_whole = self.sample_considering_mask(mask, device, batch_size) - self.fast_exposed_indices = nonzero_indices_whole.long().cpu() - self.is_fast_exposed_computed = True - return indices + if self.is_fast_exposed_computed: + chosen_indices = torch.randint(0, self.fast_exposed_indices.shape[0], (batch_size,)) + return self.fast_exposed_indices[chosen_indices].to(device) + else: + indices, nonzero_indices_whole = self.sample_considering_mask(mask, device, batch_size) + self.fast_exposed_indices = nonzero_indices_whole.long().cpu() + self.is_fast_exposed_computed = True + return indices else: if self.is_well_exposed_computed: chosen_indices = torch.randint(0, self.well_exposed_indices.shape[0], (batch_size,)) @@ -200,7 +218,7 @@ class PixelSampler: return indices - def collate_image_dataset_batch(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False, is_fast_expo: bool = False): + def collate_image_dataset_batch(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False, is_fast_expo: bool = False, only_light_sources: bool = False): """ Operates on a batch of images and samples pixels to use for generating rays. Returns a collated batch which is input to the Graph. @@ -222,8 +240,11 @@ class PixelSampler: ) else: indices = self.sample_method( - num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device, is_fast_expo=is_fast_expo - ) + num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device, + is_fast_expo=is_fast_expo, + only_light_sources=only_light_sources, + validity=batch["saturation_mask"] + ) else: if self.config.is_equirectangular: indices = self.sample_method_equirectangular( @@ -347,14 +368,20 @@ class PixelSampler: self.well_image_batch[key] = image_batch[key][well_ones] self.fast_image_batch[key] = image_batch[key][fast_ones] well_pixel_batch = self.collate_image_dataset_batch( - self.well_image_batch, int(1.0 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image + self.well_image_batch, int(0.2 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image ) pixel_batch = well_pixel_batch #.copy() fast_pixel_batch = self.collate_image_dataset_batch( - self.fast_image_batch, int(0.0 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image, is_fast_expo=True + self.fast_image_batch, int(0.3 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image, is_fast_expo=True ) for key, value in fast_pixel_batch.items(): pixel_batch[key] = torch.cat((pixel_batch[key], value), 0) + + light_pixel_batch = self.collate_image_dataset_batch( + self.fast_image_batch, int(0.5 * self.num_rays_per_batch), keep_full_image=self.config.keep_full_image, is_fast_expo=True, only_light_sources=True + ) + for key, value in light_pixel_batch.items(): + pixel_batch[key] = torch.cat((pixel_batch[key], value), 0) else: pixel_batch = self.collate_image_dataset_batch( image_batch, self.num_rays_per_batch, keep_full_image=self.config.keep_full_image diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 25ad3a965..5a549369b 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -252,8 +252,9 @@ class NerfactoField(Field): selector = ((positions > 0.0) & (positions < 1.0)).all(dim=-1) positions = positions * selector[..., None] self._sample_locations = positions - if not self._sample_locations.requires_grad: - self._sample_locations.requires_grad = True + # if not self._sample_locations.requires_grad: + # self._sample_locations.requires_grad = True + # import pdb; pdb.set_trace() positions_flat = positions.view(-1, 3) h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1) density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) @@ -270,7 +271,6 @@ class NerfactoField(Field): self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None ) -> Dict[FieldHeadNames, Tensor]: assert density_embedding is not None - # with torch.no_grad(): outputs = {} if ray_samples.camera_indices is None: raise AttributeError("Camera indices are not provided.") diff --git a/nerfstudio/lantern/config.py b/nerfstudio/lantern/config.py index 74650961f..4e5f72c9c 100644 --- a/nerfstudio/lantern/config.py +++ b/nerfstudio/lantern/config.py @@ -27,13 +27,14 @@ def Get_lantern_config(): pipeline=VanillaPipelineConfig( datamanager=HDRVanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), - train_num_rays_per_batch=4096, - eval_num_rays_per_batch=4096, - camera_optimizer=CameraOptimizerConfig( - mode="SO3xR3", - optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2), - scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000), - ), + train_num_rays_per_batch=8192, + eval_num_rays_per_batch=8192, + # camera_optimizer=CameraOptimizerConfig( + # mode="SO3xR3", + # optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2), + # scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000), + # ), + camera_optimizer=None, ), model=LanternModelConfig(eval_num_rays_per_chunk=1 << 15), ), diff --git a/nerfstudio/lantern/dataset.py b/nerfstudio/lantern/dataset.py index bbfa71f1f..b6f9a1030 100644 --- a/nerfstudio/lantern/dataset.py +++ b/nerfstudio/lantern/dataset.py @@ -125,6 +125,10 @@ class HDRInputDataset(InputDataset): # Let's and the masks # data["mask"] = torch.logical_and(data["mask"], data["saturation_mask"]) + # image_filename = self._dataparser_outputs.image_filenames[image_idx].stem + # if image_filename.startswith('right_sfm'): + # data["mask"] = torch.zeros(data["mask"].shape) + if self._dataparser_outputs.exposures is not None: exposure_value = self._dataparser_outputs.exposures[image_idx] data["exposure"] = exposure_value diff --git a/nerfstudio/lantern/model.py b/nerfstudio/lantern/model.py index 72486b217..283e7a013 100644 --- a/nerfstudio/lantern/model.py +++ b/nerfstudio/lantern/model.py @@ -352,16 +352,26 @@ class LanternModel(NerfactoModel): # loss_dict["validity_loss_fast_exposure"] = loss_mask_fast_expo gt_rgb_w = gt_rgb.clone() + # u = 5000. + # gt_rgb_w = torch.log(1. + u * gt_rgb_w) / torch.log(torch.tensor(1.+u)) + # gt_rgb_w[(exposures_resized != 1.0).repeat(1, 3)] = 1.0# torch.clamp(gt_rgb_w[(exposures_resized != 1.0).repeat(1, 3)] / 0.004, 0, 1) # import pdb; pdb.set_trace() gt_rgb_f = gt_rgb.clone() + u = 5000. + gt_rgb_f = torch.log(1. + u * gt_rgb_f) / torch.log(torch.tensor(1.+u)) + # gt_rgb_f[(exposures_resized == 1.0).repeat(1, 3)] = gt_rgb_f[(exposures_resized == 1.0).repeat(1, 3)] * 0.004 # loss_dict["rgb_loss"] = self.rgb_loss(gt_rgb_w, pred_rgb) # loss_dict["rgb_loss_fast"] = self.rgb_loss(gt_rgb_f, pred_rgb_f) # loss_dict["rgb_loss"] = self.rgb_loss(gt_rgb_f, pred_rgb_f) - # loss_dict["rgb_loss"] = (mask_w_w * ((gt_rgb_w - pred_rgb) ** 2)).mean() - # loss_dict["rgb_loss_fast"] = (mask_f_w * ((10.0 * gt_rgb_f - pred_rgb_f) ** 2)).mean() + loss_dict["rgb_loss"] = (mask_w_w * ((gt_rgb_w - pred_rgb) ** 2)).mean() + loss_dict["rgb_loss_fast"] = (mask_f_w * ((gt_rgb_f - pred_rgb_f) ** 2)).mean() + + # RGB_fast = (pred_rgb * pred_rgb_f) * 0.004 + # loss_dict["rgb_loss_fast"] = 0.0005 * (mask_f_w * ((gt_rgb_f - RGB_fast) ** 2)).mean() + # print("rgb_loss_fast: ", loss_dict["rgb_loss_fast"]) # loss_dict["rgb_loss_fast"] = 300.0 * (mask_f_w * ((gt_rgb_f - pred_rgb_f) ** 2)).mean() # import pdb; pdb.set_trace() @@ -369,7 +379,7 @@ class LanternModel(NerfactoModel): # if weights_for_RGB_loss is not None: # loss_dict["rgb_loss"] = (weights_for_RGB_loss * ((gt_rgb - pred_rgb) ** 2)).mean() # else: - loss_dict["rgb_loss"] = self.rgb_loss(gt_rgb, pred_rgb) + # loss_dict["rgb_loss"] = self.rgb_loss(gt_rgb, pred_rgb) # print("rgb_loss: ", loss_dict["rgb_loss"], ", rgb_fast_expo_loss: ", loss_dict["rgb_fast_expo_loss"]) if self.training: @@ -378,6 +388,7 @@ class LanternModel(NerfactoModel): ) assert metrics_dict is not None and "distortion" in metrics_dict loss_dict["distortion_loss"] = self.config.distortion_loss_mult * metrics_dict["distortion"] + # loss_dict["fg"] = (torch.pow(1 - torch.sum( outputs["weights_list"][-1], dim=0), 2)).mean() if self.config.predict_normals: # orientation loss for computed normals loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean( diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py index db9e78406..a9544db5b 100644 --- a/nerfstudio/process_data/equirect_utils.py +++ b/nerfstudio/process_data/equirect_utils.py @@ -576,9 +576,9 @@ def generate_planar_projections_from_equirectangular_GT( frames = [] idx = 0 files_list = [file for file in os.listdir(frame_dir) - if os.path.isfile(os.path.join(frame_dir, file))] + if os.path.isfile(os.path.join(frame_dir, file)) and (prefix + f"_{file[:-4]}") in camera_to_worlds_panos.keys()] files_list = sorted(files_list) - + num_ims = len(files_list) with progress: for i in progress.track(files_list, description="", total=num_ims): diff --git a/nerfstudio/process_data/openSFM_to_nerfstudio.py b/nerfstudio/process_data/openSFM_to_nerfstudio.py index 3f397856c..5c07359bb 100644 --- a/nerfstudio/process_data/openSFM_to_nerfstudio.py +++ b/nerfstudio/process_data/openSFM_to_nerfstudio.py @@ -49,6 +49,13 @@ def apply_correction(coefs, img, CORRECTION_CURVE_TYPE) -> np.ndarray: return img +def distance(point_array, ref): + pow_2 = np.power(point_array - ref, 2) + sum_pow_2 = np.sum(pow_2, axis=-1) + dist = np.sqrt(sum_pow_2) + return dist + + def opensfm_to_opengl(shot_data): rvec = np.array(tuple(map(float, shot_data['rotation']))) tvec = np.array(tuple(map(float, shot_data['translation']))) @@ -139,20 +146,21 @@ class OpenSFMToNeRFStudioDataset(BaseConverterToNerfstudioDataset): if not self.skip_image_processing: openSFM_reconstruction = io.load_from_json(self.metadata) frames_names = openSFM_reconstruction[0]['shots'].keys() - - # Computing the basis trandoem matrix + frames_names = sorted(frames_names) + + # Computing the basis transform matrix rot_list = [] trn_list = [] + for fname in frames_names: if fname.startswith("left_sfm"): shot_data = openSFM_reconstruction[0]['shots'][fname] left_C2W, left_rvec, left_tvec = opensfm_to_opengl(shot_data) - + right_fname = fname.replace("left_sfm", "right_sfm") if right_fname in frames_names: shot_data = openSFM_reconstruction[0]['shots'][right_fname] right_C2W, right_rvec, right_tvec = opensfm_to_opengl(shot_data) - basis_change = np.linalg.inv(left_C2W) @ right_C2W from scipy.spatial.transform import Rotation as R @@ -163,37 +171,64 @@ class OpenSFMToNeRFStudioDataset(BaseConverterToNerfstudioDataset): else: print("The right one is not registered for " + fname) + # Filtering the data points + trn_np = np.asarray(trn_list) + trn_median = np.median(trn_np, 0) + dist = distance(trn_np, trn_median) + percentile_95 = np.percentile(dist, 95) + valid_idx = dist < percentile_95 + + trn_np = trn_np[valid_idx] + rot_np = np.asarray(rot_list)[valid_idx] + basis_change = np.zeros((4, 4)) + # basis_change = solution_w_SVD from sksurgerycore.algorithms.averagequaternions import average_quaternions - avg_quat = average_quaternions(np.asarray(rot_list)) + avg_quat = average_quaternions(rot_np) r = R.from_quat(avg_quat) basis_change[0:3, 0:3] = r.as_matrix() - basis_change[0:3, 3] = np.mean(np.asarray(trn_list), 0) + basis_change[0:3, 3] = np.mean(trn_np, 0) basis_change[3, 3] = 1.0 print("basis_change: ", basis_change) - + # Convert OpenSFM coordinates to the NeRFStudio - camera_dict = {} - camera_dict["frames"] = [] + translations = [] for fname in frames_names: # if fname.startswith("left_sfm"): if fname.startswith("left_e1"): shot_data = openSFM_reconstruction[0]['shots'][fname] - left_C2W, _, _ = opensfm_to_opengl(shot_data) - camera_dict["frames"].append( - { - "file_path": fname.split('.')[0], - "transform_matrix": left_C2W.tolist() - }) + _, _, left_tvec = opensfm_to_opengl(shot_data) + translations.append(left_tvec) + + translations_np = np.asarray(translations) + left_trn_median = np.median(translations_np, 0) + dist = distance(translations_np, left_trn_median) + percentile_95 = np.percentile(dist, 99) + valid_idx = dist < percentile_95 + # Convert OpenSFM coordinates to the NeRFStudio + camera_dict = {} + camera_dict["frames"] = [] + for idx, fname in enumerate(frames_names): + # if fname.startswith("left_sfm"): + if fname.startswith("left_e1"): right_fname = fname.replace("left_e1", "right_e2") # right_fname = fname.replace("left_sfm", "right_sfm") - right_C2W = left_C2W @ basis_change - camera_dict["frames"].append( - { - "file_path": right_fname.split('.')[0], - "transform_matrix": right_C2W.tolist() - }) + shot_data = openSFM_reconstruction[0]['shots'][fname] + left_C2W, left_rvec, left_tvec = opensfm_to_opengl(shot_data) + if valid_idx[idx]: + camera_dict["frames"].append( + { + "file_path": fname.split('.')[0], + "transform_matrix": left_C2W.tolist() + }) + + right_C2W = left_C2W @ basis_change + camera_dict["frames"].append( + { + "file_path": right_fname.split('.')[0], + "transform_matrix": right_C2W.tolist() + }) with open(self.output_dir / 'panoramic_transforms.json', 'w') as f: json.dump(camera_dict, f, indent=4) diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index bfdcf0a0c..9af596eb7 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -158,6 +158,7 @@ def _render_trajectory_video( output_image = outputs[rendered_output_name] is_depth = rendered_output_name.find("depth") != -1 is_validity = rendered_output_name.find("validity_w") != -1 or rendered_output_name.find("validity_f") != -1 + is_fast_exposure = rendered_output_name.find("rgb_fast") != -1 if is_depth: output_image = ( colormaps.apply_depth_colormap( @@ -170,6 +171,12 @@ def _render_trajectory_video( .cpu() .numpy() ) + elif is_fast_exposure: + u = 5000. + img_uncompress = torch.exp(output_image * torch.log(torch.tensor(u+1.))) - 1. + img_uncompress /= u + output_image = img_uncompress + output_image = (output_image.cpu().numpy()) elif is_validity: mask_f = torch.clip(outputs["validity_f"], 0, 1) # mask_f = torch.ones(mask_f.shape).to(mask_f.get_device()) - mask_f @@ -187,10 +194,6 @@ def _render_trajectory_video( .numpy() ) elif image_format == "exr": - # u = 10. - # output_image = torch.exp(output_image * torch.log(torch.tensor(u+1.))) - 1. - # output_image /= u - # output_image = output_image output_image = ( output_image.cpu().numpy() ) @@ -547,7 +550,7 @@ class RenderInterpolated(BaseRender): eval_num_rays_per_chunk=self.eval_num_rays_per_chunk, test_mode="test", ) - + install_checks.check_ffmpeg_installed() output_names = [] @@ -558,18 +561,37 @@ class RenderInterpolated(BaseRender): output_names.append(file_name.stem) else: assert pipeline.datamanager.train_dataset is not None - cameras = pipeline.datamanager.train_dataset.cameras - for file_name in pipeline.datamanager.train_dataset.image_filenames: - output_names.append(file_name.stem) - + poses = [] # pipeline.datamanager.train_dataset.cameras + ks_list = [] + Ks = pipeline.datamanager.train_dataset.cameras.get_intrinsics_matrices() + for idx, file_name in enumerate(pipeline.datamanager.train_dataset.image_filenames): + # if file_name.stem.startswith('right_e2'): + if file_name.stem.startswith('right_sfm'): + output_names.append(file_name.stem) + poses.append(pipeline.datamanager.train_dataset.cameras[idx].camera_to_worlds) + ks_list.append(Ks[idx, :, :]) - - seconds = self.interpolation_steps * len(cameras) / self.frame_rate - camera_path = get_interpolated_camera_path( - cameras=cameras, - steps=self.interpolation_steps, - order_poses=self.order_poses, + poses = np.stack(poses, axis=0) + poses = torch.tensor(poses, dtype=torch.float32) + ks = np.stack(ks_list, axis=0) + ks = torch.tensor(ks, dtype=torch.float32) + # import pdb; pdb.set_trace() + cameras = Cameras( + fx = ks[:, 0, 0], + fy = ks[:, 1, 1], + cx = ks[0, 0, 2], + cy = ks[0, 1, 2], + camera_type = pipeline.datamanager.train_dataset.cameras.camera_type[0], + camera_to_worlds = poses, ) + # import pdb; pdb.set_trace() + + seconds = self.interpolation_steps * len(cameras) / self.frame_rate + camera_path = cameras #get_interpolated_camera_path( + # cameras=cameras, + # steps=self.interpolation_steps, + # order_poses=self.order_poses, + # ) _render_trajectory_video( pipeline, diff --git a/nerfstudio/utils/colormaps.py b/nerfstudio/utils/colormaps.py index c8aca1e51..617075b9f 100644 --- a/nerfstudio/utils/colormaps.py +++ b/nerfstudio/utils/colormaps.py @@ -24,7 +24,7 @@ from torch import Tensor from nerfstudio.utils import colors -Colormaps = Literal["default", "turbo", "viridis", "magma", "inferno", "cividis", "gray", "pca"] +Colormaps = Literal["default", "turbo", "viridis", "magma", "inferno", "cividis", "gray", "pca", "inverse-mu"] @dataclass(frozen=True) @@ -66,13 +66,15 @@ def apply_colormap( # default for rgb images if image.shape[-1] == 3: # if(colormap_options.exposure_scale != 0 or image.max() > 1): - if True: - img_gamma_22 = image - img_gamma_22 *= 2.**(colormap_options.exposure_scale) - img_gamma_22 = torch.pow(img_gamma_22, 1./2.2) - return torch.clamp(img_gamma_22, 0., 1.) - else: - return image + if colormap_options.colormap == 'inverse-mu': + u = 5000. + img_uncompress = torch.exp(image * torch.log(torch.tensor(u+1.))) - 1. + img_uncompress /= u + image = img_uncompress + img_gamma_22 = image + img_gamma_22 *= 2.**(colormap_options.exposure_scale) + img_gamma_22 = torch.pow(img_gamma_22, 1./2.2) + return torch.clamp(img_gamma_22, 0., 1.) # rendering depth outputs if image.shape[-1] == 1 and torch.is_floating_point(image): diff --git a/nerfstudio/viewer/server/control_panel.py b/nerfstudio/viewer/server/control_panel.py index 8f2e1bc34..b55885614 100644 --- a/nerfstudio/viewer/server/control_panel.py +++ b/nerfstudio/viewer/server/control_panel.py @@ -377,7 +377,7 @@ def _get_colormap_options(dimensions: int, dtype: type) -> List[Colormaps]: """ colormap_options: List[Colormaps] = [] if dimensions == 3: - colormap_options = ["default"] + colormap_options = ["default", "inverse-mu"] if dimensions == 1 and dtype == torch.float: colormap_options = [c for c in list(get_args(Colormaps)) if c not in ("default", "pca")] if dimensions > 3: