dmlc--dgl
xiang song(charlie.song)
67cb7a43a0
* Deprecate multi-graph * Handle heterograph and edge_ids * lint * Fix * Remove multigraph in C++ end * Fix lint * Add some test and fix something * Fix * Fix * upd * Fix some test case * Fix * Fix Co-authored-by: Ubuntu <ubuntu@ip-172-31-51-214.ec2.internal> Co-authored-by: Jinjing Zhou <VoVAllen@users.noreply.github.com> Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>
285 行
11 KiB
Python
285 行
11 KiB
Python
import dgl
|
|
from mxnet import nd
|
|
import numpy as np
|
|
|
|
def bbox_improve(bbox):
|
|
'''bbox encoding'''
|
|
area = (bbox[:,2] - bbox[:,0]) * (bbox[:,3] - bbox[:,1])
|
|
return nd.concat(bbox, area.expand_dims(1))
|
|
|
|
def extract_edge_bbox(g):
|
|
'''bbox encoding'''
|
|
src, dst = g.edges(order='eid')
|
|
n = g.number_of_edges()
|
|
src_bbox = g.ndata['pred_bbox'][src.asnumpy()]
|
|
dst_bbox = g.ndata['pred_bbox'][dst.asnumpy()]
|
|
edge_bbox = nd.zeros((n, 4), ctx=g.ndata['pred_bbox'].context)
|
|
edge_bbox[:,0] = nd.stack(src_bbox[:,0], dst_bbox[:,0]).min(axis=0)
|
|
edge_bbox[:,1] = nd.stack(src_bbox[:,1], dst_bbox[:,1]).min(axis=0)
|
|
edge_bbox[:,2] = nd.stack(src_bbox[:,2], dst_bbox[:,2]).max(axis=0)
|
|
edge_bbox[:,3] = nd.stack(src_bbox[:,3], dst_bbox[:,3]).max(axis=0)
|
|
return edge_bbox
|
|
|
|
def build_graph_train(g_slice, gt_bbox, img, ids, scores, bbox, feat_ind,
|
|
spatial_feat, iou_thresh=0.5,
|
|
bbox_improvement=True, scores_top_k=50, overlap=False):
|
|
'''given ground truth and predicted bboxes, assign the label to the predicted w.r.t iou_thresh'''
|
|
# match and re-factor the graph
|
|
img_size = img.shape[2:4]
|
|
gt_bbox[:, :, 0] /= img_size[1]
|
|
gt_bbox[:, :, 1] /= img_size[0]
|
|
gt_bbox[:, :, 2] /= img_size[1]
|
|
gt_bbox[:, :, 3] /= img_size[0]
|
|
bbox[:, :, 0] /= img_size[1]
|
|
bbox[:, :, 1] /= img_size[0]
|
|
bbox[:, :, 2] /= img_size[1]
|
|
bbox[:, :, 3] /= img_size[0]
|
|
|
|
n_graph = len(g_slice)
|
|
g_pred_batch = []
|
|
for gi in range(n_graph):
|
|
g = g_slice[gi]
|
|
ctx = g.ndata['bbox'].context
|
|
inds = np.where(scores[gi, :, 0].asnumpy() > 0)[0].tolist()
|
|
if len(inds) == 0:
|
|
return None
|
|
if len(inds) > scores_top_k:
|
|
top_score_inds = scores[gi, inds, 0].asnumpy().argsort()[::-1][0:scores_top_k]
|
|
inds = np.array(inds)[top_score_inds].tolist()
|
|
|
|
n_nodes = len(inds)
|
|
roi_ind = feat_ind[gi, inds].squeeze(axis=1)
|
|
g_pred = dgl.DGLGraph()
|
|
g_pred.add_nodes(n_nodes, {'pred_bbox': bbox[gi, inds],
|
|
'node_feat': spatial_feat[gi, roi_ind],
|
|
'node_class_pred': ids[gi, inds, 0],
|
|
'node_class_logit': nd.log(scores[gi, inds, 0] + 1e-7)})
|
|
|
|
# iou matching
|
|
ious = nd.contrib.box_iou(gt_bbox[gi], g_pred.ndata['pred_bbox']).asnumpy()
|
|
H, W = ious.shape
|
|
h = H
|
|
w = W
|
|
pred_to_gt_ind = np.array([-1 for i in range(W)])
|
|
pred_to_gt_class_match = [0 for i in range(W)]
|
|
pred_to_gt_class_match_id = [0 for i in range(W)]
|
|
while h > 0 and w > 0:
|
|
ind = int(ious.argmax())
|
|
row_ind = ind // W
|
|
col_ind = ind % W
|
|
if ious[row_ind, col_ind] < iou_thresh:
|
|
break
|
|
pred_to_gt_ind[col_ind] = row_ind
|
|
gt_node_class = g.ndata['node_class'][row_ind]
|
|
pred_node_class = g_pred.ndata['node_class_pred'][col_ind]
|
|
if gt_node_class == pred_node_class:
|
|
pred_to_gt_class_match[col_ind] = 1
|
|
pred_to_gt_class_match_id[col_ind] = row_ind
|
|
ious[row_ind, :] = -1
|
|
ious[:, col_ind] = -1
|
|
h -= 1
|
|
w -= 1
|
|
|
|
n_nodes = g_pred.number_of_nodes()
|
|
triplet = []
|
|
adjmat = np.zeros((n_nodes, n_nodes))
|
|
|
|
src, dst = g.all_edges(order='eid')
|
|
eid_keys = np.column_stack([src.asnumpy(), dst.asnumpy()])
|
|
eid_dict = {}
|
|
for i, key in enumerate(eid_keys):
|
|
k = tuple(key)
|
|
if k not in eid_dict:
|
|
eid_dict[k] = [i]
|
|
else:
|
|
eid_dict[k].append(i)
|
|
ori_rel_class = g.edata['rel_class'].asnumpy()
|
|
for i in range(n_nodes):
|
|
for j in range(n_nodes):
|
|
if i != j:
|
|
if pred_to_gt_class_match[i] and pred_to_gt_class_match[j]:
|
|
sub_gt_id = pred_to_gt_class_match_id[i]
|
|
ob_gt_id = pred_to_gt_class_match_id[j]
|
|
eids = eid_dict[(sub_gt_id, ob_gt_id)]
|
|
rel_cls = ori_rel_class[eids]
|
|
n_edges_between = len(rel_cls)
|
|
for ii in range(n_edges_between):
|
|
triplet.append((i, j, rel_cls[ii]))
|
|
adjmat[i,j] = 1
|
|
else:
|
|
triplet.append((i, j, 0))
|
|
src, dst, rel_class = tuple(zip(*triplet))
|
|
rel_class = nd.array(rel_class, ctx=ctx).expand_dims(1)
|
|
g_pred.add_edges(src, dst, data={'rel_class': rel_class})
|
|
|
|
# other operations
|
|
n_nodes = g_pred.number_of_nodes()
|
|
n_edges = g_pred.number_of_edges()
|
|
if bbox_improvement:
|
|
g_pred.ndata['pred_bbox'] = bbox_improve(g_pred.ndata['pred_bbox'])
|
|
g_pred.edata['rel_bbox'] = extract_edge_bbox(g_pred)
|
|
g_pred.edata['batch_id'] = nd.zeros((n_edges, 1), ctx = ctx) + gi
|
|
|
|
# remove non-overlapping edges
|
|
if overlap:
|
|
overlap_ious = nd.contrib.box_iou(g_pred.ndata['pred_bbox'][:,0:4],
|
|
g_pred.ndata['pred_bbox'][:,0:4]).asnumpy()
|
|
cols, rows = np.where(overlap_ious <= 1e-7)
|
|
if cols.shape[0] > 0:
|
|
eids = g_pred.edge_ids(cols, rows)[2].asnumpy().tolist()
|
|
if len(eids):
|
|
g_pred.remove_edges(eids)
|
|
if g_pred.number_of_edges() == 0:
|
|
g_pred = None
|
|
g_pred_batch.append(g_pred)
|
|
|
|
if n_graph > 1:
|
|
return dgl.batch(g_pred_batch)
|
|
else:
|
|
return g_pred_batch[0]
|
|
|
|
def build_graph_validate_gt_obj(img, gt_ids, bbox, spatial_feat,
|
|
bbox_improvement=True, overlap=False):
|
|
'''given ground truth bbox and label, build graph for validation'''
|
|
n_batch = img.shape[0]
|
|
img_size = img.shape[2:4]
|
|
bbox[:, :, 0] /= img_size[1]
|
|
bbox[:, :, 1] /= img_size[0]
|
|
bbox[:, :, 2] /= img_size[1]
|
|
bbox[:, :, 3] /= img_size[0]
|
|
ctx = img.context
|
|
|
|
g_batch = []
|
|
for btc in range(n_batch):
|
|
inds = np.where(bbox[btc].sum(1).asnumpy() > 0)[0].tolist()
|
|
if len(inds) == 0:
|
|
continue
|
|
n_nodes = len(inds)
|
|
g_pred = dgl.DGLGraph()
|
|
g_pred.add_nodes(n_nodes, {'pred_bbox': bbox[btc, inds],
|
|
'node_feat': spatial_feat[btc, inds],
|
|
'node_class_pred': gt_ids[btc, inds, 0],
|
|
'node_class_logit': nd.zeros_like(gt_ids[btc, inds, 0], ctx=ctx)})
|
|
|
|
edge_list = []
|
|
for i in range(n_nodes - 1):
|
|
for j in range(i + 1, n_nodes):
|
|
edge_list.append((i, j))
|
|
src, dst = tuple(zip(*edge_list))
|
|
g_pred.add_edges(src, dst)
|
|
g_pred.add_edges(dst, src)
|
|
|
|
n_nodes = g_pred.number_of_nodes()
|
|
n_edges = g_pred.number_of_edges()
|
|
if bbox_improvement:
|
|
g_pred.ndata['pred_bbox'] = bbox_improve(g_pred.ndata['pred_bbox'])
|
|
g_pred.edata['rel_bbox'] = extract_edge_bbox(g_pred)
|
|
g_pred.edata['batch_id'] = nd.zeros((n_edges, 1), ctx = ctx) + btc
|
|
|
|
g_batch.append(g_pred)
|
|
|
|
if len(g_batch) == 0:
|
|
return None
|
|
if len(g_batch) > 1:
|
|
return dgl.batch(g_batch)
|
|
return g_batch[0]
|
|
|
|
def build_graph_validate_gt_bbox(img, ids, scores, bbox, spatial_feat, gt_ids=None,
|
|
bbox_improvement=True, overlap=False):
|
|
'''given ground truth bbox, build graph for validation'''
|
|
n_batch = img.shape[0]
|
|
img_size = img.shape[2:4]
|
|
bbox[:, :, 0] /= img_size[1]
|
|
bbox[:, :, 1] /= img_size[0]
|
|
bbox[:, :, 2] /= img_size[1]
|
|
bbox[:, :, 3] /= img_size[0]
|
|
ctx = img.context
|
|
|
|
g_batch = []
|
|
for btc in range(n_batch):
|
|
id_btc = scores[btc][:,:,0].argmax(0)
|
|
score_btc = scores[btc][:,:,0].max(0)
|
|
inds = np.where(bbox[btc].sum(1).asnumpy() > 0)[0].tolist()
|
|
if len(inds) == 0:
|
|
continue
|
|
n_nodes = len(inds)
|
|
g_pred = dgl.DGLGraph()
|
|
g_pred.add_nodes(n_nodes, {'pred_bbox': bbox[btc, inds],
|
|
'node_feat': spatial_feat[btc, inds],
|
|
'node_class_pred': id_btc,
|
|
'node_class_logit': nd.log(score_btc + 1e-7)})
|
|
|
|
edge_list = []
|
|
for i in range(n_nodes - 1):
|
|
for j in range(i + 1, n_nodes):
|
|
edge_list.append((i, j))
|
|
src, dst = tuple(zip(*edge_list))
|
|
g_pred.add_edges(src, dst)
|
|
g_pred.add_edges(dst, src)
|
|
|
|
n_nodes = g_pred.number_of_nodes()
|
|
n_edges = g_pred.number_of_edges()
|
|
if bbox_improvement:
|
|
g_pred.ndata['pred_bbox'] = bbox_improve(g_pred.ndata['pred_bbox'])
|
|
g_pred.edata['rel_bbox'] = extract_edge_bbox(g_pred)
|
|
g_pred.edata['batch_id'] = nd.zeros((n_edges, 1), ctx = ctx) + btc
|
|
|
|
g_batch.append(g_pred)
|
|
|
|
if len(g_batch) == 0:
|
|
return None
|
|
if len(g_batch) > 1:
|
|
return dgl.batch(g_batch)
|
|
return g_batch[0]
|
|
|
|
def build_graph_validate_pred(img, ids, scores, bbox, feat_ind, spatial_feat,
|
|
bbox_improvement=True, scores_top_k=50, overlap=False):
|
|
'''given predicted bbox, build graph for validation'''
|
|
n_batch = img.shape[0]
|
|
img_size = img.shape[2:4]
|
|
bbox[:, :, 0] /= img_size[1]
|
|
bbox[:, :, 1] /= img_size[0]
|
|
bbox[:, :, 2] /= img_size[1]
|
|
bbox[:, :, 3] /= img_size[0]
|
|
ctx = img.context
|
|
|
|
g_batch = []
|
|
for btc in range(n_batch):
|
|
inds = np.where(scores[btc, :, 0].asnumpy() > 0)[0].tolist()
|
|
if len(inds) == 0:
|
|
continue
|
|
if len(inds) > scores_top_k:
|
|
top_score_inds = scores[btc, inds, 0].asnumpy().argsort()[::-1][0:scores_top_k]
|
|
inds = np.array(inds)[top_score_inds].tolist()
|
|
n_nodes = len(inds)
|
|
roi_ind = feat_ind[btc, inds].squeeze(axis=1)
|
|
|
|
g_pred = dgl.DGLGraph()
|
|
g_pred.add_nodes(n_nodes, {'pred_bbox': bbox[btc, inds],
|
|
'node_feat': spatial_feat[btc, roi_ind],
|
|
'node_class_pred': ids[btc, inds, 0],
|
|
'node_class_logit': nd.log(scores[btc, inds, 0] + 1e-7)})
|
|
|
|
edge_list = []
|
|
for i in range(n_nodes - 1):
|
|
for j in range(i + 1, n_nodes):
|
|
edge_list.append((i, j))
|
|
src, dst = tuple(zip(*edge_list))
|
|
g_pred.add_edges(src, dst)
|
|
g_pred.add_edges(dst, src)
|
|
|
|
n_nodes = g_pred.number_of_nodes()
|
|
n_edges = g_pred.number_of_edges()
|
|
if bbox_improvement:
|
|
g_pred.ndata['pred_bbox'] = bbox_improve(g_pred.ndata['pred_bbox'])
|
|
g_pred.edata['rel_bbox'] = extract_edge_bbox(g_pred)
|
|
g_pred.edata['batch_id'] = nd.zeros((n_edges, 1), ctx = ctx) + btc
|
|
|
|
g_batch.append(g_pred)
|
|
|
|
if len(g_batch) == 0:
|
|
return None
|
|
if len(g_batch) > 1:
|
|
return dgl.batch(g_batch)
|
|
return g_batch[0]
|