Spaces:
Running on Zero
Running on Zero
File size: 9,698 Bytes
7405e85 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 | '''
Modified from latentSplat and pixelSplat to handle extrapolation and more context views
'''
from dataclasses import dataclass
from typing import Literal, Optional
import torch
from jaxtyping import Float, Int64
from torch import Tensor
import random
from .view_sampler import ViewSampler
def farthest_point_sample(xyz, npoint):
"""
Input:
xyz: pointcloud data, [B, N, 3]
npoint: number of samples
Return:
centroids: sampled pointcloud index, [B, npoint]
"""
device = xyz.device
B, N, C = xyz.shape
centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
distance = torch.ones(B, N).to(device) * 1e10
batch_indices = torch.arange(B, dtype=torch.long).to(device)
barycenter = torch.sum((xyz), 1)
barycenter = barycenter / xyz.shape[1]
barycenter = barycenter.view(B, 1, 3)
dist = torch.sum((xyz - barycenter) ** 2, -1)
farthest = torch.max(dist, 1)[1]
for i in range(npoint):
# print("The %d farthest pts %s " % (i, farthest))
centroids[:, i] = farthest
centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
dist = torch.sum((xyz - centroid) ** 2, -1)
mask = dist < distance
distance[mask] = dist[mask]
farthest = torch.max(distance, -1)[1]
return centroids
@dataclass
class ViewSamplerBoundedV2Cfg:
name: Literal["boundedv2"]
num_context_views: int
num_target_views: int
min_distance_between_context_views: int
max_distance_between_context_views: int
max_distance_to_context_views: int
context_gap_warm_up_steps: int
target_gap_warm_up_steps: int
initial_min_distance_between_context_views: int
initial_max_distance_between_context_views: int
initial_max_distance_to_context_views: int
extra_views_sampling_strategy: Optional[Literal["random", "farthest_point", "equal"]] = "random"
target_views_replace_sample: Optional[bool] = True
class ViewSamplerBoundedV2(ViewSampler[ViewSamplerBoundedV2Cfg]):
def schedule(self, initial: int, final: int, steps: int) -> int:
fraction = self.global_step / steps
return min(initial + int((final - initial) * fraction), final)
def sample(
self,
scene: str,
extrinsics: Float[Tensor, "view 4 4"],
intrinsics: Float[Tensor, "view 3 3"],
device: torch.device = torch.device("cpu"),
max_num_views: Optional[int] = None,
min_context_views: int = 0,
max_context_views: int = 0,
min_view_dist: int | None = None,
max_view_dist: int | None = None,
) -> tuple[
Int64[Tensor, " context_view"], # indices for context views
Int64[Tensor, " target_view"], # indices for target views
]:
num_views, _, _ = extrinsics.shape
if max_num_views is not None:
num_views = min(num_views, max_num_views)
if min_context_views > 0 and max_context_views > 0 and self.stage != "test":
random_num_views = random.randint(min_context_views, max_context_views)
else:
random_num_views = None
# Compute the context view spacing based on the current global step.
if self.stage == "test":
# When testing, always use the full gap.
max_context_gap = self.cfg.max_distance_between_context_views
min_context_gap = self.cfg.max_distance_between_context_views
elif self.cfg.context_gap_warm_up_steps > 0:
max_context_gap = self.schedule(
self.cfg.initial_max_distance_between_context_views,
self.cfg.max_distance_between_context_views,
self.cfg.context_gap_warm_up_steps,
)
min_context_gap = self.schedule(
self.cfg.initial_min_distance_between_context_views,
self.cfg.min_distance_between_context_views,
self.cfg.context_gap_warm_up_steps,
)
else:
max_context_gap = self.cfg.max_distance_between_context_views
min_context_gap = self.cfg.min_distance_between_context_views
if min_view_dist is not None and max_view_dist is not None:
# for mixed dataset training, with different sampling distance
min_context_gap = min_view_dist
max_context_gap = max_view_dist
if random_num_views is not None:
# smaller context gap accordingly
scale_factor = max(max_context_views // random_num_views, 1)
max_context_gap = max_context_gap // scale_factor
min_context_gap = min_context_gap // scale_factor
if not self.cameras_are_circular:
max_context_gap = min(
num_views - 1, max_context_gap
)
# Compute the margin from context window to target window based on the current global step
if self.stage != "test" and self.cfg.target_gap_warm_up_steps > 0:
max_target_gap = self.schedule(
self.cfg.initial_max_distance_to_context_views,
self.cfg.max_distance_to_context_views,
self.cfg.target_gap_warm_up_steps,
)
else:
max_target_gap = self.cfg.max_distance_to_context_views
# Pick the gap between the context views.
if max_context_gap < min_context_gap:
raise ValueError("Example does not have enough frames!")
context_gap = torch.randint(
min_context_gap,
max_context_gap + 1,
size=tuple(),
device=device,
).item()
# Pick the left and right context indices.
index_context_left = torch.randint(
low=0,
high=num_views if self.cameras_are_circular else num_views - context_gap,
size=tuple(),
device=device,
).item()
if self.stage == "test":
index_context_left = index_context_left * 0
index_context_right = index_context_left + context_gap
index_target_left = index_context_left - max_target_gap
index_target_right = index_context_right + max_target_gap
if not self.cameras_are_circular:
index_target_left = max(0, index_target_left)
index_target_right = min(num_views - 1, index_target_right)
# Pick the target view indices.
if self.stage == "test":
# When testing, pick all.
index_target = torch.arange(
index_target_left,
index_target_right + 1,
device=device,
)
else:
# When training or validating (visualizing), pick at random.
if self.cfg.target_views_replace_sample:
index_target = torch.randint(
index_target_left,
index_target_right + 1,
size=(self.cfg.num_target_views,),
device=device,
)
else: # sample without replacement
# similarly, ok to index
index_target_candidates = torch.arange(
index_target_left,
index_target_right + 1,
device=device,
)
indices = torch.randperm(index_target_right + 1 - index_target_left, device=device)[
: self.cfg.num_target_views
]
index_target = index_target_candidates[indices]
# Apply modulo for circular datasets.
if self.cameras_are_circular:
index_target %= num_views
index_context_right %= num_views
# If more than two context views are desired, pick extra context views between
# the left and right ones.
if random_num_views is not None:
total_num_views = random_num_views
else:
total_num_views = self.cfg.num_context_views
if total_num_views > 2:
num_extra_views = total_num_views - 2
extra_views = []
if self.cfg.extra_views_sampling_strategy == 'random':
while len(set(extra_views)) != num_extra_views:
extra_views = torch.randint(
index_context_left + 1,
index_context_right,
(num_extra_views,),
).tolist()
elif self.cfg.extra_views_sampling_strategy == 'farthest_point':
context_bounded_index = torch.arange(index_context_left, index_context_right + 1)
candidate_views_position = extrinsics[context_bounded_index, :3, -1].unsqueeze(0)
index_context_local = farthest_point_sample(
candidate_views_position, total_num_views
).squeeze(0)
# remap context index back to global scene based index
index_context = context_bounded_index[index_context_local]
index_context_left = index_context[0].item()
index_context_right = index_context[-1].item()
extra_views = index_context[1:-1].tolist()
elif self.cfg.extra_views_sampling_strategy == 'equal':
pass
# sort the index
extra_views = sorted(extra_views)
else:
extra_views = []
return (
torch.tensor((index_context_left, *extra_views, index_context_right)),
index_target,
)
@property
def num_context_views(self) -> int:
return self.cfg.num_context_views
@property
def num_target_views(self) -> int:
return self.cfg.num_target_views
|