Implement DeMemWM segment split preprocessing

algorithms/dememwm/df_video.py

@@ -26,6 +26,7 @@ import glob
 
 # Utility Functions
 _DEMEMWM_SEGMENT_KEYS = ("target", "anchor", "dynamic", "revisit")
+_DEMEMWM_STREAM_KEYS = ("anchor", "dynamic", "revisit")
 
 
 def _segment_value_to_int(value, key):
@@ -52,6 +53,9 @@ def _preprocess_dememwm_latent_batch(batch):
         key: _segment_value_to_int(batch["memory_segments"][key], key)
         for key in _DEMEMWM_SEGMENT_KEYS
     }
+    segment_lengths = dict(memory_segments)
+    target_length = segment_lengths["target"]
+    stream_lengths = {key: segment_lengths[key] for key in _DEMEMWM_STREAM_KEYS}
     memory_masks = {key: batch["memory_masks"][key] for key in _DEMEMWM_SEGMENT_KEYS}
 
     # Latent dataset batches are already VAE-encoded: B x T_all x C x H_lat x W_lat.
@@ -60,6 +64,33 @@ def _preprocess_dememwm_latent_batch(batch):
     poses = rearrange(batch["poses"], "b t d -> t b d").contiguous()
     frame_indices = rearrange(batch["frame_indices"], "b t -> t b").contiguous()
 
+    segment_slices = {}
+    start = 0
+    for key in _DEMEMWM_SEGMENT_KEYS:
+        stop = start + segment_lengths[key]
+        segment_slices[key] = slice(start, stop)
+        start = stop
+    target_slice = segment_slices["target"]
+    stream_slices = {key: segment_slices[key] for key in _DEMEMWM_STREAM_KEYS}
+    if start != latents.shape[0]:
+        raise ValueError(
+            f"memory_segments sum to {start} frames, but latent batch has {latents.shape[0]}"
+        )
+
+    sequence_tensors = {
+        "latents": latents,
+        "actions": actions,
+        "poses": poses,
+        "frame_indices": frame_indices,
+    }
+    # Packed sequence order stays [target][anchor][dynamic][revisit] in T x B layout.
+    segments = {
+        key: {name: tensor[segment_slices[key]] for name, tensor in sequence_tensors.items()}
+        for key in _DEMEMWM_SEGMENT_KEYS
+    }
+    target_tensors = segments["target"]
+    stream_tensors = {key: segments[key] for key in _DEMEMWM_STREAM_KEYS}
+
     # Keep original image H/W for later ray geometry; latent H/W is not a substitute.
     return {
         "latents": latents,
@@ -67,6 +98,15 @@ def _preprocess_dememwm_latent_batch(batch):
         "poses": poses,
         "frame_indices": frame_indices,
         "memory_segments": memory_segments,
+        "segment_lengths": segment_lengths,
+        "target_length": target_length,
+        "stream_lengths": stream_lengths,
+        "segment_slices": segment_slices,
+        "target_slice": target_slice,
+        "stream_slices": stream_slices,
+        "segments": segments,
+        "target_tensors": target_tensors,
+        "stream_tensors": stream_tensors,
         "memory_masks": memory_masks,
         "image_hw": batch["image_hw"],
     }

tests/test_dememwm_latent_dataset.py

@@ -3,6 +3,7 @@ import unittest
 from pathlib import Path
 
 import numpy as np
+import torch
 from omegaconf import OmegaConf
 
 from datasets.video.memory_selection import select_memory_indices
@@ -138,6 +139,48 @@ class MemorySelectionTests(unittest.TestCase):
 
 
 class DeMemWMLatentDatasetTests(unittest.TestCase):
+
+    def test_preprocess_splits_sequence_tensors_from_memory_segments(self):
+        from algorithms.dememwm.df_video import _preprocess_dememwm_latent_batch
+
+        batch = {
+            "latents": torch.arange(6, dtype=torch.float32).view(1, 6, 1, 1, 1),
+            "actions": (100 + torch.arange(6, dtype=torch.float32)).view(1, 6, 1),
+            "poses": (200 + torch.arange(6, dtype=torch.float32)).view(1, 6, 1).repeat(1, 1, 5),
+            "frame_indices": (10 + torch.arange(6, dtype=torch.long)).view(1, 6),
+            "memory_segments": {
+                "target": torch.tensor([2]),
+                "anchor": torch.tensor([1]),
+                "dynamic": torch.tensor([3]),
+                "revisit": torch.tensor([0]),
+            },
+            "memory_masks": {
+                "target": torch.tensor([[True, False]]),
+                "anchor": torch.tensor([[True]]),
+                "dynamic": torch.tensor([[True, False, True]]),
+                "revisit": torch.zeros((1, 0), dtype=torch.bool),
+            },
+            "image_hw": torch.tensor([[360, 640]], dtype=torch.long),
+        }
+
+        preprocessed = _preprocess_dememwm_latent_batch(batch)
+
+        self.assertEqual(preprocessed["target_length"], 2)
+        self.assertEqual(preprocessed["stream_lengths"], {"anchor": 1, "dynamic": 3, "revisit": 0})
+        self.assertEqual((preprocessed["target_slice"].start, preprocessed["target_slice"].stop), (0, 2))
+        self.assertEqual(
+            {key: (slc.start, slc.stop) for key, slc in preprocessed["stream_slices"].items()},
+            {"anchor": (2, 3), "dynamic": (3, 6), "revisit": (6, 6)},
+        )
+        self.assertEqual(preprocessed["target_tensors"]["latents"][:, 0, 0, 0, 0].tolist(), [0.0, 1.0])
+        self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["latents"][:, 0, 0, 0, 0].tolist(), [3.0, 4.0, 5.0])
+        self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["actions"][:, 0, 0].tolist(), [103.0, 104.0, 105.0])
+        self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["poses"][:, 0, 0].tolist(), [203.0, 204.0, 205.0])
+        self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["frame_indices"][:, 0].tolist(), [13, 14, 15])
+        self.assertIs(preprocessed["memory_masks"]["dynamic"], batch["memory_masks"]["dynamic"])
+        self.assertEqual(tuple(preprocessed["memory_masks"]["dynamic"].shape), (1, 3))
+        self.assertEqual(preprocessed["memory_masks"]["dynamic"].device.type, "cpu")
+
     def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)
@@ -191,7 +234,23 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual(tuple(preprocessed["poses"].shape), (9, 1, 5))
         self.assertEqual(tuple(preprocessed["frame_indices"].shape), (9, 1))
         self.assertEqual(preprocessed["memory_segments"], sample["memory_segments"])
+        self.assertEqual(preprocessed["segment_lengths"], sample["memory_segments"])
+        self.assertEqual(preprocessed["target_length"], 3)
+        self.assertEqual(preprocessed["stream_lengths"], {"anchor": 2, "dynamic": 2, "revisit": 2})
         self.assertTrue(all(isinstance(value, int) for value in preprocessed["memory_segments"].values()))
+        self.assertEqual(
+            {key: (slc.start, slc.stop) for key, slc in preprocessed["segment_slices"].items()},
+            {"target": (0, 3), "anchor": (3, 5), "dynamic": (5, 7), "revisit": (7, 9)},
+        )
+        self.assertEqual((preprocessed["target_slice"].start, preprocessed["target_slice"].stop), (0, 3))
+        self.assertEqual(
+            {key: (slc.start, slc.stop) for key, slc in preprocessed["stream_slices"].items()},
+            {"anchor": (3, 5), "dynamic": (5, 7), "revisit": (7, 9)},
+        )
+        self.assertEqual(preprocessed["target_tensors"]["frame_indices"][:, 0].tolist(), [106, 107, 108])
+        self.assertEqual(preprocessed["stream_tensors"]["anchor"]["frame_indices"][:, 0].tolist(), [100, 101])
+        self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["frame_indices"][:, 0].tolist(), [104, 105])
+        self.assertEqual(preprocessed["stream_tensors"]["revisit"]["frame_indices"].shape[0], 2)
         self.assertEqual(tuple(preprocessed["memory_masks"]["target"].shape), (1, 3))
         self.assertEqual(tuple(preprocessed["memory_masks"]["anchor"].shape), (1, 2))
         self.assertEqual(tuple(preprocessed["memory_masks"]["dynamic"].shape), (1, 2))