Commit 16c009ac authored by Gioele Ciaparrone's avatar Gioele Ciaparrone Committed by Facebook Github Bot

Added convenience arguments to infer_simple.py (#400)

Summary:
- Added command line argument to infer_simple.py to choose a custom output image format.
- Added command line argument to infer_simple.py to choose to output an image even when no object is found in it.
Closes https://github.com/facebookresearch/Detectron/pull/400

Reviewed By: ir413

Differential Revision: D8705039

Pulled By: rbgirshick

fbshipit-source-id: fbcd707821a58004971a88946a17f005545ecab4
parent f752666d
......@@ -251,7 +251,7 @@ def vis_one_image_opencv(
def vis_one_image(
im, im_name, output_dir, boxes, segms=None, keypoints=None, thresh=0.9,
kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=False,
ext='pdf'):
ext='pdf', out_when_no_box=False):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
......@@ -260,7 +260,7 @@ def vis_one_image(
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
if (boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh) and not out_when_no_box:
return
dataset_keypoints, _ = keypoint_utils.get_keypoints()
......@@ -281,9 +281,12 @@ def vis_one_image(
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
if boxes is None:
sorted_inds = [] # avoid crash when 'boxes' is None
else:
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
for i in sorted_inds:
......
......@@ -83,9 +83,22 @@ def parse_args():
default='jpg',
type=str
)
parser.add_argument(
'--always-out',
dest='out_when_no_box',
help='output image even when no object is found',
action='store_true'
)
parser.add_argument(
'im_or_folder', help='image or folder of images', default=None
)
parser.add_argument(
'--output-ext',
dest='output_ext',
help='output image file format (default: pdf)',
default='pdf',
type=str
)
if len(sys.argv) == 1:
parser.print_help()
sys.exit(1)
......@@ -115,7 +128,7 @@ def main(args):
for i, im_name in enumerate(im_list):
out_name = os.path.join(
args.output_dir, '{}'.format(os.path.basename(im_name) + '.pdf')
args.output_dir, '{}'.format(os.path.basename(im_name) + '.' + args.output_ext)
)
logger.info('Processing {} -> {}'.format(im_name, out_name))
im = cv2.imread(im_name)
......@@ -145,7 +158,9 @@ def main(args):
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2
kp_thresh=2,
ext=args.output_ext,
out_when_no_box=args.out_when_no_box
)
......
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