diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index bffb0bd03..e133d3757 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -18,6 +18,7 @@ Generic Writer class from __future__ import annotations import enum +import os from abc import abstractmethod from pathlib import Path from time import time @@ -283,7 +284,12 @@ class WandbWriter(Writer): """WandDB Writer Class""" def __init__(self, log_dir: Path): - wandb.init(project="nerfstudio-project", dir=str(log_dir), reinit=True, name=log_dir.name) + wandb.init( + project=os.environ.get("WANDB_PROJECT", "nerfstudio-project"), + dir=os.environ.get("WANDB_DIR", str(log_dir)), + name=os.environ.get("WANDB_NAME", log_dir.name), + reinit=True, + ) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: image = torch.permute(image, (2, 0, 1))