add record() method to basicNCA model, create Animator class

Author: henrykrumb

ncalab/models/basicNCA.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple, List
 
 import numpy as np
 
@@ -336,6 +336,29 @@ def predict(self, image: torch.Tensor, steps: int = 100) -> Prediction:
             prediction = self.forward(x, steps=steps)
             return prediction
 
+    def record(self, image: torch.Tensor, steps: int = 100) -> List[Prediction]:
+        """
+        Record predictions for all time steps and return the resulting
+        sequence of predictions.
+
+        :param image [torch.Tensor]: Input image, BCWH.
+
+        :returns [List[Prediction]]: List of Prediction objects.
+        """
+        assert steps >= 1
+        assert image.shape[1] <= self.num_channels
+        self.eval()
+        sequence = []
+        with torch.no_grad():
+            x = image.clone()
+            x = pad_input(x, self, noise=self.pad_noise)
+            x = self.prepare_input(x)
+            for _ in range(steps):
+                prediction = self.forward(x, steps=1)
+                sequence.append(prediction)
+                x = prediction.output_image
+            return sequence
+
     def validate(
         self, image: torch.Tensor, label: torch.Tensor, steps: int
     ) -> Optional[Tuple[Dict[str, float], Prediction]]:

ncalab/models/growingNCA.py

@@ -74,40 +74,29 @@ def validate(
         """
         return None
 
+    def make_seed(self, width: int, height: int) -> torch.Tensor:
+        x = torch.zeros((1, self.num_channels, width, height)).to(self.device)
+        # set seed in center
+        x[:, 3:, width // 2, height // 2] = 1.0
+        return x
+
     def grow(
-        self, width: int, height: int, steps: int = 100, save_steps=False
-    ) -> np.ndarray | List[np.ndarray]:
+        self, seed: torch.Tensor, steps: int = 100
+    ) -> List[np.ndarray]:
         """
-        Run the growth process and return the resulting output image.
+        Run the growth process and return the resulting output sequence.
 
-        :param width [int]: Output image width.
-        :param height [int]: Output image height.
+        :param seed [torch.Tensor]: Seed image, can be generated through make_seed.
         :param steps [int]: Number of inference steps. Defaults to 100.
 
-        :returns [np.ndarray]: Image channels of the output image.
+        :returns [List[np.ndarray]]: Sequence of output images.
         """
         with torch.no_grad():
-            # TODO make use of autostepper, if available
             self.eval()
-            x = torch.zeros((1, self.num_channels, width, height)).to(self.device)
-            # set seed in center
-            x[:, 3:, width // 2, height // 2] = 1.0
-
-            if save_steps:
-                step_outs = []
-                for _ in range(steps):
-                    prediction = self.forward(x, steps=1)  # type: ignore[assignment]
-                    step_outs.append(
-                        np.clip(
-                            prediction.image_channels.squeeze(0).detach().cpu().numpy(),
-                            0,
-                            1,
-                        )
-                    )
-                    x = prediction.output_image
-                return step_outs
-            else:
-                prediction = self.forward(x, steps=steps)  # type: ignore[assignment]
-            out_np = prediction.image_channels.detach().cpu().numpy().squeeze(0)
-            out_np = np.clip(out_np, 0, 1)
-            return out_np
+            output = []
+            sequence = self.record(seed, steps=steps)
+            for prediction in sequence:
+                output.append(
+                    np.clip(prediction.image_channels_np.squeeze(0), 0, 1)
+                )
+            return output

ncalab/visualization/animation.py

@@ -2,33 +2,65 @@
 
 import matplotlib.pyplot as plt  # type: ignore[import-untyped]
 import matplotlib.animation as animation  # type: ignore[import-untyped]
+import numpy as np
+import torch
 
 
-class NCAAnimator:
-    def __init__(self, nca, x, steps=100):
-        """ """
+class Animator:
+    def __init__(
+        self,
+        nca,
+        seed: torch.Tensor,
+        steps=100,
+        interval=100,
+        repeat=True,
+        repeat_delay=3000,
+    ):
+        nca.eval()
+
         fig, ax = plt.subplots()
         fig.set_size_inches(2, 2)
-        im = ax.imshow(x[0, :3].permute(0, 2, 3, 1).cpu(), animated=True)
+
+        # first frame is input image
+        if nca.immutable_image_channels:
+            first_frame = seed[0, -nca.num_output_channels :]
+        else:
+            first_frame = seed[0, : nca.num_image_channels]
+        first_frame = first_frame.permute(1, 2, 0).detach().cpu().numpy()
+        im = ax.imshow(
+            first_frame,
+            animated=True,
+        )
+
+        predictions = nca.record(seed, steps)
+        images = []
+        for prediction in predictions:
+            if nca.immutable_image_channels:
+                output_image = prediction.output_channels_np[0]
+            else:
+                output_image = prediction.image_channels_np[0]
+            output_image = output_image.transpose(1, 2, 0)
+            output_image = np.clip(output_image, 0, 1)
+            images.append(output_image)
+
         ax.set_axis_off()
         plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
         plt.margins(0, 0)
         plt.tight_layout()
 
         def update(i):
-            nonlocal x
-            x = nca(x)  # --> BWHC
-            im.set_array(x)
+            nonlocal images
+            im.set_array(images[i])
             return (im,)
 
         self.animation_fig = animation.FuncAnimation(
             fig,
             update,
             frames=steps,
-            interval=100,
+            interval=interval,
             blit=True,
-            repeat=True,
-            repeat_delay=3000,
+            repeat=repeat,
+            repeat_delay=repeat_delay,
         )
 
     def save(self, path: str | Path):

tasks/growing_emoji/eval_growing_emoji.py

@@ -11,14 +11,13 @@
     get_compute_device,
     print_NCALab_banner,
     fix_random_seed,
+    Animator,
 )
 
 import click
 
 import torch
 
-import matplotlib.pyplot as plt  # type: ignore[import-untyped]
-import matplotlib.animation as animation  # type: ignore[import-untyped]
 
 TASK_PATH = Path(__file__).parent
 FIGURE_PATH = TASK_PATH / "figures"
@@ -53,33 +52,12 @@ def eval_growing_emoji(gpu: bool, gpu_index: int):
             weights_only=True,
         )
     )
-    nca.eval()
 
-    images = nca.grow(48, 48, steps=100, save_steps=True)
+    seed = nca.make_seed(48, 48)
+    animator = Animator(nca, seed)
 
-    fig, ax = plt.subplots()
-    fig.set_size_inches(2, 2)
-    im = ax.imshow(images[0].transpose(1, 2, 0), animated=True)
-    ax.set_axis_off()
-    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
-    plt.margins(0, 0)
-    plt.tight_layout()
-
-    def update(i):
-        im.set_array(images[i].transpose(1, 2, 0))
-        return (im,)
-
-    animation_fig = animation.FuncAnimation(
-        fig,
-        update,
-        frames=len(images),
-        interval=10,
-        blit=True,
-        repeat=True,
-        repeat_delay=3000,
-    )
     out_path = FIGURE_PATH / "growing_emoji.gif"
-    animation_fig.save(out_path)
+    animator.save(out_path)
     click.secho(f"Done. You'll find the generated GIF in {out_path} .")