tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Support splatfacto training for metashape exports (#3122)

Browse Source 
* Added the possibility to import a metashape ply file

* FIxed double import and import order

* Ruff Format

* Ruff Format

* Ruff Format

* Ruff SUCKS

* make ply as an option and use applied_transform

* fix ruff

* fix bugs when ply doesn't exist

* fix bugs when ply doesn't exist

* no backtrace for ValueError

* hot fix

* Add return type

---------

Co-authored-by: simonbethke <simon.bethke@gmail.com>
Co-authored-by: Jianbo Ye <jianboye@amazon.com>
Co-authored-by: Brent Yi <yibrenth@gmail.com>
This commit is contained in:
J.Y 2024-05-01 21:24:33 -04:00 committed by GitHub
parent 3c4c3ae07e
commit 49f5934ab3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 46 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
nerfstudio/models/splatfacto.py
View File

@ -650,8 +650,8 @@ class SplatfactoModel(Model):
return image
@staticmethod
def get_empty_outputs(camera, background):
rgb = background.repeat(int(camera.height.item()), int(camera.width.item()), 1)
def get_empty_outputs(width: int, height: int, background: torch.Tensor) -> Dict[str, Union[torch.Tensor, List]]:
rgb = background.repeat(height, width, 1)
depth = background.new_ones(*rgb.shape[:2], 1) * 10
accumulation = background.new_zeros(*rgb.shape[:2], 1)
return {"rgb": rgb, "depth": depth, "accumulation": accumulation, "background": background}
@ -694,7 +694,7 @@ class SplatfactoModel(Model):
if self.crop_box is not None and not self.training:
crop_ids = self.crop_box.within(self.means).squeeze()
if crop_ids.sum() == 0:
return self.get_empty_outputs(camera, background)
return self.get_empty_outputs(int(camera.width.item()), int(camera.height.item()), background)
else:
crop_ids = None
camera_downscale = self._get_downscale_factor()
@ -754,7 +754,7 @@ class SplatfactoModel(Model):
camera.rescale_output_resolution(camera_downscale)
if (self.radii).sum() == 0:
return self.get_empty_outputs(camera, background)
return self.get_empty_outputs(W, H, background)
if self.config.sh_degree > 0:
viewdirs = means_crop.detach() - optimized_camera_to_world.detach()[:3, 3] # (N, 3)

24
nerfstudio/process_data/metashape_utils.py
View File

@ -17,9 +17,10 @@
import json
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Dict, List
from typing import Dict, List, Optional
import numpy as np
import open3d as o3d
from nerfstudio.process_data.process_data_utils import CAMERA_MODELS
from nerfstudio.utils.rich_utils import CONSOLE
@ -36,6 +37,7 @@ def metashape_to_json(
image_filename_map: Dict[str, Path],
xml_filename: Path,
output_dir: Path,
ply_filename: Optional[Path] = None, # type: ignore
verbose: bool = False,
) -> List[str]:
"""Convert Metashape data into a nerfstudio dataset.
@ -44,6 +46,7 @@ def metashape_to_json(
image_filename_map: Mapping of original image filenames to their saved locations.
xml_filename: Path to the metashape cameras xml file.
output_dir: Path to the output directory.
ply_filename: Path to the exported ply file.
verbose: Whether to print verbose output.
Returns:
@ -181,17 +184,34 @@ def metashape_to_json(
if component_id in component_dict:
transform = component_dict[component_id] @ transform
# Metashape camera is looking towards -Z, +X is to the right and +Y is to the top/up of the first cam
# Rotate the scene according to nerfstudio convention
transform = transform[[2, 0, 1, 3], :]
# Convert from Metashape's camera coordinate system (OpenCV) to ours (OpenGL)
transform[:, 1:3] *= -1
frame["transform_matrix"] = transform.tolist()
frames.append(frame)
data["frames"] = frames
applied_transform = np.eye(4)[:3, :]
applied_transform = applied_transform[np.array([2, 0, 1]), :]
data["applied_transform"] = applied_transform.tolist()
summary = []
if ply_filename is not None:
assert ply_filename.exists()
pc = o3d.io.read_point_cloud(str(ply_filename))
points3D = np.asarray(pc.points)
points3D = np.einsum("ij,bj->bi", applied_transform[:3, :3], points3D) + applied_transform[:3, 3]
pc.points = o3d.utility.Vector3dVector(points3D)
o3d.io.write_point_cloud(str(output_dir / "sparse_pc.ply"), pc)
data["ply_file_path"] = "sparse_pc.ply"
summary.append(f"Imported {ply_filename} as starting points")
with open(output_dir / "transforms.json", "w", encoding="utf-8") as f:
json.dump(data, f, indent=4)
summary = []
if num_skipped == 1:
summary.append(f"{num_skipped} image skipped because it was missing its camera pose.")
if num_skipped > 1:

28
nerfstudio/scripts/process_data.py
View File

@ -15,12 +15,11 @@
#!/usr/bin/env python
"""Processes a video or image sequence to a nerfstudio compatible dataset."""
import sys
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Union
from typing import Optional, Union
import numpy as np
import tyro
@ -230,12 +229,18 @@ class ProcessMetashape(BaseConverterToNerfstudioDataset, _NoDefaultProcessMetash
This script assumes that cameras have been aligned using Metashape. After alignment, it is necessary to export the
camera poses as a `.xml` file. This option can be found under `File > Export > Export Cameras`.
Additionally, the points can be exported as pointcloud under `File > Export > Export Point Cloud`. Make sure to
export the data in non-binary format and exclude the normals.
This script does the following:
1. Scales images to a specified size.
2. Converts Metashape poses into the nerfstudio format.
"""
ply: Optional[Path] = None
"""Path to the Metashape point export ply file."""
num_downscales: int = 3
"""Number of times to downscale the images. Downscales by 2 each time. For example a value of 3
will downscale the images by 2x, 4x, and 8x."""
@ -248,11 +253,17 @@ class ProcessMetashape(BaseConverterToNerfstudioDataset, _NoDefaultProcessMetash
if self.xml.suffix != ".xml":
raise ValueError(f"XML file {self.xml} must have a .xml extension")
if not self.xml.exists:
if not self.xml.exists():
raise ValueError(f"XML file {self.xml} doesn't exist")
if self.eval_data is not None:
raise ValueError("Cannot use eval_data since cameras were already aligned with Metashape.")
if self.ply is not None:
if self.ply.suffix != ".ply":
raise ValueError(f"PLY file {self.ply} must have a .ply extension")
if not self.ply.exists():
raise ValueError(f"PLY file {self.ply} doesn't exist")
self.output_dir.mkdir(parents=True, exist_ok=True)
image_dir = self.output_dir / "images"
image_dir.mkdir(parents=True, exist_ok=True)
@ -291,6 +302,7 @@ class ProcessMetashape(BaseConverterToNerfstudioDataset, _NoDefaultProcessMetash
image_filename_map=image_filename_map,
xml_filename=self.xml,
output_dir=self.output_dir,
ply_filename=self.ply,
verbose=self.verbose,
)
)
@ -335,7 +347,7 @@ class ProcessRealityCapture(BaseConverterToNerfstudioDataset, _NoDefaultProcessR
if self.csv.suffix != ".csv":
raise ValueError(f"CSV file {self.csv} must have a .csv extension")
if not self.csv.exists:
if not self.csv.exists():
raise ValueError(f"CSV file {self.csv} doesn't exist")
if self.eval_data is not None:
raise ValueError("Cannot use eval_data since cameras were already aligned with RealityCapture.")
@ -413,12 +425,12 @@ class ProcessODM(BaseConverterToNerfstudioDataset):
shots_file = self.data / "odm_report" / "shots.geojson"
reconstruction_file = self.data / "opensfm" / "reconstruction.json"
if not shots_file.exists:
if not shots_file.exists():
raise ValueError(f"shots file {shots_file} doesn't exist")
if not shots_file.exists:
if not shots_file.exists():
raise ValueError(f"cameras file {cameras_file} doesn't exist")
if not orig_images_dir.exists:
if not orig_images_dir.exists():
raise ValueError(f"Images dir {orig_images_dir} doesn't exist")
if self.eval_data is not None:
@ -522,7 +534,7 @@ def entrypoint():
tyro.extras.set_accent_color("bright_yellow")
try:
tyro.cli(Commands).main()
except RuntimeError as e:
except (RuntimeError, ValueError) as e:
CONSOLE.log("[bold red]" + e.args[0])