Simple Project Code

.gitignore

@@ -3,6 +3,9 @@ __pycache__/
 *.py[cod]
 *$py.class
 
+*.jpg
+*.png
+
 # C extensions
 *.so
 

easy_viz_cnn/colors.py

@@ -0,0 +1,50 @@
+from itertools import cycle
+from typing import NamedTuple
+
+
+class RGB(NamedTuple):
+    red: float
+    green: float
+    blue: float
+    alpha: float = 1.0
+
+
+DARK_LIGHT = cycle([RGB(0.4, 0.4, 0.4), RGB(0.8, 0.8, 0.8)])
+BLUE_GRAY = cycle([RGB(0.654, 0.780, 0.905), RGB(0.8, 0.8, 0.8)])
+GREEN_GRAY = cycle([RGB(0.5, 0.8, 0.5), RGB(0.8, 0.8, 0.8)])
+ORANGE_BROWN = cycle([RGB(0.9, 0.6, 0.4), RGB(0.6, 0.4, 0.2)])
+PURPLE_PINK = cycle([RGB(0.7, 0.5, 0.9), RGB(0.9, 0.7, 0.7)])
+TEAL_GRAY = cycle([RGB(0.5, 0.8, 0.8), RGB(0.8, 0.8, 0.8)])
+RED_BEIGE = cycle([RGB(0.9, 0.4, 0.4), RGB(0.9, 0.8, 0.7)])
+SUNSET = cycle([
+    RGB(0.9, 0.6, 0.3),  # Orange
+    RGB(0.8, 0.4, 0.4),  # Red
+    RGB(0.9, 0.7, 0.7),  # Pink
+    RGB(0.7, 0.5, 0.9)   # Purple
+])
+OCEAN = cycle([
+    RGB(0.2, 0.6, 0.8),  # Light Blue
+    RGB(0.0, 0.5, 0.5),  # Teal
+    RGB(0.4, 0.8, 0.6),  # Light Green
+    RGB(0.2, 0.4, 0.7)   # Darker Blue
+])
+FOREST = cycle([
+    RGB(0.3, 0.7, 0.3),  # Green
+    RGB(0.4, 0.3, 0.2),  # Brown
+    RGB(0.1, 0.4, 0.2),  # Dark Green
+    RGB(0.1, 0.2, 0.4)   # Dark Blue
+])
+PASTEL = cycle([
+    RGB(0.9, 0.7, 0.8),  # Pastel Pink
+    RGB(0.7, 0.8, 0.9),  # Pastel Blue
+    RGB(0.9, 0.9, 0.7),  # Pastel Yellow
+    RGB(0.7, 0.9, 0.8)   # Pastel Green
+])
+RAINBOW = cycle([
+    RGB(1.0, 0.0, 0.0),  # Red
+    RGB(1.0, 0.5, 0.0),  # Orange
+    RGB(1.0, 1.0, 0.0),  # Yellow
+    RGB(0.0, 1.0, 0.0),  # Green
+    RGB(0.0, 0.0, 1.0),  # Blue
+    RGB(0.5, 0.0, 1.0)   # Purple
+])

easy_viz_cnn/layers.py

@@ -0,0 +1,180 @@
+from typing import TypeVar
+
+
+T = TypeVar("T", int, tuple[int, int])
+
+
+class Layer:
+    """Base class for all layers in the CNN model."""
+
+    def __str__(self) -> str:
+        raise NotImplementedError("Subclasses should implement this!")
+
+    def output_shape(self, input_size):
+        raise NotImplementedError("Subclasses should implement this!")
+
+
+class InputLayer(Layer):
+    """Represents an input layer in a CNN."""
+
+    def __init__(self, input_size):
+        self.input_size = input_size
+
+    def __str__(self) -> str:
+        return f"Input({self.input_size})"
+
+    def output_shape(self, *args, **kwargs):
+        return self.input_size
+
+
+class Conv1DLayer(Layer):
+    """
+    Represents a 1D convolutional layer in a CNN.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        Number of output channels.
+    kernel_size : int
+        Size of the kernel.
+    stride : int, optional
+        Stride of the convolution. Default is 1.
+    padding : int, optional
+        Padding of the convolution. Default is 0.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+    ):
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+
+    def __str__(self) -> str:
+        return f"Conv1D({self.in_channels}, {self.out_channels}, {self.kernel_size}, {self.stride}, {self.padding})"
+
+    def output_shape(self, input_size):
+        output_size = (input_size - self.kernel_size + 2 * self.padding) // self.stride + 1
+        return (output_size, self.out_channels)
+
+
+class Conv2DLayer(Layer):
+    """
+    Represents a 2D convolutional layer in a CNN.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        Number of output channels.
+    kernel_size : int or tuple[int, int]
+        Size of the kernel.
+    stride : int or tuple[int, int], optional
+        Stride of the convolution. Default is 1.
+    padding : int or tuple[int, int], optional
+        Padding of the convolution. Default is 0.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, int],
+        stride: int | tuple[int, int] = 1,
+        padding: int | tuple[int, int] = 0,
+    ):
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        if isinstance(kernel_size, int):
+            self.kernel_size = (kernel_size, kernel_size)
+        else:
+            self.kernel_size = kernel_size
+
+        if isinstance(stride, int):
+            self.stride = (stride, stride)
+        else:
+            self.stride = stride
+
+        if isinstance(padding, int):
+            self.padding = (padding, padding)
+        else:
+            self.padding = padding
+
+    def __str__(self) -> str:
+        return f"Conv2D({self.in_channels}, {self.out_channels}, {self.kernel_size}, {self.stride}, {self.padding})"
+
+    def output_shape(self, input_size):
+        output_height = (input_size[0] - self.kernel_size[0] + 2 * self.padding[0]) // self.stride[
+            0
+        ] + 1
+        output_width = (input_size[1] - self.kernel_size[1] + 2 * self.padding[0]) // self.stride[
+            0
+        ] + 1
+        return (output_height, output_width, self.out_channels)
+
+
+class PoolLayer(Layer):
+    """
+    Represents a pooling layer in a CNN.
+
+    Parameters
+    ----------
+    kernel_size : int
+        Size of the kernel.
+    stride : int, optional
+        Stride of the pooling. Default is the same as the kernel size.
+    pool_type : str, optional
+        Type of the pooling. Can be "max" or "avg". Default is "max".
+    """
+
+    def __init__(
+        self,
+        kernel_size: int,
+        stride: int = None,
+        pool_type: str = "max",
+    ):
+        self.kernel_size = kernel_size
+        self.stride = stride or kernel_size
+        self.pool_type = pool_type
+
+    def __str__(self) -> str:
+        return f"{self.pool_type.capitalize()}Pool({self.kernel_size}, {self.stride})"
+
+    def output_shape(self, input_size):
+        output_height = (input_size[0] - self.kernel_size) // self.stride + 1
+        output_width = (input_size[1] - self.kernel_size) // self.stride + 1
+        return (output_height, output_width, input_size[2])
+
+
+class FCLayer(Layer):
+    """
+    Represents a fully connected (dense) layer in a CNN.
+
+    Parameters
+    ----------
+    in_features : int
+        Number of input features.
+    out_features : int
+        Number of output features.
+    """
+
+    def __init__(self, in_features: int, out_features: int):
+        self.in_features = in_features
+        self.out_features = out_features
+
+    def __str__(self) -> str:
+        return f"FC({self.in_features}, {self.out_features})"
+
+    def output_shape(self, input_size: int) -> tuple[int, int]:
+        return (1, self.out_features)

easy_viz_cnn/models.py

@@ -0,0 +1,76 @@
+from matplotlib.figure import Figure
+import matplotlib.pyplot as plt
+
+from easy_viz_cnn.layers import Layer
+
+
+class SimpleCNNModel:
+    """
+    Represents a simple CNN model composed of various layers.
+
+    Parameters
+    ----------
+    layers : list[Layer]
+        List of layers in the model.
+    """
+
+    def __init__(self, layers: list[Layer]):
+        self.layers = layers
+
+    def __str__(self) -> str:
+        return " -> ".join([str(layer) for layer in self.layers])
+
+    def features(self):
+        input_size = None
+        for layer in self.layers:
+            input_size = layer.output_shape(input_size)
+            yield input_size
+
+    def create_figure(
+        self,
+        title: str = "CNN Architecture",
+        title_fontsize: int = 12,
+        layer_color: str = "lightblue",
+        layer_fontsize: int = 10,
+    ) -> Figure:
+        """
+        Plot the architecture of the CNN model.
+
+        Parameters
+        ----------
+        title : str, optional
+            Title of the plot. Default is "CNN Architecture".
+        title_fontsize : int, optional
+            Font size of the title. Default is 12.
+        layer_color : str, optional
+            Background color of the layer boxes. Default is "lightblue".
+        layer_fontsize : int, optional
+            Font size of the layer labels. Default is 10.
+        """
+
+        fig, ax = plt.subplots(figsize=(8, 6))
+        ax.set_title(title, fontsize=title_fontsize, pad=5)
+        ax.axis("off")
+
+        x_offset = 0.5
+        y_offset = 0.80  # Start lower to fit more layers
+        layer_height = 0.12  # Larger layer height
+        layer_spacing = layer_height + 0.02  # Tight spacing between layers
+
+        for layer in self.layers:
+            ax.text(
+                x_offset,
+                y_offset,
+                str(layer),
+                ha="center",
+                va="center",
+                fontsize=layer_fontsize,
+                bbox=dict(
+                    facecolor=layer_color,
+                    edgecolor="black",
+                    boxstyle=f"round,pad={layer_height * 5}",
+                ),
+            )
+            y_offset -= layer_spacing
+
+        return fig