diff --git a/manim_ml/gridded_rectangle.py b/manim_ml/gridded_rectangle.py new file mode 100644 index 0000000..1068f98 --- /dev/null +++ b/manim_ml/gridded_rectangle.py @@ -0,0 +1,197 @@ +from manim import * +import numpy as np + +class CornersRectangle(Rectangle): + """Rectangle with functionality for getting the corner coordinates""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.corners = VGroup( + *[Dot(corner_point) for corner_point in self.get_vertices()] + ) + self.corners.set_fill(opacity=0.0) + self.add(self.corners) + + def get_corners_dict(self): + """Returns a dictionary of the corners""" + return { + "top_left": self.corners[3], + "top_right": self.corners[0], + "bottom_left": self.corners[2], + "bottom_right": self.corners[1], + } + +class GriddedRectangle(VGroup): + """Rectangle object with grid lines""" + + def __init__(self, center, color=WHITE, height=2.0, width=4.0, + mark_paths_closed=True, close_new_points=True, + grid_xstep=None, grid_ystep=None, grid_stroke_width=0.0, #DEFAULT_STROKE_WIDTH/2, + grid_stroke_color=None, grid_stroke_opacity=None, **kwargs): + super().__init__() + # Fields + self.center = center + self.mark_paths_closed = mark_paths_closed + self.close_new_points = close_new_points + self.grid_xstep = grid_xstep + self.grid_ystep = grid_ystep + self.grid_stroke_width = grid_stroke_width + self.grid_stroke_color = grid_stroke_color + self.grid_stroke_opacity = grid_stroke_opacity + self.rotation_angles = [0, 0, 0] + # Make inner_rectangle + self.inner_rectangle = Rectangle( + width=width, + height=height, + stroke_opacity=0.0, + stroke_width=0.0 + ) + print(self.inner_rectangle.get_vertices()) + # self.inner_rectangle = Polygon( + """ + points = [ + self.center + np.array([width / 2, height / 2, 0]), + self.center + np.array([width / 2, -1*(height / 2), 0]), + self.center + np.array([-1 * (width / 2), -1 * (height / 2), 0]), + self.center + np.array([-1 * (width / 2), height / 2, 0]), + ] + self.inner_rectangle = Polygram( + points, + stroke_opacity=0.0, + stroke_width=0.0 + ) + """ + self.add(self.inner_rectangle) + # Make outline rectangle + self.outline_rectangle = SurroundingRectangle( + self.inner_rectangle, + color=color, + buff=0.0, + **kwargs + ) + self.add(self.outline_rectangle) + # Move to center + self.move_to(self.center) + # Setup Object + # TODO re-implement gridded rectangle + # self.grid_lines = self.make_grid_lines() + # self.add(self.grid_lines) + # Make dots for the corners + # Make outer corner dots + self.outer_corners = VGroup( + *[Dot(corner_point) for corner_point in self.outline_rectangle.get_vertices()] + ) + + self.outer_corners.set_fill(opacity=0.0) + self.add(self.outer_corners) + # Make inner corner dots + self.inner_corners = VGroup( + *[Dot(corner_point) for corner_point in self.inner_rectangle.get_vertices()] + ) + self.inner_corners.set_fill(opacity=0.0) + self.add(self.inner_corners) + + def get_corners_dict(self, inner_rectangle=False): + """Returns a dictionary of the corners""" + if inner_rectangle: + return { + "top_left": self.inner_corners[3], + "top_right": self.inner_corners[0], + "bottom_left": self.inner_corners[2], + "bottom_right": self.inner_corners[1], + } + else: + return { + "top_left": self.outer_corners[3], + "top_right": self.outer_corners[0], + "bottom_left": self.outer_corners[2], + "bottom_right": self.outer_corners[1], + } + + def make_grid_lines(self): + """Make grid lines in rectangle""" + grid_lines = VGroup() + width = self.width + height = self.width + + v = self.inner_rectangle.get_vertices() + if self.grid_xstep is not None: + grid_xstep = abs(self.grid_xstep) + count = int(width / grid_xstep) + grid = VGroup( + *( + Line( + v[1] + i * grid_xstep * RIGHT, + v[1] + i * grid_xstep * RIGHT + height * DOWN, + color=self.color, + stroke_width=self.grid_stroke_width + ) + for i in range(1, count) + ) + ) + grid_lines.add(grid) + + if self.grid_ystep is not None: + grid_ystep = abs(self.grid_ystep) + count = int(height / grid_ystep) + grid = VGroup( + *( + Line( + v[1] + i * grid_ystep * DOWN, + v[1] + i * grid_ystep * DOWN + width * RIGHT, + color=self.color, + stroke_width = self.grid_stroke_width + ) + for i in range(1, count) + ) + ) + grid_lines.add(grid) + + return grid_lines + + def rotate_about_origin(self, angle, axis=OUT, axes=[]): + self.rotation_angles[np.nonzero(axis)[0][0]] = angle + return super().rotate_about_origin(angle, axis, axes) + + def get_normal_vector(self): + """Gets the vector normal to main rectangle face""" + # Get three corner points + corner_1 = self.rectangle.get_top() + corner_2 = self.rectangle.get_left() + corner_3 = self.rectangle.get_right() + # Make vectors from them + a = corner_1 - corner_3 + b = corner_1 - corner_2 + # Compute cross product + normal_vector = np.cross(b, a) + normal_vector /= np.linalg.norm(normal_vector) + + return normal_vector + + def get_rotation_axis_and_angle(self): + """Gets the angle of rotation necessary to rotate something from the default z-axis to the rectangle""" + def unit_vector(vector): + """ Returns the unit vector of the vector. """ + return vector / np.linalg.norm(vector) + + def angle_between(v1, v2): + """ Returns the angle in radians between vectors 'v1' and 'v2':: + + >>> angle_between((1, 0, 0), (0, 1, 0)) + 1.5707963267948966 + >>> angle_between((1, 0, 0), (1, 0, 0)) + 0.0 + >>> angle_between((1, 0, 0), (-1, 0, 0)) + 3.141592653589793 + """ + v1_u = unit_vector(v1) + v2_u = unit_vector(v2) + return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0)) + + normal_vector = self.get_normal_vector() + z_axis = Z_AXIS + # Get angle between normal vector and z axis + axis = np.cross(normal_vector, z_axis) + angle = angle_between(normal_vector, z_axis) + + return axis, angle \ No newline at end of file diff --git a/manim_ml/neural_network/layers/__init__.py b/manim_ml/neural_network/layers/__init__.py index 224a1c1..46edf6a 100644 --- a/manim_ml/neural_network/layers/__init__.py +++ b/manim_ml/neural_network/layers/__init__.py @@ -1,7 +1,7 @@ from .convolutional3d_to_convolutional3d import Convolutional3DToConvolutional3D from .convolutional2d_to_convolutional2d import Convolutional2DToConvolutional2D -from .convolutional_3d import Convolutional3DLayer -from .convolutional_2d import Convolutional2DLayer +from .convolutional3d import Convolutional3DLayer +from .convolutional2d import Convolutional2DLayer from .feed_forward_to_vector import FeedForwardToVector from .paired_query_to_feed_forward import PairedQueryToFeedForward from .embedding_to_feed_forward import EmbeddingToFeedForward diff --git a/manim_ml/neural_network/layers/convolutional_2d.py b/manim_ml/neural_network/layers/convolutional2d.py similarity index 87% rename from manim_ml/neural_network/layers/convolutional_2d.py rename to manim_ml/neural_network/layers/convolutional2d.py index efdf43f..24ed9d6 100644 --- a/manim_ml/neural_network/layers/convolutional_2d.py +++ b/manim_ml/neural_network/layers/convolutional2d.py @@ -1,19 +1,17 @@ from manim import * from matplotlib import animation -from xarray import align from manim_ml.neural_network.layers.parent_layers import VGroupNeuralNetworkLayer class Convolutional2DLayer(VGroupNeuralNetworkLayer): def __init__(self, feature_map_height, feature_map_width, filter_width, filter_height, - stride=1, cell_width=0.5, pixel_width=0.5, feature_map_color=BLUE, filter_color=ORANGE, + stride=1, cell_width=0.5, feature_map_color=BLUE, filter_color=ORANGE, **kwargs): super(VGroupNeuralNetworkLayer, self).__init__(**kwargs) self.feature_map_height = feature_map_height self.feature_map_width = feature_map_width self.filter_width = filter_width self.filter_height = filter_height - self.pixel_width = pixel_width self.feature_map_color = feature_map_color self.filter_color = filter_color self.stride = stride diff --git a/manim_ml/neural_network/layers/convolutional2d_to_convolutional2d.py b/manim_ml/neural_network/layers/convolutional2d_to_convolutional2d.py index 79dee83..e3c98f0 100644 --- a/manim_ml/neural_network/layers/convolutional2d_to_convolutional2d.py +++ b/manim_ml/neural_network/layers/convolutional2d_to_convolutional2d.py @@ -1,6 +1,6 @@ from cv2 import line from manim import * -from manim_ml.neural_network.layers.convolutional_2d import Convolutional2DLayer +from manim_ml.neural_network.layers.convolutional2d import Convolutional2DLayer from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer class Convolutional2DToConvolutional2D(ConnectiveLayer): diff --git a/manim_ml/neural_network/layers/convolutional3d.py b/manim_ml/neural_network/layers/convolutional3d.py new file mode 100644 index 0000000..3086ea0 --- /dev/null +++ b/manim_ml/neural_network/layers/convolutional3d.py @@ -0,0 +1,96 @@ +from manim import * +from manim_ml.neural_network.layers.parent_layers import ThreeDLayer, VGroupNeuralNetworkLayer +from manim_ml.gridded_rectangle import GriddedRectangle +import numpy as np + +class Convolutional3DLayer(VGroupNeuralNetworkLayer, ThreeDLayer): + """Handles rendering a convolutional layer for a nn""" + + def __init__(self, num_feature_maps, feature_map_width, feature_map_height, + filter_width, filter_height, cell_width=0.2, filter_spacing=0.1, color=BLUE, + pulse_color=ORANGE, filter_color=ORANGE, stride=1, stroke_width=2.0, **kwargs): + super(VGroupNeuralNetworkLayer, self).__init__(**kwargs) + self.num_feature_maps = num_feature_maps + self.feature_map_height = feature_map_height + self.filter_color = filter_color + self.feature_map_width = feature_map_width + self.filter_width = filter_width + self.filter_height = filter_height + self.cell_width = cell_width + self.filter_spacing = filter_spacing + self.color = color + self.pulse_color = pulse_color + self.stride = stride + self.stroke_width = stroke_width + # Make the feature maps + self.feature_maps = self.construct_feature_maps() + self.add(self.feature_maps) + + def construct_feature_maps(self): + """Creates the neural network layer""" + # Draw rectangles that are filled in with opacity + feature_maps = VGroup() + for filter_index in range(self.num_feature_maps): + rectangle = GriddedRectangle( + center=[0, 0, filter_index * self.filter_spacing], # Center coordinate + color=self.color, + height=self.feature_map_height * self.cell_width, + width=self.feature_map_width * self.cell_width, + fill_color=self.color, + fill_opacity=0.2, + stroke_color=self.color, + stroke_width=self.stroke_width, + # grid_xstep=self.cell_width, + # grid_ystep=self.cell_width, + # grid_stroke_width=DEFAULT_STROKE_WIDTH/2 + ) + # Rotate about z axis + rectangle.rotate_about_origin( + 90 * DEGREES, + np.array([0, 1, 0]) + ) + feature_maps.add(rectangle) + + return feature_maps + + def make_forward_pass_animation( + self, + run_time=5, + corner_pulses=False, + layer_args={}, + **kwargs + ): + """Convolution forward pass animation""" + # Note: most of this animation is done in the Convolution3DToConvolution3D layer + print(f"Corner pulses: {corner_pulses}") + if corner_pulses: + passing_flashes = [] + for line in self.corner_lines: + pulse = ShowPassingFlash( + line.copy() + .set_color(self.pulse_color) + .set_stroke(opacity=1.0), + time_width=0.5, + run_time=run_time, + rate_func=rate_functions.linear + ) + passing_flashes.append(pulse) + + # per_filter_run_time = run_time / len(self.feature_maps) + # Make animation group + animation_group = AnimationGroup( + *passing_flashes, + # filter_flashes + ) + else: + animation_group = AnimationGroup() + + return animation_group + + def scale(self, scale_factor, **kwargs): + self.cell_width *= scale_factor + super().scale(scale_factor, **kwargs) + + @override_animation(Create) + def _create_override(self, **kwargs): + return FadeIn(self.feature_maps) diff --git a/manim_ml/neural_network/layers/convolutional3d_to_convolutional3d.py b/manim_ml/neural_network/layers/convolutional3d_to_convolutional3d.py index 6cc110c..832a9ea 100644 --- a/manim_ml/neural_network/layers/convolutional3d_to_convolutional3d.py +++ b/manim_ml/neural_network/layers/convolutional3d_to_convolutional3d.py @@ -1,66 +1,287 @@ -from cProfile import run from manim import * -from manim_ml.neural_network.layers.convolutional_3d import Convolutional3DLayer -from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer +from manim_ml.neural_network.layers.convolutional3d import Convolutional3DLayer +from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer, ThreeDLayer +from manim_ml.gridded_rectangle import GriddedRectangle, CornersRectangle -class Convolutional3DToConvolutional3D(ConnectiveLayer): +class Filters(VGroup): + """Group for showing a collection of filters connecting two layers""" + + def __init__(self, input_layer, output_layer, line_color=ORANGE, stroke_width=2.0): + super().__init__() + self.input_layer = input_layer + self.output_layer = output_layer + self.line_color = line_color + self.stroke_width = stroke_width + # Make the filter + self.input_rectangles = self.make_input_feature_map_rectangles() + self.add(self.input_rectangles) + self.output_rectangles = self.make_output_feature_map_rectangles() + self.add(self.output_rectangles) + self.connective_lines = self.make_connective_lines() + self.add(self.connective_lines) + + def make_input_feature_map_rectangles(self): + rectangles = [] + + rectangle_width = self.input_layer.filter_width * self.input_layer.cell_width + rectangle_height = self.input_layer.filter_height * self.input_layer.cell_width + filter_color = self.input_layer.filter_color + + for index, feature_map in enumerate(self.input_layer.feature_maps): + rectangle = GriddedRectangle( + center=feature_map.get_center(), + width=rectangle_width, + height=rectangle_height, + fill_color=filter_color, + stroke_color=filter_color, + fill_opacity=0.2, + z_index=2, + stroke_width=self.stroke_width, + ) + # Center on feature map + # rectangle.move_to(feature_map.get_center()) + # Rotate so it is in the yz plane + rectangle.rotate( + 90 * DEGREES, + axis=[0, 1, 0] + ) + # Get the feature map top left corner + feature_map_top_left = feature_map.get_corners_dict(inner_rectangle=True)["top_left"] + rectangle_top_left = rectangle.get_corners_dict()["top_left"] + # Move the rectangle to the corner location + rectangle.next_to( + feature_map_top_left, + submobject_to_align=rectangle_top_left, + buff=0.0 + ) + + rectangles.append(rectangle) + + feature_map_rectangles = VGroup(*rectangles) + + return feature_map_rectangles + + def make_output_feature_map_rectangles(self): + rectangles = [] + + rectangle_width = self.output_layer.cell_width + rectangle_height = self.output_layer.cell_width + filter_color = self.output_layer.filter_color + + for index, feature_map in enumerate(self.output_layer.feature_maps): + rectangle = GriddedRectangle( + center=feature_map.get_center(), + width=rectangle_width, + height=rectangle_height, + fill_color=filter_color, + stroke_color=filter_color, + fill_opacity=0.2, + stroke_width=self.stroke_width, + z_index=2, + ) + # Center on feature map + # rectangle.move_to(feature_map.get_center()) + # Rotate so it is in the yz plane + rectangle.rotate( + 90 * DEGREES, + axis=[0, 1, 0] + ) + # Get the feature map top left corner + feature_map_top_left = feature_map.get_corners_dict(inner_rectangle=True)["top_left"] + rectangle_top_left = rectangle.get_corners_dict()["top_left"] + # Move the rectangle to the corner location + rectangle.next_to( + feature_map_top_left, + submobject_to_align=rectangle_top_left, + buff=0.0 + ) + + rectangles.append(rectangle) + + feature_map_rectangles = VGroup(*rectangles) + + return feature_map_rectangles + + def make_connective_lines(self): + """Lines connecting input filter with output node""" + + corner_names = ["top_left", "bottom_left", "top_right", "bottom_right"] + + def make_input_connective_lines(): + """Makes connective lines between the corners of the input filters""" + first_input_rectangle = self.input_rectangles[0] + last_input_rectangle = self.input_rectangles[-1] + # Get the corner dots for each rectangle + first_input_corners = first_input_rectangle.get_corners_dict() + last_input_corners = last_input_rectangle.get_corners_dict() + # Iterate through each corner and make the lines + lines = [] + for corner_name in corner_names: + line = Line( + first_input_corners[corner_name].get_center(), + last_input_corners[corner_name].get_center(), + color=self.line_color, + stroke_width=self.stroke_width + ) + lines.append(line) + + return VGroup(*lines) + + def make_output_connective_lines(): + """Makes connective lines between the corners of the output filters""" + first_output_rectangle = self.output_rectangles[0] + last_output_rectangle = self.output_rectangles[-1] + # Get the corner dots for each rectangle + first_output_corners = first_output_rectangle.get_corners_dict() + last_output_corners = last_output_rectangle.get_corners_dict() + # Iterate through each corner and make the lines + lines = [] + for corner_name in corner_names: + line = Line( + first_output_corners[corner_name].get_center(), + last_output_corners[corner_name].get_center(), + color=self.line_color, + stroke_width=self.stroke_width + ) + lines.append(line) + + return VGroup(*lines) + + def make_input_to_output_connective_lines(): + """Make connective lines between last input filter and first output filter""" + last_input_rectangle = self.input_rectangles[-1] + first_output_rectangle = self.output_rectangles[0] + # Get the corner dots for each rectangle + last_input_corners = last_input_rectangle.get_corners_dict() + first_output_corners = first_output_rectangle.get_corners_dict() + # Iterate through each corner and make the lines + lines = [] + for corner_name in corner_names: + line = Line( + last_input_corners[corner_name].get_center(), + first_output_corners[corner_name].get_center(), + color=self.line_color, + stroke_width=self.stroke_width + ) + lines.append(line) + + return VGroup(*lines) + + input_lines = make_input_connective_lines() + output_lines = make_output_connective_lines() + input_output_lines = make_input_to_output_connective_lines() + + connective_lines = VGroup( + *input_lines, + *output_lines, + *input_output_lines + ) + + return connective_lines + +class Convolutional3DToConvolutional3D(ConnectiveLayer, ThreeDLayer): """Feed Forward to Embedding Layer""" input_class = Convolutional3DLayer output_class = Convolutional3DLayer - def __init__(self, input_layer, output_layer, color=WHITE, pulse_color=RED, - **kwargs): - super().__init__(input_layer, output_layer, input_class=Convolutional3DLayer, output_class=Convolutional3DLayer, - **kwargs) + def __init__(self, input_layer: Convolutional3DLayer, output_layer: Convolutional3DLayer, + color=WHITE, filter_opacity=0.3, line_color=WHITE, + pulse_color=ORANGE, **kwargs): + super().__init__(input_layer, output_layer, input_class=Convolutional3DLayer, + output_class=Convolutional3DLayer, **kwargs) self.color = color + self.filter_color = self.input_layer.filter_color + self.filter_width = self.input_layer.filter_width + self.filter_height = self.input_layer.filter_height + self.feature_map_width = self.input_layer.feature_map_width + self.feature_map_height = self.input_layer.feature_map_height + self.num_input_feature_maps = self.input_layer.num_feature_maps + self.num_output_feature_maps = self.output_layer.num_feature_maps + self.cell_width = self.input_layer.cell_width + self.stride = self.input_layer.stride + self.filter_opacity = filter_opacity + self.line_color = line_color self.pulse_color = pulse_color - self.lines = self.make_lines() - self.add(self.lines) - - def make_lines(self): - """Make lines connecting the input and output layers""" - lines = VGroup() - # Get the first and last rectangle - input_rectangle = self.input_layer.rectangles[-1] - output_rectangle = self.output_layer.rectangles[0] - input_vertices = input_rectangle.get_vertices() - output_vertices = output_rectangle.get_vertices() - # Go through each vertex - for vertex_index in range(len(input_vertices)): - # Make a line - line = Line( - start=input_vertices[vertex_index], - end=output_vertices[vertex_index], - color=self.color, - stroke_opacity=0.0 - ) - lines.add(line) - - return lines - - def make_forward_pass_animation(self, layer_args={}, run_time=1.5, **kwargs): - """Forward pass animation from conv to conv""" - animations = [] - # Go thorugh the lines - for line in self.lines: - pulse = ShowPassingFlash( - line.copy() - .set_color(self.pulse_color) - .set_stroke(opacity=1.0), - time_width=0.5, - run_time=run_time - ) - animations.append(pulse) - # Make animation group + def make_filter_propagation_animation(self): + """Make filter propagation animation""" + # TODO implement this + raise NotImplementedError() + # Deprecated code + old_z_index = self.filter_lines.z_index + lines_copy = self.filter_lines.copy().set_color(ORANGE).set_z_index(old_z_index + 1) animation_group = AnimationGroup( - *animations, - run_time=run_time + Create(lines_copy, lag_ratio=0.0), + # FadeOut(self.filter_lines), + FadeOut(lines_copy), + lag_ratio=1.0 ) return animation_group + def make_forward_pass_animation(self, layer_args={}, run_time=10.5, **kwargs): + """Forward pass animation from conv2d to conv2d""" + animations = [] + # Create the filters, output nodes (feature map square), and lines + filters = Filters(self.input_layer, self.output_layer) + self.add(filters) + # Make animations for creating the filters, output_nodes, and filter_lines + # TODO decide if I want to create the filters at the start of a conv + # animation or have them there by default + # animations.append( + # Create(filters) + # ) + # Make shift amounts + right_shift = np.array([0, self.input_layer.cell_width, 0])# * 1.55 + left_shift = np.array([0, -1*self.input_layer.cell_width, 0])# * 1.55 + up_shift = np.array([0, 0, -1*self.input_layer.cell_width])# * 1.55 + down_shift = np.array([0, 0, self.input_layer.cell_width])# * 1.55 + # Make filter shifting animations + num_y_moves = int((self.feature_map_height - self.filter_height) / self.stride) + num_x_moves = int((self.feature_map_width - self.filter_width) / self.stride) + for y_move in range(num_y_moves): + # Go right num_x_moves + for x_move in range(num_x_moves): + # Shift right + shift_animation = ApplyMethod( + filters.shift, + self.stride * right_shift + ) + # shift_animation = self.animate.shift(right_shift) + animations.append(shift_animation) + + # Go back left num_x_moves and down one + shift_amount = self.stride * num_x_moves * left_shift + self.stride * down_shift + # Make the animation + shift_animation = ApplyMethod( + filters.shift, + shift_amount + ) + animations.append(shift_animation) + # Do last row move right + for x_move in range(num_x_moves): + # Shift right + shift_animation = ApplyMethod( + filters.shift, + self.stride * right_shift + ) + # shift_animation = self.animate.shift(right_shift) + animations.append(shift_animation) + + # Remove filters + return Succession( + *animations, + lag_ratio=1.0 + ) + + def set_z_index(self, z_index, family=False): + """Override set_z_index""" + super().set_z_index(4) + + def scale(self, scale_factor, **kwargs): + self.cell_width *= scale_factor + super().scale(scale_factor, **kwargs) + @override_animation(Create) def _create_override(self, **kwargs): return AnimationGroup() - diff --git a/manim_ml/neural_network/layers/convolutional_3d.py b/manim_ml/neural_network/layers/convolutional_3d.py deleted file mode 100644 index 80b1402..0000000 --- a/manim_ml/neural_network/layers/convolutional_3d.py +++ /dev/null @@ -1,122 +0,0 @@ -from manim import * -from manim_ml.neural_network.layers.parent_layers import VGroupNeuralNetworkLayer - -class Convolutional3DLayer(VGroupNeuralNetworkLayer): - """Handles rendering a convolutional layer for a nn""" - - def __init__(self, num_filters, filter_width, filter_height, filter_spacing=0.1, color=BLUE, - pulse_color=ORANGE, **kwargs): - super(VGroupNeuralNetworkLayer, self).__init__(**kwargs) - self.num_filters = num_filters - self.filter_width = filter_width - self.filter_height = filter_height - self.filter_spacing = filter_spacing - self.color = color - self.pulse_color = pulse_color - - self._construct_layer(num_filters=self.num_filters, filter_width=self.filter_width, filter_height=self.filter_height) - - def _construct_layer(self, num_filters=5, filter_width=4, filter_height=4): - """Creates the neural network layer""" - # Make axes, but hide the lines - axes = ThreeDAxes( - tips=False, - x_length=1, - y_length=1, - x_axis_config={ - "include_ticks": False, - "stroke_width": 0.0 - }, - y_axis_config={ - "include_ticks": False, - "stroke_width": 0.0 - }, - z_axis_config={ - "include_ticks": False, - "stroke_width": 0.0 - } - ) - self.add(axes) - # Set the camera angle so that the - # self.set_camera_orientation(phi=75 * DEGREES, theta=30 * DEGREES) - # Draw rectangles that are filled in with opacity - self.rectangles = VGroup() - for filter_index in range(num_filters): - rectangle = Rectangle( - color=self.color, - height=filter_height, - width=filter_width, - fill_color=self.color, - fill_opacity=0.2, - stroke_color=WHITE, - ) - rectangle.rotate_about_origin((80 - filter_index*0.5) * DEGREES, np.array([0, 1, 0])) # Rotate about z axis - rectangle.rotate_about_origin(15 * DEGREES, np.array([1, 0, 0])) # Rotate about x axis - rectangle.shift(np.array([filter_index*self.filter_spacing, filter_height*0.5, -3])) - - self.rectangles.add(rectangle) - - self.add(self.rectangles) - - self.corner_lines = self.make_filter_corner_lines() - self.add(self.corner_lines) - - def make_filter_corner_lines(self): - """Make filter corner lines""" - corner_lines = VGroup() - - first_rectangle = self.rectangles[0] - last_rectangle = self.rectangles[-1] - first_vertices = first_rectangle.get_vertices() - last_vertices = last_rectangle.get_vertices() - for vertex_index in range(len(first_vertices)): - # Make a line - line = Line( - start=first_vertices[vertex_index], - end=last_vertices[vertex_index], - color=WHITE, - stroke_opacity=0.0 - ) - corner_lines.add(line) - - return corner_lines - - def make_forward_pass_animation(self, run_time=5, layer_args={}, **kwargs): - """Convolution forward pass animation""" - passing_flashes = [] - for line in self.corner_lines: - pulse = ShowPassingFlash( - line.copy() - .set_color(self.pulse_color) - .set_stroke(opacity=1.0), - time_width=0.5, - run_time=run_time, - rate_func=rate_functions.linear - ) - passing_flashes.append(pulse) - - per_filter_run_time = run_time / len(self.rectangles) - filter_flashes = [] - for filter in self.rectangles: - single_flash = Succession( - ApplyMethod(filter.set_color, self.pulse_color, run_time=per_filter_run_time/4), - Wait(per_filter_run_time/2), - ApplyMethod(filter.set_color, self.color, run_time=per_filter_run_time/4), - ApplyMethod(filter.set_stroke_color, WHITE, run_time=0.0) - ) - filter_flashes.append(single_flash) - - filter_flashes = Succession( - *filter_flashes, - ) - # Make animation group - animation_group = AnimationGroup( - *passing_flashes, - filter_flashes - ) - - return animation_group - - @override_animation(Create) - def _create_override(self, **kwargs): - return FadeIn(self.rectangles) diff --git a/manim_ml/neural_network/layers/parent_layers.py b/manim_ml/neural_network/layers/parent_layers.py index 6ef782c..baaf228 100644 --- a/manim_ml/neural_network/layers/parent_layers.py +++ b/manim_ml/neural_network/layers/parent_layers.py @@ -35,6 +35,12 @@ class VGroupNeuralNetworkLayer(NeuralNetworkLayer): def _create_override(self): return super()._create_override() +class ThreeDLayer(ABC): + """Abstract class for 3D layers""" + + def __init__(self): + pass + class ConnectiveLayer(VGroupNeuralNetworkLayer): """Forward pass animation for a given pair of layers""" diff --git a/manim_ml/neural_network/neural_network.py b/manim_ml/neural_network/neural_network.py index 0a51ca5..170499f 100644 --- a/manim_ml/neural_network/neural_network.py +++ b/manim_ml/neural_network/neural_network.py @@ -14,16 +14,18 @@ import textwrap from manim_ml.neural_network.layers.embedding import EmbeddingLayer from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer -from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer +from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer, ThreeDLayer from manim_ml.neural_network.layers.util import get_connective_layer from manim_ml.list_group import ListGroup from manim_ml.neural_network.neural_network_transformations import InsertLayer, RemoveLayer class NeuralNetwork(Group): + """Neural Network Visualization Container Class""" def __init__(self, input_layers, edge_color=WHITE, layer_spacing=0.2, animation_dot_color=RED, edge_width=2.5, dot_radius=0.03, - title=" "): + title=" ", camera=None, camera_phi=-70 * DEGREES, + camera_theta=-80 * DEGREES): super(Group, self).__init__() self.input_layers = ListGroup(*input_layers) self.edge_width = edge_width @@ -33,6 +35,11 @@ class NeuralNetwork(Group): self.dot_radius = dot_radius self.title_text = title self.created = False + self.camera = camera + # Set the camera orientation for 3D Layers + if not self.camera is None: + self.camera.set_phi(camera_phi) + self.camera.set_theta(camera_theta) # TODO take layer_node_count [0, (1, 2), 0] # and make it have explicit distinct subspaces self._place_layers() @@ -59,7 +66,8 @@ class NeuralNetwork(Group): current_layer = self.input_layers[layer_index] current_layer.move_to(previous_layer) # TODO Temp fix - if isinstance(current_layer, EmbeddingLayer) or isinstance(previous_layer, EmbeddingLayer): + if isinstance(current_layer, EmbeddingLayer) \ + or isinstance(previous_layer, EmbeddingLayer): shift_vector = np.array([(previous_layer.get_width()/2 + current_layer.get_width()/2 - 0.2), 0, 0]) else: shift_vector = np.array([(previous_layer.get_width()/2 + current_layer.get_width()/2) + self.layer_spacing, 0, 0]) @@ -71,6 +79,11 @@ class NeuralNetwork(Group): all_layers = ListGroup() for layer_index in range(len(self.input_layers) - 1): current_layer = self.input_layers[layer_index] + # Make the layer fixed in frame if its not 3D + if not isinstance(current_layer, ThreeDLayer): + self.camera.add_fixed_orientation_mobjects(current_layer) + self.camera.add_fixed_in_frame_mobjects(current_layer) + # Add the layer to the list of layers all_layers.add(current_layer) next_layer = self.input_layers[layer_index + 1] # Check if layer is actually a nested NeuralNetwork @@ -83,7 +96,16 @@ class NeuralNetwork(Group): # Find connective layer with correct layer pair connective_layer = get_connective_layer(current_layer, next_layer) connective_layers.add(connective_layer) + # Make the layer fixed in frame if its not 3D + if not isinstance(current_layer, ThreeDLayer): + self.camera.add_fixed_orientation_mobjects(connective_layer) + self.camera.add_fixed_in_frame_mobjects(connective_layer) + # Add the layer to the list of layers all_layers.add(connective_layer) + # Check if final layer is a 3D layer + if not isinstance(self.input_layers[-1], ThreeDLayer): + self.camera.add_fixed_orientation_mobjects(self.input_layers[-1]) + self.camera.add_fixed_in_frame_mobjects(self.input_layers[-1]) # Add final layer all_layers.add(self.input_layers[-1]) # Handle layering @@ -114,11 +136,11 @@ class NeuralNetwork(Group): return animation_group - def make_forward_pass_animation(self, run_time=10, passing_flash=True, layer_args={}, + def make_forward_pass_animation(self, run_time=None, passing_flash=True, layer_args={}, **kwargs): """Generates an animation for feed forward propagation""" all_animations = [] - per_layer_runtime = run_time/len(self.all_layers) + per_layer_runtime = run_time / len(self.all_layers) if not run_time is None else None for layer_index, layer in enumerate(self.all_layers): # Get the layer args if isinstance(layer, ConnectiveLayer): @@ -148,7 +170,6 @@ class NeuralNetwork(Group): # Make the animation group animation_group = Succession( *all_animations, - run_time=run_time, lag_ratio=1.0 ) diff --git a/tests/test_3d_camera_move.py b/tests/test_3d_camera_move.py new file mode 100644 index 0000000..f04527e --- /dev/null +++ b/tests/test_3d_camera_move.py @@ -0,0 +1,41 @@ +from manim import * +from PIL import Image + +from manim_ml.neural_network.layers.convolutional3d import Convolutional3DLayer +from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer +from manim_ml.neural_network.layers.image import ImageLayer +from manim_ml.neural_network.neural_network import NeuralNetwork + +# Make the specific scene +config.pixel_height = 1200 +config.pixel_width = 1900 +config.frame_height = 6.0 +config.frame_width = 6.0 + +class CombinedScene(ThreeDScene): + def construct(self): + image = Image.open('../assets/mnist/digit.jpeg') + numpy_image = np.asarray(image) + # Make nn + nn = NeuralNetwork( + [ + ImageLayer(numpy_image, height=1.4), + Convolutional3DLayer(1, 5, 5, 3, 3, filter_spacing=0.2), + Convolutional3DLayer(2, 3, 3, 1, 1, filter_spacing=0.2), + FeedForwardLayer(3, rectangle_stroke_width=4, node_stroke_width=4), + ], + layer_spacing=0.5, + camera=self.camera + ) + + nn.scale(1.3) + # Center the nn + nn.move_to(ORIGIN) + self.add(nn) + # Play animation + forward_pass = nn.make_forward_pass_animation( + corner_pulses=False + ) + self.play( + forward_pass + ) diff --git a/tests/test_convolutional_3d_layer.py b/tests/test_convolutional_3d_layer.py index abe5e92..b504fe3 100644 --- a/tests/test_convolutional_3d_layer.py +++ b/tests/test_convolutional_3d_layer.py @@ -1,7 +1,7 @@ from manim import * from PIL import Image -from manim_ml.neural_network.layers.convolutional_3d import Convolutional3DLayer +from manim_ml.neural_network.layers.convolutional3d import Convolutional3DLayer from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer from manim_ml.neural_network.layers.image import ImageLayer from manim_ml.neural_network.neural_network import NeuralNetwork @@ -22,34 +22,64 @@ class SingleConvolutionalLayerScene(ThreeDScene): self.set_camera_orientation(phi=280*DEGREES, theta=-10*DEGREES, gamma=90*DEGREES) # self.play(nn.make_forward_pass_animation(run_time=5)) +class Simple3DConvScene(ThreeDScene): + + def construct(self): + """ + TODO + - [X] Make grid lines for the CNN filters + - [ ] Make Scanning filter effect + - [ ] Have filter box go accross each input feature map + - [ ] Make filter lines effect + - [ ] Make flowing animation down filter lines + """ + # Make nn + layers = [ + Convolutional3DLayer(1, 5, 5, 5, 5, feature_map_height=3, filter_width=3, filter_height=3), + Convolutional3DLayer(1, 3, 3, 1, 1, feature_map_width=3, filter_width=3, filter_height=3), + ] + nn = NeuralNetwork(layers) + # Center the nn + nn.move_to(ORIGIN) + self.add(nn) + # Play animation + # self.set_camera_orientation(phi=280*DEGREES, theta=-10*DEGREES, gamma=90*DEGREES) + self.play(nn.make_forward_pass_animation(run_time=30)) + # Make the specific scene config.pixel_height = 1200 config.pixel_width = 1900 -config.frame_height = 12.0 -config.frame_width = 12.0 +config.frame_height = 6.0 +config.frame_width = 6.0 -class CombinedScene(ThreeDScene, Scene): +class CombinedScene(ThreeDScene): def construct(self): image = Image.open('../assets/mnist/digit.jpeg') numpy_image = np.asarray(image) # Make nn - nn = NeuralNetwork([ - ImageLayer(numpy_image, height=1.4), - Convolutional3DLayer(3, 3, 3, filter_spacing=0.2), - Convolutional3DLayer(5, 2, 2, filter_spacing=0.2), - Convolutional3DLayer(10, 2, 1, filter_spacing=0.2), - FeedForwardLayer(3, rectangle_stroke_width=4, node_stroke_width=4).scale(2), - FeedForwardLayer(1, rectangle_stroke_width=4, node_stroke_width=4).scale(2) - ], layer_spacing=0.2) + nn = NeuralNetwork( + [ + ImageLayer(numpy_image, height=1.4), + Convolutional3DLayer(1, 5, 5, 3, 3, filter_spacing=0.2), + Convolutional3DLayer(2, 3, 3, 1, 1, filter_spacing=0.2), + FeedForwardLayer(3, rectangle_stroke_width=4, node_stroke_width=4), + FeedForwardLayer(1, rectangle_stroke_width=4, node_stroke_width=4) + ], + layer_spacing=0.5, + camera=self.camera + ) nn.scale(1.3) # Center the nn nn.move_to(ORIGIN) - self.play(Create(nn)) + self.add(nn) # Play animation # self.set_camera_orientation(phi=280* DEGREES, theta=-20*DEGREES, gamma=90 * DEGREES) # self.begin_ambient_camera_rotation() - forward_pass = nn.make_forward_pass_animation(run_time=10) + forward_pass = nn.make_forward_pass_animation( + run_time=10, + corner_pulses=False + ) print(forward_pass) self.play( forward_pass diff --git a/tests/test_embedding_layer.py b/tests/test_embedding_layer.py index b24cf81..0de916e 100644 --- a/tests/test_embedding_layer.py +++ b/tests/test_embedding_layer.py @@ -108,4 +108,4 @@ class QueryEmbeddingNNScene(Scene): } } ) - ) \ No newline at end of file + ) diff --git a/tests/test_flow.py b/tests/test_flow.py new file mode 100644 index 0000000..a6e1f13 --- /dev/null +++ b/tests/test_flow.py @@ -0,0 +1,6 @@ +from manim_ml.flow.flow import * + +class TestScene(Scene): + + def construct(self): + self.add(Rectangle()) \ No newline at end of file diff --git a/tests/test_opengl_shader.py b/tests/test_opengl_shader.py new file mode 100644 index 0000000..20fdfdd --- /dev/null +++ b/tests/test_opengl_shader.py @@ -0,0 +1,87 @@ +import manim.utils.opengl as opengl +from manim import * +from manim.opengl import * # type: ignore + +class InlineShaderExample(Scene): + def construct(self): + config["background_color"] = "#333333" + + c = Circle(fill_opacity=0.7).shift(UL) + self.add(c) + + shader = Shader( + self.renderer.context, + source={ + "vertex_shader": """ + #version 330 + in vec4 in_vert; + in vec4 in_color; + out vec4 v_color; + uniform mat4 u_model_view_matrix; + uniform mat4 u_projection_matrix; + void main() { + v_color = in_color; + vec4 camera_space_vertex = u_model_view_matrix * in_vert; + vec4 clip_space_vertex = u_projection_matrix * camera_space_vertex; + gl_Position = clip_space_vertex; + } + """, + "fragment_shader": """ + #version 330 + in vec4 v_color; + out vec4 frag_color; + void main() { + frag_color = v_color; + } + + void main () { + // Previously, you'd have rendered your complete scene into a texture + // bound to "fullScreenTexture." + vec4 rValue = texture2D(fullscreenTexture, gl_TexCoords[0] - rOffset); + vec4 gValue = texture2D(fullscreenTexture, gl_TexCoords[0] - gOffset); + vec4 bValue = texture2D(fullscreenTexture, gl_TexCoords[0] - bOffset); + + // Combine the offset colors. + gl_FragColor = vec4(rValue.r, gValue.g, bValue.b, 1.0); + } + """, + }, + ) + shader.set_uniform("u_model_view_matrix", opengl.view_matrix()) + shader.set_uniform( + "u_projection_matrix", + opengl.orthographic_projection_matrix(), + ) + + attributes = np.zeros( + 6, + dtype=[ + ("in_vert", np.float32, (4,)), + ("in_color", np.float32, (4,)), + ], + ) + attributes["in_vert"] = np.array( + [ + [-1, -1, 0, 1], + [-1, 1, 0, 1], + [1, 1, 0, 1], + [-1, -1, 0, 1], + [1, -1, 0, 1], + [1, 1, 0, 1], + ], + ) + attributes["in_color"] = np.array( + [ + [0, 0, 1, 1], + [0, 0, 1, 1], + [0, 0, 1, 1], + [0, 0, 1, 1], + [0, 0, 1, 1], + [0, 0, 1, 1], + ], + ) + mesh = Mesh(shader, attributes) + self.add(mesh) + + self.wait(5) + # self.embed_2() \ No newline at end of file diff --git a/tests/test_surrounding_rectangle.py b/tests/test_surrounding_rectangle.py new file mode 100644 index 0000000..5f8d4b7 --- /dev/null +++ b/tests/test_surrounding_rectangle.py @@ -0,0 +1,10 @@ +from manim import * + +class SurroundingRectangleTest(Scene): + + def construct(self): + rectangle = Rectangle(width=1, height=1, color=WHITE, fill_opacity=1.0) + self.add(rectangle) + + surrounding_rectangle = SurroundingRectangle(rectangle, color=GREEN, buff=0.0) + self.add(surrounding_rectangle) \ No newline at end of file