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)
def my_draw_graph_pred(im, boxes, cls_score, rel_score, gt_to_pred, roidb,
im_num):
"""
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: 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_scores = np.max(rel_score, 2)
sorteds = np.argsort(-rel_pred_scores, axis=None)
indexes = np.unravel_index(np.argsort(-rel_pred_scores, axis=None),
rel_pred_scores.shape)
rel_pred = []
all_rels = []
import cv2
good_det = np.argsort(-np.max(cls_score, 1), axis=None)
im2 = im #.copy()[:, :, (2, 1, 0)].astype(np.uint8) #im #cv2.imread('/home/gay/ext_lib/scene-graph-TF-release/data/scannet/images/1.jpg')w
for n in range(min(len(cls_pred), 10)):
lidx = cls_pred[good_det[n]]
b = boxes[good_det[n], :]
cv2.rectangle(im2, (int(b[0]), int(b[1])), (int(b[2]), int(b[3])),
(255, 0, 0), 5)
cv2.rectangle(im2, (int(b[4]), int(b[5])), (int(b[6]), int(b[7])),
(255, 0, 0), 5)
cv2.rectangle(im2, (int(b[8]), int(b[9])), (int(b[10]), int(b[11])),
(255, 0, 0), 5)
cv2.rectangle(im2, (int(b[12]), int(b[13])), (int(b[14]), int(b[15])),
(255, 0, 0), 5)
cv2.rectangle(im2, (int(b[16]), int(b[17])), (int(b[18]), int(b[19])),
(255, 0, 0), 5)
l = cls_pred[n]
cv2.putText(im2, cfg.ind_to_class[lidx],
(int(b[0]) + 40, int(b[1]) + 40), cv2.FONT_HERSHEY_PLAIN,
1, (0, 0, 255), 2)
#cv2.imshow('inage',im)
cv2.imwrite('visual' + str(im_num) + '.png', im2)
#cv2.waitKey()
for n in range(min(len(indexes[0]), 7)):
i = indexes[0][n]
j = indexes[1][n]
#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,
save_f='graph' + str(im_num) + '.png')