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

Merge 50b2f0349f into 50e0e3c70c

Browse Source
This commit is contained in:
basebala 2026-07-09 16:56:40 +00:00 committed by GitHub
commit cd2c34a4f3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 34 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
Dockerfile
View File

@ -81,11 +81,12 @@ RUN git clone https://github.com/colmap/colmap.git && \
# Upgrade pip and install dependencies.
# pip install torch==2.2.2 torchvision==0.17.2 --index-url https://download.pytorch.org/whl/cu118 && \
RUN pip install --no-cache-dir --upgrade pip 'setuptools<70.0.0' && \
RUN pip install --no-cache-dir --upgrade pip setuptools wheel && \
pip install --no-cache-dir 'setuptools<70.0.0' && \
pip install --no-cache-dir torch==2.1.2+cu118 torchvision==0.16.2+cu118 'numpy<2.0.0' --extra-index-url https://download.pytorch.org/whl/cu118 && \
git clone --branch master --recursive https://github.com/cvg/Hierarchical-Localization.git /opt/hloc && \
cd /opt/hloc && git checkout v1.4 && python3.10 -m pip install --no-cache-dir . && cd ~ && \
TCNN_CUDA_ARCHITECTURES="${CUDA_ARCHITECTURES}" pip install --no-cache-dir "git+https://github.com/NVlabs/tiny-cuda-nn.git@b3473c81396fe927293bdfd5a6be32df8769927c#subdirectory=bindings/torch" && \
TCNN_CUDA_ARCHITECTURES="${CUDA_ARCHITECTURES}" pip install --no-cache-dir --no-build-isolation "git+https://github.com/NVlabs/tiny-cuda-nn.git@b3473c81396fe927293bdfd5a6be32df8769927c#subdirectory=bindings/torch" && \
pip install --no-cache-dir pycolmap==0.6.1 pyceres==2.1 omegaconf==2.3.0
# Install gsplat and nerfstudio.
@ -97,8 +98,10 @@ RUN pip install --no-cache-dir --upgrade pip 'setuptools<70.0.0' && \
COPY --from=source /tmp/nerfstudio/ /tmp/nerfstudio
RUN export TORCH_CUDA_ARCH_LIST="$(echo "$CUDA_ARCHITECTURES" | tr ';' '\n' | awk '$0 > 70 {print substr($0,1,1)"."substr($0,2)}' | tr '\n' ' ' | sed 's/ $//')" && \
export MAX_JOBS=4 && \
pip install --no-cache-dir torch==2.1.2+cu118 torchvision==0.16.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118 && \
pip install --upgrade setuptools && \
GSPLAT_VERSION="$(sed -n 's/.*gsplat==\s*\([^," '"'"']*\).*/\1/p' /tmp/nerfstudio/pyproject.toml)" && \
pip install --no-cache-dir git+https://github.com/nerfstudio-project/gsplat.git@v${GSPLAT_VERSION} && \
pip install --no-cache-dir --no-build-isolation git+https://github.com/nerfstudio-project/gsplat.git@v${GSPLAT_VERSION} && \
pip install --no-cache-dir /tmp/nerfstudio 'numpy<2.0.0' && \
rm -rf /tmp/nerfstudio

7
nerfstudio/data/dataparsers/nerfstudio_dataparser.py
View File

@ -31,6 +31,7 @@ from nerfstudio.data.utils.dataparsers_utils import (
get_train_eval_split_all,
get_train_eval_split_filename,
get_train_eval_split_fraction,
get_train_eval_split_indices,
get_train_eval_split_interval,
)
from nerfstudio.utils.io import load_from_json
@ -59,7 +60,7 @@ class NerfstudioDataParserConfig(DataParserConfig):
"""The method to use to center the poses."""
auto_scale_poses: bool = True
"""Whether to automatically scale the poses to fit in +/- 1 bounding box."""
eval_mode: Literal["fraction", "filename", "interval", "all"] = "fraction"
eval_mode: Literal["fraction", "filename", "interval", "all", "indices"] = "fraction"
"""
The method to use for splitting the dataset into train and eval.
Fraction splits based on a percentage for train and the remaining for eval.
@ -77,6 +78,8 @@ class NerfstudioDataParserConfig(DataParserConfig):
"""Replace the unknown pixels with this color. Relevant if you have a mask but still sample everywhere."""
load_3D_points: bool = False
"""Whether to load the 3D points from the colmap reconstruction."""
eval_image_indices: Tuple[int, ...] = (0,)
"""Specifies the image indices to use during eval; if None, uses all."""
@dataclass
@ -212,6 +215,8 @@ class Nerfstudio(DataParser):
i_train, i_eval = get_train_eval_split_filename(image_filenames)
elif self.config.eval_mode == "interval":
i_train, i_eval = get_train_eval_split_interval(image_filenames, self.config.eval_interval)
elif self.config.eval_mode == "indices":
i_train, i_eval = get_train_eval_split_indices(image_filenames, self.config.eval_image_indices)
elif self.config.eval_mode == "all":
CONSOLE.log(
"[yellow] Be careful with '--eval-mode=all'. If using camera optimization, the cameras may diverge in the current implementation, giving unpredictable results."

20
nerfstudio/data/utils/dataparsers_utils.py
View File

@ -99,3 +99,23 @@ def get_train_eval_split_all(image_filenames: List) -> Tuple[np.ndarray, np.ndar
i_train = i_all
i_eval = i_all
return i_train, i_eval
def get_train_eval_split_indices(
image_filenames: List, eval_image_indices: Tuple[int, ...]
) -> Tuple[np.ndarray, np.ndarray]:
"""
Get the train/eval split based on specified indices in the config.
Args:
image_filenames: list of image filenames
eval_image_indices: Tuple of indices to use for evaluation.
"""
for idx in eval_image_indices:
if idx >= len(image_filenames) or idx < 0:
raise ValueError(f"Eval index {idx} is out of bounds for the number of images {len(image_filenames)}.")
all_indices = set(range(len(image_filenames)))
eval_indices = set(eval_image_indices)
train_indices = all_indices - eval_indices
return np.array(sorted(train_indices)), np.array(sorted(eval_indices))

4
tests/process_data/test_process_images.py
View File

@ -59,7 +59,7 @@ def test_process_images_skip_colmap(tmp_path: Path):
id=1,
xyz=np.array([0, 0, 0]),
rgb=np.array([0, 0, 0]),
error=np.array([0]),
error=0.0,
image_ids=np.array([1]),
point2D_idxs=np.array([0]),
),
@ -144,7 +144,7 @@ def test_process_images_recursively_skip_colmap(tmp_path: Path):
id=1,
xyz=np.array([0, 0, 0]),
rgb=np.array([0, 0, 0]),
error=np.array([0]),
error=0.0,
image_ids=np.array([1]),
point2D_idxs=np.array([0]),
),