# ========================================================= 
# Image Colorization using Transfer Learning (ResNet50) 
# ========================================================= 
import tensorflow as tf 
from tensorflow.keras import layers, Model 
from tensorflow.keras.applications import ResNet50 
import numpy as np 
import cv2 
import os 
from skimage.color import rgb2lab, lab2rgb 
import matplotlib.pyplot as plt 
from sklearn.model_selection import train_test_split 
# ========================================================= 
# PARAMETERS 
# ========================================================= 
IMG_SIZE = 224 
BATCH_SIZE = 8 
EPOCHS = 20 
DATASET_PATH = "dataset/"   # folder containing RGB images 
# ========================================================= 
# PREPROCESSING FUNCTION 
# ========================================================= 
def preprocess_image(img): 
img = cv2.resize(img, (IMG_SIZE, IMG_SIZE)) 
img = img.astype("float32") / 255 
 
    lab = rgb2lab(img) 
 
    L = lab[:,:,0] / 100 
    ab = lab[:,:,1:] / 128 
 
    L = np.expand_dims(L, axis=-1) 
 
    return L, ab 
 
 
# ========================================================= 
# LOAD DATASET 
# ========================================================= 
 
X_L = [] 
Y_ab = [] 
 
for file in os.listdir(DATASET_PATH): 
 
    path = os.path.join(DATASET_PATH, file) 
 
    img = cv2.imread(path) 
 
    if img is None: 
        continue 
 
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 
 
    L, ab = preprocess_image(img) 
 
    X_L.append(L) 
    Y_ab.append(ab) 
 
X_L = np.array(X_L) 
Y_ab = np.array(Y_ab) 
 
print("Dataset shape:", X_L.shape) 
 
# ========================================================= 
# TRAIN TEST SPLIT 
# ========================================================= 
 
X_train, X_test, y_train, y_test = train_test_split( 
    X_L, 
    Y_ab, 
    test_size=0.2, 
    random_state=42 
) 
 
# ========================================================= 
# BUILD ENCODER (TRANSFER LEARNING) 
# ========================================================= 
 
def build_encoder(): 
 
    base_model = ResNet50( 
        weights="imagenet", 
        include_top=False, 
        input_shape=(IMG_SIZE, IMG_SIZE, 3) 
    ) 
 
    for layer in base_model.layers: 
        layer.trainable = False 
 
    return base_model 
 
 
# ========================================================= 
# BUILD COLORIZATION MODEL 
# ========================================================= 
 
def build_model(): 
 
    encoder = build_encoder() 
 
    input_l = layers.Input(shape=(IMG_SIZE, IMG_SIZE, 1)) 
 
    # convert grayscale to 3 channels 
    x = layers.Concatenate()([input_l, input_l, input_l]) 
 
    features = encoder(x) 
 
    # Decoder 
    x = layers.Conv2D(256, 3, padding='same', activation='relu')(features) 
    x = layers.UpSampling2D((2,2))(x) 
 
    x = layers.Conv2D(128, 3, padding='same', activation='relu')(x) 
    x = layers.UpSampling2D((2,2))(x) 
 
    x = layers.Conv2D(64, 3, padding='same', activation='relu')(x) 
    x = layers.UpSampling2D((2,2))(x) 
 
    x = layers.Conv2D(32, 3, padding='same', activation='relu')(x) 
    x = layers.UpSampling2D((2,2))(x) 
 
    x = layers.Conv2D(16, 3, padding='same', activation='relu')(x) 
 
    output_ab = layers.Conv2D(2, 1, activation='tanh')(x) 
 
    model = Model(inputs=input_l, outputs=output_ab) 
 
    return model 
 
 
# ========================================================= 
# CREATE MODEL 
# ========================================================= 
 
model = build_model() 
 
model.compile( 
    optimizer='adam', 
    loss='mse' 
) 
 
model.summary() 
 
# ========================================================= 
# TRAIN MODEL 
# ========================================================= 
 
history = model.fit( 
    X_train, 
    y_train, 
    validation_data=(X_test, y_test), 
    epochs=EPOCHS, 
    batch_size=BATCH_SIZE 
) 
 
# ========================================================= 
# POST PROCESSING 
# ========================================================= 
 
def reconstruct_image(L, ab): 
 
    L = L * 100 
    ab = ab * 128 
 
    lab = np.concatenate([L, ab], axis=-1) 
 
    rgb = lab2rgb(lab) 
 
    return rgb 
 
 
# ========================================================= 
# TEST MODEL 
# ========================================================= 
 
pred_ab = model.predict(X_test) 
 
index = 0 
 
pred_img = reconstruct_image(X_test[index], pred_ab[index]) 
plt.figure(figsize=(8,4)) 
plt.subplot(1,2,1) 
plt.title("Grayscale") 
plt.imshow(X_test[index].squeeze(), cmap='gray') 
plt.axis("off") 
plt.subplot(1,2,2) 
plt.title("Colorized") 
plt.imshow(pred_img) 
plt.axis("off") 
plt.show() 
# ========================================================= 
# SAVE MODEL 
# ========================================================= 
model.save("colorization_resnet_model.h5") 
print("Model saved successfully!") 