Composable Nueral Network Skelatan structure

manim_ml/decision_tree/decision_tree.py

@@ -0,0 +1,26 @@
+"""
+    Module for visualizing decision trees in Manim. 
+    It parses a decision tree classifier from sklearn. 
+"""
+from manim import *
+from manim_ml.one_to_one_sync import OneToOneSync
+
+class LeafNode(VGroup):
+    pass
+
+class SplitNode(VGroup):
+    pass
+
+class DecisionTreeDiagram(Graph):
+    """Decision Tree Digram Class for Manim"""
+    pass
+
+class DecisionTreeEmbedding():
+    """Embedding for the decision tree"""
+    pass
+
+class DecisionTreeContainer(OneToOneSync):
+    """Connects the DecisionTreeDiagram to the DecisionTreeEmbedding"""
+
+    def __init__(self):
+        pass

manim_ml/image.py

@@ -0,0 +1,23 @@
+from manim import *
+import numpy as np
+
+class GrayscaleImageMobject(ImageMobject):
+    """Mobject for creating images in Manim from numpy arrays"""
+
+    def __init__(self, numpy_image, height=2.3):
+        self.numpy_image = numpy_image
+
+        assert len(np.shape(self.numpy_image)) == 2 
+        input_image = self.numpy_image[None, :, :]
+        # Convert grayscale to rgb version of grayscale
+        input_image = np.repeat(input_image, 3, axis=0)
+        input_image = np.rollaxis(input_image, 0, start=3)
+
+        super().__init__(input_image, image_mode="RGB")
+
+        self.set_resampling_algorithm(RESAMPLING_ALGORITHMS["nearest"])
+        self.scale_to_fit_height(height)
+
+    @override_animation(Create)
+    def create(self, run_time=2):
+        return FadeIn(self)

manim_ml/manifold.py

@@ -0,0 +1,3 @@
+"""
+    Visaulization of a latent Manifold
+"""

manim_ml/neural_network/connective_layers.py

@@ -0,0 +1,101 @@
+"""
+    Layers that describe the connections between user layers.
+"""
+from manim import *
+from manim_ml.neural_network.layers import NeuralNetworkLayer
+from abc import ABC, abstractmethod
+
+class ConnectiveLayer(NeuralNetworkLayer):
+    """Forward pass animation for a given pair of layers"""
+
+    @abstractmethod
+    def __init__(self, input_layer, output_layer):
+        super(NeuralNetworkLayer, self).__init__()
+        self.input_layer = input_layer
+        self.output_layer = output_layer
+
+    @abstractmethod
+    def make_forward_pass_animation(self):
+        pass
+
+class FeedForwardToFeedForward(ConnectiveLayer):
+    """Layer for connecting FeedForward layer to FeedForwardLayer"""
+
+    def __init__(self, input_layer, output_layer, passing_flash=True,
+                dot_radius=0.05, animation_dot_color=RED, edge_color=WHITE,
+                edge_width=0.5):
+        super().__init__(input_layer, output_layer)
+        self.passing_flash = passing_flash
+        self.edge_color = edge_color
+        self.dot_radius = dot_radius
+        self.animation_dot_color = animation_dot_color
+        self.edge_width = edge_width
+
+        self.edges = self.construct_edges()
+        self.add(self.edges)
+
+    def construct_edges(self):
+        # Go through each node in the two layers and make a connecting line
+        edges = []
+        for node_i in self.input_layer.node_group:
+            for node_j in self.output_layer.node_group:
+                line = Line(node_i.get_center(), node_j.get_center(), 
+                            color=self.edge_color, stroke_width=self.edge_width)
+                edges.append(line)
+
+        edges = Group(*edges)
+        return edges
+
+    def make_forward_pass_animation(self, run_time=1):
+        """Animation for passing information from one FeedForwardLayer to the next"""
+        path_animations = []
+        dots = []
+        for edge in self.edges:
+            dot = Dot(color=self.animation_dot_color, fill_opacity=1.0, radius=self.dot_radius)
+            # Handle layering
+            dot.set_z_index(1)
+            # Add to dots group
+            dots.append(dot)
+            # Make the animation
+            if self.passing_flash:
+                print("passing flash")
+                anim = ShowPassingFlash(edge.copy().set_color(self.animation_dot_color), time_width=0.2, run_time=3)
+            else:
+                anim = MoveAlongPath(dot, edge, run_time=run_time, rate_function=sigmoid)
+            path_animations.append(anim)
+
+        if not self.passing_flash:
+            dots = Group(*dots)
+            self.add(dots)
+
+        path_animations = AnimationGroup(*path_animations)
+
+        return path_animations
+
+class ImageToFeedForward(ConnectiveLayer):
+    """Image Layer to FeedForward layer"""
+
+    def __init__(self, input_layer, output_layer, animation_dot_color=RED,
+                dot_radius=0.05):
+        self.animation_dot_color = animation_dot_color
+        self.dot_radius = dot_radius
+        # Input assumed to be ImageLayer
+        # Output assumed to be FeedForwardLayer
+        super().__init__(input_layer, output_layer)
+
+    def make_forward_pass_animation(self):
+        """Makes dots diverge from the given location and move to the feed forward nodes decoder"""
+        animations = []
+        image_mobject = self.input_layer.image_mobject
+        # Move the dots to the centers of each of the nodes in the FeedForwardLayer
+        image_location  = image_mobject.get_center()
+        for node in self.output_layer.node_group:
+            new_dot = Dot(image_location, radius=self.dot_radius, color=self.animation_dot_color)
+            per_node_succession = Succession(
+                Create(new_dot),
+                new_dot.animate.move_to(node.get_center()),
+            )
+            animations.append(per_node_succession)
+
+        animation_group = AnimationGroup(*animations)
+        return animation_group

manim_ml/neural_network/feed_forward.py

@@ -1,17 +0,0 @@
-from manim import *
-from manim_ml.neural_network.layers import FeedForwardLayer
-from manim_ml.neural_network.neural_network import NeuralNetwork
-
-class FeedForwardNeuralNetwork(NeuralNetwork):
-    """NeuralNetwork with just feed forward layers"""
-
-    def __init__(self, layer_node_count, node_radius=0.08, 
-                node_color=BLUE, **kwargs):
-
-        # construct layers
-        layers = []
-        for num_nodes in layer_node_count:
-            layer = FeedForwardLayer(num_nodes, node_color=node_color, node_radius=node_radius)
-            layers.append(layer)
-        # call super class
-        super().__init__(layers, **kwargs)

manim_ml/neural_network/layers.py

@@ -1,79 +1,15 @@
+from typing import overload
 from manim import *
 from abc import ABC, abstractmethod
+from manim_ml.image import GrayscaleImageMobject
 
-class NeuralNetworkLayer(ABC, VGroup):
+class NeuralNetworkLayer(ABC, Group):
     """Abstract Neural Network Layer class"""
 
     @abstractmethod
     def make_forward_pass_animation(self):
         pass
 
-class ConnectiveLayer(NeuralNetworkLayer):
-    """Forward pass animation for a given pair of layers"""
-
-    @abstractmethod
-    def __init__(self, input_layer, output_layer):
-        super(NeuralNetworkLayer, self).__init__()
-        self.input_layer = input_layer
-        self.output_layer = output_layer
-
-    @abstractmethod
-    def make_forward_pass_animation(self):
-        pass
-
-class FeedForwardToFeedForward(ConnectiveLayer):
-
-    def __init__(self, input_layer, output_layer, passing_flash=True,
-                dot_radius=0.05, animation_dot_color=RED, edge_color=WHITE,
-                edge_width=0.5):
-        super().__init__(input_layer, output_layer)
-        self.passing_flash = passing_flash
-        self.edge_color = edge_color
-        self.dot_radius = dot_radius
-        self.animation_dot_color = animation_dot_color
-        self.edge_width = edge_width
-
-        self.edges = self.construct_edges()
-        self.add(self.edges)
-
-    def construct_edges(self):
-        # Go through each node in the two layers and make a connecting line
-        edges = []
-        for node_i in self.input_layer.node_group:
-            for node_j in self.output_layer.node_group:
-                line = Line(node_i.get_center(), node_j.get_center(), 
-                            color=self.edge_color, stroke_width=self.edge_width)
-                edges.append(line)
-
-        edges = VGroup(*edges)
-        return edges
-
-    def make_forward_pass_animation(self, run_time=1):
-        """Animation for passing information from one FeedForwardLayer to the next"""
-        path_animations = []
-        dots = []
-        for edge in self.edges:
-            dot = Dot(color=self.animation_dot_color, fill_opacity=1.0, radius=self.dot_radius)
-            # Handle layering
-            dot.set_z_index(1)
-            # Add to dots group
-            dots.append(dot)
-            # Make the animation
-            if self.passing_flash:
-                print("passing flash")
-                anim = ShowPassingFlash(edge.copy().set_color(self.animation_dot_color), time_width=0.2, run_time=3)
-            else:
-                anim = MoveAlongPath(dot, edge, run_time=run_time, rate_function=sigmoid)
-            path_animations.append(anim)
-
-        if not self.passing_flash:
-            dots = VGroup(*dots)
-            self.add(dots)
-
-        path_animations = AnimationGroup(*path_animations)
-
-        return path_animations
-
 class FeedForwardLayer(NeuralNetworkLayer):
     """Handles rendering a layer for a neural network"""
 
@@ -94,7 +30,7 @@ class FeedForwardLayer(NeuralNetworkLayer):
         self.rectangle_fill_color = rectangle_fill_color
         self.animation_dot_color = animation_dot_color
 
-        self.node_group = VGroup()
+        self.node_group = Group()
 
         self._construct_neural_network_layer()
 
@@ -127,3 +63,24 @@ class FeedForwardLayer(NeuralNetworkLayer):
 
         return succession
  
+class ImageLayer(NeuralNetworkLayer):
+    """Image Layer for Neural Network"""
+
+    def __init__(self, numpy_image, height=1.5):
+        super().__init__()
+        self.numpy_image = numpy_image
+        if len(np.shape(self.numpy_image)) == 2:
+            # Assumed Grayscale
+            self.image_mobject = GrayscaleImageMobject(self.numpy_image, height=height)
+        elif len(np.shape(self.numpy_image)) == 3:
+            # Assumed RGB
+            self.image_mobject = ImageMobject(self.numpy_image)
+
+        self.add(self.image_mobject)
+
+    def make_forward_pass_animation(self):
+        return Create(self.image_mobject)
+
+    @property
+    def width(self):
+        return self.image_mobject.width

manim_ml/neural_network/neural_network.py

@@ -11,44 +11,50 @@ Example:
 """
 from manim import *
 from matplotlib import animation
-from manim_ml.neural_network.layers import FeedForwardToFeedForward, FeedForwardLayer
+from numpy import isin
+from manim_ml.neural_network.layers import FeedForwardLayer, ImageLayer
+from manim_ml.neural_network.connective_layers import FeedForwardToFeedForward, ImageToFeedForward
 
-class NeuralNetwork(VGroup):
+class NeuralNetwork(Group):
 
     def __init__(self, layers, edge_color=WHITE, layer_spacing=0.8,
-                    animation_dot_color=RED, edge_width=1.5, dot_radius=0.05):
+                    animation_dot_color=RED, edge_width=1.5, dot_radius=0.03):
         super().__init__()
-        self.layers = layers
+        self.layers = Group(*layers)
         self.edge_width = edge_width
         self.edge_color = edge_color
         self.layer_spacing = layer_spacing
         self.animation_dot_color = animation_dot_color
         self.dot_radius = dot_radius
-
         # TODO take layer_node_count [0, (1, 2), 0] 
         # and make it have explicit distinct subspaces
-        self.layers = self._construct_layers()
+        self._place_layers()
         self.connective_layers = self._construct_connective_layers()
+        # Center the whole diagram by default
+        self.all_group = Group(self.layers, self.connective_layers)
+        self.all_group.move_to(ORIGIN)
 
         self.add(self.connective_layers)
         self.add(self.layers)
 
-    def _construct_layers(self):
+    def _place_layers(self):
         """Creates the neural network"""
-        layers = VGroup()
-        # Create each layer
-        for layer_index, layer in enumerate(self.layers):
+        # TODO implement more sophisticated custom layouts
+        for layer_index in range(1, len(self.layers)):
+            previous_layer = self.layers[layer_index - 1]
+            current_layer = self.layers[layer_index]
             # Manage spacing
-            layer.move_to([self.layer_spacing * layer_index, 0, 0])
+            # Default: half each width times 2
+            spacing = config.frame_width * 0.05 + (previous_layer.width / 2 + current_layer.width / 2)
+            current_layer.move_to(previous_layer.get_center())
+            current_layer.shift(np.array([spacing, 0, 0]))
             # Add layer to VGroup
-            layers.add(layer)
         # Handle layering
-        layers.set_z_index(2)
-        return layers
+        self.layers.set_z_index(2)
 
     def _construct_connective_layers(self):
         """Draws connecting lines between layers"""
-        connective_layers = VGroup()
+        connective_layers = Group()
         for layer_index in range(len(self.layers) - 1):
             current_layer = self.layers[layer_index]
             next_layer = self.layers[layer_index + 1]
@@ -59,6 +65,10 @@ class NeuralNetwork(VGroup):
                                                     edge_width=self.edge_width)
 
                 connective_layers.add(edge_layer)
+            elif isinstance(current_layer, ImageLayer) \
+                and isinstance(next_layer, FeedForwardLayer):
+                image_to_feedforward = ImageToFeedForward(current_layer, next_layer, dot_radius=self.dot_radius)
+                connective_layers.add(image_to_feedforward)
             else:
                 raise Exception(f"Unimplemented connection for layer types: {type(current_layer)} and {type(next_layer)}")
 
@@ -66,7 +76,7 @@ class NeuralNetwork(VGroup):
         connective_layers.set_z_index(0)
         return connective_layers
 
-    def make_forward_pass_animation(self, run_time=2, passing_flash=True):
+    def make_forward_pass_animation(self, run_time=5, passing_flash=True):
         """Generates an animation for feed forward propogation"""
         all_animations = []
 
@@ -89,7 +99,7 @@ class NeuralNetwork(VGroup):
 class FeedForwardNeuralNetwork(NeuralNetwork):
     """NeuralNetwork with just feed forward layers"""
 
-    def __init__(self, layer_node_count, node_radius=1.0, 
+    def __init__(self, layer_node_count, node_radius=0.08, 
                 node_color=BLUE, **kwargs):
         # construct layers
         layers = []

manim_ml/one_to_one_sync.py

@@ -0,0 +1,10 @@
+"""
+    Module for handling syncing two animations one to one.
+    The goal here is to zip up two classes and their respective animations,
+    and create a joint class with the same animations that runs the animations
+    for both classes at the same time. This way we can connect two isomorphic
+    views of the same concept and visualize them at the same time. 
+"""
+
+class OneToOneSync():
+    pass

manim_ml/oracle_guidance.py

@@ -1,17 +0,0 @@
-"""
-    Here is a animated explanatory figure for the "ORACLE GUIDED IMAGE SYNTHESIS WITH RELATIVE
-QUERIES" paper. 
-"""
-from manim import *
-import neural_network
-
-class OracleGuidanceScene(Scene):
-    """
-        Oracle Guidance Scene object
-    """
-
-    def construct_neural_networks(self):
-        """Constructs the encoder and decoder networks"""
-
-    def construct(self):
-        pass

tests/test_neural_network.py

@@ -1,7 +1,8 @@
 from manim import *
-from manim_ml.neural_network.layers import FeedForwardLayer
-from manim_ml.neural_network.neural_network import NeuralNetwork
-from manim_ml.neural_network.feed_forward import FeedForwardNeuralNetwork
+from manim_ml.neural_network.layers import FeedForwardLayer, ImageLayer
+from manim_ml.neural_network.neural_network import NeuralNetwork, FeedForwardNeuralNetwork
+from PIL import Image
+import numpy as np
 
 config.pixel_height = 720
 config.pixel_width = 1280
@@ -20,7 +21,11 @@ class NeuralNetworkScene(Scene):
 
     def construct(self):
         # Make the Layer object
-        layers = [FeedForwardLayer(3), FeedForwardLayer(5), FeedForwardLayer(3)]
+        layers = [
+            FeedForwardLayer(3), 
+            FeedForwardLayer(5), 
+            FeedForwardLayer(3)
+        ]
         nn = NeuralNetwork(layers)
         nn.move_to(ORIGIN)
         # Make Animation
@@ -29,6 +34,26 @@ class NeuralNetworkScene(Scene):
 
         self.play(forward_propagation_animation)
 
+class ImageNeuralNetworkScene(Scene):
+
+    def construct(self):
+        image = Image.open('images/image.jpeg')
+        numpy_image = np.asarray(image)
+        # Make nn
+        layers = [
+            ImageLayer(numpy_image, height=1.0),
+            FeedForwardLayer(3), 
+            FeedForwardLayer(5),
+            FeedForwardLayer(3),
+            FeedForwardLayer(3)
+        ]
+        nn = NeuralNetwork(layers)
+        # Center the nn
+        nn.move_to(ORIGIN)
+        self.add(nn)
+        # Play animation
+        self.play(nn.make_forward_pass_animation())
+
 if __name__ == "__main__":
     """Render all scenes"""
     # Feed Forward Neural Network