Remove stale DeMemWM raw-video helpers

.exp_artifact/dememwm_remaining_fix_plan.md

@@ -119,7 +119,7 @@ Fix DeMemWM nested split filtering
 
 ## Substep 2: Remove Stale Raw-Video Helper Path From DeMemWM Algorithm
 
-Status: `[ ]`
+Status: `[x]`
 
 Bug:
 

algorithms/dememwm/df_video.py

@@ -1,7 +1,6 @@
 from collections.abc import Mapping
 import os
 import random
-import math
 import numpy as np
 import torch
 import torch.distributed as dist
@@ -9,7 +8,6 @@ import torch.nn.functional as F
 import torchvision.transforms.functional as TF
 from torchvision.transforms import InterpolationMode
 from PIL import Image
-from packaging import version as pver
 from einops import rearrange
 from tqdm import tqdm
 from omegaconf import DictConfig, open_dict
@@ -317,247 +315,6 @@ def _apply_memory_route_masks(frame_memory_masks, query_noise_levels, cfg, diffu
     return routed
 
 
-def euler_to_rotation_matrix(pitch, yaw):
-    """
-    Convert pitch and yaw angles (in radians) to a 3x3 rotation matrix.
-    Supports batch input.
-
-    Args:
-        pitch (torch.Tensor): Pitch angles in radians.
-        yaw (torch.Tensor): Yaw angles in radians.
-
-    Returns:
-        torch.Tensor: Rotation matrix of shape (batch_size, 3, 3).
-    """
-    cos_pitch, sin_pitch = torch.cos(pitch), torch.sin(pitch)
-    cos_yaw, sin_yaw = torch.cos(yaw), torch.sin(yaw)
-
-    R_pitch = torch.stack([
-        torch.ones_like(pitch), torch.zeros_like(pitch), torch.zeros_like(pitch),
-        torch.zeros_like(pitch), cos_pitch, -sin_pitch,
-        torch.zeros_like(pitch), sin_pitch, cos_pitch
-    ], dim=-1).reshape(-1, 3, 3)
-
-    R_yaw = torch.stack([
-        cos_yaw, torch.zeros_like(yaw), sin_yaw,
-        torch.zeros_like(yaw), torch.ones_like(yaw), torch.zeros_like(yaw),
-        -sin_yaw, torch.zeros_like(yaw), cos_yaw
-    ], dim=-1).reshape(-1, 3, 3)
-
-    return torch.matmul(R_yaw, R_pitch)
-
-
-def euler_to_camera_to_world_matrix(pose):
-    """
-    Convert (x, y, z, pitch, yaw) to a 4x4 camera-to-world transformation matrix using torch.
-    Supports both (5,) and (f, b, 5) shaped inputs.
-
-    Args:
-        pose (torch.Tensor): Pose tensor of shape (5,) or (f, b, 5).
-
-    Returns:
-        torch.Tensor: Camera-to-world transformation matrix of shape (4, 4).
-    """
-
-    origin_dim = pose.ndim
-    if origin_dim == 1:
-        pose = pose.unsqueeze(0).unsqueeze(0)  # Convert (5,) -> (1, 1, 5)
-    elif origin_dim == 2:
-        pose = pose.unsqueeze(0)
-
-    x, y, z, pitch, yaw = pose[..., 0], pose[..., 1], pose[..., 2], pose[..., 3], pose[..., 4]
-    pitch, yaw = torch.deg2rad(pitch), torch.deg2rad(yaw)
-
-    # Compute rotation matrix (batch mode)
-    R = euler_to_rotation_matrix(pitch, yaw)  # Shape (f*b, 3, 3)
-
-    # Create the 4x4 transformation matrix
-    eye = torch.eye(4, dtype=torch.float32, device=pose.device)
-    camera_to_world = eye.repeat(R.shape[0], 1, 1)  # Shape (f*b, 4, 4)
-
-    # Assign rotation
-    camera_to_world[:, :3, :3] = R
-
-    # Assign translation
-    camera_to_world[:, :3, 3] = torch.stack([x.reshape(-1), y.reshape(-1), z.reshape(-1)], dim=-1)
-
-    # Reshape back to (f, b, 4, 4) if needed
-    if origin_dim == 3:
-        return camera_to_world.view(pose.shape[0], pose.shape[1], 4, 4)
-    elif origin_dim == 2:
-        return camera_to_world.view(pose.shape[0], 4, 4)
-    else:
-        return camera_to_world.squeeze(0).squeeze(0)  # Convert (1,1,4,4) -> (4,4)
-
-def is_inside_fov_3d_hv(points, center, center_pitch, center_yaw, fov_half_h, fov_half_v):
-    """
-    Check whether points are within a given 3D field of view (FOV) 
-    with separately defined horizontal and vertical ranges.
-
-    The center view direction is specified by pitch and yaw (in degrees).
-
-    :param points: (N, B, 3) Sample point coordinates
-    :param center: (3,) Center coordinates of the FOV
-    :param center_pitch: Pitch angle of the center view (in degrees)
-    :param center_yaw: Yaw angle of the center view (in degrees)
-    :param fov_half_h: Horizontal half-FOV angle (in degrees)
-    :param fov_half_v: Vertical half-FOV angle (in degrees)
-    :return: Boolean tensor (N, B), indicating whether each point is inside the FOV
-    """
-    # Compute vectors relative to the center
-    vectors = points - center  # shape (N, B, 3)
-    x = vectors[..., 0]
-    y = vectors[..., 1]
-    z = vectors[..., 2]
-    
-    # Compute horizontal angle (yaw): measured with respect to the z-axis as the forward direction,
-    # and the x-axis as left-right, resulting in a range of -180 to 180 degrees.
-    azimuth = torch.atan2(x, z) * (180 / math.pi)
-    
-    # Compute vertical angle (pitch): measured with respect to the horizontal plane,
-    # resulting in a range of -90 to 90 degrees.
-    elevation = torch.atan2(y, torch.sqrt(x**2 + z**2)) * (180 / math.pi)
-    
-    # Compute the angular difference from the center view (handling circular angle wrap-around)
-    diff_azimuth = (azimuth - center_yaw).abs() % 360
-    diff_elevation = (elevation - center_pitch).abs() % 360
-    
-    # Adjust values greater than 180 degrees to the shorter angular difference
-    diff_azimuth = torch.where(diff_azimuth > 180, 360 - diff_azimuth, diff_azimuth)
-    diff_elevation = torch.where(diff_elevation > 180, 360 - diff_elevation, diff_elevation)
-    
-    # Check if both horizontal and vertical angles are within their respective FOV limits
-    return (diff_azimuth < fov_half_h) & (diff_elevation < fov_half_v)
-    
-def generate_points_in_sphere(n_points, radius):
-    # Sample three independent uniform distributions
-    samples_r = torch.rand(n_points)       # For radius distribution
-    samples_phi = torch.rand(n_points)     # For azimuthal angle phi
-    samples_u = torch.rand(n_points)       # For polar angle theta
-
-    # Apply cube root to ensure uniform volumetric distribution
-    r = radius * torch.pow(samples_r, 1/3)
-    # Azimuthal angle phi uniformly distributed in [0, 2π]
-    phi = 2 * math.pi * samples_phi
-    # Convert u to theta to ensure cos(theta) is uniformly distributed
-    theta = torch.acos(1 - 2 * samples_u)
-
-    # Convert spherical coordinates to Cartesian coordinates
-    x = r * torch.sin(theta) * torch.cos(phi)
-    y = r * torch.sin(theta) * torch.sin(phi)
-    z = r * torch.cos(theta)
-
-    points = torch.stack((x, y, z), dim=1)
-    return points
-
-def tensor_max_with_number(tensor, number):
-    number_tensor = torch.tensor(number, dtype=tensor.dtype, device=tensor.device)
-    result = torch.max(tensor, number_tensor)
-    return result
-
-def custom_meshgrid(*args):
-    # ref: https://pytorch.org/docs/stable/generated/torch.meshgrid.html?highlight=meshgrid#torch.meshgrid
-    if pver.parse(torch.__version__) < pver.parse('1.10'):
-        return torch.meshgrid(*args)
-    else:
-        return torch.meshgrid(*args, indexing='ij')
-    
-def camera_to_world_to_world_to_camera(camera_to_world: torch.Tensor) -> torch.Tensor:
-    """
-    Convert Camera-to-World matrices to World-to-Camera matrices for a tensor with shape (f, b, 4, 4).
-
-    Args:
-        camera_to_world (torch.Tensor): A tensor of shape (f, b, 4, 4), where:
-            f = number of frames,
-            b = batch size.
-
-    Returns:
-        torch.Tensor: A tensor of shape (f, b, 4, 4) representing the World-to-Camera matrices.
-    """
-    # Ensure input is a 4D tensor
-    assert camera_to_world.ndim == 4 and camera_to_world.shape[2:] == (4, 4), \
-        "Input must be of shape (f, b, 4, 4)"
-    
-    # Extract the rotation (R) and translation (T) parts
-    R = camera_to_world[:, :, :3, :3]  # Shape: (f, b, 3, 3)
-    T = camera_to_world[:, :, :3, 3]   # Shape: (f, b, 3)
-    
-    # Initialize an identity matrix for the output
-    world_to_camera = torch.eye(4, device=camera_to_world.device).unsqueeze(0).unsqueeze(0)
-    world_to_camera = world_to_camera.repeat(camera_to_world.size(0), camera_to_world.size(1), 1, 1)  # Shape: (f, b, 4, 4)
-    
-    # Compute the rotation (transpose of R)
-    world_to_camera[:, :, :3, :3] = R.transpose(2, 3)
-    
-    # Compute the translation (-R^T * T)
-    world_to_camera[:, :, :3, 3] = -torch.matmul(R.transpose(2, 3), T.unsqueeze(-1)).squeeze(-1)
-    
-    return world_to_camera.to(camera_to_world.dtype)
-
-def convert_to_plucker(poses, curr_frame, focal_length, image_width, image_height):
-
-    intrinsic = np.asarray([focal_length * image_width,
-                                focal_length * image_height,
-                                0.5 * image_width,
-                                0.5 * image_height], dtype=np.float32)
-
-    c2ws = get_relative_pose(poses, zero_first_frame_scale=curr_frame)
-    c2ws = rearrange(c2ws, "t b m n -> b t m n")
-
-    K = torch.as_tensor(intrinsic, device=poses.device, dtype=poses.dtype).repeat(c2ws.shape[0],c2ws.shape[1],1)  # [B, F, 4]
-    plucker_embedding = ray_condition(K, c2ws, image_height, image_width, device=c2ws.device)
-    plucker_embedding = rearrange(plucker_embedding, "b t h w d -> t b h w d").contiguous()
-
-    return plucker_embedding
-
-
-def get_relative_pose(abs_c2ws, zero_first_frame_scale):
-    abs_w2cs = camera_to_world_to_world_to_camera(abs_c2ws)
-    target_cam_c2w = torch.tensor([
-        [1, 0, 0, 0],
-        [0, 1, 0, 0],
-        [0, 0, 1, 0],
-        [0, 0, 0, 1]
-    ]).to(abs_c2ws.device).to(abs_c2ws.dtype)
-    abs2rel = target_cam_c2w @ abs_w2cs[zero_first_frame_scale]
-    ret_poses = [abs2rel @ abs_c2w for abs_c2w in abs_c2ws]
-    ret_poses = torch.stack(ret_poses)
-    return ret_poses
-
-def ray_condition(K, c2w, H, W, device):
-    # c2w: B, V, 4, 4
-    # K: B, V, 4
-
-    B = K.shape[0]
-
-    j, i = custom_meshgrid(
-        torch.linspace(0, H - 1, H, device=device, dtype=c2w.dtype),
-        torch.linspace(0, W - 1, W, device=device, dtype=c2w.dtype),
-    )
-    i = i.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5  # [B, HxW]
-    j = j.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5  # [B, HxW]
-
-    fx, fy, cx, cy = K.chunk(4, dim=-1)  # B,V, 1
-
-    zs = torch.ones_like(i, device=device, dtype=c2w.dtype)  # [B, HxW]
-    xs = -(i - cx) / fx * zs
-    ys = -(j - cy) / fy * zs 
-
-    zs = zs.expand_as(ys)
-
-    directions = torch.stack((xs, ys, zs), dim=-1)  # B, V, HW, 3
-    directions = directions / directions.norm(dim=-1, keepdim=True)  # B, V, HW, 3
-
-    rays_d = directions @ c2w[..., :3, :3].transpose(-1, -2)  # B, V, 3, HW
-    rays_o = c2w[..., :3, 3]  # B, V, 3
-    rays_o = rays_o[:, :, None].expand_as(rays_d)  # B, V, 3, HW
-    # c2w @ dirctions
-    rays_dxo = torch.linalg.cross(rays_o, rays_d)
-    plucker = torch.cat([rays_dxo, rays_d], dim=-1)
-    plucker = plucker.reshape(B, c2w.shape[1], H, W, 6)  # B, V, H, W, 6
-
-    return plucker
-
 def random_transform(tensor):
     """
     Apply the same random translation, rotation, and scaling to all frames in the batch.
@@ -891,35 +648,12 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         self.validation_step_outputs.clear()
 
     def _preprocess_batch(self, batch):
-        if isinstance(batch, Mapping):
-            return _preprocess_dememwm_latent_batch(batch)
-
-        xs, conditions, pose_conditions, frame_index = batch
-
-        if self.action_cond_dim:
-            conditions = torch.cat([torch.zeros_like(conditions[:, :1]), conditions[:, 1:]], 1)
-            conditions = rearrange(conditions, "b t d -> t b d").contiguous()
-        else:
-            raise NotImplementedError("Only support external cond.")
-
-        pose_conditions = rearrange(pose_conditions, "b t d -> t b d").contiguous()
-        c2w_mat = euler_to_camera_to_world_matrix(pose_conditions)
-        xs = rearrange(xs, "b t c ... -> t b c ...").contiguous()
-        frame_index = rearrange(frame_index, "b t -> t b").contiguous()
-
-        return xs, conditions, pose_conditions, c2w_mat, frame_index
-    
-    def encode(self, x):
-        # vae encoding
-        T = x.shape[0]
-        H, W = x.shape[-2:]
-        scaling_factor = 0.07843137255
-
-        x = rearrange(x, "t b c h w -> (t b) c h w")
-        with torch.no_grad():
-            x = self.vae.encode(x * 2 - 1).mean * scaling_factor
-        x = rearrange(x, "(t b) (h w) c -> t b c h w", t=T, h=H // self.vae.patch_size, w=W // self.vae.patch_size)
-        return x
+        if not isinstance(batch, Mapping):
+            raise TypeError(
+                "DeMemWM requires the latent dict batch contract "
+                "from video_minecraft_dememwm_latent; raw tuple batches are unsupported."
+            )
+        return _preprocess_dememwm_latent_batch(batch)
 
     def decode(self, x):
         total_frames = x.shape[0]
@@ -930,116 +664,6 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         x = rearrange(x, "(t b) c h w-> t b c h w", t=total_frames)
         return x
 
-    def _generate_condition_indices(self, curr_frame, memory_condition_length, xs_pred, pose_conditions, frame_idx, horizon):
-        """
-        Generate indices for condition similarity based on the current frame and pose conditions.
-        """
-        if curr_frame < memory_condition_length:
-            random_idx = [i for i in range(curr_frame)] + [0] * (memory_condition_length - curr_frame)
-            random_idx = np.repeat(np.array(random_idx)[:, None], xs_pred.shape[1], -1)
-        else:
-            # Generate points in a sphere and filter based on field of view
-            num_samples = 10000
-            radius = 30
-            points = generate_points_in_sphere(num_samples, radius).to(pose_conditions.device)
-            points = points[:, None].repeat(1, pose_conditions.shape[1], 1)
-            points += pose_conditions[curr_frame, :, :3][None]
-            fov_half_h = torch.tensor(105 / 2, device=pose_conditions.device)
-            fov_half_v = torch.tensor(75 / 2, device=pose_conditions.device)
-
-            # in_fov1 = is_inside_fov_3d_hv(
-            #     points, pose_conditions[curr_frame, :, :3],
-            #     pose_conditions[curr_frame, :, -2], pose_conditions[curr_frame, :, -1],
-            #     fov_half_h, fov_half_v
-            # )
-
-            in_fov1 = torch.stack([
-                is_inside_fov_3d_hv(points, pc[:, :3], pc[:, -2], pc[:, -1], fov_half_h, fov_half_v)
-                for pc in pose_conditions[curr_frame:curr_frame+horizon]
-            ])
-
-            in_fov1 = torch.sum(in_fov1, 0) > 0
-
-            # Compute overlap ratios and select indices
-            in_fov_list = torch.stack([
-                is_inside_fov_3d_hv(points, pc[:, :3], pc[:, -2], pc[:, -1], fov_half_h, fov_half_v)
-                for pc in pose_conditions[:curr_frame]
-            ])
-
-            random_idx = []
-            for _ in range(memory_condition_length):
-                overlap_ratio = ((in_fov1.bool() & in_fov_list).sum(1)) / in_fov1.sum()
-                
-                confidence = overlap_ratio + (curr_frame - frame_idx[:curr_frame]) / curr_frame * (-0.2)
-
-                if len(random_idx) > 0:
-                    confidence[torch.cat(random_idx)] = -1e10
-                _, r_idx = torch.topk(confidence, k=1, dim=0)
-                random_idx.append(r_idx[0])
-
-                # choice 1: directly remove overlapping region
-                occupied_mask = in_fov_list[r_idx[0, range(in_fov1.shape[-1])], :, range(in_fov1.shape[-1])].permute(1,0)
-                in_fov1 = in_fov1 & ~occupied_mask
-
-                # choice 2: apply similarity filter 
-                # cos_sim = F.cosine_similarity(xs_pred.to(r_idx.device)[r_idx[:, range(in_fov1.shape[1])], 
-                #     range(in_fov1.shape[1])], xs_pred.to(r_idx.device)[:curr_frame], dim=2)
-                # cos_sim = cos_sim.mean((-2,-1))
-
-                # mask_sim = cos_sim>0.9
-                # in_fov_list = in_fov_list & ~mask_sim[:,None].to(in_fov_list.device)
-
-            random_idx = torch.stack(random_idx).cpu()
-
-        return random_idx
-
-    def _prepare_conditions(self, 
-                            start_frame, curr_frame, horizon, conditions, 
-                            pose_conditions, c2w_mat, frame_idx, random_idx,
-                            image_width, image_height):
-        """
-        Prepare input conditions and pose conditions for sampling.
-        """
-
-        padding = torch.zeros((len(random_idx),) + conditions.shape[1:], device=conditions.device, dtype=conditions.dtype)
-        input_condition = torch.cat([conditions[start_frame:curr_frame + horizon], padding], dim=0)
-
-        batch_size = conditions.shape[1]
-
-        if self.use_plucker:
-            if self.relative_embedding:
-                frame_idx_list = []
-                input_pose_condition = []
-                for i in range(start_frame, curr_frame + horizon):
-                    input_pose_condition.append(convert_to_plucker(torch.cat([c2w_mat[i:i+1],c2w_mat[random_idx[:,range(batch_size)], range(batch_size)]]).clone(), 0, focal_length=self.focal_length,
-                                                image_width=image_width, image_height=image_height).to(conditions.dtype))
-                    frame_idx_list.append(torch.cat([frame_idx[i:i+1]-frame_idx[i:i+1], frame_idx[random_idx[:,range(batch_size)], range(batch_size)]-frame_idx[i:i+1]]))
-                input_pose_condition = torch.cat(input_pose_condition)
-                frame_idx_list = torch.cat(frame_idx_list)
-
-            else:
-                input_pose_condition = torch.cat([c2w_mat[start_frame : curr_frame + horizon], c2w_mat[random_idx[:,range(batch_size)], range(batch_size)]], dim=0).clone()
-                input_pose_condition = convert_to_plucker(input_pose_condition, 0, focal_length=self.focal_length)
-                frame_idx_list = None
-        else:
-            input_pose_condition = torch.cat([pose_conditions[start_frame : curr_frame + horizon], pose_conditions[random_idx[:,range(batch_size)], range(batch_size)]], dim=0).clone()
-            frame_idx_list = None
-
-        return input_condition, input_pose_condition, frame_idx_list
-
-    def _prepare_noise_levels(self, scheduling_matrix, m, curr_frame, batch_size, memory_condition_length):
-        """
-        Prepare noise levels for the current sampling step.
-        """
-        from_noise_levels = np.concatenate((np.zeros((curr_frame,), dtype=np.int64), scheduling_matrix[m]))[:, None].repeat(batch_size, axis=1)
-        to_noise_levels = np.concatenate((np.zeros((curr_frame,), dtype=np.int64), scheduling_matrix[m + 1]))[:, None].repeat(batch_size, axis=1)
-        if memory_condition_length:
-            from_noise_levels = np.concatenate([from_noise_levels, np.zeros((memory_condition_length, from_noise_levels.shape[-1]), dtype=np.int32)], axis=0)
-            to_noise_levels = np.concatenate([to_noise_levels, np.zeros((memory_condition_length, from_noise_levels.shape[-1]), dtype=np.int32)], axis=0)
-        from_noise_levels = torch.from_numpy(from_noise_levels).to(self.device)
-        to_noise_levels = torch.from_numpy(to_noise_levels).to(self.device)
-        return from_noise_levels, to_noise_levels
-
     def validation_step(self, batch, batch_idx, namespace="validation") -> STEP_OUTPUT:
         """
         Perform a single validation step.

tests/test_dememwm_latent_dataset.py

@@ -572,9 +572,6 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
                 self.horizons.append(horizon)
                 return np.stack([np.full(horizon, level, dtype=np.int64) for level in (2, 1, 0)])
 
-            def _generate_condition_indices(self, *args, **kwargs):
-                raise AssertionError("latent dict validation must not use the legacy memory_condition_length path")
-
             def encode(self, xs):
                 raise AssertionError("latent dict validation must not VAE-encode inputs")