def draw_graph_pred(im, boxes, cls_score, rel_score, gt_to_pred, roidb, idx):
"""
Draw a predicted scene graph. To keep the graph interpretable, only draw
the node and edge predictions that have correspounding ground truth
labels.
args:
im: image
boxes: prediceted boxes
cls_score: object classification scores
rel_score: relation classification scores
gt_to_pred: a mapping from ground truth box indices to predicted box indices
idx: for saving
roidb: roidb
"""
gt_relations = roidb['gt_relations']
im = im[:, :, (2, 1, 0)].copy()
cls_pred = np.argmax(cls_score, 1)
rel_pred_mat = np.argmax(rel_score, 2)
rel_pred = []
all_rels = []
for i in xrange(rel_pred_mat.shape[0]):
for j in xrange(rel_pred_mat.shape[1]):
# find graph predictions (nodes and edges) that have
# correspounding ground truth annotations
# ignore nodes that have no edge connections
for rel in gt_relations:
if rel[0] not in gt_to_pred or rel[1] not in gt_to_pred:
continue
# discard duplicate grounding
if [i, j] in all_rels:
continue
if i == gt_to_pred[rel[0]] and j == gt_to_pred[rel[1]]:
rel_pred.append([i, j, rel_pred_mat[i, j], 1])
all_rels.append([i, j])
rel_pred = np.array(rel_pred)
if rel_pred.size == 0:
return
# indices of predicted boxes
pred_inds = rel_pred[:, :2].ravel()
# draw graph predictions
graph_dict = draw_scene_graph(cls_pred, pred_inds, rel_pred)
viz_scene_graph(im, boxes, cls_pred, pred_inds, rel_pred, preprocess=False)
out_boxes = []
for box, cls in zip(boxes[pred_inds], cls_pred[pred_inds]):
out_boxes.append(box[cls * 4:(cls + 1) * 4].tolist())
graph_dict['boxes'] = out_boxes
do_save = 'y'
if do_save == 'y':
scipy.misc.imsave('cherry/im_%i.png' % idx, im)
fn = open('cherry/graph_%i.json' % idx, 'w+')
json.dump(graph_dict, fn)
print(idx)
def draw_graph_pred(im, boxes, cls_score, rel_score, gt_to_pred, roidb, iter_num, save_dir, use_gt=False):
"""
Draw a predicted scene graph. To keep the graph interpretable, only draw
the node and edge predictions that have corresponding ground truth
labels.
args:
im: image
boxes: predicted boxes
cls_score: object classification scores
rel_score: relation classification scores
gt_to_pred: a mapping from ground truth box indices to predicted box indices
idx: for saving
roidb: roidb
"""
gt_relations = roidb['gt_relations']
im = im[:, :, (2, 1, 0)].copy()
cls_pred = np.argmax(cls_score, 1)
rel_pred_mat = np.argmax(rel_score, 2)
rel_pred = []
all_rels = []
for i in xrange(rel_pred_mat.shape[0]):
for j in xrange(rel_pred_mat.shape[1]):
# find graph predictions (nodes and edges) that have
# corresponding ground truth annotations
# ignore nodes that have no edge connections
for rel in gt_relations:
if rel[0] not in gt_to_pred or rel[1] not in gt_to_pred:
continue
# discard duplicate grounding
if [i, j] in all_rels:
continue
if i == gt_to_pred[rel[0]] and j == gt_to_pred[rel[1]]:
rel_pred.append([i, j, rel_pred_mat[i,j], 1])
all_rels.append([i, j])
rel_pred = np.array(rel_pred)
if use_gt:
cls_pred = roidb['gt_classes']
rel_pred = gt_relations
if rel_pred.size == 0:
return
# indices of predicted boxes
pred_inds = rel_pred[:, :2].ravel()
# draw graph predictions
graph_dict = draw_scene_graph(cls_pred, pred_inds, rel_pred, iter_num, save_dir)
viz_scene_graph(im, boxes, cls_pred, iter_num, save_dir, pred_inds, rel_pred, preprocess=False)
# Obsolete: Uncomment this to get the scene graph predictions
"""
def draw_graph_pred_no_truth(im, boxes, cls_score, rel_score):
"""
Draw a predicted scene graph. To keep the graph interpretable, only draw
the node and edge predictions that have correspounding ground truth
labels.
args:
im: image
boxes: prediceted boxes
cls_score: object classification scores
rel_score: relation classification scores
gt_to_pred: a mapping from ground truth box indices to predicted box indices
idx: for saving
roidb: roidb
"""
#gt_relations = roidb['gt_relations']
im = im[:, :, (2, 1, 0)].copy()
cls_pred = np.argmax(cls_score, 1)
rel_pred_mat = np.argmax(rel_score, 2)
rel_pred = []
all_rels = []
for i in xrange(rel_pred_mat.shape[0]):
for j in xrange(rel_pred_mat.shape[1]):
# discard duplicate grounding
if [i, j] in all_rels:
continue
rel_pred.append([i, j, rel_pred_mat[i, j], 1])
all_rels.append([i, j])
rel_pred = np.array(rel_pred)
if rel_pred.size == 0:
return
# indices of predicted boxes
pred_inds = rel_pred[:, :2].ravel()
# draw graph predictions
graph_dict = draw_scene_graph(cls_pred, pred_inds, rel_pred)
viz_scene_graph(im, boxes, cls_pred, pred_inds, rel_pred, preprocess=False)