tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code

Fix Nerfacto + Aria (#2950)

Browse Source
This commit is contained in:
Brent Yi 2024-02-26 10:16:54 -08:00 committed by GitHub
parent 4dcbfc55fb
commit c491e3e101
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 41 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
nerfstudio/data/pixel_samplers.py
View File

@ -178,22 +178,50 @@ class PixelSampler:
if isinstance(mask, torch.Tensor) and not self.config.ignore_mask:
indices = self.sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device)
else:
rand_samples = torch.rand((batch_size, 3), device=device)
# convert random samples tto radius and theta
radii = self.config.fisheye_crop_radius * torch.sqrt(rand_samples[:, 1])
theta = 2.0 * torch.pi * rand_samples[:, 2]
# Rejection sampling.
valid: Optional[torch.Tensor] = None
indices = None
while True:
samples_needed = batch_size if valid is None else int(batch_size - torch.sum(valid).item())
# convert radius and theta to x and y between -radii and radii
x = radii * torch.cos(theta)
y = radii * torch.sin(theta)
# Check if done!
if samples_needed == 0:
break
# Multiply by the batch size and height/width to get pixel indices.
indices = torch.floor(
torch.stack([rand_samples[:, 0], y, x], dim=1)
* torch.tensor([num_images, image_height // 2, image_width // 2], device=device)
+ torch.tensor([0, image_height // 2, image_width // 2], device=device)
).long()
rand_samples = torch.rand((samples_needed, 2), device=device)
# Convert random samples to radius and theta.
radii = self.config.fisheye_crop_radius * torch.sqrt(rand_samples[:, 0])
theta = 2.0 * torch.pi * rand_samples[:, 1]
# Convert radius and theta to x and y.
x = (radii * torch.cos(theta) + image_width // 2).long()
y = (radii * torch.sin(theta) + image_height // 2).long()
sampled_indices = torch.stack(
[torch.randint(0, num_images, size=(samples_needed,), device=device), y, x], dim=-1
)
# Update indices.
if valid is None:
indices = sampled_indices
valid = (
(sampled_indices[:, 1] >= 0)
& (sampled_indices[:, 1] < image_height)
& (sampled_indices[:, 2] >= 0)
& (sampled_indices[:, 2] < image_width)
)
else:
assert indices is not None
not_valid = ~valid
indices[not_valid, :] = sampled_indices
valid[not_valid] = (
(sampled_indices[:, 1] >= 0)
& (sampled_indices[:, 1] < image_height)
& (sampled_indices[:, 2] >= 0)
& (sampled_indices[:, 2] < image_width)
)
assert indices is not None
assert indices.shape == (batch_size, 3)
return indices
def collate_image_dataset_batch(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False):