|
|
from glob import glob |
|
|
from tqdm import tqdm, trange |
|
|
import numpy as np |
|
|
import mmcv |
|
|
import pickle |
|
|
from sklearn.cluster import KMeans |
|
|
|
|
|
def load_and_process_traj(ann_file): |
|
|
|
|
|
data_infos = mmcv.load(ann_file, file_format='pkl') |
|
|
ego_trajs = [] |
|
|
map_locs = [] |
|
|
|
|
|
for data in tqdm(data_infos): |
|
|
if np.sum(np.array(data["gt_fut_bbox_sdc_mask"][0, :8],dtype=np.float32))==8: |
|
|
traj = data["gt_fut_bbox_sdc_lidar"][0, :8, :2] |
|
|
ego_trajs.append(traj) |
|
|
map_locs.append(data['map_location']) |
|
|
|
|
|
|
|
|
with open('/cpfs01/user/liuhaochen/AlgEngine_nuplan/traj_vocab/full_traj.pkl', 'wb') as writer: |
|
|
pickle.dump([ego_trajs, map_locs], writer) |
|
|
|
|
|
def process_kmeans_vocab(traj_file, anchors=4096, loc=''): |
|
|
with open(traj_file, 'rb') as reader: |
|
|
[traj_data, locs] = pickle.load(reader) |
|
|
print(traj_data[0].shape) |
|
|
if loc=='': |
|
|
end_p = np.array([traj[-1,:2] for traj in traj_data]) |
|
|
else: |
|
|
end_p = [] |
|
|
for traj,l in zip(traj_data, locs): |
|
|
if l==loc: |
|
|
end_p.append(traj[-1,:2]) |
|
|
end_p = np.array(end_p) |
|
|
|
|
|
print(end_p.shape) |
|
|
|
|
|
kmeans = KMeans(n_clusters=anchors, verbose=True) |
|
|
kmeans.fit(end_p) |
|
|
print('fit end') |
|
|
centroids = kmeans.cluster_centers_ |
|
|
|
|
|
with open(f'/cpfs01/user/liuhaochen/AlgEngine_nuplan/traj_vocab/vocab_center_{anchors}.pkl', 'wb') as writer: |
|
|
pickle.dump(centroids, writer) |
|
|
|
|
|
print('processing the representative trajs...') |
|
|
rep_traj = [] |
|
|
rep_loc = [] |
|
|
for i in trange(centroids.shape[0]): |
|
|
centroid = centroids[i] |
|
|
dist_arg = np.argmin(np.linalg.norm(end_p - centroid[np.newaxis, :2], axis=1)) |
|
|
rep_traj.append(traj_data[dist_arg]) |
|
|
rep_loc.append(locs[dist_arg]) |
|
|
|
|
|
with open(f'/cpfs01/user/liuhaochen/AlgEngine_nuplan/traj_vocab/vocab_traj_{anchors}.pkl', 'wb') as writer: |
|
|
pickle.dump([rep_traj, rep_loc], writer) |
|
|
|
|
|
import matplotlib.pyplot as plt |
|
|
def visualization(traj_file, kmeans_files, suffix='4096'): |
|
|
with open(traj_file, 'rb') as reader: |
|
|
[traj_data, loc_data] = pickle.load(reader) |
|
|
|
|
|
with open(kmeans_files, 'rb') as reader: |
|
|
k_means_data = pickle.load(reader) |
|
|
|
|
|
sing_mask = np.array(loc_data)=='singapore' |
|
|
plt.figure() |
|
|
plt.scatter(k_means_data[:, 0], k_means_data[:, 1], c='orange', |
|
|
marker='*', s=5, zorder=3) |
|
|
plt.scatter(k_means_data[sing_mask, 0], k_means_data[sing_mask, 1], c='red', |
|
|
marker='*', s=5, zorder=4) |
|
|
|
|
|
for traj in tqdm(traj_data): |
|
|
plt.plot(traj[:, 0], traj[:, 1], color='navy', |
|
|
alpha=0.6, zorder=1) |
|
|
|
|
|
plt.savefig(f'/cpfs01/user/liuhaochen/AlgEngine_nuplan/traj_vocab/vocab_{suffix}.png') |
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
visualization( |
|
|
'/cpfs01/user/liuhaochen/AlgEngine_nuplan/traj_vocab/vocab_traj_8192.pkl', |
|
|
'/cpfs01/user/liuhaochen/AlgEngine_nuplan/traj_vocab/vocab_center_8192.pkl', |
|
|
'8192' |
|
|
) |
|
|
|
|
|
visualization( |
|
|
'/cpfs01/user/liuhaochen/AlgEngine_nuplan/traj_vocab/vocab_traj_4096.pkl', |
|
|
'/cpfs01/user/liuhaochen/AlgEngine_nuplan/traj_vocab/vocab_center_4096.pkl', |
|
|
'4096' |
|
|
) |
|
|
|
|
|
|
|
|
|
