Commit e7f3a882 authored by Yukun Chen's avatar Yukun Chen Committed by Yuxin Wu

fix dp extraction bug for visualize(). (#934)

* fix dp extraction bug for visualize().

* gt_masks is actually not needed here.
parent e119ef75
......@@ -349,11 +349,7 @@ def visualize(model, model_path, nr_visualize=100, output_dir='output'):
utils.fs.mkdir_p(output_dir)
with tqdm.tqdm(total=nr_visualize) as pbar:
for idx, dp in itertools.islice(enumerate(df), nr_visualize):
img = dp[0]
if cfg.MODE_MASK:
gt_boxes, gt_labels, gt_masks = dp[-3:]
else:
gt_boxes, gt_labels = dp[-2:]
img, gt_boxes, gt_labels = dp['image'], dp['gt_boxes'], dp['gt_labels']
rpn_boxes, rpn_scores, all_scores, \
final_boxes, final_scores, final_labels = pred(img, gt_boxes, gt_labels)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment