#%%
import tensorflow as tf
import imageio
from skimage import color
from skimage.transform import resize
import matplotlib.pyplot as plt
import numpy as np

# 모델 로드
model = tf.keras.models.load_model('color.keras')

# 이미지 로드 및 전처리
harry = imageio.imread('./harry.jpg')
original_shape = harry.shape # 원본 이미지 사이즈 저장
harry_resized = resize(harry, (256, 256, 3))
harry_gray = np.mean(harry_resized, axis=-1, keepdims=True)
harry_gray_reshaped = harry_gray.reshape((1, 256, 256, 1))

# 모델 예측
test = model.predict(harry_gray_reshaped)

# 예측 결과 크기 조정
test_resized = resize(test[0], original_shape)

# 결과 표시
plt.imshow(test_resized)
plt.show()
# %%