def construct_graph_object(self):
"""
Constructs the entire Graph object to serve as input to the MPN, and stores it in self.graph_obj,
"""
# Load Appearance Data
reid_embeddings, node_feats = self._load_appearance_data()
# Determine graph connectivity (i.e. edges) and compute edge features
edge_ixs = self._get_edge_ixs(reid_embeddings)
edge_feats_dict = compute_edge_feats_dict(edge_ixs = edge_ixs, det_df = self.graph_df,
fps = self.seq_info_dict['fps'],
use_cuda = self.inference_mode)
edge_feats = [edge_feats_dict[feat_names] for feat_names in self.dataset_params['edge_feats_to_use'] if feat_names in edge_feats_dict]
edge_feats = torch.stack(edge_feats).T
# Compute embeddings distances. Pairwise distance computation might create out of memmory errors, hence we batch it
emb_dists = []
for i in range(0, edge_ixs[0].shape[0], 50000):
emb_dists.append(F.pairwise_distance(reid_embeddings[edge_ixs[0][i:i + 50000]],
reid_embeddings[edge_ixs[1][i:i + 50000]]).view(-1, 1))
emb_dists = torch.cat(emb_dists, dim=0)
# Add embedding distances to edge features if needed
if 'emb_dist' in self.dataset_params['edge_feats_to_use']:
edge_feats = torch.cat((edge_feats, emb_dists), dim = 1)
self.graph_obj = Graph(x = node_feats,
edge_attr = torch.cat((edge_feats, edge_feats), dim = 0),
edge_index = torch.cat((edge_ixs, torch.stack((edge_ixs[1], edge_ixs[0]))), dim=1))
if self.inference_mode:
self.graph_obj.reid_emb_dists = torch.cat((emb_dists, emb_dists))
self.graph_obj.to(torch.device("cuda" if torch.cuda.is_available() and self.inference_mode else "cpu"))
def _get_edge_ixs(self, reid_embeddings):
"""
Constructs graph edges by taking pairs of nodes with valid time connections (not in same frame, not too far
apart in time) and perhaps taking KNNs according to reid embeddings.
Args:
reid_embeddings: torch.tensor with shape (num_nodes, reid_embeds_dim)
Returns:
torch.tensor withs shape (2, num_edges)
"""
edge_ixs = get_time_valid_conn_ixs(
frame_num=torch.from_numpy(self.graph_df.frame.values),
max_frame_dist=self.max_frame_dist,
use_cuda=self.inference_mode
and self.graph_df["frame_path"].iloc[0].find("MOT17-03") == -1,
)
edge_feats_dict = None
if (
"max_feet_vel" in self.dataset_params
and self.dataset_params["max_feet_vel"] is not None
): # New parameter. We do graph pruning based on feet velocity
# print("VELOCITY PRUNING")
edge_feats_dict = compute_edge_feats_dict(
edge_ixs=edge_ixs,
det_df=self.graph_df,
fps=self.seq_info_dict["fps"],
use_cuda=self.inference_mode,
)
feet_vel = torch.sqrt(
edge_feats_dict["norm_feet_x_dists"] ** 2
+ edge_feats_dict["norm_feet_y_dists"] ** 2
)
vel_mask = feet_vel < self.dataset_params["max_feet_vel"]
edge_ixs = edge_ixs.T[vel_mask].T
for feat_name, feat_vals in edge_feats_dict.items():
edge_feats_dict[feat_name] = feat_vals[vel_mask]
# During inference, top k nns must not be done here, as it is computed independently for sequence chunks
if not self.inference_mode and self.dataset_params["top_k_nns"] is not None:
reid_pwise_dist = F.pairwise_distance(
reid_embeddings[edge_ixs[0]], reid_embeddings[edge_ixs[1]]
)
k_nns_mask = get_knn_mask(
pwise_dist=reid_pwise_dist,
edge_ixs=edge_ixs,
num_nodes=self.graph_df.shape[0],
top_k_nns=self.dataset_params["top_k_nns"],
reciprocal_k_nns=self.dataset_params["reciprocal_k_nns"],
symmetric_edges=False,
use_cuda=self.inference_mode,
)
edge_ixs = edge_ixs.T[k_nns_mask].T
if edge_feats_dict is not None:
for feat_name, feat_vals in edge_feats_dict.items():
edge_feats_dict[feat_name] = feat_vals[k_nns_mask]
return edge_ixs, edge_feats_dict
def construct_graph_object(self, set_reid_embeddings, set_node_feats,
start_frame):
"""
Constructs the entire Graph object to serve as input to the MPN, and stores it in self.graph_obj,
"""
if set_reid_embeddings is None or set_node_feats is None:
# Load Appearance Data
reid_embeddings, node_feats = self._load_appearance_data()
else:
reid_embeddings = []
node_feats = []
for ix in range(len(self.graph_df)):
row = self.graph_df.iloc[ix]
row_id = int(row["id"])
row_frame = int(row["frame"]) + start_frame
reid_embeddings.append(set_reid_embeddings[row_id][row_frame])
node_feats.append(set_node_feats[row_id][row_frame])
reid_embeddings = torch.stack(reid_embeddings).cuda()
node_feats = torch.stack(node_feats).cuda()
# Determine graph connectivity (i.e. edges) and compute edge features
edge_ixs = self._get_edge_ixs(reid_embeddings)
edge_feats_dict = compute_edge_feats_dict(
edge_ixs=edge_ixs,
det_df=self.graph_df,
fps=self.seq_info_dict["fps"],
use_cuda=self.inference_mode,
)
edge_feats = [
edge_feats_dict[feat_names]
for feat_names in self.dataset_params["edge_feats_to_use"]
if feat_names in edge_feats_dict
]
edge_feats = torch.stack(edge_feats).T
# Compute embeddings distances. Pairwise distance computation might create out of memmory errors, hence we batch it
emb_dists = []
for i in range(0, edge_ixs[0].shape[0], 50000):
emb_dists.append(
F.pairwise_distance(
reid_embeddings[edge_ixs[0][i:i + 50000]],
reid_embeddings[edge_ixs[1][i:i + 50000]],
).view(-1, 1))
emb_dists = torch.cat(emb_dists, dim=0)
# Add embedding distances to edge features if needed
if "emb_dist" in self.dataset_params["edge_feats_to_use"]:
edge_feats = torch.cat((edge_feats, emb_dists), dim=1)
self.graph_obj = Graph(
x=node_feats,
edge_attr=torch.cat((edge_feats, edge_feats), dim=0),
edge_index=torch.cat(
(edge_ixs, torch.stack((edge_ixs[1], edge_ixs[0]))), dim=1),
)
if self.inference_mode:
self.graph_obj.reid_emb_dists = torch.cat((emb_dists, emb_dists))
self.graph_obj.to(
torch.device("cuda" if torch.cuda.is_available()
and self.inference_mode else "cpu"))
def construct_graph_object(self):
"""
Constructs the entire Graph object to serve as input to the MPN, and stores it in self.graph_obj,
"""
# Load Appearance Data
reid_embeddings, node_feats = self._load_appearance_data()
joints = self._get_joints()
if "joints" in self.dataset_params["node_feats_to_use"]:
node_feats = torch.cat(
(node_feats, joints.reshape((joints.shape[0], -1)).float()), dim=1
)
if "bb_positions" in self.dataset_params["node_feats_to_use"]:
positions = torch.zeros((node_feats.shape[0], 4), dtype=torch.float32)
positions[:, 0] = torch.tensor(self.graph_df.bb_left.array)
positions[:, 1] = torch.tensor(self.graph_df.bb_top.array)
positions[:, 2] = torch.tensor(self.graph_df.bb_right.array)
positions[:, 3] = torch.tensor(self.graph_df.bb_bot.array)
if self.inference_mode:
positions = positions.cuda()
node_feats = torch.cat((node_feats, positions), dim=1)
# TODO: Add other options for incorporating joints
# Determine graph connectivity (i.e. edges) and compute edge features
edge_ixs, edge_feats_dict = self._get_edge_ixs(reid_embeddings)
# print(edge_ixs.shape)
if edge_feats_dict is None:
edge_feats_dict = compute_edge_feats_dict(
edge_ixs=edge_ixs,
det_df=self.graph_df,
fps=self.seq_info_dict["fps"],
use_cuda=self.inference_mode,
)
edge_feats = [
edge_feats_dict[feat_names]
for feat_names in self.dataset_params["edge_feats_to_use"]
if feat_names in edge_feats_dict
]
edge_feats = torch.stack(edge_feats).T
# print("Edge features", edge_feats.shape)
# Compute embeddings distances. Pairwise distance computation might create out of memmory errors, hence we batch it
emb_dists = []
for i in range(0, edge_ixs.shape[1], 50000):
emb_dists.append(
F.pairwise_distance(
reid_embeddings[edge_ixs[0][i : i + 50000]],
reid_embeddings[edge_ixs[1][i : i + 50000]],
).view(-1, 1)
)
emb_dists = torch.cat(emb_dists, dim=0)
# Add embedding distances to edge features if needed
if "emb_dist" in self.dataset_params["edge_feats_to_use"]:
edge_feats = torch.cat((edge_feats, emb_dists), dim=1)
self.graph_obj = Graph(
x=node_feats,
edge_attr=torch.cat((edge_feats, edge_feats), dim=0),
edge_index=torch.cat(
(edge_ixs, torch.stack((edge_ixs[1], edge_ixs[0]))), dim=1
),
)
if self.inference_mode:
self.graph_obj.reid_emb_dists = torch.cat((emb_dists, emb_dists))
self.graph_obj.to(
torch.device(
"cuda" if torch.cuda.is_available() and self.inference_mode else "cpu"
)
)