Prepare DeMemWM pose and image metadata

algorithms/dememwm/df_video.py

@@ -48,6 +48,16 @@ def _segment_value_to_int(value, key):
     return int(value)
 
 
+def _normalize_dememwm_image_hw(image_hw, batch_size):
+    image_hw = image_hw if torch.is_tensor(image_hw) else torch.as_tensor(image_hw)
+    image_hw = image_hw.to(dtype=torch.long)
+    if image_hw.ndim == 1:
+        image_hw = image_hw.unsqueeze(0)
+    if image_hw.shape[0] == 1 and batch_size != 1:
+        image_hw = image_hw.expand(batch_size, -1)
+    return image_hw.contiguous()
+
+
 def _preprocess_dememwm_latent_batch(batch):
     memory_segments = {
         key: _segment_value_to_int(batch["memory_segments"][key], key)
@@ -61,8 +71,10 @@ def _preprocess_dememwm_latent_batch(batch):
     # Latent dataset batches are already VAE-encoded: B x T_all x C x H_lat x W_lat.
     latents = rearrange(batch["latents"], "b t c ... -> t b c ...").contiguous()
     actions = rearrange(batch["actions"], "b t d -> t b d").contiguous()
+    # Dataset poses are the frame-memory geometry metadata in packed T_all x B x 5 order.
     poses = rearrange(batch["poses"], "b t d -> t b d").contiguous()
     frame_indices = rearrange(batch["frame_indices"], "b t -> t b").contiguous()
+    image_hw = _normalize_dememwm_image_hw(batch["image_hw"], latents.shape[1])
 
     segment_slices = {}
     start = 0
@@ -104,6 +116,7 @@ def _preprocess_dememwm_latent_batch(batch):
         "actions": actions,
         "action_conditions": action_conditions,
         "poses": poses,
+        "frame_memory_pose": poses,
         "frame_indices": frame_indices,
         "memory_segments": memory_segments,
         "segment_lengths": segment_lengths,
@@ -116,7 +129,7 @@ def _preprocess_dememwm_latent_batch(batch):
         "target_tensors": target_tensors,
         "stream_tensors": stream_tensors,
         "memory_masks": memory_masks,
-        "image_hw": batch["image_hw"],
+        "image_hw": image_hw,
     }
 
 

tests/test_dememwm_latent_dataset.py

@@ -179,10 +179,44 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual(preprocessed["target_tensors"]["action_conditions"][:, 0, 0].tolist(), [0.0, 101.0])
         self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["action_conditions"][:, 0, 0].tolist(), [0.0, 0.0, 0.0])
         self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["poses"][:, 0, 0].tolist(), [203.0, 204.0, 205.0])
+        self.assertEqual(preprocessed["frame_memory_pose"][:, 0, 0].tolist(), [200.0, 201.0, 202.0, 203.0, 204.0, 205.0])
+        self.assertIs(preprocessed["frame_memory_pose"], preprocessed["poses"])
         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")
+        self.assertEqual(preprocessed["image_hw"].tolist(), [[360, 640]])
+        self.assertEqual(preprocessed["image_hw"].dtype, torch.long)
+
+    def test_preprocess_expands_shared_image_hw_metadata(self):
+        from algorithms.dememwm.df_video import _preprocess_dememwm_latent_batch
+
+        batch = {
+            "latents": torch.zeros((2, 1, 1, 1, 1), dtype=torch.float32),
+            "actions": torch.zeros((2, 1, 1), dtype=torch.float32),
+            "poses": torch.arange(10, dtype=torch.float32).view(2, 1, 5),
+            "frame_indices": torch.zeros((2, 1), dtype=torch.long),
+            "memory_segments": {
+                "target": torch.tensor([1, 1]),
+                "anchor": torch.tensor([0, 0]),
+                "dynamic": torch.tensor([0, 0]),
+                "revisit": torch.tensor([0, 0]),
+            },
+            "memory_masks": {
+                "target": torch.ones((2, 1), dtype=torch.bool),
+                "anchor": torch.zeros((2, 0), dtype=torch.bool),
+                "dynamic": torch.zeros((2, 0), dtype=torch.bool),
+                "revisit": torch.zeros((2, 0), dtype=torch.bool),
+            },
+            "image_hw": torch.tensor([720, 1280], dtype=torch.long),
+        }
+
+        preprocessed = _preprocess_dememwm_latent_batch(batch)
+
+        self.assertEqual(tuple(preprocessed["frame_memory_pose"].shape), (1, 2, 5))
+        self.assertEqual(preprocessed["frame_memory_pose"][:, :, 0].tolist(), [[0.0, 5.0]])
+        self.assertEqual(tuple(preprocessed["image_hw"].shape), (2, 2))
+        self.assertEqual(preprocessed["image_hw"].tolist(), [[720, 1280], [720, 1280]])
 
     def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):
         with tempfile.TemporaryDirectory() as tmp:
@@ -240,6 +274,8 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertTrue(torch.equal(preprocessed["action_conditions"][3:], torch.zeros_like(preprocessed["actions"][3:])))
         self.assertTrue(torch.equal(preprocessed["actions"][3:], torch.ones_like(preprocessed["actions"][3:])))
         self.assertEqual(tuple(preprocessed["poses"].shape), (9, 1, 5))
+        self.assertIs(preprocessed["frame_memory_pose"], preprocessed["poses"])
+        self.assertEqual(preprocessed["frame_memory_pose"][:7, 0, 0].tolist(), [0.0, 1.0, 2.0, -6.0, -5.0, -2.0, -1.0])
         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"])