segmentation/showexample.py at master · lmycross/segmentation · 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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import torch
from torchvision import transforms
from torch.autograd import Variable
import tiramisu_nobias
import imageio


def decode_segmap(temp):
    Sky = [128, 128, 128]
    Building = [128, 0, 0]
    Pole = [192, 192, 128]
    Road = [128, 64, 128]
    Pavement = [60, 40, 222]
    Tree = [128, 128, 0]
    SignSymbol = [192, 128, 128]
    Fence = [64, 64, 128]
    Car = [64, 0, 128]
    Pedestrian = [64, 64, 0]
    Bicyclist = [0, 128, 192]
    Unlabelled = [0, 0, 0]


    label_colours = np.array([Sky, Building, Pole, Road,
                                  Pavement, Tree, SignSymbol, Fence, Car,
                                  Pedestrian, Bicyclist, Unlabelled])
    r = temp.copy()
    g = temp.copy()
    b = temp.copy()
    for l in range(12):
        r[temp == l] = label_colours[l, 0]
        g[temp == l] = label_colours[l, 1]
        b[temp == l] = label_colours[l, 2]

    rgb = np.zeros((temp.shape[0], temp.shape[1], 3))
    rgb[:, :, 0] = r
    rgb[:, :, 1] = g
    rgb[:, :, 2] = b

    return rgb

label_num=11
model = tiramisu_nobias.FCDenseNet103(label_num)
model.load_state_dict(torch.load('C:/Users/mchiwml4/pycode/net-data/segmentation/camvid/tiramisu2/net_params65.pth'))

model=model
model.eval()


truthimagedir='C:/Users/mchiwml4/github/SegNet-Tutorial/CamVid/testannot/Seq05VD_f00660.png'

truth=np.asarray(Image.open(truthimagedir))
truth = decode_segmap(truth)

truth=np.uint8(truth)


imagedir='C:/Users/mchiwml4/github/SegNet-Tutorial/CamVid/test/Seq05VD_f00660.png'
inputs = imageio.imread(imagedir)
transform1 = transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.41189488770418226, 0.4251328066237724, 0.432670702070482),(0.3070734955953852, 0.3116110784489235, 0.3072184293428751))
                ])
img=transform1(inputs)
img=img.unsqueeze(0)
img= Variable(img,requires_grad=False)

outputs=model(img)

pred = outputs.data.max(1)[1].numpy()
pred=pred.squeeze(0)
pre = decode_segmap(pred)
pre=np.uint8(pre)


plt.figure()
plt.imshow(inputs)
plt.axis('off')
plt.figure()
plt.imshow(truth)
plt.axis('off')
plt.figure()
plt.imshow(pre)
plt.axis('off')
plt.show()