Skip to content
GitLab
Commit ab48995a authored by Kunming Li's avatar Kunming Li
Browse files

First commit

parent ca681ad5
Expand all Hide whitespace changes
Inline Side-by-side
.ipynb_checkpoints/Untitled-checkpoint.ipynb 0 → 100644
View file @ ab48995a
%% Cell type:code id: tags:
``` python
#!/usr/bin/env python
import rospy
from std_msgs.msg import String
from zio_obstacle_msgs.msg import ObjectsSequenceStamped
import os
import math
import sys
import torch
import numpy as np
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import pickle
import argparse
import glob
import torch.distributions.multivariate_normal as torchdist
from utils import *
from metrics import *
from model import social_stgcnn
import copy
### time scale
### none
### deepsort tracker test
class PedBuffer(object):
def __init__(self):
# frame id, ped id, x,y
self.time_stamp = 0
self.ped_data = []
self.graph_data = []
self.skip_num = 4
def append(self, data):
#### time filter: need set time scale
self.time_stamp += 1
# add seuqnece
ped_num = len(data.objects_sequence[0].objects)
for i in range(0,ped_num):
ped_id = data.objects_sequence[0].objects[i].id
ped_x = data.objects_sequence[0].objects[0].pose.pose.position.x
ped_y = data.objects_sequence[0].objects[0].pose.pose.position.y
self.ped_data.append([self.time_stamp,ped_id,ped_x,ped_y])
#self.ped_data = np.array(self.ped_data)
def seq_to_graph(self,seq_,norm_lap_matr = True):
seq_ = seq_.squeeze()
#seq_rel = seq_rel.squeeze()
seq_len = seq_.shape[2]
max_nodes = seq_.shape[0]
V = np.zeros((seq_len,max_nodes,seq_.shape[1]))
A = np.zeros((seq_len,max_nodes,max_nodes))
for s in range(seq_len):
step_ = seq_[:,:,s]
for h in range(len(step_)):
V[s,h,:] = step_[h]
A[s,h,h] = 1
for k in range(h+1,len(step_)):
l2_norm = anorm(step_[h],step_[k])
A[s,h,k] = l2_norm
A[s,k,h] = l2_norm
if norm_lap_matr:
G = nx.from_numpy_matrix(A[s,:,:])
A[s,:,:] = nx.normalized_laplacian_matrix(G).toarray()
return torch.from_numpy(V).type(torch.float),\
torch.from_numpy(A).type(torch.float)
def create_sequence(self):
### set 8 observe timestep, change it to 4.
if self.time_stamp>= 8*self.skip_num:
time_start = int(self.time_stamp - 8 * self.skip_num)
time_end = int(self.time_stamp*4)
tmp_data = np.array(self.ped_data)
ped_seq = tmp_data[tmp_data[:,0]== time_start,:]
for i in range(time_start,time_end,self.skip_num):
fol_seq = tmp_data[tmp_data[:,0]== i,:]
ped_seq = np.concatenate((ped_seq,fol_seq),axis = 0)
ped_id_list = np.unique(ped_seq[:,1])
ped_traj_all = np.empty([1,2,8])
#ped_traj_rel_all = np.zeros([len(ped_id_list),2,8])
for i in ped_id_list:
ped_traj = ped_seq[ped_seq[:,1]==i,2:].T
if ped_traj.shape[1]!=8*self.skip_num:
print("'----------------------------------")
print("this trajectory is ignored")
print("-----------------------------------")
else:
ped_traj = ped_traj[np.newaxis,:,:]
#print(ped_traj.shape)
ped_traj_all= np.append(ped_traj_all,ped_traj,axis=0)
ped_traj_rel_all = ped_traj_all[:,:,1:] - ped_traj_all[:,:,:-1]
#print(ped_traj_rel_all.shape)
if ped_traj_all.shape[0]>1:
V_obs, A = self.seq_to_graph(ped_traj_all)
else:
V_obs = None
A = None
return V_obs, ped_traj_all, A
class Predictor(object):
def __init__(self):
exp_path = '/home/kunmingli/Desktop/Social-STGCNN/checkpoint/social-stgcnn_att_all_812/'
model_path = exp_path+'/val_best.pth'
args_path = exp_path+'/args.pkl'
with open(args_path,'rb') as f:
args = pickle.load(f)
stats= exp_path+'/constant_metrics.pkl'
with open(stats,'rb') as f:
cm = pickle.load(f)
self.model = social_stgcnn(n_stgcnn =args.n_stgcnn,n_txpcnn=args.n_txpcnn,
output_feat=args.output_size,seq_len=args.obs_seq_len,
kernel_size=args.kernel_size,pred_seq_len=args.pred_seq_len).cuda()
self.model.load_state_dict(torch.load(model_path))
self.pred_traj_type = 'traj'
self.num_sam = 20
def att_gcn(self, V_obs, A_obs):
V_obs = torch.unsqueeze(V_obs,0)
V_obs_tmp =V_obs.permute(0,3,1,2)
V_pred,_ = self.model(V_obs_tmp.cuda(),A_obs.cuda().squeeze())
return V_pred
def prediction_result(self,V_pred,obs_traj):
V_pred = V_pred.permute(0,2,3,1)
V_pred = V_pred.squeeze()
V_x = seq_to_nodes(ob_traj)
sx = torch.exp(V_pred[:,:,2]) #sx
sy = torch.exp(V_pred[:,:,3]) #sy
corr = torch.tanh(V_pred[:,:,4]) #corr
cov = torch.zeros(V_pred.shape[0],V_pred.shape[1],2,2).cuda()
cov[:,:,0,0]= sx*sx
cov[:,:,0,1]= corr*sx*sy
cov[:,:,1,0]= corr*sx*sy
cov[:,:,1,1]= sy*sy
mean = V_pred[:,:,0:2]
mvnormal = torchdist.MultivariateNormal(mean,cov)
res = []
for i in range(0, self.num_sam):
V_pred = mvnormal.sample()
V_pred_rel_to_abs = nodes_rel_to_nodes_abs(V_pred.data.cpu().numpy().squeeze().copy(),
V_x[-1,:,:].copy())
res.append(V_pred_rel_to_abs)
return res
class Publisher(object):
def __init__(self):