diff --git a/detect.py b/detect.py index d52f8139..b2acb320 100755 --- a/detect.py +++ b/detect.py @@ -18,8 +18,8 @@ parser.add_argument('-txt_out', type=bool, default=False) parser.add_argument('-cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path') parser.add_argument('-class_path', type=str, default='data/coco.names', help='path to class label file') -parser.add_argument('-conf_thres', type=float, default=0.8, help='object confidence threshold') -parser.add_argument('-nms_thres', type=float, default=0.5, help='iou threshold for non-maximum suppression') +parser.add_argument('-conf_thres', type=float, default=0.99, help='object confidence threshold') +parser.add_argument('-nms_thres', type=float, default=0.45, help='iou threshold for non-maximum suppression') parser.add_argument('-batch_size', type=int, default=1, help='size of the batches') parser.add_argument('-img_size', type=int, default=32 * 13, help='size of each image dimension') opt = parser.parse_args() @@ -33,7 +33,8 @@ def detect(opt): # Load model model = Darknet(opt.cfg, opt.img_size) - weights_path = 'checkpoints/yolov3.weights' + #weights_path = 'checkpoints/yolov3.weights' + weights_path = 'checkpoints/latest.pt' if weights_path.endswith('.weights'): # saved in darknet format load_weights(model, weights_path) else: # endswith('.pt'), saved in pytorch format @@ -130,7 +131,7 @@ def detect(opt): if opt.plot_flag: # Add the bbox to the plot - label = '%s %.2f' % (classes[int(cls_pred)], cls_conf) if cls_conf > 0.05 else None + label = '%s %.2f' % (classes[int(cls_pred)], conf) color = bbox_colors[int(np.where(unique_classes == int(cls_pred))[0])] plot_one_box([x1, y1, x2, y2], img, label=label, color=color, line_thickness=3)