diff --git a/README.md b/README.md index 9c62965..3b1c446 100644 --- a/README.md +++ b/README.md @@ -116,6 +116,7 @@ To evaluate a trained network: comming soon. ## Demo +### Image ```Shell python demo.py --threshold 0.5 --iou_threshold 0.5 --score --weight checkpoint_VOC_efficientdet-d1_34.pth --file_name demo.png ``` @@ -126,11 +127,16 @@ Output:

-## Webcam Demo +### Webcam Demo You can use a webcam in a real-time demo by running: ```Shell -python demo.py --threshold 0.5 --iou_threshold 0.5 --cam --score --weight checkpoint_VOC_efficientdet-d1_34.pth +python demo.py --threshold 0.5 --iou_threshold 0.5 --score --weight checkpoint_VOC_efficientdet-d1_34.pth --cam +``` + +### Video Demo +```Shell +python demo.py --threshold 0.5 --iou_threshold 0.5 --score --weight checkpoint_VOC_efficientdet-d1_34.pth --cam /path/to/video ``` ## Performance diff --git a/demo.py b/demo.py index 7219ef4..1f9759d 100644 --- a/demo.py +++ b/demo.py @@ -25,7 +25,7 @@ parser.add_argument('-w', '--weight', default='./weights/voc0712.pth', type=str, help='Weight model path') parser.add_argument('-c', '--cam', - action="store_true", help='Use camera') + help='Use camera') parser.add_argument('-f', '--file_name', default='pic.jpg', help='Image path') parser.add_argument('--num_class', default=21, type=int, @@ -84,6 +84,8 @@ def process(self, file_name=None, img=None, show=False): bbox_scores = list() colors = list() for j in range(scores.shape[0]): + if scores[j] < args.threshold: + continue bbox = transformed_anchors[[j], :][0].data.cpu().numpy() x1 = int(bbox[0]*origin_img.shape[1]/self.size_image[1]) y1 = int(bbox[1]*origin_img.shape[0]/self.size_image[0]) @@ -129,8 +131,11 @@ def process(self, file_name=None, img=None, show=False): else: return origin_img - def camera(self): - cap = cv2.VideoCapture(0) + def camera(self, video_path=None): + if video_path is not None: + cap = cv2.VideoCapture(video_path) + else: + cap = cv2.VideoCapture(0) if not cap.isOpened(): print("Unable to open camera") exit(-1) @@ -174,6 +179,6 @@ def camera(self): detect = Detect(weights=args.weight) print('cam: ', args.cam) if args.cam: - detect.camera() + detect.camera(args.cam) else: detect.process(file_name=args.file_name, show=True)