dmlc--dgl
Tong He
cbee427839
* add working scripts * add frcnn training script * remove redundent files * refactor validation computation, will optimize sgdet and training * validation finally finished * f-rcnn training * test reldn * rm file * update reldn training * data preprocess to h5 * temp * use coco json * fix conflict * new obj dataset for detection * update training * before cleanup * remove abundant files * add arg parse to train * cleanup code file * update * fix * add readme * add ipynb as demo * add demo pic * update readme * add demo script * improve paths * improve readme * add docstrings * fix args description * update readme * add models from s3 * update README Co-authored-by: Minjie Wang <minjie.wang@nyu.edu>
47 行
1.7 KiB
Python
47 行
1.7 KiB
Python
import numpy as np
|
|
import gluoncv as gcv
|
|
from matplotlib import pyplot as plt
|
|
|
|
def plot_sg(img, preds, obj_classes, rel_classes, topk=1):
|
|
'''visualization of generated scene graph'''
|
|
size = img.shape[0:2]
|
|
box_scale = np.array([size[1], size[0], size[1], size[0]])
|
|
topk = min(topk, preds.shape[0])
|
|
ax = gcv.utils.viz.plot_image(img)
|
|
for i in range(topk):
|
|
rel = int(preds[i, 2])
|
|
src = int(preds[i, 3])
|
|
dst = int(preds[i, 4])
|
|
src_name = obj_classes[src]
|
|
dst_name = obj_classes[dst]
|
|
rel_name = rel_classes[rel]
|
|
src_bbox = preds[i, 5:9] * box_scale
|
|
dst_bbox = preds[i, 9:13] * box_scale
|
|
|
|
src_center = np.array([(src_bbox[0] + src_bbox[2]) / 2, (src_bbox[1] + src_bbox[3]) / 2])
|
|
dst_center = np.array([(dst_bbox[0] + dst_bbox[2]) / 2, (dst_bbox[1] + dst_bbox[3]) / 2])
|
|
rel_center = (src_center + dst_center) / 2
|
|
|
|
line_x = np.array([(src_bbox[0] + src_bbox[2]) / 2, (dst_bbox[0] + dst_bbox[2]) / 2])
|
|
line_y = np.array([(src_bbox[1] + src_bbox[3]) / 2, (dst_bbox[1] + dst_bbox[3]) / 2])
|
|
|
|
ax.plot(line_x, line_y,
|
|
linewidth=3.0, alpha=0.7, color=plt.cm.cool(rel))
|
|
|
|
ax.text(src_center[0], src_center[1],
|
|
'{:s}'.format(src_name),
|
|
bbox=dict(alpha=0.5),
|
|
fontsize=12, color='white')
|
|
ax.text(dst_center[0], dst_center[1],
|
|
'{:s}'.format(dst_name),
|
|
bbox=dict(alpha=0.5),
|
|
fontsize=12, color='white')
|
|
ax.text(rel_center[0], rel_center[1],
|
|
'{:s}'.format(rel_name),
|
|
bbox=dict(alpha=0.5),
|
|
fontsize=12, color='white')
|
|
return ax
|
|
|
|
plot_sg(img, preds, 2)
|
|
|