UNet/test.py at master · Beokane/UNet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
from tqdm import tqdm
from PIL import Image

from archs import UNet
import params

config = params.parse_args()

def image_label_transform(image, label):
        image_transform = transforms.Compose([
            transforms.Resize((config.input_height, config.input_width)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[
                                 0.229, 0.224, 0.225])
        ])
        label_transform = transforms.Compose([
            transforms.Resize((config.input_height, config.input_width)),
            transforms.ToTensor(),
            transforms.Lambda(lambda x: ((x > 0) & (x < 255)).float())
        ])
        return image_transform(image), label_transform(label)

def main():
    # 定义数据预处理和加载数据
    test_image_transform = transforms.Compose([
        transforms.Resize((config.input_height, config.input_width)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    test_label_transform = transforms.Compose([
        transforms.Resize((config.input_height, config.input_width)),
        transforms.ToTensor(),
        transforms.Lambda(lambda x: ((x > 0) & (x < 255)).float())
    ])

    test_dataset = datasets.VOCSegmentation(
        root='./data', year='2007', image_set='test', download=False, transform=test_image_transform, target_transform=test_label_transform)
    test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=4)
    model = UNet(num_classes=1, input_channels=3)
    model.load_state_dict(torch.load('./models/model.pth'), strict=False)
    model = model.cuda()
    model.eval()
    os.makedirs('./output', exist_ok=True)
    test_filenames = test_dataset.images
    with torch.no_grad():
        with tqdm(test_loader, desc='Testing', unit='batch') as pbar:
            for i, (image, label) in enumerate(pbar):
                image = image.to('cuda')
                output = model(image)
                probs = torch.sigmoid(output)
                preds = (probs > 0.5).float()
                preds_numpy = preds.squeeze(1).cpu().numpy()
                label_numpy = label.squeeze(1).cpu().numpy()
                for j in range(preds_numpy.shape[0]):
                    pred_image = Image.fromarray((preds_numpy[j] * 255).astype('uint8'))
                    label_image = Image.fromarray((label_numpy[j] * 255).astype('uint8'))
                    file_name = os.path.basename(test_filenames[i])
                    pred_image.save(f'./output/{file_name}_pred.png')
                    label_image.save(f'./output/{file_name}_label.png')
                torch.cuda.empty_cache()
        print('\nfinish test')

if __name__ == '__main__':
    main()