test.py

import torch
from tqdm import tqdm
from torch.utils.data import DataLoader
from sklearn.metrics import jaccard_score
from utils_data import Dataset_Seg
from utils_model import UNet
from pathlib import Path
from utils import set_seed

def compute_iou_on_test(model, test_loader, device, threshold=0.5):
    model.eval()
    preds = []
    targets = []

    with torch.no_grad():
        for batch in tqdm(test_loader, desc="Evaluating IoU", unit="img"):
            images, true_masks = batch
            images = images.to(device=device, dtype=torch.float32, memory_format=torch.channels_last)
            true_masks = true_masks.to(device=device, dtype=torch.long)

            output = model(images)

            if model.n_classes == 1:
                probs = torch.sigmoid(output.squeeze(1))
                pred_masks = (probs > threshold).long()
                preds.extend(pred_masks.cpu().numpy().flatten())
                targets.extend(true_masks.cpu().numpy().flatten())
            else:
                probs = torch.softmax(output, dim=1)
                pred_masks = probs.argmax(dim=1)
                preds.extend(pred_masks.cpu().numpy().flatten())
                targets.extend(true_masks.cpu().numpy().flatten())

    iou = jaccard_score(targets, preds, average='binary' if model.n_classes == 1 else 'macro')
    print(f"IoU: {iou:.4f}")
    return iou


if __name__ == '__main__': 
    
    path_to_data = '/content/drive/MyDrive/archive'

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

    # Model
    model = UNet(n_channels=3, n_classes=1, bilinear=False)
    model.to(device=device)
    model.load_state_dict(torch.load(Path(path_to_data) / 'checkpoint_best.pth'), map_location=device)

    test_data = Dataset_Seg(path_to_data, 'test')
    loader_args = dict(batch_size=1, pin_memory=True)
    test_loader = DataLoader(test_data, shuffle=False, drop_last=False, **loader_args)

    iou_score = compute_iou_on_test(model, test_loader, device)
    print("Final IoU on test set:", iou_score)