-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest.py
308 lines (250 loc) · 13.6 KB
/
test.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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
import torch
import torch.nn as nn
# import image_utils
import utils.image_utils
from utils.image_utils import sampling, image2world,get_patch_grad
from utils.kmeans import kmeans
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import cv2
def torch_multivariate_gaussian_heatmap(coordinates, H, W, dist, sigma_factor, ratio, device, rot=False):
"""
Create Gaussian Kernel for CWS
"""
ax = torch.linspace(0, H, H, device=device) - coordinates[1]
ay = torch.linspace(0, W, W, device=device) - coordinates[0]
xx, yy = torch.meshgrid([ax, ay])
meshgrid = torch.stack([yy, xx], dim=-1)
radians = torch.atan2(dist[0], dist[1])
c, s = torch.cos(radians), torch.sin(radians)
R = torch.Tensor([[c, s], [-s, c]]).to(device)
if rot:
R = torch.matmul(torch.Tensor([[0, -1], [1, 0]]).to(device), R)
dist_norm = dist.square().sum(-1).sqrt() + 5 # some small padding to avoid division by zero
conv = torch.Tensor([[dist_norm / sigma_factor / ratio, 0], [0, dist_norm / sigma_factor]]).to(device)
conv = torch.square(conv)
T = torch.matmul(R, conv)
T = torch.matmul(T, R.T)
kernel = (torch.matmul(meshgrid, torch.inverse(T)) * meshgrid).sum(-1)
kernel = torch.exp(-0.5 * kernel)
return kernel / kernel.sum()
def evaluate(model, val_loader, val_images, num_goals, num_traj, obs_len, batch_size, device, input_template, waypoints, resize, temperature,exp_name, use_TTST=False, use_CWS=False, rel_thresh=0.002, CWS_params=None, dataset_name=None, homo_mat=None, mode='val'):
"""
:param model: torch model
:param val_loader: torch dataloader
:param val_images: dict with keys: scene_name value: preprocessed image as torch.Tensor
:param num_goals: int, number of goals
:param num_traj: int, number of trajectories per goal
:param obs_len: int, observed timesteps
:param batch_size: int, batch_size
:param device: torch device
:param input_template: torch.Tensor, heatmap template
:param waypoints: number of waypoints
:param resize: resize factor
:param temperature: float, temperature to control peakiness of heatmap
:param use_TTST: bool
:param use_CWS: bool
:param rel_thresh: float
:param CWS_params: dict
:param dataset_name: ['sdd','ind','eth']
:param params: dict with hyperparameters
:param homo_mat: dict with homography matrix
:param mode: ['val', 'test']
:return: val_ADE, val_FDE for one epoch
"""
model.eval()
val_ADE = []
val_FDE = []
counter = 0
with torch.no_grad():
# outer loop, for loop over each scene as scenes have different image size and to calculate segmentation only once
for trajectory, meta, scene in val_loader:
# Get scene image and apply semantic segmentation
scene_image = val_images[scene].to(device).unsqueeze(0)
for i in range(0, len(trajectory), batch_size):
# Create Heatmaps for past and ground-truth future trajectories
_, _, H, W = scene_image.shape
observed = trajectory[i:i+batch_size, :obs_len, :].to(device)
observed_map = get_patch_grad(observed, H, W,device=device)
observed_map = torch.stack(observed_map).reshape([-1, obs_len, H, W])
gt_future = trajectory[i:i+batch_size, obs_len:].to(device)
semantic_image = scene_image.expand(observed_map.shape[0], -1, -1, -1)
# Forward pass
# Calculate features
feature_input = torch.cat([semantic_image, observed_map], dim=1)
features = model.pred_feature(feature_input)
# Predict goal and waypoint probability distributions
pred_waypoint_map = model.pred_end(features,w=W,h=H)
pred_waypoint_map = pred_waypoint_map[:, waypoints]
pred_waypoint_map_sigmoid = pred_waypoint_map / temperature
pred_waypoint_map_sigmoid = model.sigmoid(pred_waypoint_map_sigmoid)
################################################ TTST ##################################################
if dataset_name =="sdd":
sample_scene = ((1-semantic_image[:, 3])/2)+(semantic_image[:, 1]/2)
else:
sample_scene = semantic_image[:, 1]
if use_TTST:
# TTST Begin
# sample a large amount of goals to be clustered
goal_samples = sampling(pred_waypoint_map_sigmoid[:, -1:], num_samples=10000, replacement=True, rel_threshold=rel_thresh,scene=sample_scene)
goal_samples = goal_samples.permute(2, 0, 1, 3)
num_clusters = num_goals - 1
goal_samples_softargmax = model.softargmax(pred_waypoint_map[:, -1:]) # first sample is softargmax sample
# Iterate through all person/batch_num, as this k-Means implementation doesn't support batched clustering
goal_samples_list = []
for person in range(goal_samples.shape[1]):
goal_sample = goal_samples[:, person, 0]
# Actual k-means clustering, Outputs:
# cluster_ids_x - Information to which cluster_idx each point belongs to
# cluster_centers - list of centroids, which are our new goal samples
cluster_ids_x, cluster_centers = kmeans(X=goal_sample, num_clusters=num_clusters, distance='euclidean', device=device, tqdm_flag=False, tol=0.001, iter_limit=1000)
goal_samples_list.append(cluster_centers)
goal_samples = torch.stack(goal_samples_list).permute(1, 0, 2).unsqueeze(2)
goal_samples = torch.cat([goal_samples_softargmax.unsqueeze(0), goal_samples], dim=0)
# TTST End
else:
goal_samples = sampling(pred_waypoint_map_sigmoid[:, -1:],scene=sample_scene, num_samples=num_goals)
goal_samples = goal_samples.permute(2, 0, 1, 3)
# Predict waypoints:
# in case len(waypoints) == 1, so only goal is needed (goal counts as one waypoint in this implementation)
if len(waypoints) == 1:
waypoint_samples = goal_samples
if use_CWS and len(waypoints) > 1:
sigma_factor = CWS_params['sigma_factor']
ratio = CWS_params['ratio']
rot = CWS_params['rot']
goal_samples = goal_samples.repeat(num_traj, 1, 1, 1) # repeat K_a times
last_observed = trajectory[i:i+batch_size, obs_len-1].to(device) # [N, 2]
waypoint_samples_list = [] # in the end this should be a list of [K, N, # waypoints, 2] waypoint coordinates
for g_num, waypoint_samples in enumerate(goal_samples.squeeze(2)):
waypoint_list = [] # for each K sample have a separate list
waypoint_list.append(waypoint_samples)
for waypoint_num in reversed(range(len(waypoints)-1)):
distance = last_observed - waypoint_samples
gaussian_heatmaps = []
traj_idx = g_num // num_goals # idx of trajectory for the same goal
for dist, coordinate in zip(distance, waypoint_samples): # for each person
length_ratio = 1 / (waypoint_num + 2)
gauss_mean = coordinate + (dist * length_ratio) # Get the intermediate point's location using CV model
sigma_factor_ = sigma_factor - traj_idx
gaussian_heatmaps.append(torch_multivariate_gaussian_heatmap(gauss_mean, H, W, dist, sigma_factor_, ratio, device, rot))
gaussian_heatmaps = torch.stack(gaussian_heatmaps) # [N, H, W]
waypoint_map_before = pred_waypoint_map_sigmoid[:, waypoint_num]
waypoint_map = waypoint_map_before * gaussian_heatmaps
# normalize waypoint map
waypoint_map = (waypoint_map.flatten(1) / waypoint_map.flatten(1).sum(-1, keepdim=True)).view_as(waypoint_map)
# For first traj samples use softargmax
if g_num // num_goals == 0:
# Softargmax
waypoint_samples = model.softargmax(waypoint_map.unsqueeze(0))
waypoint_samples = waypoint_samples.squeeze(0)
else:
waypoint_samples = sampling(waypoint_map.unsqueeze(1), num_samples=1, rel_threshold=0.05,scene=sample_scene)
waypoint_samples = waypoint_samples.permute(2, 0, 1, 3)
waypoint_samples = waypoint_samples.squeeze(2).squeeze(0)
waypoint_list.append(waypoint_samples)
waypoint_list = waypoint_list[::-1]
waypoint_list = torch.stack(waypoint_list).permute(1, 0, 2) # permute back to [N, # waypoints, 2]
waypoint_samples_list.append(waypoint_list)
waypoint_samples = torch.stack(waypoint_samples_list)
# CWS End
# If not using CWS, and we still need to sample waypoints (i.e., not only goal is needed)
elif not use_CWS and len(waypoints) > 1:
waypoint_samples = sampling(pred_waypoint_map_sigmoid[:, :-1], scene=sample_scene,num_samples=num_goals * num_traj)
waypoint_samples = waypoint_samples.permute(2, 0, 1, 3)
goal_samples = goal_samples.repeat(num_traj, 1, 1, 1) # repeat K_a times
waypoint_samples = torch.cat([waypoint_samples, goal_samples], dim=2)
future_samples = []
for waypoint in waypoint_samples:
# waypoint_map = get_patch(input_template, waypoint.reshape(-1, 2).cpu().numpy(), H, W)
waypoint_map = get_patch_grad(waypoint.reshape(-1, 2).cpu().numpy(), H, W,device=device)
waypoint_map = torch.stack(waypoint_map).reshape([-1, len(waypoints), H, W])
waypoint_maps_downsampled = nn.AvgPool2d(kernel_size=2, stride=2)(waypoint_map)
pred_traj_map = model.pred_traj(features,waypoint_maps_downsampled,w=W,h=H)
pred_traj = model.softargmax(pred_traj_map)
future_samples.append(pred_traj)
future_samples = torch.stack(future_samples)
#畫圖#####################################################################################################
draw_fig(future_samples=future_samples,observed=observed,gt_future=gt_future,semantic_image=sample_scene,scene_name=scene,exp_name=exp_name)
#畫圖#####################################################################################################
# waypoint_map = get_patch_grad(waypoint_samples.reshape(-1, 2).cpu().numpy(), H, W,device=device)
# # waypoint_map = get_patch(input_template, waypoint_samples.reshape(-1, 2).cpu().numpy(), H, W)
# waypoint_map = torch.stack(waypoint_map).reshape([-1, len(waypoints), H, W])
# waypoint_map_downsampled = nn.AvgPool2d(kernel_size=2, stride=2)(waypoint_map)
# # plt.imshow(waypoint_map[0,0].cpu())
# # plt.axis('off')
# # plt.savefig(f"/home/ycchen/Desktop/ycchen_stuff/bbbb.png")
# pred_traj_map = model.pred_traj(features,waypoint_map_downsampled,w=W,h=H)
# pred_traj_map = model.sigmoid(pred_traj_map)
# pred_traj = model.softargmax(pred_traj_map)
# pred_traj_map_soft = model.softmax(pred_traj_map)
# future_samples = pred_traj
# future_samples = torch.stack(future_samples).unsqueeze(0)
# fig, axes = plt.subplots(1, 12, figsize=(12, 6))
# for i in range(12):
# axes[i].imshow(pred_traj_map[0, i].cpu())
# axes[i].scatter(x=gt_future[0,i,0].cpu(),y=gt_future[0,i,1].cpu(),c='r',s=1)
# axes[i].scatter(x=pred_traj[0,i,0].cpu(),y=pred_traj[0,i,1].cpu(),c='b',s=1)
# axes[i].set_title(f"Pred: {i}")
# axes[i].axis('off')
# plt.savefig(f"/home/ycchen/Desktop/ycchen_stuff/MyModel/save_image/pred.png")
# plt.clf()
# converts ETH/UCY pixel coordinates back into world-coordinates
# converts ETH/UCY pixel coordinates back into world-coordinates
if dataset_name == 'eth':
waypoint_samples = image2world(waypoint_samples, scene, homo_mat, resize)
future_samples = image2world(future_samples, scene, homo_mat, resize)
# pred_traj = image2world(pred_traj, scene, homo_mat, resize)
gt_future = image2world(gt_future, scene, homo_mat, resize)
gt_goal = gt_future[:, -1:]
val_FDE.append(((((gt_goal - waypoint_samples[:, :, -1:])) ** 2).sum(dim=3) ** 0.5).min(dim=0)[0])
val_ADE.append(((((gt_future - future_samples)) ** 2).sum(dim=3) ** 0.5).mean(dim=2).min(dim=0)[0])
else:
gt_goal = gt_future[:, -1:]
val_FDE.append(((((gt_goal - waypoint_samples[:, :, -1:])/ resize) ** 2).sum(dim=3) ** 0.5).min(dim=0)[0])
val_ADE.append(((((gt_future - future_samples)/ resize) ** 2).sum(dim=3) ** 0.5).mean(dim=2).min(dim=0)[0])
val_ADE = torch.cat(val_ADE).mean()
val_FDE = torch.cat(val_FDE).mean()
return val_ADE.item(), val_FDE.item()
def draw_fig(future_samples,gt_future,observed,semantic_image,scene_name,exp_name):
colors = cm.rainbow(np.linspace(0, 1, future_samples.shape[0]))
if "sdd" in exp_name:
frame = cv2.imread(f"./data/SDD/test/{scene_name}/reference.jpg")
else:
name = exp_name.split("_")[0]
frame = cv2.imread(f"/home/ycchen/Desktop/ycchen_stuff/MyModel/data/refenence_map/{name}.png")
mask = np.uint8(255 * semantic_image[1].cpu())
mask = cv2.applyColorMap(mask, cv2.COLORMAP_JET)
# 找場景圖
frames = {"frame": frame}
utils.image_utils.resize(frames, factor=0.25, seg_mask=True)
utils.image_utils.pad(frames)
frame = frames["frame"]
# frame = cv2.resize(frame, (192,160), interpolation=cv2.INTER_NEAREST)
frame = cv2.addWeighted(frame, 0.7, mask, 0.3, 0)
plt.imshow(frame)
# plt.imshow(semantic_image[0, 0].cpu())
predd = []
agg = 0
for i, c in enumerate(colors):
traj_plt = plt.scatter(x=future_samples[i, agg, :-1, 0].cpu(), y=future_samples[i, agg, :-1, 1].cpu(), color=c)
fin_plt = plt.scatter(x=future_samples[i, agg, -1, 0].cpu(), y=future_samples[i, agg, -1, 1].cpu(), marker="*",
color=c)
predd.append(traj_plt)
# predd.append(fin_plt)
gt_past = plt.scatter(x=observed[agg, :, 0].cpu(), y=observed[agg, :, 1].cpu(), marker='x', c="b")
gt_fur = plt.scatter(x=gt_future[0, :, 0].cpu(), y=gt_future[0, :, 1].cpu(), marker='x', c="g")
# gt_end = plt.scatter(x=gt_future[0, -1, 0].cpu(), y=gt_future[0, -1, 1].cpu(), marker='x', c="g")
predd.append(gt_past)
predd.append(gt_fur)
# predd.append(gt_end)
plt.legend(predd,
["pred_1 trajectory", "pred_2 trajectory", "pred_3 trajectory", "pred_4 trajectory", "pred_5 trajectory",
"observed trajectory","future trajectory"],loc=2)
# plt.legend(predd,
# ["pred_1 endpoint", "pred_2 endpoint", "pred_3 endpoint", "pred_4 endpoint", "pred_5 endpoint",
# "future endpoint"])
plt.axis('off')
plt.savefig(f"/home/ycchen/Desktop/ycchen_stuff/MyModel/save_image/{exp_name}_pred_real.png")
plt.clf()