Manim Morphological Dilation Animation

dilation.py

"""
Manim animation showing morphological dilation. Input image and kernel can be modified and the
output will still look good, although you may need to move/translate it to be viewable.
Mostly written by Claude, I needed to modify it a bit.

`manim -pqh dilation.py`
"""

from manim import *
import numpy as np

class MorphologicalDilation(Scene):
    def construct(self):
        # first sample input
        input_image = [
            [0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 1, 0, 0],
            [0, 0, 1, 0, 0, 0, 1, 0, 0],
            [0, 0, 0, 1, 0, 1, 0, 0, 0],
            [0, 0, 0, 1, 1, 1, 0, 0, 0],
            [0, 0, 0 ,0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 1, 0],
        ]
        # vertical line input
        # input_image = [
        #     [0, 0, 0, 0, 0, 0, 0, 0, 0],
        #     [0, 0, 0, 0, 1, 0, 0, 0, 0],
        #     [0, 0, 0, 0, 1, 0, 0, 0, 0],
        #     [0, 0, 0, 0, 1, 0, 0, 0, 0],
        #     [0, 0, 0, 0, 1, 0, 0, 0, 0],
        #     [0, 0, 0 ,0, 1, 0, 0, 0, 0],
        #     [0, 0, 0, 0, 0, 0, 0, 0, 0],
        # ]
        kernel_size = (3, 1)  # (width, height)
        
        # Grid position parameters
        input_grid_position = [-5.7, 1, 0]  # Moved left side
        output_grid_position = [1, 1, 0]  # Moved right side

        # You can modify these parameters as needed
        pixel_size = 0.6
        grid_color = BLUE
        bg_color = BLACK
        kernel_color = YELLOW
        max_color = GREEN
        output_color = WHITE
        text_color = BLUE
        input_fill_color = WHITE
        
        # Convert image to numpy array for easier manipulation
        input_array = np.array(input_image)
        
        # Calculate dimensions
        rows, cols = input_array.shape
        k_width, k_height = kernel_size
        
        # Calculate output dimensions
        output_rows = rows
        output_cols = cols
        
        # Create output array (initialized with zeros)
        output_array = np.zeros((output_rows, output_cols))
        
        # Create title
        title = Text("Morphological Dilation (3x1 Kernel)").scale(0.8).to_edge(UP)
        self.play(Write(title))
        
        # Helper function to create grid
        def create_grid(array, color_map, pos):
            grid = VGroup()
            values = VGroup()
            
            for i in range(len(array)):
                for j in range(len(array[0])):
                    # Create square
                    square = Square(side_length=pixel_size)
                    square.set_stroke(grid_color, 2)
                    
                    # Position the square
                    square.move_to([j * pixel_size + pos[0], -i * pixel_size + pos[1], 0])
                    
                    # Fill color based on value
                    value = array[i][j]
                    fill_color = color_map(value)
                    square.set_fill(fill_color, opacity=0.8 if value > 0 else 0.1)
                    
                    # Add value text with blue color
                    # change text color based on background
                    if fill_color == bg_color:
                        text_color = WHITE
                    else:
                        text_color = BLACK
                    value_text = Text(str(int(value)), color=text_color).scale(0.3)
                    value_text.move_to(square.get_center())
                    
                    grid.add(square)
                    values.add(value_text)
            
            return grid, values
        
        # Create input grid
        def input_color_map(value):
            return input_fill_color if value > 0 else bg_color
        
        input_grid, input_values = create_grid(
            input_array, 
            input_color_map, 
            input_grid_position
        )
        
        input_label = Text("Input Image").scale(0.6)
        input_label.next_to(input_grid, UP * 1.5, buff=0.5)
        
        # Create output grid
        def output_color_map(value):
            return output_color if value > 0 else bg_color
        
        output_grid, output_values = create_grid(
            output_array, 
            output_color_map, 
            output_grid_position
        )
        
        output_label = Text("Output Image").scale(0.6)
        output_label.next_to(output_grid, UP * 1.5, buff=0.5)
        
        # Show input and output grids
        self.play(
            FadeIn(input_grid), 
            FadeIn(input_values),
            FadeIn(input_label),
            FadeIn(output_grid), 
            FadeIn(output_label)
        )
        
        
        # Create animation for dilation process
        def get_kernel_grid(i, j):
            kernel_grid = VGroup()
            # check if any 1s in the kernel area
            kernel_area = input_array[max(0, i - pad_height):min(rows, i + pad_height + 1), 
                                       max(0, j - pad_width):min(cols, j + pad_width + 1)]
            if np.any(kernel_area > 0):
                fill_color = GREEN
            else:
                fill_color = RED
                
            # Create squares for kernel
            for ki in range(k_height):
                for kj in range(k_width):
                    # Get corresponding position in the input image
                    ni = i - k_height // 2 + ki
                    nj = j - k_width // 2 + kj
                    
                    # Create square for kernel position
                    square = Square(side_length=pixel_size)
                    square.set_stroke(kernel_color, 2)
                    
                    # Position the square relative to input grid
                    square.move_to([
                        nj * pixel_size + input_grid_position[0], 
                        -ni * pixel_size + input_grid_position[1], 
                        0
                    ])
                    
                    # Set fill color - only if the position is valid


                    square.set_fill(fill_color, opacity=0.5)

                    
                    kernel_grid.add(square)
            
            return kernel_grid
        
        # Padding for valid kernel positions
        pad_height, pad_width = k_height // 2, k_width // 2
        
        for i in range(rows):
            for j in range(cols):
                # Create kernel grid at current position
                kernel_grid = get_kernel_grid(i, j)
                
                # Find the maximum in the kernel neighborhood
                max_val = 0
                max_pos = None
                
                for ki in range(k_height):
                    for kj in range(k_width):
                        ni = i - k_height // 2 + ki
                        nj = j - k_width // 2 + kj
                        
                        # Check if position is valid
                        if 0 <= ni < rows and 0 <= nj < cols:
                            if input_array[ni][nj] > max_val:
                                max_val = input_array[ni][nj]
                                max_pos = (ni, nj)
                
                # Update output array
                output_array[i][j] = max_val
                
                # Create max indicator
                max_indicator = None
                if max_pos:
                    ni, nj = max_pos
                    max_indicator = Square(side_length=pixel_size)
                    max_indicator.set_stroke(max_color, 4)
                    max_indicator.move_to([
                        nj * pixel_size + input_grid_position[0], 
                        -ni * pixel_size + input_grid_position[1], 
                        0
                    ])
                
                # Update output cell
                output_cell = output_grid[i * cols + j]
                updated_cell = output_cell.copy()
                updated_cell.set_fill(output_color_map(max_val), opacity=0.8 if max_val > 0 else 0.1)
                
                # Update output value text
                output_val_text = output_values[i * cols + j]
                text_color = BLACK if max_val > 0 else WHITE
                updated_val_text = Text(str(int(max_val)), color=text_color).scale(0.3)
                updated_val_text.move_to(output_cell.get_center())
                
                # show kernel and max indicator
                kernel_animations = [FadeIn(kernel_grid)]
                if max_indicator:
                    kernel_animations.append(FadeIn(max_indicator))
                self.play(*kernel_animations, run_time=0.1)
                
                # uncomment for more breathing room between animations
                self.wait(0.1)
                
                # update output cell
                output_animations = [
                    Transform(output_cell, updated_cell),
                    Transform(output_val_text, updated_val_text)
                ]
                self.play(*output_animations, run_time=0.05)
                
                # Remove kernel and max indicator
                removal_animations = [FadeOut(kernel_grid)]
                if max_indicator:
                    removal_animations.append(FadeOut(max_indicator))
                
                self.play(*removal_animations, run_time=0.1)
        
        self.wait(2)

if __name__ == "__main__":
    scene = MorphologicalDilation()
    scene.render()