forked from DQSSSSS/VectorNet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
visualization.py
70 lines (60 loc) · 1.99 KB
/
visualization.py
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
import torch
from configure import device
from dataloader import load_test, load_train
import matplotlib.pyplot as plt
vectorNet = torch.load('VectorNet-test.model').to(device)
def viz(x, y):
data = x.permute(1, 0, 2) # [vNumber, batch size, len]
id = data[0, :, 0].long()
print(id.item())
batchSize, len = data.shape[1], data.shape[2]
j = 1
for i in range(1, data.shape[0]):
if i == data.shape[0] - 1 or \
data[i, 0, len - 1] != data[i + 1, 0, len - 1]:
listX = []
listY = []
listX.append(data[j, 0, 0])
listY.append(data[j, 0, 1])
while j <= i:
listX.append(data[j, 0, 2])
listY.append(data[j, 0, 3])
j += 1
if data[i, 0, -1] == id.item():
plt.plot(listX, listY, 'r', linewidth=3) # agent
print(listX)
print(listY)
print('---------')
elif data[i, 0, 4] == 1:
plt.plot(listX, listY, 'black', linewidth=3) # others
else:
plt.plot(listX, listY, 'b') # lane
listX = [0]
listY = [0]
for i in range(0, y.shape[1], 2):
listX.append(y[0, i])
listY.append(y[0, i + 1])
print(listX)
print(listY)
print('---------')
plt.plot(listX, listY, 'g', linewidth=3) # ground truth
x = x.to(device)
myPredict = vectorNet(x)
listX = [0]
listY = [0]
for i in range(0, y.shape[1], 2):
listX.append(myPredict[0, i])
listY.append(myPredict[0, i + 1])
print(listX)
print(listY)
print('---------')
plt.plot(listX, listY, 'yellow', linewidth=3) # predict
plt.show()
if __name__ == '__main__':
data = load_train()
dataset = torch.utils.data.DataLoader(data, batch_size=1)
for data, y in dataset:
offset = data[:, -1, :] # [0, 0, 0, 0, 0, maxX, maxY, ..., 0]
data = data[:, 0:data.shape[1] - 1, :]
viz(data, y)
# exit(0)