Hardcode FOV sampler constants

datasets/video/memory_selection.py

@@ -9,6 +9,10 @@ import torch
 
 
 SEGMENT_KEYS = ("anchor", "dynamic", "revisit")
+_FOV_NUM_POINTS = 10000
+_FOV_RADIUS = 30.0
+_FOV_HALF_H = 105.0 / 2.0
+_FOV_HALF_V = 75.0 / 2.0
 
 
 def cfg_get(cfg, key: str, default=None):
@@ -54,19 +58,14 @@ def _pose_directions(poses: torch.Tensor) -> torch.Tensor:
     return directions / torch.linalg.vector_norm(directions, dim=-1, keepdim=True).clamp_min(1e-6)
 
 
-def _sample_points_in_sphere(center: torch.Tensor, cfg=None) -> torch.Tensor:
-    num_points = max(1, int(cfg_get(cfg, "fov_num_points", 10000)))
-    radius = float(cfg_get(cfg, "fov_radius", cfg_get(cfg, "fov_spatial_radius", 30.0)))
-    seed = int(cfg_get(cfg, "fov_sample_seed", 0))
+def _sample_points_in_sphere(center: torch.Tensor) -> torch.Tensor:
     device = center.device
     dtype = center.dtype
-    generator = torch.Generator(device=device)
-    generator.manual_seed(seed)
 
-    samples_r = torch.rand(num_points, device=device, dtype=dtype, generator=generator)
-    samples_phi = torch.rand(num_points, device=device, dtype=dtype, generator=generator)
-    samples_u = torch.rand(num_points, device=device, dtype=dtype, generator=generator)
-    r = radius * torch.pow(samples_r, 1.0 / 3.0)
+    samples_r = torch.rand(_FOV_NUM_POINTS, device=device, dtype=dtype)
+    samples_phi = torch.rand(_FOV_NUM_POINTS, device=device, dtype=dtype)
+    samples_u = torch.rand(_FOV_NUM_POINTS, device=device, dtype=dtype)
+    r = _FOV_RADIUS * torch.pow(samples_r, 1.0 / 3.0)
     phi = 2.0 * torch.pi * samples_phi
     theta = torch.acos(1.0 - 2.0 * samples_u)
 
@@ -81,14 +80,12 @@ def _sample_points_in_sphere(center: torch.Tensor, cfg=None) -> torch.Tensor:
     return center[None, :3] + offsets
 
 
-def _inside_fov_points(points: torch.Tensor, poses: torch.Tensor, cfg=None) -> torch.Tensor:
+def _inside_fov_points(points: torch.Tensor, poses: torch.Tensor) -> torch.Tensor:
     points = points.to(dtype=torch.float32)
     poses = _as_pose_tensor(poses).to(device=points.device, dtype=torch.float32)
     if points.numel() == 0 or poses.numel() == 0:
         return torch.zeros((poses.shape[0], points.shape[0]), device=points.device, dtype=torch.bool)
 
-    fov_half_h = float(cfg_get(cfg, "fov_half_h", 105.0 / 2.0))
-    fov_half_v = float(cfg_get(cfg, "fov_half_v", 75.0 / 2.0))
     vectors = points[None, :, :] - poses[:, None, :3]
     x = vectors[..., 0]
     y = vectors[..., 1]
@@ -97,18 +94,18 @@ def _inside_fov_points(points: torch.Tensor, poses: torch.Tensor, cfg=None) -> t
     elevation = torch.rad2deg(torch.atan2(y, torch.sqrt(x * x + z * z)))
     yaw_diff = _wrap_degrees(azimuth - poses[:, None, 4])
     pitch_diff = _wrap_degrees(elevation - poses[:, None, 3])
-    return (yaw_diff < fov_half_h) & (pitch_diff < fov_half_v)
+    return (yaw_diff < _FOV_HALF_H) & (pitch_diff < _FOV_HALF_V)
 
 
-def _target_fov_points(target_poses: torch.Tensor, cfg=None) -> torch.Tensor:
-    """Sample a WorldMem-style sphere and keep points visible in the target FOV."""
+def _target_fov_points(target_poses: torch.Tensor) -> torch.Tensor:
+    """Sample a fixed-radius sphere and keep points visible in the target FOV."""
 
     target_poses = _as_pose_tensor(target_poses)
     if target_poses.shape[0] == 0:
         return torch.zeros((0, 3), device=target_poses.device, dtype=target_poses.dtype)
 
-    points = _sample_points_in_sphere(target_poses[0], cfg)
-    target_inside = _inside_fov_points(points, target_poses, cfg).any(dim=0)
+    points = _sample_points_in_sphere(target_poses[0])
+    target_inside = _inside_fov_points(points, target_poses).any(dim=0)
     return points[target_inside]
 
 
@@ -122,12 +119,11 @@ def _pose_preselect(candidates: np.ndarray, poses_t: torch.Tensor, target_positi
 
     candidates_t = torch.as_tensor(candidates, device=poses_t.device, dtype=torch.long)
     targets_t = torch.as_tensor(target_positions, device=poses_t.device, dtype=torch.long)
-    radius = float(cfg_get(cfg, "fov_radius", cfg_get(cfg, "fov_spatial_radius", 30.0)))
     candidate_poses = poses_t.index_select(0, candidates_t)
     target_poses = poses_t.index_select(0, targets_t)
     candidate_forward = _pose_directions(candidate_poses)
     target_forward = _pose_directions(target_poses)
-    translation_norm = torch.linalg.vector_norm(candidate_poses[:, None, :3] - target_poses[None, :, :3], dim=-1) / max(radius, 1e-6)
+    translation_norm = torch.linalg.vector_norm(candidate_poses[:, None, :3] - target_poses[None, :, :3], dim=-1) / _FOV_RADIUS
     dot = (candidate_forward[:, None, :] * target_forward[None, :, :]).sum(dim=-1).clamp(-1.0, 1.0)
     angular = torch.acos(dot) / torch.pi
     pose_distance = (translation_norm + angular).min(dim=1).values
@@ -138,7 +134,7 @@ def _pose_preselect(candidates: np.ndarray, poses_t: torch.Tensor, target_positi
 
 def _candidate_fov_masks(poses_t: torch.Tensor, candidates: np.ndarray, target_positions: np.ndarray, cfg=None) -> Tuple[torch.Tensor, torch.Tensor]:
     targets_t = torch.as_tensor(target_positions, device=poses_t.device, dtype=torch.long)
-    points = _target_fov_points(poses_t.index_select(0, targets_t), cfg)
+    points = _target_fov_points(poses_t.index_select(0, targets_t))
     if len(candidates) == 0 or points.shape[0] == 0:
         return (
             torch.zeros((len(candidates), 0), device=poses_t.device, dtype=torch.bool),
@@ -152,7 +148,7 @@ def _candidate_fov_masks(poses_t: torch.Tensor, candidates: np.ndarray, target_p
     score_parts = []
     for start in range(0, len(candidates), chunk_size):
         chunk = candidates_t[start : start + chunk_size]
-        inside = _inside_fov_points(points, poses_t.index_select(0, chunk), cfg)
+        inside = _inside_fov_points(points, poses_t.index_select(0, chunk))
         inside_parts.append(inside)
         score_parts.append(inside.float().mean(dim=1))
     return torch.cat(inside_parts, dim=0), torch.cat(score_parts, dim=0)
@@ -223,8 +219,8 @@ def _select_by_pose_similarity(
     candidate_poses = poses_t.index_select(0, candidates_t)
     target_poses = poses_t.index_select(0, targets_t)
 
-    # Training follows WorldMem's cheap pose-similarity sampling, not the 10k-point FOV selector.
-    spatial_scale = max(float(cfg_get(cfg, "pose_similarity_radius", cfg_get(cfg, "fov_radius", 30.0))), 1e-6)
+    # Training uses cheap pose-similarity sampling, not the 10k-point FOV selector.
+    spatial_scale = max(float(cfg_get(cfg, "pose_similarity_radius", _FOV_RADIUS)), 1e-6)
     angle_scale = max(float(cfg_get(cfg, "pose_similarity_angle_scale", 180.0)), 1e-6)
     spatial_delta = torch.linalg.vector_norm(candidate_poses[:, None, :3] - target_poses[None, :, :3], dim=-1)
     angle_delta = _wrap_degrees(candidate_poses[:, None, 3:] - target_poses[None, :, 3:]).mean(dim=-1)

tests/test_dememwm_latent_dataset.py

@@ -26,8 +26,6 @@ def _selection_cfg(**overrides):
         "local_context_exclusion_frames": 2,
         "plucker_moment_radius": 30.0,
         "anchor_diverse_selection": False,
-        "fov_num_points": 512,
-        "fov_sample_seed": 0,
         "pose_preselect_topk": 64,
         "candidate_chunk_size": 64,
     }