test.py

import os
import numpy as np
import tensorflow as tf
from tqdm import tqdm
import cv2

from data import load_data, tf_dataset
from architecture import build_unet

H=256
W=256
num_classes=3

if __name__ == "__main__":
    """Seeding the environment"""
    np.random.seed(42)
    tf.random.set_seed(42)

    """To call the dataset"""
    path = "E:/MIC_Projects/U-Net_Implementation/"
    (train_x, train_y), (valid_x, valid_y), (test_x, test_y) = load_data(path)
    print(f"Dataset: Train:{len(train_x)} - Valid:{len(valid_x)} - Test:{len(test_x)}")

    """Model"""
    model=tf.keras.models.load_model("model.h5")

    """Saving the masks"""
    for x,y in tqdm(zip(test_x,test_y), total = len(test_x)):
        name = x.split("/")[-1]

        ## Read image
        x = cv2.imread(x, cv2.IMREAD_COLOR)
        x = cv2.resize(x, (W, H))
        x = x / 255.0
        x = x.astype(np.float32)

        ## Read mask
        y = cv2.imread(y, cv2.IMREAD_GRAYSCALE)
        y = cv2.resize(y, (W, H))  ## (256, 256)
        y = y - 1
        y = np.expand_dims(y, axis=-1)  ## (256, 256, 1)
        y = y * (255 / num_classes)
        y = y.astype(np.int32)
        y = np.concatenate([y, y, y], axis=2)

        ## Prediction
        p = model.predict(np.expand_dims(x, axis=0))[0]
        p = np.argmax(p, axis=-1)#to extract maximun index value
        p = np.expand_dims(p, axis=-1)
        p = p * (255 / num_classes)
        p = p.astype(np.int32)
        p = np.concatenate([p, p, p], axis=2)

        x = x * 255.0
        x = x.astype(np.int32)

        h, w, _ = x.shape
        line = np.ones((h, 10, 3)) * 255

        # print(x.shape, line.shape, y.shape, line.shape, p.shape)

        final_image = np.concatenate([x, line, y, line, p], axis=1)
        cv2.imwrite(f"results/{name}", final_image)