Test script? about pretrained-backbones-unet

[mberkay0]

You can use a script as shown below.

from backbones_unet.utils.dataset import SemanticSegmentationDataset
from torch.utils.data import DataLoader

test_img_path = './images'
# dataloader 
test_dataset = SemanticSegmentationDataset(test_img_path)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, pin_memory=True)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# probability map threshold
threshold = 0.5

# predictions
preds = []
for imgs in test_loader:
    with torch.no_grad():
        pred = model(imgs.to(device))
        pred[:, :][pred[:, :] >= threshold] = 1
        pred[:, :][pred[:, :] < threshold] = 0
    preds.append(pred)

# predicted masks
preds = torch.mean(torch.stack(preds, dim=0), dim=0).permute((0,2,3,1)).cpu().detach()

[bc-bytes]

Thank you for the quick reply!