Skip to content

Commit

Permalink
updates
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher committed Sep 5, 2018
1 parent 0bfc4bc commit af9864d
Showing 1 changed file with 5 additions and 4 deletions.
9 changes: 5 additions & 4 deletions detect.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@

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()
Expand All @@ -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
Expand Down Expand Up @@ -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)

Expand Down

0 comments on commit af9864d

Please sign in to comment.