Add direct render+export and notifications

nerfstudio/scripts/exporter.py

@@ -57,6 +57,8 @@ class Exporter:
     """Path to the config YAML file."""
     output_dir: Path
     """Path to the output directory."""
+    complete: bool = False
+    """Set to True when export is finished."""
 
 
 def validate_pipeline(normal_method: str, normal_output_name: str, pipeline: Pipeline) -> None:
@@ -74,7 +76,10 @@ def validate_pipeline(normal_method: str, normal_output_name: str, pipeline: Pip
         pixel_area = torch.ones_like(origins[..., :1])
         camera_indices = torch.zeros_like(origins[..., :1])
         ray_bundle = RayBundle(
-            origins=origins, directions=directions, pixel_area=pixel_area, camera_indices=camera_indices
+            origins=origins,
+            directions=directions,
+            pixel_area=pixel_area,
+            camera_indices=camera_indices,
         )
         outputs = pipeline.model(ray_bundle)
         if normal_output_name not in outputs:
@@ -135,7 +140,12 @@ def main(self) -> None:
         # Increase the batchsize to speed up the evaluation.
         assert isinstance(
             pipeline.datamanager,
-            (VanillaDataManager, ParallelDataManager, FullImageDatamanager, RandomCamerasDataManager),
+            (
+                VanillaDataManager,
+                ParallelDataManager,
+                FullImageDatamanager,
+                RandomCamerasDataManager,
+            ),
         )
         assert pipeline.datamanager.train_pixel_sampler is not None
         pipeline.datamanager.train_pixel_sampler.num_rays_per_batch = self.num_rays_per_batch
@@ -153,7 +163,7 @@ def main(self) -> None:
             estimate_normals=estimate_normals,
             rgb_output_name=self.rgb_output_name,
             depth_output_name=self.depth_output_name,
-            normal_output_name=self.normal_output_name if self.normal_method == "model_output" else None,
+            normal_output_name=(self.normal_output_name if self.normal_method == "model_output" else None),
             crop_obb=crop_obb,
             std_ratio=self.std_ratio,
         )
@@ -180,6 +190,8 @@ def main(self) -> None:
         print("\033[A\033[A")
         CONSOLE.print("[bold green]:white_check_mark: Saving Point Cloud")
 
+        self.complete = True
+
 
 @dataclass
 class ExportTSDFMesh(Exporter):
@@ -241,18 +253,21 @@ def main(self) -> None:
         if self.texture_method == "nerf":
             # load the mesh from the tsdf export
             mesh = get_mesh_from_filename(
-                str(self.output_dir / "tsdf_mesh.ply"), target_num_faces=self.target_num_faces
+                str(self.output_dir / "tsdf_mesh.ply"),
+                target_num_faces=self.target_num_faces,
             )
             CONSOLE.print("Texturing mesh with NeRF")
             texture_utils.export_textured_mesh(
                 mesh,
                 pipeline,
                 self.output_dir,
-                px_per_uv_triangle=self.px_per_uv_triangle if self.unwrap_method == "custom" else None,
+                px_per_uv_triangle=(self.px_per_uv_triangle if self.unwrap_method == "custom" else None),
                 unwrap_method=self.unwrap_method,
                 num_pixels_per_side=self.num_pixels_per_side,
             )
 
+        self.complete = True
+
 
 @dataclass
 class ExportPoissonMesh(Exporter):
@@ -316,7 +331,12 @@ def main(self) -> None:
         # Increase the batchsize to speed up the evaluation.
         assert isinstance(
             pipeline.datamanager,
-            (VanillaDataManager, ParallelDataManager, FullImageDatamanager, RandomCamerasDataManager),
+            (
+                VanillaDataManager,
+                ParallelDataManager,
+                FullImageDatamanager,
+                RandomCamerasDataManager,
+            ),
         )
         assert pipeline.datamanager.train_pixel_sampler is not None
         pipeline.datamanager.train_pixel_sampler.num_rays_per_batch = self.num_rays_per_batch
@@ -336,7 +356,7 @@ def main(self) -> None:
             estimate_normals=estimate_normals,
             rgb_output_name=self.rgb_output_name,
             depth_output_name=self.depth_output_name,
-            normal_output_name=self.normal_output_name if self.normal_method == "model_output" else None,
+            normal_output_name=(self.normal_output_name if self.normal_method == "model_output" else None),
             crop_obb=crop_obb,
             std_ratio=self.std_ratio,
         )
@@ -366,18 +386,21 @@ def main(self) -> None:
         if self.texture_method == "nerf":
             # load the mesh from the poisson reconstruction
             mesh = get_mesh_from_filename(
-                str(self.output_dir / "poisson_mesh.ply"), target_num_faces=self.target_num_faces
+                str(self.output_dir / "poisson_mesh.ply"),
+                target_num_faces=self.target_num_faces,
             )
             CONSOLE.print("Texturing mesh with NeRF")
             texture_utils.export_textured_mesh(
                 mesh,
                 pipeline,
                 self.output_dir,
-                px_per_uv_triangle=self.px_per_uv_triangle if self.unwrap_method == "custom" else None,
+                px_per_uv_triangle=(self.px_per_uv_triangle if self.unwrap_method == "custom" else None),
                 unwrap_method=self.unwrap_method,
                 num_pixels_per_side=self.num_pixels_per_side,
             )
 
+        self.complete = True
+
 
 @dataclass
 class ExportMarchingCubesMesh(Exporter):
@@ -439,11 +462,13 @@ def main(self) -> None:
             mesh,
             pipeline,
             self.output_dir,
-            px_per_uv_triangle=self.px_per_uv_triangle if self.unwrap_method == "custom" else None,
+            px_per_uv_triangle=(self.px_per_uv_triangle if self.unwrap_method == "custom" else None),
             unwrap_method=self.unwrap_method,
             num_pixels_per_side=self.num_pixels_per_side,
         )
 
+        self.complete = True
+
 
 @dataclass
 class ExportCameraPoses(Exporter):
@@ -460,7 +485,10 @@ def main(self) -> None:
         assert isinstance(pipeline, VanillaPipeline)
         train_frames, eval_frames = collect_camera_poses(pipeline)
 
-        for file_name, frames in [("transforms_train.json", train_frames), ("transforms_eval.json", eval_frames)]:
+        for file_name, frames in [
+            ("transforms_train.json", train_frames),
+            ("transforms_eval.json", eval_frames),
+        ]:
             if len(frames) == 0:
                 CONSOLE.print(f"[bold yellow]No frames found for {file_name}. Skipping.")
                 continue
@@ -629,6 +657,8 @@ def main(self) -> None:
 
         ExportGaussianSplat.write_ply(str(filename), count, map_to_tensors)
 
+        self.complete = True
+
 
 Commands = tyro.conf.FlagConversionOff[
     Union[

nerfstudio/scripts/render.py

@@ -75,9 +75,10 @@ def _render_trajectory_video(
     depth_near_plane: Optional[float] = None,
     depth_far_plane: Optional[float] = None,
     colormap_options: colormaps.ColormapOptions = colormaps.ColormapOptions(),
-    render_nearest_camera=False,
+    render_nearest_camera: bool = False,
     check_occlusions: bool = False,
-) -> None:
+    kill_flag: List[bool] = [False],
+) -> bool:
     """Helper function to create a video of the spiral trajectory.
 
     Args:
@@ -137,6 +138,9 @@ def _render_trajectory_video(
 
         with progress:
             for camera_idx in progress.track(range(cameras.size), description=""):
+                if kill_flag[0]:
+                    return False
+
                 obb_box = None
                 if crop_data is not None:
                     obb_box = crop_data.obb
@@ -202,9 +206,13 @@ def _render_trajectory_video(
                 for rendered_output_name in rendered_output_names:
                     if rendered_output_name not in outputs:
                         CONSOLE.rule("Error", style="red")
-                        CONSOLE.print(f"Could not find {rendered_output_name} in the model outputs", justify="center")
                         CONSOLE.print(
-                            f"Please set --rendered_output_name to one of: {outputs.keys()}", justify="center"
+                            f"Could not find {rendered_output_name} in the model outputs",
+                            justify="center",
+                        )
+                        CONSOLE.print(
+                            f"Please set --rendered_output_name to one of: {outputs.keys()}",
+                            justify="center",
                         )
                         sys.exit(1)
                     output_image = outputs[rendered_output_name]
@@ -256,10 +264,17 @@ def _render_trajectory_video(
                 render_image = np.concatenate(render_image, axis=1)
                 if output_format == "images":
                     if image_format == "png":
-                        media.write_image(output_image_dir / f"{camera_idx:05d}.png", render_image, fmt="png")
+                        media.write_image(
+                            output_image_dir / f"{camera_idx:05d}.png",
+                            render_image,
+                            fmt="png",
+                        )
                     if image_format == "jpeg":
                         media.write_image(
-                            output_image_dir / f"{camera_idx:05d}.jpg", render_image, fmt="jpeg", quality=jpeg_quality
+                            output_image_dir / f"{camera_idx:05d}.jpg",
+                            render_image,
+                            fmt="jpeg",
+                            quality=jpeg_quality,
                         )
                 if output_format == "video":
                     if writer is None:
@@ -287,7 +302,15 @@ def _render_trajectory_video(
         table.add_row("Video", str(output_filename))
     else:
         table.add_row("Images", str(output_image_dir))
-    CONSOLE.print(Panel(table, title="[bold][green]:tada: Render Complete :tada:[/bold]", expand=False))
+    CONSOLE.print(
+        Panel(
+            table,
+            title="[bold][green]:tada: Render Complete :tada:[/bold]",
+            expand=False,
+        )
+    )
+
+    return True
 
 
 def insert_spherical_metadata_into_file(
@@ -432,6 +455,11 @@ class BaseRender:
     """If true, checks line-of-sight occlusions when computing camera distance and rejects cameras not visible to each other"""
     camera_idx: Optional[int] = None
     """Index of the training camera to render."""
+    kill_flag: List[bool] = field(default_factory=lambda: [False])
+    """Stop execution of render if set to True."""
+
+    def kill(self) -> None:
+        self.kill_flag[0] = True
 
 
 @dataclass
@@ -442,6 +470,8 @@ class RenderCameraPath(BaseRender):
     """Filename of the camera path to render."""
     output_format: Literal["images", "video"] = "video"
     """How to save output data."""
+    complete: bool = True
+    """Set to True when render is finished."""
 
     def main(self) -> None:
         """Main function."""
@@ -485,7 +515,7 @@ def main(self) -> None:
         if self.camera_idx is not None:
             camera_path.metadata = {"cam_idx": self.camera_idx}
 
-        _render_trajectory_video(
+        self.complete = _render_trajectory_video(
             pipeline,
             camera_path,
             output_filename=self.output_path,
@@ -501,6 +531,7 @@ def main(self) -> None:
             colormap_options=self.colormap_options,
             render_nearest_camera=self.render_nearest_camera,
             check_occlusions=self.check_occlusions,
+            kill_flag=self.kill_flag,
         )
 
         if (
@@ -536,6 +567,7 @@ def main(self) -> None:
                 colormap_options=self.colormap_options,
                 render_nearest_camera=self.render_nearest_camera,
                 check_occlusions=self.check_occlusions,
+                kill_flag=self.kill_flag,
             )
 
             self.output_path = Path(str(left_eye_path.parent)[:-5] + ".mp4")
@@ -639,6 +671,7 @@ def main(self) -> None:
             colormap_options=self.colormap_options,
             render_nearest_camera=self.render_nearest_camera,
             check_occlusions=self.check_occlusions,
+            kill_flag=self.kill_flag,
         )
 
 
@@ -694,6 +727,7 @@ def main(self) -> None:
             colormap_options=self.colormap_options,
             render_nearest_camera=self.render_nearest_camera,
             check_occlusions=self.check_occlusions,
+            kill_flag=self.kill_flag,
         )
 
 
@@ -731,7 +765,10 @@ def main(self):
 
         def update_config(config: TrainerConfig) -> TrainerConfig:
             data_manager_config = config.pipeline.datamanager
-            assert isinstance(data_manager_config, (VanillaDataManagerConfig, FullImageDatamanagerConfig))
+            assert isinstance(
+                data_manager_config,
+                (VanillaDataManagerConfig, FullImageDatamanagerConfig),
+            )
             data_manager_config.eval_num_images_to_sample_from = -1
             data_manager_config.eval_num_times_to_repeat_images = -1
             if isinstance(data_manager_config, VanillaDataManagerConfig):
@@ -741,7 +778,11 @@ def update_config(config: TrainerConfig) -> TrainerConfig:
                 data_manager_config.data = self.data
             if self.downscale_factor is not None:
                 assert hasattr(data_manager_config.dataparser, "downscale_factor")
-                setattr(data_manager_config.dataparser, "downscale_factor", self.downscale_factor)
+                setattr(
+                    data_manager_config.dataparser,
+                    "downscale_factor",
+                    self.downscale_factor,
+                )
             return config
 
         config, pipeline, _, _ = eval_setup(
@@ -806,10 +847,12 @@ def update_config(config: TrainerConfig) -> TrainerConfig:
                         if rendered_output_name not in all_outputs:
                             CONSOLE.rule("Error", style="red")
                             CONSOLE.print(
-                                f"Could not find {rendered_output_name} in the model outputs", justify="center"
+                                f"Could not find {rendered_output_name} in the model outputs",
+                                justify="center",
                             )
                             CONSOLE.print(
-                                f"Please set --rendered-output-name to one of: {all_outputs}", justify="center"
+                                f"Please set --rendered-output-name to one of: {all_outputs}",
+                                justify="center",
                             )
                             sys.exit(1)
 
@@ -876,7 +919,10 @@ def update_config(config: TrainerConfig) -> TrainerConfig:
                             media.write_image(output_path.with_suffix(".png"), output_image, fmt="png")
                         elif self.image_format == "jpeg":
                             media.write_image(
-                                output_path.with_suffix(".jpg"), output_image, fmt="jpeg", quality=self.jpeg_quality
+                                output_path.with_suffix(".jpg"),
+                                output_image,
+                                fmt="jpeg",
+                                quality=self.jpeg_quality,
                             )
                         else:
                             raise ValueError(f"Unknown image format {self.image_format}")
@@ -889,7 +935,13 @@ def update_config(config: TrainerConfig) -> TrainerConfig:
         )
         for split in self.split.split("+"):
             table.add_row(f"Outputs {split}", str(self.output_path / split))
-        CONSOLE.print(Panel(table, title="[bold][green]:tada: Render on split {} Complete :tada:[/bold]", expand=False))
+        CONSOLE.print(
+            Panel(
+                table,
+                title="[bold][green]:tada: Render on split {} Complete :tada:[/bold]",
+                expand=False,
+            )
+        )
 
 
 Commands = tyro.conf.FlagConversionOff[