def write_traj_map(self, trajmap: TrajectoryMapManager):
full_map = trajmap.traj_map
half_size = 16 # half of map size, in map size
if not self.rotate:
center_pos = trajmap.real2map(self.traj_train[-1])
else:
center_pos = trajmap.real2map(self.traj_original[self.obs_length])
original_map = cv2.resize(
full_map[np.maximum(center_pos[0] - 2 * half_size, 0):np.
minimum(center_pos[0] + 2 * half_size, full_map.shape[0]),
np.maximum(center_pos[1] - 2 * half_size, 0):np.
minimum(center_pos[1] +
2 * half_size, full_map.shape[1]), ],
(4 * half_size, 4 * half_size))
final_map = original_map[half_size:3 * half_size,
half_size:3 * half_size]
if self.reverse:
final_map = np.flip(original_map[half_size:3 * half_size,
half_size:3 * half_size])
if self.rotate:
final_map = cv2.warpAffine(
original_map,
cv2.getRotationMatrix2D(
(2 * half_size, 2 * half_size),
self.rotate,
1,
),
(4 * half_size, 4 * half_size),
)
final_map = final_map[half_size:3 * half_size,
half_size:3 * half_size]
self.traj_map = final_map
def prepare_test_agents_batch(self,
agents_batch: dict,
test_on_neighbors=False):
# create trajectory map for each batch
if not type(self.given_maps_when_test) == np.ndarray:
traj_maps = [
TrajectoryMapManager(agents_batch[batch_index])
for batch_index in agents_batch
]
else:
traj_maps = self.given_maps_when_test
print('Using given maps')
# write traj map and save batch order
test_index = dict()
for batch_index, traj_map in zip(agents_batch, traj_maps):
total_count = 0
if not batch_index in test_index:
test_index[batch_index] = []
for agent_index, _ in enumerate(agents_batch[batch_index]):
agents_batch[batch_index][agent_index].write_traj_map(traj_map)
start_count = total_count
total_count += 1
if test_on_neighbors:
agents_batch[batch_index][
agent_index].write_traj_map_for_neighbors(traj_map)
nei_len = agents_batch[batch_index][
agent_index].neighbor_number
total_count += nei_len
test_index[batch_index].append(
[i for i in range(start_count, total_count)])
return agents_batch, test_index
def write_traj_map_for_neighbors(self, trajmap: TrajectoryMapManager):
self.traj_map_neighbors = []
full_map = trajmap.traj_map
half_size = 16
for nei_traj in self.get_neighbor_traj():
center_pos = trajmap.real2map(nei_traj[-1, :])
original_map = cv2.resize(
full_map[np.maximum(center_pos[0] - 2 * half_size, 0):np.
minimum(center_pos[0] +
2 * half_size, full_map.shape[0]),
np.maximum(center_pos[1] - 2 * half_size, 0):np.
minimum(center_pos[1] +
2 * half_size, full_map.shape[1]), ],
(4 * half_size, 4 * half_size))
final_map = original_map[half_size:3 * half_size,
half_size:3 * half_size]
self.traj_map_neighbors.append(final_map)
def sample_data(self,
data_manager,
person_index,
add_noise=False,
reverse=False,
rotate=False,
desc='Calculate agent data',
use_time_bar=True,
random_sample=False,
sample_start=0.0,
given_trajmap=False,
return_trajmap=False):
"""
Sample training data from data_manager.
`random_sample`: 为0到1的正数时表示随机取样百分比，为-1到0的负数时表示按照数据集时间顺序百分比取样的终点，此时0～1正数`sample_start`表示起点
return: a list of Agent_Part
"""
agents = []
if person_index == 'auto':
if random_sample > 0 and random_sample < 1:
if USE_SEED:
random.seed(SEED)
person_index = random.sample(
[i for i in range(data_manager.person_number)],
int(data_manager.person_number * random_sample),
)
elif random_sample == 0 or random_sample >= 1 or random_sample < -1:
person_index = range(data_manager.person_number)
elif random_sample < 0 and random_sample >= -1:
person_index = [
i for i in range(
(data_manager.person_number *
np.abs(sample_start)).astype(int), # start index
(data_manager.person_number *
np.abs(random_sample)).astype(int), # end index
)
]
if use_time_bar:
itera = tqdm(person_index, desc=desc)
else:
itera = person_index
for person in itera:
agent_current = data_manager.agent_data[person]
start_frame = agent_current.start_frame
end_frame = agent_current.end_frame
for frame_point in range(start_frame, end_frame, self.args.step):
if frame_point + self.total_frames > end_frame:
break
# type: Agent_Part
sample_agent = data_manager.get_trajectory(
person,
frame_point,
frame_point + self.obs_frames,
frame_point + self.total_frames,
calculate_social=self.args.calculate_social,
normalization=self.args.normalization,
add_noise=add_noise,
reverse=reverse,
rotate=rotate,
)
agents.append(sample_agent)
if not given_trajmap:
traj_trajmap = TrajectoryMapManager(agents)
for index in range(len(agents)):
agents[index].write_traj_map(traj_trajmap)
if return_trajmap:
return agents, traj_trajmap
else:
return agents
else:
for index in range(len(agents)):
agents[index].write_traj_map(given_trajmap)
return agents