glenn-jocher commited on
Commit
a9dc0c2
·
1 Parent(s): bfd51f6

update check_img_size() for model strides

Browse files
Files changed (2) hide show
  1. detect.py +1 -1
  2. test.py +1 -1
detect.py CHANGED
@@ -25,6 +25,7 @@ def detect(save_img=False):
25
  # torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning
26
  # model.fuse()
27
  model.to(device).eval()
 
28
  if half:
29
  model.half() # to FP16
30
 
@@ -152,7 +153,6 @@ if __name__ == '__main__':
152
  parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
153
  parser.add_argument('--augment', action='store_true', help='augmented inference')
154
  opt = parser.parse_args()
155
- opt.img_size = check_img_size(opt.img_size)
156
  print(opt)
157
 
158
  with torch.no_grad():
 
25
  # torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning
26
  # model.fuse()
27
  model.to(device).eval()
28
+ imgsz = check_img_size(imgsz, s=model.model[-1].stride.max()) # check img_size
29
  if half:
30
  model.half() # to FP16
31
 
 
153
  parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
154
  parser.add_argument('--augment', action='store_true', help='augmented inference')
155
  opt = parser.parse_args()
 
156
  print(opt)
157
 
158
  with torch.no_grad():
test.py CHANGED
@@ -34,6 +34,7 @@ def test(data,
34
  torch_utils.model_info(model)
35
  model.fuse()
36
  model.to(device)
 
37
 
38
  # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
39
  # if device.type != 'cpu' and torch.cuda.device_count() > 1:
@@ -242,7 +243,6 @@ if __name__ == '__main__':
242
  parser.add_argument('--merge', action='store_true', help='use Merge NMS')
243
  parser.add_argument('--verbose', action='store_true', help='report mAP by class')
244
  opt = parser.parse_args()
245
- opt.img_size = check_img_size(opt.img_size)
246
  opt.save_json = opt.save_json or opt.data.endswith('coco.yaml')
247
  opt.data = check_file(opt.data) # check file
248
  print(opt)
 
34
  torch_utils.model_info(model)
35
  model.fuse()
36
  model.to(device)
37
+ imgsz = check_img_size(imgsz, s=model.model[-1].stride.max()) # check img_size
38
 
39
  # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
40
  # if device.type != 'cpu' and torch.cuda.device_count() > 1:
 
243
  parser.add_argument('--merge', action='store_true', help='use Merge NMS')
244
  parser.add_argument('--verbose', action='store_true', help='report mAP by class')
245
  opt = parser.parse_args()
 
246
  opt.save_json = opt.save_json or opt.data.endswith('coco.yaml')
247
  opt.data = check_file(opt.data) # check file
248
  print(opt)