diff --git a/fcos/bin/fcos b/fcos/bin/fcos index e2d309ab..c34167e7 100644 --- a/fcos/bin/fcos +++ b/fcos/bin/fcos @@ -23,36 +23,37 @@ def pretty_print(bbox_results): )) -parser = argparse.ArgumentParser(description="FCOS Object Detector") -parser.add_argument( - "input_image", - help="path or url to an input image", -) -args = parser.parse_args() +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="FCOS Object Detector") + parser.add_argument( + "input_image", + help="path or url to an input image", + ) + args = parser.parse_args() -fcos = FCOS( - model_name="fcos_syncbn_bs32_c128_MNV2_FPN_1x", - nms_thresh=0.6, - cpu_only=not torch.cuda.is_available() # if you do not have GPUs, please set cpu_only as True -) + fcos = FCOS( + model_name="fcos_syncbn_bs32_c128_MNV2_FPN_1x", + nms_thresh=0.6, + cpu_only=not torch.cuda.is_available() # if you do not have GPUs, please set cpu_only as True + ) -im = io.imread(args.input_image) -assert im.shape[-1] == 3, "only 3-channel images are supported" + im = io.imread(args.input_image) + assert im.shape[-1] == 3, "only 3-channel images are supported" -# convert from RGB to BGR because fcos assumes the BRG input image -im = im[..., ::-1].copy() + # convert from RGB to BGR because fcos assumes the BRG input image + im = im[..., ::-1].copy() -# resize image to have its shorter size == 800 -f = 800.0 / float(min(im.shape[:2])) -im = cv2.resize(im, (0, 0), fx=f, fy=f) + # resize image to have its shorter size == 800 + f = 800.0 / float(min(im.shape[:2])) + im = cv2.resize(im, (0, 0), fx=f, fy=f) -start_time = time.time() + start_time = time.time() -bbox_results = fcos.detect(im) + bbox_results = fcos.detect(im) -inference_time = time.time() - start_time + inference_time = time.time() - start_time -pretty_print(bbox_results) -print("Predicted in {:.2f} seconds.".format(inference_time)) -print("Press any key to exit...") -fcos.show_bboxes(im, bbox_results) + pretty_print(bbox_results) + print("Predicted in {:.2f} seconds.".format(inference_time)) + print("Press any key to exit...") + fcos.show_bboxes(im, bbox_results)