import numpy as np
import matplotlib.pyplot as plt

# 입력 이미지 정의
input_image = np.array([
    [0, 1, 2],
    [3, 4, 5],
    [6, 7, 8]
])

# 커널(필터) 정의
kernel = np.array([
    [0, 1],
    [2, 3]
])

# 출력 이미지 초기화 (2x2)
output_shape = (input_image.shape[0] - kernel.shape[0] + 1, 
                input_image.shape[1] - kernel.shape[1] + 1)
output_image = np.zeros(output_shape)

# 컨볼루션 수행
for i in range(output_shape[0]):
    for j in range(output_shape[1]):
        region = input_image[i:i + kernel.shape[0], j:j + kernel.shape[1]]
        output_image[i, j] = np.sum( region * kernel )
    pass
pass

# 시각화
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
fs = 24       # 텍스트 폰트 크기
tick_fs = 18  # 눈금 폰트 크기
tick_color = 'gray'  # 눈금 색상

# 입력 이미지 시각화
axes[0].imshow(input_image, cmap='Blues', interpolation='none')
axes[0].set_title('Input', fontsize=18)
axes[0].set_xticks(range(input_image.shape[1]))
axes[0].set_yticks(range(input_image.shape[0]))
axes[0].tick_params(axis='both', labelsize=tick_fs, colors=tick_color)
for i, j in np.ndindex(input_image.shape):
    axes[0].text(j, i, f'{input_image[i, j]}', fontsize=fs, ha='center', va='center')

# 커널 시각화
axes[1].imshow(kernel, cmap='Blues', interpolation='none')
axes[1].set_title('Kernel', fontsize=18)
axes[1].set_xticks(range(kernel.shape[1]))
axes[1].set_yticks(range(kernel.shape[0]))
axes[1].tick_params(axis='both', labelsize=tick_fs, colors=tick_color)
for i, j in np.ndindex(kernel.shape):
    axes[1].text(j, i, f'{kernel[i, j]}', fontsize=fs, ha='center', va='center')

# 출력 이미지 시각화
axes[2].imshow(output_image, cmap='Blues', interpolation='none')
axes[2].set_title('Output', fontsize=18)
axes[2].set_xticks(range(output_image.shape[1]))
axes[2].set_yticks(range(output_image.shape[0]))
axes[2].tick_params(axis='both', labelsize=tick_fs, colors=tick_color)
for i, j in np.ndindex(output_image.shape):
    axes[2].text(j, i, f'{int(output_image[i, j])}', fontsize=fs, ha='center', va='center')

plt.tight_layout()
plt.show()