Test.py

import torch
import lightning as L
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader,Dataset
import numpy as np
from Model.model import My_model
from BUSI import *
from UDIAT import *

device = 'cuda' if torch.cuda.is_available() else 'cpu'

def iou_score(y_pred, y_true):
    smooth = 1e-5
    y_pred = torch.sigmoid(y_pred)

    y_pred = y_pred.data.cpu().numpy()
    y_true = y_true.data.cpu().numpy()

    y_pred = y_pred > 0.5
    y_true = y_true > 0.5
    intersection = (y_pred & y_true).sum()
    union = (y_pred | y_true).sum()

    return (intersection + smooth) / (union + smooth)

def dice_score(y_pred, y_true, smooth=0.):

    y_pred = torch.sigmoid(y_pred)

    y_pred = y_pred.view(-1)
    y_true = y_true.view(-1)
    intersection = (y_pred * y_true).sum()

    return (2. * intersection + smooth) / (y_pred.sum() + y_true.sum() + smooth)

model = My_model().to(device)

class Segmentor(L.LightningModule):
    def __init__(self, model=model):
        super().__init__()
        self.model = model

    def forward(self, x):
        return self.model(x)

    def test_step(self, batch, batch_idx):
        image, y_true = batch
        y_pred = self.model(image)

        # Calculate Dice and IoU scores
        dice = dice_score(y_pred, y_true)
        iou = iou_score(y_pred, y_true)

        # Compute TP, FP, FN, TN for Precision, Recall, and F-Score
        y_pred_bin = (y_pred > 0.5).float()  # Threshold predictions to binary

        TP = (y_pred_bin * y_true).sum(dim=(1, 2, 3))
        FP = ((y_pred_bin == 1) & (y_true == 0)).sum(dim=(1, 2, 3))
        FN = ((y_pred_bin == 0) & (y_true == 1)).sum(dim=(1, 2, 3))

        # Precision, Recall, and F-Score
        precision = TP / (TP + FP + 1e-8)
        recall = TP / (TP + FN + 1e-8)
        f_score = 2 * (precision * recall) / (precision + recall + 1e-8)

        # Average metrics over the batch
        precision_mean = precision.mean().item()
        recall_mean = recall.mean().item()
        f_score_mean = f_score.mean().item()

        # Log all metrics
        metrics = {
            "Test Dice": dice,
            "Test IoU": iou,
            "Test Precision": precision_mean,
            "Test Recall": recall_mean,
            "Test F-Score": f_score_mean,
        }
        self.log_dict(metrics, prog_bar=True)
        return metrics
    
def visualize_prediction_VT(model, dataset):
  model.eval()
  plt.figure(figsize=(6, 2*5), layout='compressed')
  x1, y1 = dataset[20]
  x2, y2 = dataset[21]
  x3, y3 = dataset[22]
  x4, y4 = dataset[23]
  x5, y5 = dataset[24]
  x = torch.cat((x1.unsqueeze(0),x2.unsqueeze(0),x3.unsqueeze(0),x4.unsqueeze(0),x5.unsqueeze(0))).cuda()
  y = torch.cat((y1.unsqueeze(0),y2.unsqueeze(0),y3.unsqueeze(0),y4.unsqueeze(0),y5.unsqueeze(0))).cuda()
  # x = x.unsqueeze(0)
  # y = yk.unsqueeze(1)

  y_pred = model(x).data.squeeze()
  y_pred[y_pred>0.5] = 1
  y_pred[y_pred<=0.5] = 0
  y_pred = y_pred.cpu().numpy()
  for i in range(5):
            # convert torch to numpy
            xa = x[i]
            xa = xa.permute(1, 2, 0).cpu().numpy()
            ya = y[i]
            ya = ya.permute(1, 2, 0).cpu().numpy()
                    # visualization
            plt.subplot(5, 3, 3*i + 1)
            plt.title(f"Image {i+1}")
            plt.imshow(xa)
            plt.axis('off')

            plt.subplot(5, 3, 3*i + 2)
            plt.title(f"Ground truth {i+1}")
            plt.imshow(ya, cmap='gray')
            plt.axis('off')

            plt.subplot(5, 3, 3*i + 3)
            plt.title(f"SK {i+1}")
            plt.imshow(y_pred[i], cmap='gray')
            plt.axis('off')

CHECKPOINT_PATH = '/content/drive/MyDrive/Weights/Ultrasound/Version53/BUSI/ckpt0.8288.ckpt'

# Prediction
trainer = L.Trainer()
segmentor = Segmentor.load_from_checkpoint(CHECKPOINT_PATH, model = model)
segmentor.eval()
trainer.test(segmentor, testloader_busi)

device = 'cuda'
# Load the Segmentor LightningModule directly from the checkpoint.
# The 'model' argument provides the class to instantiate the actual model inside Segmentor.
segmentor_loaded = Segmentor.load_from_checkpoint(CHECKPOINT_PATH, model=My_model().to(device))

# Extract the UNet model instance with loaded weights from the Segmentor
model_for_visualization = segmentor_loaded.model.to(device)

visualize_prediction_VT(model=model_for_visualization, dataset=BUSILoader(x_test, y_test, transform=test_transform))