joegue / ca-net Goto Github PK

    import os
    import cv2
    import time
    import torch
    import imageio
    import numpy as np
    from glob import glob
    from tqdm import tqdm
    from operator import add
    from sklearn.metrics import (
        jaccard_score, f1_score, recall_score, precision_score, accuracy_score, fbeta_score)
    from binary import assd
    from parameters import *
    from msu import U_Net
    from utils import create_dir, seeding, make_channel_last
    from data import load_data
    from crf import apply_crf

    def calculate_metrics(y_true, y_pred):
      y_true = y_true.cpu().numpy()
      y_pred = y_pred.cpu().numpy()

      y_pred = y_pred > 0.5
      y_pred = y_pred.reshape(-1)
      y_pred = y_pred.astype(np.uint8)

      y_true = y_true > 0.5
      y_true = y_true.reshape(-1)
      y_true = y_true.astype(np.uint8)

    ## Score
    score_jaccard = jaccard_score(y_true, y_pred, average='binary')
    score_f1 = f1_score(y_true, y_pred, average='binary')
    score_recall = recall_score(y_true, y_pred, average='binary')
    score_precision = precision_score(y_true, y_pred, average='binary', zero_division=1)
    score_acc = accuracy_score(y_true, y_pred)
    score_fbeta = fbeta_score(y_true, y_pred, beta=1.0, average='binary', zero_division=1)

    return [score_jaccard, score_f1, score_recall, score_precision, score_acc, score_fbeta]

    if __name__ == "__main__":

    """ Load dataset """
    test_x = sorted(glob(os.path.join(root_path, "test/images", "*.png")))
    test_y = sorted(glob(os.path.join(root_path, "test/masks", "*.png")))

    """ Load the checkpoint """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = U_Net()
    model = model.to(device)
    model.load_state_dict(torch.load('checkpoint.pth', map_location=device))
    model.eval()

    """ Testing """
    metrics_score = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

    for i, (x, y) in enumerate(zip(test_x, test_y)):
        name = y.split("/")[-1].split(".")[0]

        ## Image
        image = cv2.imread(x, cv2.IMREAD_GRAYSCALE)
        image1 = cv2.resize(image, size)
        ori_img1 = image1
        image1 = np.expand_dims(image1, axis=0)
        #image1 = np.transpose(image1, (2, 0, 1))
        image1 = image1/255.0
        image1 = np.expand_dims(image1, axis=0)
        image1 = image1.astype(np.float32)
        image1 = torch.from_numpy(image1)
        image1 = image1.to(device)

        ## Mask
        mask = cv2.imread(y, cv2.IMREAD_GRAYSCALE)
        mask1 = cv2.resize(mask, size)
        ori_mask1 = mask1
        mask1 = np.expand_dims(mask1, axis=0)
        mask1 = mask1/255.0
        mask1 = np.expand_dims(mask1, axis=0)
        mask1 = mask1.astype(np.float32)
        mask1 = torch.from_numpy(mask1)
        mask1 = mask1.to(device)

        with torch.no_grad():
            pred_y1 = torch.sigmoid(model(image1))

            """ Evaluation metrics """
            score = calculate_metrics(mask1, pred_y1)
            metrics_score = list(map(add, metrics_score, score))

    jaccard = metrics_score[0]/len(test_x)
    f1 = metrics_score[1]/len(test_x)
    recall = metrics_score[2]/len(test_x)
    precision = metrics_score[3]/len(test_x)
    acc = metrics_score[4]/len(test_x)
    f2 = metrics_score[5]/len(test_x)

    print(f"Jaccard: {jaccard:1.4f} - F1: {f1:1.4f} - Recall: {recall:1.4f} - Precision: {precision:1.4f} - Acc: {acc:1.4f} - F2: {f2:1.4f}")