近期复现了QCNet:
https://github.com/ZikangZhou/QCNet
在源代码的基础上增加了轨迹预测可视化部分的代码，主要还是得了解Argoverse2 数据集及相关的API，可以参考这篇博客：
https://blog.csdn.net/m0_56423263/article/details/134593815
环境安装：
我的nvidia-smi:CUDA 11.3;nvcc -V:11.8
python3.8.16
torch2.0.0+cu118
torch-cluster1.6.3+pt20cu118
torch-scatter2.1.2+pt20cu118
pytorch-lightning2.1.2
pip install torch==2.0.0 torchvision==0.15.1 torchaudio==2.0.1 --index-url https://download.pytorch.org/whl/cu118
pip install torch-cluster -f https://data.pyg.org/whl/torch-2.0.0+cu118.html
pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu118.html
### 我的cudnn和torch 不匹配 让他直接用自己的cudnn
export LD_LIBRARY_PATH=/path/to/your/venv/lib:$LD_LIBRARY_PATH
qcnet.py
# Copyright (c) 2023, Zikang Zhou. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from itertools import chain
from itertools import compress
from pathlib import Path
from typing import Optional
import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.data import Batch
from torch_geometric.data import HeteroData
from losses import MixtureNLLLoss
from losses import NLLLoss
from metrics import Brier
from metrics import MR
from metrics import minADE
from metrics import minAHE
from metrics import minFDE
from metrics import minFHE
from modules import QCNetDecoder
from modules import QCNetEncoder
from visual import plot_single_vehicle
from test_zy import run_generate_scenario_visualizations
try:
from av2.datasets.motion_forecasting.eval.submission import ChallengeSubmission
except ImportError:
ChallengeSubmission = object
class QCNet(pl.LightningModule):
def __init__(self,
dataset: str,
input_dim: int,
hidden_dim: int,
output_dim: int,
output_head: bool,
num_historical_steps: int,
num_future_steps: int,
num_modes: int,
num_recurrent_steps: int,
num_freq_bands: int,
num_map_layers: int,
num_agent_layers: int,
num_dec_layers: int,
num_heads: int,
head_dim: int,
dropout: float,
pl2pl_radius: float,
time_span: Optional[int],
pl2a_radius: float,
a2a_radius: float,
num_t2m_steps: Optional[int],
pl2m_radius: float,
a2m_radius: float,
lr: float,
weight_decay: float,
T_max: int,
submission_dir: str,
submission_file_name: str,
**kwargs) -> None:
super(QCNet, self).__init__()
submission_dir='/mnt/ve_share2/zy/QCNet/submission'
self.save_hyperparameters()
self.dataset = dataset
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.output_dim = output_dim
self.output_head = output_head
self.num_historical_steps = num_historical_steps
self.num_future_steps = num_future_steps
self.num_modes = num_modes
self.num_recurrent_steps = num_recurrent_steps
self.num_freq_bands = num_freq_bands
self.num_map_layers = num_map_layers
self.num_agent_layers = num_agent_layers
self.num_dec_layers = num_dec_layers
self.num_heads = num_heads
self.head_dim = head_dim
self.dropout = dropout
self.pl2pl_radius = pl2pl_radius
self.time_span = time_span
self.pl2a_radius = pl2a_radius
self.a2a_radius = a2a_radius
self.num_t2m_steps = num_t2m_steps
self.pl2m_radius = pl2m_radius
self.a2m_radius = a2m_radius
self.lr = lr
self.weight_decay = weight_decay
self.T_max = T_max
self.submission_dir ='/mnt/ve_share2/zy/QCNet/submission' #submission_dir
self.submission_file_name = submission_file_name
self.encoder = QCNetEncoder(
dataset=dataset,
input_dim=input_dim,
hidden_dim=hidden_dim,
num_historical_steps=num_historical_steps,
pl2pl_radius=pl2pl_radius,
time_span=time_span,
pl2a_radius=pl2a_radius,
a2a_radius=a2a_radius,
num_freq_bands=num_freq_bands,
num_map_layers=num_map_layers,
num_agent_layers=num_agent_layers,
num_heads=num_heads,
head_dim=head_dim,
dropout=dropout,
)
self.decoder = QCNetDecoder(
dataset=dataset,
input_dim=input_dim,
hidden_dim=hidden_dim,
output_dim=output_dim,
output_head=output_head,
num_historical_steps=num_historical_steps,
num_future_steps=num_future_steps,
num_modes=num_modes,
num_recurrent_steps=num_recurrent_steps,
num_t2m_steps=num_t2m_steps,
pl2m_radius=pl2m_radius,
a2m_radius=a2m_radius,
num_freq_bands=num_freq_bands,
num_layers=num_dec_layers,
num_heads=num_heads,
head_dim=head_dim,
dropout=dropout,
)
self.reg_loss = NLLLoss(component_distribution=['laplace'] * output_dim + ['von_mises'] * output_head,
reduction='none')
self.cls_loss = MixtureNLLLoss(component_distribution=['laplace'] * output_dim + ['von_mises'] * output_head,
reduction='none')
self.Brier = Brier(max_guesses=6)
self.minADE = minADE(max_guesses=6)
self.minAHE = minAHE(max_guesses=6)
self.minFDE = minFDE(max_guesses=6)
self.minFHE = minFHE(max_guesses=6)
self.MR = MR(max_guesses=6)
self.test_predictions = dict()
def forward(self, data: HeteroData):
scene_enc = self.encoder(data)
pred = self.decoder(data, scene_enc)
return pred
def training_step(self,
data,
batch_idx):
if isinstance(data, Batch):
data['agent']['av_index'] += data['agent']['ptr'][:-1]
reg_mask = data['agent']['predict_mask'][:, self.num_historical_steps:]
cls_mask = data['agent']['predict_mask'][:, -1]
pred = self(data)
if self.output_head:
traj_propose = torch.cat([pred['loc_propose_pos'][..., :self.output_dim],
pred['loc_propose_head'],
pred['scale_propose_pos'][..., :self.output_dim],
pred['conc_propose_head']], dim=-1)
traj_refine = torch.cat([pred['loc_refine_pos'][..., :self.output_dim],
pred['loc_refine_head'],
pred['scale_refine_pos'][..., :self.output_dim],
pred['conc_refine_head']], dim=-1)
else:
traj_propose = torch.cat([pred['loc_propose_pos'][..., :self.output_dim],
pred['scale_propose_pos'][..., :self.output_dim]], dim=-1)
traj_refine = torch.cat([pred['loc_refine_pos'][..., :self.output_dim],
pred['scale_refine_pos'][..., :self.output_dim]], dim=-1)
pi = pred['pi']
gt = torch.cat([data['agent']['target'][..., :self.output_dim], data['agent']['target'][..., -1:]], dim=-1)
l2_norm = (torch.norm(traj_propose[..., :self.output_dim] -
gt[..., :self.output_dim].unsqueeze(1), p=2, dim=-1) * reg_mask.unsqueeze(1)).sum(dim=-1)
best_mode = l2_norm.argmin(dim=-1)
traj_propose_best = traj_propose[torch.arange(traj_propose.size(0)), best_mode]
traj_refine_best = traj_refine[torch.arange(traj_refine.size(0)), best_mode]
reg_loss_propose = self.reg_loss(traj_propose_best,
gt[..., :self.output_dim + self.output_head]).sum(dim=-1) * reg_mask
reg_loss_propose = reg_loss_propose.sum(dim=0) / reg_mask.sum(dim=0).clamp_(min=1)
reg_loss_propose = reg_loss_propose.mean()
reg_loss_refine = self.reg_loss(traj_refine_best,
gt[..., :self.output_dim + self.output_head]).sum(dim=-1) * reg_mask
reg_loss_refine = reg_loss_refine.sum(dim=0) / reg_mask.sum(dim=0).clamp_(min=1)
reg_loss_refine = reg_loss_refine.mean()
cls_loss = self.cls_loss(pred=traj_refine[:, :, -1:].detach(),
target=gt[:, -1:, :self.output_dim + self.output_head],
prob=pi,
mask=reg_mask[:, -1:]) * cls_mask
cls_loss = cls_loss.sum() / cls_mask.sum().clamp_(min=1)
self.log('train_reg_loss_propose', reg_loss_propose, prog_bar=False, on_step=True, on_epoch=True, batch_size=1)
self.log('train_reg_loss_refine', reg_loss_refine, prog_bar=False, on_step=True, on_epoch=True, batch_size=1)
self.log('train_cls_loss', cls_loss, prog_bar=False, on_step=True, on_epoch=True, batch_size=1)
loss = reg_loss_propose + reg_loss_refine + cls_loss
return loss
def validation_step(self,
data,
batch_idx):
if isinstance(data, Batch):
data['agent']['av_index'] += data['agent']['ptr'][:-1]
reg_mask = data['agent']['predict_mask'][:, self.num_historical_steps:]
cls_mask = data['agent']['predict_mask'][:, -1]
pred = self(data)
if self.output_head:
traj_propose = torch.cat([pred['loc_propose_pos'][..., :self.output_dim],
pred['loc_propose_head'],
pred['scale_propose_pos'][..., :self.output_dim],
pred['conc_propose_head']], dim=-1)
traj_refine = torch.cat([pred['loc_refine_pos'][..., :self.output_dim],
pred['loc_refine_head'],
pred['scale_refine_pos'][..., :self.output_dim],
pred['conc_refine_head']], dim=-1)
else:
traj_propose = torch.cat([pred['loc_propose_pos'][..., :self.output_dim],
pred['scale_propose_pos'][..., :self.output_dim]], dim=-1)
traj_refine = torch.cat([pred['loc_refine_pos'][..., :self.output_dim],
pred['scale_refine_pos'][..., :self.output_dim]], dim=-1)
pi = pred['pi'] #N,6
####zy visualization
eval_mask = data['agent']['category'] == 3
traj_past = data['agent']['position'][eval_mask, :self.num_historical_steps, :2].cpu()
origin_eval = data['agent']['position'][eval_mask, self.num_historical_steps - 1]
theta_eval = data['agent']['heading'][eval_mask, self.num_historical_steps - 1]
cos, sin = theta_eval.cos(), theta_eval.sin()
rot_mat = torch.zeros(eval_mask.sum(), 2, 2, device=self.device)
rot_mat[:, 0, 0] = cos
rot_mat[:, 0, 1] = sin
rot_mat[:, 1, 0] = -sin
rot_mat[:, 1, 1] = cos
traj_eval = torch.matmul(traj_refine[eval_mask, :, :, :2],
rot_mat.unsqueeze(1)) + origin_eval[:, :2].reshape(-1, 1, 1, 2)
traj_eval=traj_eval.cpu()
gt_eval = data['agent']['position'][eval_mask, self.num_historical_steps:, :2].cpu()
base_path = '/mnt/ve_share2/zy/Argoverse_2_Motion_Forecasting_Dataset/raw/val/{}'
full_path = base_path.format(data.scenario_id[0])
save_path='/mnt/ve_share2/zy/QCNet/zy'
run_generate_scenario_visualizations(full_path,save_path,1,'first',True,traj_eval[0,:,:,:])
plot_single_vehicle(traj_past.numpy(),gt_eval.numpy(),traj_eval[0,:,:,:].numpy(),data.scenario_id[0])
gt = torch.cat([data['agent']['target'][..., :self.output_dim], data['agent']['target'][..., -1:]], dim=-1)
l2_norm = (torch.norm(traj_propose[..., :self.output_dim] -
gt[..., :self.output_dim].unsqueeze(1), p=2, dim=-1) * reg_mask.unsqueeze(1)).sum(dim=-1)
best_mode = l2_norm.argmin(dim=-1)
traj_propose_best = traj_propose[torch.arange(traj_propose.size(0)), best_mode] #N,60,4
traj_refine_best = traj_refine[torch.arange(traj_refine.size(0)), best_mode]
reg_loss_propose = self.reg_loss(traj_propose_best,
gt[..., :self.output_dim + self.output_head]).sum(dim=-1) * reg_mask #N,60
reg_loss_propose = reg_loss_propose.sum(dim=0) / reg_mask.sum(dim=0).clamp_(min=1)
reg_loss_propose = reg_loss_propose.mean()
reg_loss_refine = self.reg_loss(traj_refine_best,
gt[..., :self.output_dim + self.output_head]).sum(dim=-1) * reg_mask
reg_loss_refine = reg_loss_refine.sum(dim=0) / reg_mask.sum(dim=0).clamp_(min=1)
reg_loss_refine = reg_loss_refine.mean()
cls_loss = self.cls_loss(pred=traj_refine[:, :, -1:].detach(),##N,60,1,4
target=gt[:, -1:, :self.output_dim + self.output_head],#N,1,2
prob=pi, #N,6
mask=reg_mask[:, -1:]) * cls_mask
cls_loss = cls_loss.sum() / cls_mask.sum().clamp_(min=1)
self.log('val_reg_loss_propose', reg_loss_propose, prog_bar=True, on_step=False, on_epoch=True, batch_size=1,
sync_dist=True)
self.log('val_reg_loss_refine', reg_loss_refine, prog_bar=True, on_step=False, on_epoch=True, batch_size=1,
sync_dist=True)
self.log('val_cls_loss', cls_loss, prog_bar=True, on_step=False, on_epoch=True, batch_size=1, sync_dist=True)
if self.dataset == 'argoverse_v2':
eval_mask = data['agent']['category'] == 3
else:
raise ValueError('{} is not a valid dataset'.format(self.dataset))
valid_mask_eval = reg_mask[eval_mask]
traj_eval = traj_refine[eval_mask, :, :, :self.output_dim + self.output_head]
if not self.output_head:
traj_2d_with_start_pos_eval = torch.cat([traj_eval.new_zeros((traj_eval.size(0), self.num_modes, 1, 2)),
traj_eval[..., :2]], dim=-2)
motion_vector_eval = traj_2d_with_start_pos_eval[:, :, 1:] - traj_2d_with_start_pos_eval[:, :, :-1]
head_eval = torch.atan2(motion_vector_eval[..., 1], motion_vector_eval[..., 0])
traj_eval = torch.cat([traj_eval, head_eval.unsqueeze(-1)], dim=-1)
pi_eval = F.softmax(pi[eval_mask], dim=-1)
gt_eval = gt[eval_mask]
self.Brier.update(pred=traj_eval[..., :self.output_dim], target=gt_eval[..., :self.output_dim], prob=pi_eval,
valid_mask=valid_mask_eval)
self.minADE.update(pred=traj_eval[..., :self.output_dim], target=gt_eval[..., :self.output_dim], prob=pi_eval,
valid_mask=valid_mask_eval)
self.minAHE.update(pred=traj_eval, target=gt_eval, prob=pi_eval, valid_mask=valid_mask_eval)
self.minFDE.update(pred=traj_eval[..., :self.output_dim], target=gt_eval[..., :self.output_dim], prob=pi_eval,
valid_mask=valid_mask_eval)
self.minFHE.update(pred=traj_eval, target=gt_eval, prob=pi_eval, valid_mask=valid_mask_eval)
self.MR.update(pred=traj_eval[..., :self.output_dim], target=gt_eval[..., :self.output_dim], prob=pi_eval,
valid_mask=valid_mask_eval)
self.log('val_Brier', self.Brier, prog_bar=True, on_step=False, on_epoch=True, batch_size=gt_eval.size(0))
self.log('val_minADE', self.minADE, prog_bar=True, on_step=False, on_epoch=True, batch_size=gt_eval.size(0))
self.log('val_minAHE', self.minAHE, prog_bar=True, on_step=False, on_epoch=True, batch_size=gt_eval.size(0))
self.log('val_minFDE', self.minFDE, prog_bar=True, on_step=False, on_epoch=True, batch_size=gt_eval.size(0))
self.log('val_minFHE', self.minFHE, prog_bar=True, on_step=False, on_epoch=True, batch_size=gt_eval.size(0))
self.log('val_MR', self.MR, prog_bar=True, on_step=False, on_epoch=True, batch_size=gt_eval.size(0))
def test_step(self,
data,
batch_idx):
if isinstance(data, Batch):
data['agent']['av_index'] += data['agent']['ptr'][:-1]
pred = self(data)
if self.output_head:
traj_refine = torch.cat([pred['loc_refine_pos'][..., :self.output_dim],
pred['loc_refine_head'],
pred['scale_refine_pos'][..., :self.output_dim],
pred['conc_refine_head']], dim=-1)
else:
traj_refine = torch.cat([pred['loc_refine_pos'][..., :self.output_dim],
pred['scale_refine_pos'][..., :self.output_dim]], dim=-1)
pi = pred['pi']
if self.dataset == 'argoverse_v2':
eval_mask = data['agent']['category'] == 3
else:
raise ValueError('{} is not a valid dataset'.format(self.dataset))
origin_eval = data['agent']['position'][eval_mask, self.num_historical_steps - 1]
theta_eval = data['agent']['heading'][eval_mask, self.num_historical_steps - 1]
cos, sin = theta_eval.cos(), theta_eval.sin()
rot_mat = torch.zeros(eval_mask.sum(), 2, 2, device=self.device)
rot_mat[:, 0, 0] = cos
rot_mat[:, 0, 1] = sin
rot_mat[:, 1, 0] = -sin
rot_mat[:, 1, 1] = cos
traj_eval = torch.matmul(traj_refine[eval_mask, :, :, :2],
rot_mat.unsqueeze(1)) + origin_eval[:, :2].reshape(-1, 1, 1, 2)
pi_eval = F.softmax(pi[eval_mask], dim=-1)
traj_eval = traj_eval.cpu().numpy()
pi_eval = pi_eval.cpu().numpy()
if self.dataset == 'argoverse_v2':
eval_id = list(compress(list(chain(*data['agent']['id'])), eval_mask))
if isinstance(data, Batch):
for i in range(data.num_graphs):
self.test_predictions[data['scenario_id'][i]] = (pi_eval[i], {eval_id[i]: traj_eval[i]})
else:
self.test_predictions[data['scenario_id']] = (pi_eval[0], {eval_id[0]: traj_eval[0]})
else:
raise ValueError('{} is not a valid dataset'.format(self.dataset))
def on_test_end(self):
if self.dataset == 'argoverse_v2':
ChallengeSubmission(self.test_predictions).to_parquet(
Path(self.submission_dir) / f'{self.submission_file_name}.parquet')
else:
raise ValueError('{} is not a valid dataset'.format(self.dataset))
def configure_optimizers(self):
decay = set()
no_decay = set()
whitelist_weight_modules = (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.MultiheadAttention, nn.LSTM,
nn.LSTMCell, nn.GRU, nn.GRUCell)
blacklist_weight_modules = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.LayerNorm, nn.Embedding)
for module_name, module in self.named_modules():
for param_name, param in module.named_parameters():
full_param_name = '%s.%s' % (module_name, param_name) if module_name else param_name
if 'bias' in param_name:
no_decay.add(full_param_name)
elif 'weight' in param_name:
if isinstance(module, whitelist_weight_modules):
decay.add(full_param_name)
elif isinstance(module, blacklist_weight_modules):
no_decay.add(full_param_name)
elif not ('weight' in param_name or 'bias' in param_name):
no_decay.add(full_param_name)
param_dict = {param_name: param for param_name, param in self.named_parameters()}
inter_params = decay & no_decay
union_params = decay | no_decay
assert len(inter_params) == 0
assert len(param_dict.keys() - union_params) == 0
optim_groups = [
{"params": [param_dict[param_name] for param_name in sorted(list(decay))],
"weight_decay": self.weight_decay},
{"params": [param_dict[param_name] for param_name in sorted(list(no_decay))],
"weight_decay": 0.0},
]
optimizer = torch.optim.AdamW(optim_groups, lr=self.lr, weight_decay=self.weight_decay)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=self.T_max, eta_min=0.0)
return [optimizer], [scheduler]
@staticmethod
def add_model_specific_args(parent_parser):
parser = parent_parser.add_argument_group('QCNet')
parser.add_argument('--dataset', type=str, required=True)
parser.add_argument('--input_dim', type=int, default=2)
parser.add_argument('--hidden_dim', type=int, default=128)
parser.add_argument('--output_dim', type=int, default=2)
parser.add_argument('--output_head', action='store_true')
parser.add_argument('--num_historical_steps', type=int, required=True)
parser.add_argument('--num_future_steps', type=int, required=True)
parser.add_argument('--num_modes', type=int, default=6)
parser.add_argument('--num_recurrent_steps', type=int, required=True)
parser.add_argument('--num_freq_bands', type=int, default=64)
parser.add_argument('--num_map_layers', type=int, default=1)
parser.add_argument('--num_agent_layers', type=int, default=2)
parser.add_argument('--num_dec_layers', type=int, default=2)
parser.add_argument('--num_heads', type=int, default=8)
parser.add_argument('--head_dim', type=int, default=16)
parser.add_argument('--dropout', type=float, default=0.1)
parser.add_argument('--pl2pl_radius', type=float, required=True)
parser.add_argument('--time_span', type=int, default=None)
parser.add_argument('--pl2a_radius', type=float, required=True)
parser.add_argument('--a2a_radius', type=float, required=True)
parser.add_argument('--num_t2m_steps', type=int, default=None)
parser.add_argument('--pl2m_radius', type=float, required=True)
parser.add_argument('--a2m_radius', type=float, required=True)
parser.add_argument('--lr', type=float, default=5e-4)
parser.add_argument('--weight_decay', type=float, default=1e-4)
parser.add_argument('--T_max', type=int, default=64)
parser.add_argument('--submission_dir', type=str, default='/mnt/ve_share2/zy/QCNet/submission')
parser.add_argument('--submission_file_name', type=str, default='submission')
return parent_parser
图像版本：
import matplotlib.pyplot as plt
import numpy as np
import sys
# sys.path.append("/mnt/ve_share2/zy/HIVT")
# from argoverse_api.argoverse.map_representation.map_api import ArgoverseMap
from av2.map.map_api import ArgoverseStaticMap
from typing import Callable, Dict, List, Optional, Tuple, Union
from typing import Final, List, Optional, Sequence, Set, Tuple
from av2.utils.typing import NDArrayFloat, NDArrayInt
from pathlib import Path
_LANE_SEGMENT_COLOR: Final[str] = "#E0E0E0"
_DRIVABLE_AREA_COLOR: Final[str] = "#7A7A7A"
'''
过去轨迹 (Past Trajectory):
color="#ECA154"：这是过去轨迹的颜色，用浅橙色表示。
地面真实轨迹 (Ground Trut ):
color="#d33e4c"：这是地面真实轨迹的颜色，用深红色表示。
预测轨迹 (Forecasted Trajectory):
color="#007672"：这是预测轨迹的颜色，用深绿色表示。
箭头和终点标记的颜色:
箭头和终点标记的颜色与相应轨迹的颜色一致，分别是橙色、红色和绿色。
'''
def plot_single_vehicle(
sample_past_trajectory: np.ndarray, # 过去轨迹的坐标数组，形状为 (1, 50, 2)
sample_groundtruth: np.ndarray, # 地面真实轨迹的坐标数组，形状为 (1, 60, 2)
sample_forecasted_trajectories: List[np.ndarray], # 预测轨迹的列表，每个数组形状为 (6,60, 2)
scenario_id,
):
plt.figure()
min_x = min(
np.min(sample_past_trajectory[:, :, 0]),
np.min(sample_groundtruth[:, :, 0]),
np.min([np.min(traj[:, 0]) for traj in sample_forecasted_trajectories])
)
max_x = max(
np.max(sample_past_trajectory[:, :, 0]),
np.max(sample_groundtruth[:, :, 0]),
np.max([np.max(traj[:, 0]) for traj in sample_forecasted_trajectories])
)
min_y = min(
np.min(sample_past_trajectory[:, :, 1]),
np.min(sample_groundtruth[:, :, 1]),
np.min([np.min(traj[:, 1]) for traj in sample_forecasted_trajectories])
)
max_y = max(
np.max(sample_past_trajectory[:, :, 1]),
np.max(sample_groundtruth[:, :, 1]),
np.max([np.max(traj[:, 1]) for traj in sample_forecasted_trajectories])
)
x_buffer = 5
y_buffer = 5
plt.plot(
sample_past_trajectory[0, :, 0],
sample_past_trajectory[0, :, 1],
color="#ECA154",
label="Past Trajectory",
alpha=1,
linewidth=2,
zorder=15,
ls="--"
)
plt.plot(
sample_groundtruth[0, :, 0],
sample_groundtruth[0, :, 1],
color="#d33e4c",
label="Ground Truth",
alpha=1,
linewidth=2,
zorder=20,
ls="--"
)
for i, sample_forecasted_trajectory in enumerate(sample_forecasted_trajectories):
plt.plot(
sample_forecasted_trajectory[:, 0],
sample_forecasted_trajectory[:, 1],
color="#007672",
label=f"Forecasted Trajectory {i + 1}",
alpha=1,
linewidth=2,
zorder=20,
ls="--"
)
# Plot the end marker for forecasted trajectories
plt.arrow(
sample_forecasted_trajectory[-2, 0],
sample_forecasted_trajectory[-2, 1],
sample_forecasted_trajectory[-1, 0] - sample_forecasted_trajectory[-2, 0],
sample_forecasted_trajectory[-1, 1] - sample_forecasted_trajectory[-2, 1],
color="#007672",
label="Forecasted Trajectory",
alpha=1,
linewidth=3,
zorder=25,
head_width=0.3,
head_length=0.3
)
# Plot the end marker for history
plt.arrow(
sample_past_trajectory[0, -2, 0],
sample_past_trajectory[0, -2, 1],
sample_past_trajectory[0, -1, 0] - sample_past_trajectory[0, -2, 0],
sample_past_trajectory[0, -1, 1] - sample_past_trajectory[0, -2, 1],
color="#ECA154",
label="Past Trajectory",
alpha=1,
linewidth=2.5,
zorder=25,
head_width=0.1,
)
# Plot the end marker for ground truth
plt.arrow(
sample_groundtruth[0, -2, 0],
sample_groundtruth[0, -2, 1],
sample_groundtruth[0, -1, 0] - sample_groundtruth[0, -2, 0],
sample_groundtruth[0, -1, 1] - sample_groundtruth[0, -2, 1],
color="#d33e4c",
label="Ground Truth",
alpha=1,
linewidth=3,
zorder=30,
head_width=0.1,
)
static_map_path=f"/mnt/ve_share2/zy/Argoverse_2_Motion_Forecasting_Dataset/raw/val/{scenario_id}/log_map_archive_{scenario_id}.json"
static_map_path=Path(static_map_path)
static_map = ArgoverseStaticMap.from_json(static_map_path)
_plot_static_map_elements(static_map)
plt.xlim(min_x - x_buffer, max_x + x_buffer)
plt.ylim(min_y - y_buffer, max_y + y_buffer)
plt.savefig(f'/mnt/ve_share2/zy/QCNet/zy/{scenario_id}.png')
def _plot_static_map_elements(
static_map: ArgoverseStaticMap, show_ped_xings: bool = False
) -> None:
"""Plot all static map elements associated with an Argoverse scenario.
Args:
static_map: Static map containing elements to be plotted.
show_ped_xings: Configures whether pedestrian crossings should be plotted.
"""
# Plot drivable areas
for drivable_area in static_map.vector_drivable_areas.values():
_plot_polygons([drivable_area.xyz], alpha=0.5, color=_DRIVABLE_AREA_COLOR)
# Plot lane segments
for lane_segment in static_map.vector_lane_segments.values():
_plot_polylines(
[
lane_segment.left_lane_boundary.xyz,
lane_segment.right_lane_boundary.xyz,
],
line_width=0.5,
color=_LANE_SEGMENT_COLOR,
)
# Plot pedestrian crossings
if show_ped_xings:
for ped_xing in static_map.vector_pedestrian_crossings.values():
_plot_polylines(
[ped_xing.edge1.xyz, ped_xing.edge2.xyz],
alpha=1.0,
color=_LANE_SEGMENT_COLOR,
)
def _plot_polylines(
polylines: Sequence[NDArrayFloat],
*,
style: str = "-",
line_width: float = 1.0,
alpha: float = 1.0,
color: str = "r",
) -> None:
"""Plot a group of polylines with the specified config.
Args:
polylines: Collection of (N, 2) polylines to plot.
style: Style of the line to plot (e.g. `-` for solid, `--` for dashed)
line_width: Desired width for the plotted lines.
alpha: Desired alpha for the plotted lines.
color: Desired color for the plotted lines.
"""
for polyline in polylines:
plt.plot(
polyline[:, 0],
polyline[:, 1],
style,
linewidth=line_width,
color=color,
alpha=alpha,
)
def _plot_polygons(
polygons: Sequence[NDArrayFloat], *, alpha: float = 1.0, color: str = "r"
) -> None:
"""Plot a group of filled polygons with the specified config.
Args:
polygons: Collection of polygons specified by (N,2) arrays of vertices.
alpha: Desired alpha for the polygon fill.
color: Desired color for the polygon.
"""
for polygon in polygons:
plt.fill(polygon[:, 0], polygon[:, 1], color=color, alpha=alpha)
生成视频版本
# <Copyright 2022, Argo AI, LLC. Released under the MIT license.>
"""Visualization utils for Argoverse MF scenarios."""
import io
import math
from pathlib import Path
from typing import Final, List, Optional, Sequence, Set, Tuple
import cv2
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Rectangle
from PIL import Image as img
from PIL.Image import Image
from av2.datasets.motion_forecasting.data_schema import (
ArgoverseScenario,
ObjectType,
TrackCategory,
)
from av2.map.map_api import ArgoverseStaticMap
from av2.utils.typing import NDArrayFloat, NDArrayInt
_PlotBounds = Tuple[float, float, float, float]
# Configure constants
_OBS_DURATION_TIMESTEPS: Final[int] = 50
_PRED_DURATION_TIMESTEPS: Final[int] = 60
_ESTIMATED_VEHICLE_LENGTH_M: Final[float] = 4.0
_ESTIMATED_VEHICLE_WIDTH_M: Final[float] = 2.0
_ESTIMATED_CYCLIST_LENGTH_M: Final[float] = 2.0
_ESTIMATED_CYCLIST_WIDTH_M: Final[float] = 0.7
_PLOT_BOUNDS_BUFFER_M: Final[float] = 30.0
_DRIVABLE_AREA_COLOR: Final[str] = "#7A7A7A"
_LANE_SEGMENT_COLOR: Final[str] = "#E0E0E0"
_DEFAULT_ACTOR_COLOR: Final[str] = "#D3E8EF"
_FOCAL_AGENT_COLOR: Final[str] = "#ECA25B"
_ZY: Final[str] = "#d33e4c"
_AV_COLOR: Final[str] = "#007672"
_BOUNDING_BOX_ZORDER: Final[
int
] = 100 # Ensure actor bounding boxes are plotted on top of all map elements
_STATIC_OBJECT_TYPES: Set[ObjectType] = {
ObjectType.STATIC,
ObjectType.BACKGROUND,
ObjectType.CONSTRUCTION,
ObjectType.RIDERLESS_BICYCLE,
}
def visualize_scenario(
scenario: ArgoverseScenario,
scenario_static_map: ArgoverseStaticMap,
save_path: Path,
predicted_trajectories,
) -> None:
"""Build dynamic visualization for all tracks and the local map associated with an Argoverse scenario.
Note: This function uses OpenCV to create a MP4 file using the MP4V codec.
Args:
scenario: Argoverse scenario to visualize.
scenario_static_map: Local static map elements associated with `scenario`.
save_path: Path where output MP4 video should be saved.
"""
predicted_trajectories=predicted_trajectories
# Build each frame for the video
frames: List[Image] = []
plot_bounds: _PlotBounds = (0, 0, 0, 0)
for timestep in range(_OBS_DURATION_TIMESTEPS + _PRED_DURATION_TIMESTEPS):
_, ax = plt.subplots()
# Plot static map elements and actor tracks
_plot_static_map_elements(scenario_static_map)
cur_plot_bounds = _plot_actor_tracks(ax, scenario, timestep,predicted_trajectories)
if cur_plot_bounds:
plot_bounds = cur_plot_bounds
# Set map bounds to capture focal trajectory history (with fixed buffer in all directions)
plt.xlim(
plot_bounds[0] - _PLOT_BOUNDS_BUFFER_M,
plot_bounds[1] + _PLOT_BOUNDS_BUFFER_M,
)
plt.ylim(
plot_bounds[2] - _PLOT_BOUNDS_BUFFER_M,
plot_bounds[3] + _PLOT_BOUNDS_BUFFER_M,
)
plt.gca().set_aspect("equal", adjustable="box")
# Minimize plot margins and make axes invisible
plt.gca().set_axis_off()
plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
plt.margins(0, 0)
plt.gca().xaxis.set_major_locator(plt.NullLocator())
plt.gca().yaxis.set_major_locator(plt.NullLocator())
# Save plotted frame to in-memory buffer
buf = io.BytesIO()
plt.savefig(buf, format="png")
plt.close()
buf.seek(0)
frame = img.open(buf)
frames.append(frame)
# Write buffered frames to MP4V-encoded video
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
vid_path = str(save_path.parents[0] / f"{save_path.stem}.mp4")
video = cv2.VideoWriter(vid_path, fourcc, fps=10, frameSize=frames[0].size)
for i in range(len(frames)):
frame_temp = frames[i].copy()
video.write(cv2.cvtColor(np.array(frame_temp), cv2.COLOR_RGB2BGR))
video.release()
def _plot_static_map_elements(
static_map: ArgoverseStaticMap, show_ped_xings: bool = False
) -> None:
"""Plot all static map elements associated with an Argoverse scenario.
Args:
static_map: Static map containing elements to be plotted.
show_ped_xings: Configures whether pedestrian crossings should be plotted.
"""
# Plot drivable areas
for drivable_area in static_map.vector_drivable_areas.values():
_plot_polygons([drivable_area.xyz], alpha=0.5, color=_DRIVABLE_AREA_COLOR)
# Plot lane segments
for lane_segment in static_map.vector_lane_segments.values():
_plot_polylines(
[
lane_segment.left_lane_boundary.xyz,
lane_segment.right_lane_boundary.xyz,
],
line_width=0.5,
color=_LANE_SEGMENT_COLOR,
)
# Plot pedestrian crossings
if show_ped_xings:
for ped_xing in static_map.vector_pedestrian_crossings.values():
_plot_polylines(
[ped_xing.edge1.xyz, ped_xing.edge2.xyz],
alpha=1.0,
color=_LANE_SEGMENT_COLOR,
)
def _plot_actor_tracks(
ax: plt.Axes, scenario: ArgoverseScenario, timestep: int,predicted_trajectories
) -> Optional[_PlotBounds]:
"""Plot all actor tracks (up to a particular time step) associated with an Argoverse scenario.
Args:
ax: Axes on which actor tracks should be plotted.
scenario: Argoverse scenario for which to plot actor tracks.
timestep: Tracks are plotted for all actor data up to the specified time step.
Returns:
track_bounds: (x_min, x_max, y_min, y_max) bounds for the extent of actor tracks.
"""
track_bounds = None
for track in scenario.tracks:
# Get timesteps for which actor data is valid
actor_timesteps: NDArrayInt = np.array(
[
object_state.timestep
for object_state in track.object_states
if object_state.timestep <= timestep
]
)
if actor_timesteps.shape[0] < 1 or actor_timesteps[-1] != timestep:
continue
# Get actor trajectory and heading history
actor_trajectory: NDArrayFloat = np.array(
[
list(object_state.position)
for object_state in track.object_states
if object_state.timestep <= timestep
]
)
actor_headings: NDArrayFloat = np.array(
[
object_state.heading
for object_state in track.object_states
if object_state.timestep <= timestep
]
)
# Plot polyline for focal agent location history
track_color = _DEFAULT_ACTOR_COLOR
if track.category == TrackCategory.FOCAL_TRACK:
x_min, x_max = actor_trajectory[:, 0].min(), actor_trajectory[:, 0].max()
y_min, y_max = actor_trajectory[:, 1].min(), actor_trajectory[:, 1].max()
track_bounds = (x_min, x_max, y_min, y_max)
track_color = _FOCAL_AGENT_COLOR
_plot_polylines([actor_trajectory], color=track_color, line_width=2)
##zy
if predicted_trajectories is not None and timestep > _OBS_DURATION_TIMESTEPS-1:
for i in range(predicted_trajectories.shape[0]):
predicted_trajectory = predicted_trajectories[i]
_plot_polylines([predicted_trajectory], style="--", color=_ZY, line_width=1)
elif track.track_id == "AV":
track_color = _AV_COLOR
elif track.object_type in _STATIC_OBJECT_TYPES:
continue
# Plot bounding boxes for all vehicles and cyclists
if track.object_type == ObjectType.VEHICLE:
_plot_actor_bounding_box(
ax,
actor_trajectory[-1],
actor_headings[-1],
track_color,
(_ESTIMATED_VEHICLE_LENGTH_M, _ESTIMATED_VEHICLE_WIDTH_M),
)
elif (
track.object_type == ObjectType.CYCLIST
or track.object_type == ObjectType.MOTORCYCLIST
):
_plot_actor_bounding_box(
ax,
actor_trajectory[-1],
actor_headings[-1],
track_color,
(_ESTIMATED_CYCLIST_LENGTH_M, _ESTIMATED_CYCLIST_WIDTH_M),
)
else:
plt.plot(
actor_trajectory[-1, 0],
actor_trajectory[-1, 1],
"o",
color=track_color,
markersize=4,
)
return track_bounds
def _plot_polylines(
polylines: Sequence[NDArrayFloat],
*,
style: str = "-",
line_width: float = 1.0,
alpha: float = 1.0,
color: str = "r",
) -> None:
"""Plot a group of polylines with the specified config.
Args:
polylines: Collection of (N, 2) polylines to plot.
style: Style of the line to plot (e.g. `-` for solid, `--` for dashed)
line_width: Desired width for the plotted lines.
alpha: Desired alpha for the plotted lines.
color: Desired color for the plotted lines.
"""
for polyline in polylines:
plt.plot(
polyline[:, 0],
polyline[:, 1],
style,
linewidth=line_width,
color=color,
alpha=alpha,
)
def _plot_polygons(
polygons: Sequence[NDArrayFloat], *, alpha: float = 1.0, color: str = "r"
) -> None:
"""Plot a group of filled polygons with the specified config.
Args:
polygons: Collection of polygons specified by (N,2) arrays of vertices.
alpha: Desired alpha for the polygon fill.
color: Desired color for the polygon.
"""
for polygon in polygons:
plt.fill(polygon[:, 0], polygon[:, 1], color=color, alpha=alpha)
def _plot_actor_bounding_box(
ax: plt.Axes,
cur_location: NDArrayFloat,
heading: float,
color: str,
bbox_size: Tuple[float, float],
) -> None:
"""Plot an actor bounding box centered on the actor's current location.
Args:
ax: Axes on which actor bounding box should be plotted.
cur_location: Current location of the actor (2,).
heading: Current heading of the actor (in radians).
color: Desired color for the bounding box.
bbox_size: Desired size for the bounding box (length, width).
"""
(bbox_length, bbox_width) = bbox_size
# Compute coordinate for pivot point of bounding box
d = np.hypot(bbox_length, bbox_width)
theta_2 = math.atan2(bbox_width, bbox_length)
pivot_x = cur_location[0] - (d / 2) * math.cos(heading + theta_2)
pivot_y = cur_location[1] - (d / 2) * math.sin(heading + theta_2)
vehicle_bounding_box = Rectangle(
(pivot_x, pivot_y),
bbox_length,
bbox_width,
np.degrees(heading),
color=color,
zorder=_BOUNDING_BOX_ZORDER,
)
ax.add_patch(vehicle_bounding_box)
调用上述文件
# <Copyright 2022, Argo AI, LLC. Released under the MIT license.>
"""Script to generate dynamic visualizations from a directory of Argoverse scenarios."""
from enum import Enum, unique
from pathlib import Path
from random import choices
from typing import Final
import click
from joblib import Parallel, delayed
from rich.progress import track
from av2.datasets.motion_forecasting import scenario_serialization
from zyvisual import visualize_scenario
# from zyvisual import visualize_zy
# from av2.datasets.motion_forecasting.viz.scenario_visualization import (
# visualize_scenario,
# )
from av2.map.map_api import ArgoverseStaticMap
_DEFAULT_N_JOBS: Final[int] = -2 # Use all but one CPUs
@unique
class SelectionCriteria(str, Enum):
"""Valid criteria used to select Argoverse scenarios for visualization."""
FIRST: str = "first"
RANDOM: str = "random"
def generate_scenario_visualizations(
argoverse_scenario_dir: Path,
viz_output_dir: Path,
num_scenarios: int,
selection_criteria: SelectionCriteria,
*,
debug: bool = False,
predicted_trajectories,
) -> None:
"""Generate and save dynamic visualizations for selected scenarios within `argoverse_scenario_dir`.
Args:
argoverse_scenario_dir: Path to local directory where Argoverse scenarios are stored.
viz_output_dir: Path to local directory where generated visualizations should be saved.
num_scenarios: Maximum number of scenarios for which to generate visualizations.
selection_criteria: Controls how scenarios are selected for visualization.
debug: Runs preprocessing in single-threaded mode when enabled.
"""
Path(viz_output_dir).mkdir(parents=True, exist_ok=True)
all_scenario_files = sorted(argoverse_scenario_dir.rglob("*.parquet"))
scenario_file_list = (
all_scenario_files[:num_scenarios]
if selection_criteria == SelectionCriteria.FIRST
else choices(all_scenario_files, k=num_scenarios)
) # Ignoring type here because type of "choice" is partially unknown.
# Build inner function to generate visualization for a single scenario.
def generate_scenario_visualization(scenario_path: Path,predicted_trajectories) -> None:
"""Generate and save dynamic visualization for a single Argoverse scenario.
NOTE: This function assumes that the static map is stored in the same directory as the scenario file.
Args:
scenario_path: Path to the parquet file corresponding to the Argoverse scenario to visualize.
"""
scenario_id = scenario_path.stem.split("_")[-1]
static_map_path = (
scenario_path.parents[0] / f"log_map_archive_{scenario_id}.json"
)
viz_save_path = viz_output_dir / f"{scenario_id}.mp4"
scenario = scenario_serialization.load_argoverse_scenario_parquet(scenario_path)
static_map = ArgoverseStaticMap.from_json(static_map_path)
# visualize_zy(scenario, static_map, viz_save_path,predicted_trajectories=predicted_trajectories)
visualize_scenario(scenario, static_map, viz_save_path,predicted_trajectories=predicted_trajectories)
# Generate visualization for each selected scenario in parallel (except if running in debug mode)
if debug:
for scenario_path in track(scenario_file_list):
generate_scenario_visualization(scenario_path,predicted_trajectories)
else:
Parallel(n_jobs=_DEFAULT_N_JOBS)(
delayed(generate_scenario_visualization)(scenario_path)
for scenario_path in track(scenario_file_list)
)
def run_generate_scenario_visualizations(
argoverse_scenario_dir: str,
viz_output_dir: str,
num_scenarios: int,
selection_criteria: str,
debug: bool,
predicted_trajectories= None,
) -> None:
"""Click entry point for generation of Argoverse scenario visualizations."""
generate_scenario_visualizations(
Path(argoverse_scenario_dir),
Path(viz_output_dir),
num_scenarios,
SelectionCriteria(selection_criteria.lower()),
debug=debug,
predicted_trajectories=predicted_trajectories,
)
# run_generate_scenario_visualizations(
# '/mnt/ve_share2/zy/Argoverse_2_Motion_Forecasting_Dataset/raw/val/f93736af-279f-42ad-80e7-20a42164fec8',
# "/mnt/ve_share2/zy/QCNet/zy",
# 1,
# 'first',
# True,
# predicted_trajectories
# )
结果
2442
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
