diff --git a/threestudio/models/renderers/nerf_volume_renderer.py b/threestudio/models/renderers/nerf_volume_renderer.py index ab27e476..f8a1caca 100644 --- a/threestudio/models/renderers/nerf_volume_renderer.py +++ b/threestudio/models/renderers/nerf_volume_renderer.py @@ -9,7 +9,11 @@ from threestudio.models.geometry.base import BaseImplicitGeometry from threestudio.models.materials.base import BaseMaterial from threestudio.models.renderers.base import VolumeRenderer +<<<<<<< HEAD from threestudio.utils.ops import chunk_batch, validate_empty_rays +======= +from threestudio.utils.ops import chunk_batch +>>>>>>> TextMesh coarse stage implementation (#121) from threestudio.utils.typing import * diff --git a/threestudio/models/renderers/neus_volume_renderer.py b/threestudio/models/renderers/neus_volume_renderer.py index ca53ce25..180fc907 100644 --- a/threestudio/models/renderers/neus_volume_renderer.py +++ b/threestudio/models/renderers/neus_volume_renderer.py @@ -104,7 +104,9 @@ def forward( n_rays = rays_o_flatten.shape[0] def alpha_fn(t_starts, t_ends, ray_indices): - ray_indices, t_starts, t_ends = validate_empty_rays(ray_indices, t_starts, t_ends) + ray_indices, t_starts, t_ends = validate_empty_rays( + ray_indices, t_starts, t_ends + ) t_starts, t_ends = t_starts[..., None], t_ends[..., None] t_origins = rays_o_flatten[ray_indices] t_positions = (t_starts + t_ends) / 2.0 @@ -153,7 +155,9 @@ def alpha_fn(t_starts, t_ends, ray_indices): cone_angle=0.0, ) - ray_indices, t_starts_, t_ends_ = validate_empty_rays(ray_indices, t_starts_, t_ends_) + ray_indices, t_starts_, t_ends_ = validate_empty_rays( + ray_indices, t_starts_, t_ends_ + ) ray_indices = ray_indices.long() t_starts, t_ends = t_starts_[..., None], t_ends_[..., None] t_origins = rays_o_flatten[ray_indices] @@ -163,7 +167,6 @@ def alpha_fn(t_starts, t_ends, ray_indices): positions = t_origins + t_dirs * t_positions t_intervals = t_ends - t_starts - if self.training: geo_out = self.geometry(positions, output_normal=True) rgb_fg_all = self.material( diff --git a/threestudio/utils/saving.py b/threestudio/utils/saving.py index f68ac954..b233d151 100644 --- a/threestudio/utils/saving.py +++ b/threestudio/utils/saving.py @@ -9,12 +9,12 @@ import numpy as np import torch import trimesh +import wandb from matplotlib import cm from matplotlib.colors import LinearSegmentedColormap from PIL import Image, ImageDraw from pytorch_lightning.loggers import WandbLogger -import wandb from threestudio.models.mesh import Mesh from threestudio.utils.typing import *