Zengyf-CVer commited on
Commit
112bf3b
β€’
1 Parent(s): 7cea19b

app update

Browse files
.gitignore CHANGED
@@ -65,6 +65,7 @@
65
 
66
  !requirements.txt
67
  !.pre-commit-config.yaml
 
68
 
69
  test.py
70
  test*.py
 
65
 
66
  !requirements.txt
67
  !.pre-commit-config.yaml
68
+ !data/*
69
 
70
  test.py
71
  test*.py
data/coco128.yaml ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # YOLOv5 πŸš€ by Ultralytics, GPL-3.0 license
2
+ # COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
3
+ # Example usage: python train.py --data coco128.yaml
4
+ # parent
5
+ # β”œβ”€β”€ yolov5
6
+ # └── datasets
7
+ # └── coco128 ← downloads here (7 MB)
8
+
9
+
10
+ # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
11
+ path: ../datasets/coco128 # dataset root dir
12
+ train: images/train2017 # train images (relative to 'path') 128 images
13
+ val: images/train2017 # val images (relative to 'path') 128 images
14
+ test: # test images (optional)
15
+
16
+ # Classes
17
+ nc: 80 # number of classes
18
+ names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
19
+ 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
20
+ 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
21
+ 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
22
+ 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
23
+ 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
24
+ 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
25
+ 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
26
+ 'hair drier', 'toothbrush'] # class names
27
+
28
+
29
+ # Download script/URL (optional)
30
+ download: https://ultralytics.com/assets/coco128.zip
export.py CHANGED
@@ -67,9 +67,9 @@ if platform.system() != 'Windows':
67
  from models.experimental import attempt_load
68
  from models.yolo import Detect
69
  from utils.dataloaders import LoadImages
70
- from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, check_version, colorstr,
71
- file_size, print_args, url2file)
72
- from utils.torch_utils import select_device
73
 
74
 
75
  def export_formats():
@@ -152,13 +152,12 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
152
  # Simplify
153
  if simplify:
154
  try:
155
- check_requirements(('onnx-simplifier',))
 
156
  import onnxsim
157
 
158
  LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
159
- model_onnx, check = onnxsim.simplify(model_onnx,
160
- dynamic_input_shape=dynamic,
161
- input_shapes={'images': list(im.shape)} if dynamic else None)
162
  assert check, 'assert check failed'
163
  onnx.save(model_onnx, f)
164
  except Exception as e:
@@ -217,8 +216,9 @@ def export_coreml(model, im, file, int8, half, prefix=colorstr('CoreML:')):
217
  return None, None
218
 
219
 
220
- def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
221
  # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
 
222
  try:
223
  assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
224
  try:
@@ -231,11 +231,11 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
231
  if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
232
  grid = model.model[-1].anchor_grid
233
  model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
234
- export_onnx(model, im, file, 12, train, False, simplify) # opset 12
235
  model.model[-1].anchor_grid = grid
236
  else: # TensorRT >= 8
237
  check_version(trt.__version__, '8.0.0', hard=True) # require tensorrt>=8.0.0
238
- export_onnx(model, im, file, 13, train, False, simplify) # opset 13
239
  onnx = file.with_suffix('.onnx')
240
 
241
  LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
@@ -264,6 +264,14 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
264
  for out in outputs:
265
  LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}')
266
 
 
 
 
 
 
 
 
 
267
  LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine in {f}')
268
  if builder.platform_has_fast_fp16 and half:
269
  config.set_flag(trt.BuilderFlag.FP16)
@@ -363,7 +371,7 @@ def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=c
363
  converter.optimizations = [tf.lite.Optimize.DEFAULT]
364
  if int8:
365
  from models.tf import representative_dataset_gen
366
- dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data
367
  converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)
368
  converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
369
  converter.target_spec.supported_types = []
@@ -402,7 +410,7 @@ def export_edgetpu(file, prefix=colorstr('Edge TPU:')):
402
  f = str(file).replace('.pt', '-int8_edgetpu.tflite') # Edge TPU model
403
  f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model
404
 
405
- cmd = f"edgetpu_compiler -s -o {file.parent} {f_tfl}"
406
  subprocess.run(cmd.split(), check=True)
407
 
408
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
@@ -447,7 +455,7 @@ def export_tfjs(file, prefix=colorstr('TensorFlow.js:')):
447
  LOGGER.info(f'\n{prefix} export failure: {e}')
448
 
449
 
450
- @torch.no_grad()
451
  def run(
452
  data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
453
  weights=ROOT / 'yolov5s.pt', # weights path
@@ -461,7 +469,7 @@ def run(
461
  keras=False, # use Keras
462
  optimize=False, # TorchScript: optimize for mobile
463
  int8=False, # CoreML/TF INT8 quantization
464
- dynamic=False, # ONNX/TF: dynamic axes
465
  simplify=False, # ONNX: simplify model
466
  opset=12, # ONNX: opset version
467
  verbose=False, # TensorRT: verbose log
@@ -492,6 +500,8 @@ def run(
492
  # Checks
493
  imgsz *= 2 if len(imgsz) == 1 else 1 # expand
494
  assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}'
 
 
495
 
496
  # Input
497
  gs = int(max(model.stride)) # grid size (max stride)
@@ -519,7 +529,7 @@ def run(
519
  if jit:
520
  f[0] = export_torchscript(model, im, file, optimize)
521
  if engine: # TensorRT required before ONNX
522
- f[1] = export_engine(model, im, file, train, half, simplify, workspace, verbose)
523
  if onnx or xml: # OpenVINO requires ONNX
524
  f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify)
525
  if xml: # OpenVINO
@@ -578,7 +588,7 @@ def parse_opt():
578
  parser.add_argument('--keras', action='store_true', help='TF: use Keras')
579
  parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
580
  parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
581
- parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
582
  parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
583
  parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
584
  parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
 
67
  from models.experimental import attempt_load
68
  from models.yolo import Detect
69
  from utils.dataloaders import LoadImages
70
+ from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, check_version, check_yaml,
71
+ colorstr, file_size, print_args, url2file)
72
+ from utils.torch_utils import select_device, smart_inference_mode
73
 
74
 
75
  def export_formats():
 
152
  # Simplify
153
  if simplify:
154
  try:
155
+ cuda = torch.cuda.is_available()
156
+ check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
157
  import onnxsim
158
 
159
  LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
160
+ model_onnx, check = onnxsim.simplify(model_onnx)
 
 
161
  assert check, 'assert check failed'
162
  onnx.save(model_onnx, f)
163
  except Exception as e:
 
216
  return None, None
217
 
218
 
219
+ def export_engine(model, im, file, train, half, dynamic, simplify, workspace=4, verbose=False):
220
  # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
221
+ prefix = colorstr('TensorRT:')
222
  try:
223
  assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
224
  try:
 
231
  if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
232
  grid = model.model[-1].anchor_grid
233
  model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
234
+ export_onnx(model, im, file, 12, train, dynamic, simplify) # opset 12
235
  model.model[-1].anchor_grid = grid
236
  else: # TensorRT >= 8
237
  check_version(trt.__version__, '8.0.0', hard=True) # require tensorrt>=8.0.0
238
+ export_onnx(model, im, file, 13, train, dynamic, simplify) # opset 13
239
  onnx = file.with_suffix('.onnx')
240
 
241
  LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
 
264
  for out in outputs:
265
  LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}')
266
 
267
+ if dynamic:
268
+ if im.shape[0] <= 1:
269
+ LOGGER.warning(f"{prefix}WARNING: --dynamic model requires maximum --batch-size argument")
270
+ profile = builder.create_optimization_profile()
271
+ for inp in inputs:
272
+ profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)
273
+ config.add_optimization_profile(profile)
274
+
275
  LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine in {f}')
276
  if builder.platform_has_fast_fp16 and half:
277
  config.set_flag(trt.BuilderFlag.FP16)
 
371
  converter.optimizations = [tf.lite.Optimize.DEFAULT]
372
  if int8:
373
  from models.tf import representative_dataset_gen
374
+ dataset = LoadImages(check_dataset(check_yaml(data))['train'], img_size=imgsz, auto=False)
375
  converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)
376
  converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
377
  converter.target_spec.supported_types = []
 
410
  f = str(file).replace('.pt', '-int8_edgetpu.tflite') # Edge TPU model
411
  f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model
412
 
413
+ cmd = f"edgetpu_compiler -s -d -k 10 --out_dir {file.parent} {f_tfl}"
414
  subprocess.run(cmd.split(), check=True)
415
 
416
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
 
455
  LOGGER.info(f'\n{prefix} export failure: {e}')
456
 
457
 
458
+ @smart_inference_mode()
459
  def run(
460
  data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
461
  weights=ROOT / 'yolov5s.pt', # weights path
 
469
  keras=False, # use Keras
470
  optimize=False, # TorchScript: optimize for mobile
471
  int8=False, # CoreML/TF INT8 quantization
472
+ dynamic=False, # ONNX/TF/TensorRT: dynamic axes
473
  simplify=False, # ONNX: simplify model
474
  opset=12, # ONNX: opset version
475
  verbose=False, # TensorRT: verbose log
 
500
  # Checks
501
  imgsz *= 2 if len(imgsz) == 1 else 1 # expand
502
  assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}'
503
+ if optimize:
504
+ assert device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu'
505
 
506
  # Input
507
  gs = int(max(model.stride)) # grid size (max stride)
 
529
  if jit:
530
  f[0] = export_torchscript(model, im, file, optimize)
531
  if engine: # TensorRT required before ONNX
532
+ f[1] = export_engine(model, im, file, train, half, dynamic, simplify, workspace, verbose)
533
  if onnx or xml: # OpenVINO requires ONNX
534
  f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify)
535
  if xml: # OpenVINO
 
588
  parser.add_argument('--keras', action='store_true', help='TF: use Keras')
589
  parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
590
  parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
591
+ parser.add_argument('--dynamic', action='store_true', help='ONNX/TF/TensorRT: dynamic axes')
592
  parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
593
  parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
594
  parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
models/common.py CHANGED
@@ -25,7 +25,7 @@ from utils.dataloaders import exif_transpose, letterbox
25
  from utils.general import (LOGGER, check_requirements, check_suffix, check_version, colorstr, increment_path,
26
  make_divisible, non_max_suppression, scale_coords, xywh2xyxy, xyxy2xywh)
27
  from utils.plots import Annotator, colors, save_one_box
28
- from utils.torch_utils import copy_attr, time_sync
29
 
30
 
31
  def autopad(k, p=None): # kernel, padding
@@ -305,7 +305,7 @@ class Concat(nn.Module):
305
 
306
  class DetectMultiBackend(nn.Module):
307
  # YOLOv5 MultiBackend class for python inference on various backends
308
- def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False):
309
  # Usage:
310
  # PyTorch: weights = *.pt
311
  # TorchScript: *.torchscript
@@ -331,7 +331,7 @@ class DetectMultiBackend(nn.Module):
331
  names = yaml.safe_load(f)['names']
332
 
333
  if pt: # PyTorch
334
- model = attempt_load(weights if isinstance(weights, list) else w, device=device)
335
  stride = max(int(model.stride.max()), 32) # model stride
336
  names = model.module.names if hasattr(model, 'module') else model.names # get class names
337
  model.half() if fp16 else model.float()
@@ -384,19 +384,24 @@ class DetectMultiBackend(nn.Module):
384
  logger = trt.Logger(trt.Logger.INFO)
385
  with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
386
  model = runtime.deserialize_cuda_engine(f.read())
 
387
  bindings = OrderedDict()
388
  fp16 = False # default updated below
 
389
  for index in range(model.num_bindings):
390
  name = model.get_binding_name(index)
391
  dtype = trt.nptype(model.get_binding_dtype(index))
392
- shape = tuple(model.get_binding_shape(index))
 
 
 
 
 
 
393
  data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(device)
394
  bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr()))
395
- if model.binding_is_input(index) and dtype == np.float16:
396
- fp16 = True
397
  binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
398
- context = model.create_execution_context()
399
- batch_size = bindings['images'].shape[0]
400
  elif coreml: # CoreML
401
  LOGGER.info(f'Loading {w} for CoreML inference...')
402
  import coremltools as ct
@@ -441,6 +446,8 @@ class DetectMultiBackend(nn.Module):
441
  output_details = interpreter.get_output_details() # outputs
442
  elif tfjs:
443
  raise Exception('ERROR: YOLOv5 TF.js inference is not supported')
 
 
444
  self.__dict__.update(locals()) # assign all variables to self
445
 
446
  def forward(self, im, augment=False, visualize=False, val=False):
@@ -464,7 +471,13 @@ class DetectMultiBackend(nn.Module):
464
  im = im.cpu().numpy() # FP32
465
  y = self.executable_network([im])[self.output_layer]
466
  elif self.engine: # TensorRT
467
- assert im.shape == self.bindings['images'].shape, (im.shape, self.bindings['images'].shape)
 
 
 
 
 
 
468
  self.binding_addrs['images'] = int(im.data_ptr())
469
  self.context.execute_v2(list(self.binding_addrs.values()))
470
  y = self.bindings['output'].data
@@ -550,6 +563,9 @@ class AutoShape(nn.Module):
550
  self.dmb = isinstance(model, DetectMultiBackend) # DetectMultiBackend() instance
551
  self.pt = not self.dmb or model.pt # PyTorch model
552
  self.model = model.eval()
 
 
 
553
 
554
  def _apply(self, fn):
555
  # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
@@ -562,7 +578,7 @@ class AutoShape(nn.Module):
562
  m.anchor_grid = list(map(fn, m.anchor_grid))
563
  return self
564
 
565
- @torch.no_grad()
566
  def forward(self, imgs, size=640, augment=False, profile=False):
567
  # Inference from various sources. For height=640, width=1280, RGB images example inputs are:
568
  # file: imgs = 'data/images/zidane.jpg' # str or PosixPath
 
25
  from utils.general import (LOGGER, check_requirements, check_suffix, check_version, colorstr, increment_path,
26
  make_divisible, non_max_suppression, scale_coords, xywh2xyxy, xyxy2xywh)
27
  from utils.plots import Annotator, colors, save_one_box
28
+ from utils.torch_utils import copy_attr, smart_inference_mode, time_sync
29
 
30
 
31
  def autopad(k, p=None): # kernel, padding
 
305
 
306
  class DetectMultiBackend(nn.Module):
307
  # YOLOv5 MultiBackend class for python inference on various backends
308
+ def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True):
309
  # Usage:
310
  # PyTorch: weights = *.pt
311
  # TorchScript: *.torchscript
 
331
  names = yaml.safe_load(f)['names']
332
 
333
  if pt: # PyTorch
334
+ model = attempt_load(weights if isinstance(weights, list) else w, device=device, inplace=True, fuse=fuse)
335
  stride = max(int(model.stride.max()), 32) # model stride
336
  names = model.module.names if hasattr(model, 'module') else model.names # get class names
337
  model.half() if fp16 else model.float()
 
384
  logger = trt.Logger(trt.Logger.INFO)
385
  with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
386
  model = runtime.deserialize_cuda_engine(f.read())
387
+ context = model.create_execution_context()
388
  bindings = OrderedDict()
389
  fp16 = False # default updated below
390
+ dynamic = False
391
  for index in range(model.num_bindings):
392
  name = model.get_binding_name(index)
393
  dtype = trt.nptype(model.get_binding_dtype(index))
394
+ if model.binding_is_input(index):
395
+ if -1 in tuple(model.get_binding_shape(index)): # dynamic
396
+ dynamic = True
397
+ context.set_binding_shape(index, tuple(model.get_profile_shape(0, index)[2]))
398
+ if dtype == np.float16:
399
+ fp16 = True
400
+ shape = tuple(context.get_binding_shape(index))
401
  data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(device)
402
  bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr()))
 
 
403
  binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
404
+ batch_size = bindings['images'].shape[0] # if dynamic, this is instead max batch size
 
405
  elif coreml: # CoreML
406
  LOGGER.info(f'Loading {w} for CoreML inference...')
407
  import coremltools as ct
 
446
  output_details = interpreter.get_output_details() # outputs
447
  elif tfjs:
448
  raise Exception('ERROR: YOLOv5 TF.js inference is not supported')
449
+ else:
450
+ raise Exception(f'ERROR: {w} is not a supported format')
451
  self.__dict__.update(locals()) # assign all variables to self
452
 
453
  def forward(self, im, augment=False, visualize=False, val=False):
 
471
  im = im.cpu().numpy() # FP32
472
  y = self.executable_network([im])[self.output_layer]
473
  elif self.engine: # TensorRT
474
+ if self.dynamic and im.shape != self.bindings['images'].shape:
475
+ i_in, i_out = (self.model.get_binding_index(x) for x in ('images', 'output'))
476
+ self.context.set_binding_shape(i_in, im.shape) # reshape if dynamic
477
+ self.bindings['images'] = self.bindings['images']._replace(shape=im.shape)
478
+ self.bindings['output'].data.resize_(tuple(self.context.get_binding_shape(i_out)))
479
+ s = self.bindings['images'].shape
480
+ assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"
481
  self.binding_addrs['images'] = int(im.data_ptr())
482
  self.context.execute_v2(list(self.binding_addrs.values()))
483
  y = self.bindings['output'].data
 
563
  self.dmb = isinstance(model, DetectMultiBackend) # DetectMultiBackend() instance
564
  self.pt = not self.dmb or model.pt # PyTorch model
565
  self.model = model.eval()
566
+ if self.pt:
567
+ m = self.model.model.model[-1] if self.dmb else self.model.model[-1] # Detect()
568
+ m.inplace = False # Detect.inplace=False for safe multithread inference
569
 
570
  def _apply(self, fn):
571
  # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
 
578
  m.anchor_grid = list(map(fn, m.anchor_grid))
579
  return self
580
 
581
+ @smart_inference_mode()
582
  def forward(self, imgs, size=640, augment=False, profile=False):
583
  # Inference from various sources. For height=640, width=1280, RGB images example inputs are:
584
  # file: imgs = 'data/images/zidane.jpg' # str or PosixPath
models/experimental.py CHANGED
@@ -89,8 +89,6 @@ def attempt_load(weights, device=None, inplace=True, fuse=True):
89
  if t is Detect and not isinstance(m.anchor_grid, list):
90
  delattr(m, 'anchor_grid')
91
  setattr(m, 'anchor_grid', [torch.zeros(1)] * m.nl)
92
- elif t is Conv:
93
- m._non_persistent_buffers_set = set() # torch 1.6.0 compatibility
94
  elif t is nn.Upsample and not hasattr(m, 'recompute_scale_factor'):
95
  m.recompute_scale_factor = None # torch 1.11.0 compatibility
96
 
 
89
  if t is Detect and not isinstance(m.anchor_grid, list):
90
  delattr(m, 'anchor_grid')
91
  setattr(m, 'anchor_grid', [torch.zeros(1)] * m.nl)
 
 
92
  elif t is nn.Upsample and not hasattr(m, 'recompute_scale_factor'):
93
  m.recompute_scale_factor = None # torch 1.11.0 compatibility
94
 
models/yolo.py CHANGED
@@ -7,6 +7,7 @@ Usage:
7
  """
8
 
9
  import argparse
 
10
  import os
11
  import platform
12
  import sys
@@ -49,7 +50,7 @@ class Detect(nn.Module):
49
  self.anchor_grid = [torch.zeros(1)] * self.nl # init anchor grid
50
  self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,2)
51
  self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv
52
- self.inplace = inplace # use in-place ops (e.g. slice assignment)
53
 
54
  def forward(self, x):
55
  z = [] # inference output
@@ -75,12 +76,12 @@ class Detect(nn.Module):
75
 
76
  return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
77
 
78
- def _make_grid(self, nx=20, ny=20, i=0):
79
  d = self.anchors[i].device
80
  t = self.anchors[i].dtype
81
  shape = 1, self.na, ny, nx, 2 # grid shape
82
  y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
83
- if check_version(torch.__version__, '1.10.0'): # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility
84
  yv, xv = torch.meshgrid(y, x, indexing='ij')
85
  else:
86
  yv, xv = torch.meshgrid(y, x)
@@ -259,10 +260,8 @@ def parse_model(d, ch): # model_dict, input_channels(3)
259
  for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args
260
  m = eval(m) if isinstance(m, str) else m # eval strings
261
  for j, a in enumerate(args):
262
- try:
263
  args[j] = eval(a) if isinstance(a, str) else a # eval strings
264
- except NameError:
265
- pass
266
 
267
  n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain
268
  if m in (Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
 
7
  """
8
 
9
  import argparse
10
+ import contextlib
11
  import os
12
  import platform
13
  import sys
 
50
  self.anchor_grid = [torch.zeros(1)] * self.nl # init anchor grid
51
  self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,2)
52
  self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv
53
+ self.inplace = inplace # use inplace ops (e.g. slice assignment)
54
 
55
  def forward(self, x):
56
  z = [] # inference output
 
76
 
77
  return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
78
 
79
+ def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, '1.10.0')):
80
  d = self.anchors[i].device
81
  t = self.anchors[i].dtype
82
  shape = 1, self.na, ny, nx, 2 # grid shape
83
  y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
84
+ if torch_1_10: # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility
85
  yv, xv = torch.meshgrid(y, x, indexing='ij')
86
  else:
87
  yv, xv = torch.meshgrid(y, x)
 
260
  for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args
261
  m = eval(m) if isinstance(m, str) else m # eval strings
262
  for j, a in enumerate(args):
263
+ with contextlib.suppress(NameError):
264
  args[j] = eval(a) if isinstance(a, str) else a # eval strings
 
 
265
 
266
  n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain
267
  if m in (Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
utils/autoanchor.py CHANGED
@@ -10,7 +10,7 @@ import torch
10
  import yaml
11
  from tqdm import tqdm
12
 
13
- from utils.general import LOGGER, colorstr, emojis
14
 
15
  PREFIX = colorstr('AutoAnchor: ')
16
 
@@ -45,9 +45,9 @@ def check_anchors(dataset, model, thr=4.0, imgsz=640):
45
  bpr, aat = metric(anchors.cpu().view(-1, 2))
46
  s = f'\n{PREFIX}{aat:.2f} anchors/target, {bpr:.3f} Best Possible Recall (BPR). '
47
  if bpr > 0.98: # threshold to recompute
48
- LOGGER.info(emojis(f'{s}Current anchors are a good fit to dataset βœ…'))
49
  else:
50
- LOGGER.info(emojis(f'{s}Anchors are a poor fit to dataset ⚠️, attempting to improve...'))
51
  na = m.anchors.numel() // 2 # number of anchors
52
  try:
53
  anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
@@ -62,7 +62,7 @@ def check_anchors(dataset, model, thr=4.0, imgsz=640):
62
  s = f'{PREFIX}Done βœ… (optional: update model *.yaml to use these anchors in the future)'
63
  else:
64
  s = f'{PREFIX}Done ⚠️ (original anchors better than new anchors, proceeding with original anchors)'
65
- LOGGER.info(emojis(s))
66
 
67
 
68
  def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
 
10
  import yaml
11
  from tqdm import tqdm
12
 
13
+ from utils.general import LOGGER, colorstr
14
 
15
  PREFIX = colorstr('AutoAnchor: ')
16
 
 
45
  bpr, aat = metric(anchors.cpu().view(-1, 2))
46
  s = f'\n{PREFIX}{aat:.2f} anchors/target, {bpr:.3f} Best Possible Recall (BPR). '
47
  if bpr > 0.98: # threshold to recompute
48
+ LOGGER.info(f'{s}Current anchors are a good fit to dataset βœ…')
49
  else:
50
+ LOGGER.info(f'{s}Anchors are a poor fit to dataset ⚠️, attempting to improve...')
51
  na = m.anchors.numel() // 2 # number of anchors
52
  try:
53
  anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
 
62
  s = f'{PREFIX}Done βœ… (optional: update model *.yaml to use these anchors in the future)'
63
  else:
64
  s = f'{PREFIX}Done ⚠️ (original anchors better than new anchors, proceeding with original anchors)'
65
+ LOGGER.info(s)
66
 
67
 
68
  def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
utils/autobatch.py CHANGED
@@ -8,7 +8,7 @@ from copy import deepcopy
8
  import numpy as np
9
  import torch
10
 
11
- from utils.general import LOGGER, colorstr, emojis
12
  from utils.torch_utils import profile
13
 
14
 
@@ -62,5 +62,5 @@ def autobatch(model, imgsz=640, fraction=0.9, batch_size=16):
62
  b = batch_sizes[max(i - 1, 0)] # select prior safe point
63
 
64
  fraction = np.polyval(p, b) / t # actual fraction predicted
65
- LOGGER.info(emojis(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) βœ…'))
66
  return b
 
8
  import numpy as np
9
  import torch
10
 
11
+ from utils.general import LOGGER, colorstr
12
  from utils.torch_utils import profile
13
 
14
 
 
62
  b = batch_sizes[max(i - 1, 0)] # select prior safe point
63
 
64
  fraction = np.polyval(p, b) / t # actual fraction predicted
65
+ LOGGER.info(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) βœ…')
66
  return b
utils/dataloaders.py CHANGED
@@ -3,6 +3,7 @@
3
  Dataloaders and dataset utils
4
  """
5
 
 
6
  import glob
7
  import hashlib
8
  import json
@@ -55,13 +56,10 @@ def get_hash(paths):
55
  def exif_size(img):
56
  # Returns exif-corrected PIL size
57
  s = img.size # (width, height)
58
- try:
59
  rotation = dict(img._getexif().items())[orientation]
60
  if rotation in [6, 8]: # rotation 270 or 90
61
  s = (s[1], s[0])
62
- except Exception:
63
- pass
64
-
65
  return s
66
 
67
 
@@ -91,6 +89,13 @@ def exif_transpose(image):
91
  return image
92
 
93
 
 
 
 
 
 
 
 
94
  def create_dataloader(path,
95
  imgsz,
96
  batch_size,
@@ -130,13 +135,17 @@ def create_dataloader(path,
130
  nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers]) # number of workers
131
  sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)
132
  loader = DataLoader if image_weights else InfiniteDataLoader # only DataLoader allows for attribute updates
 
 
133
  return loader(dataset,
134
  batch_size=batch_size,
135
  shuffle=shuffle and sampler is None,
136
  num_workers=nw,
137
  sampler=sampler,
138
  pin_memory=True,
139
- collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn), dataset
 
 
140
 
141
 
142
  class InfiniteDataLoader(dataloader.DataLoader):
@@ -469,7 +478,7 @@ class LoadImagesAndLabels(Dataset):
469
  [cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove items
470
  labels, shapes, self.segments = zip(*cache.values())
471
  self.labels = list(labels)
472
- self.shapes = np.array(shapes, dtype=np.float64)
473
  self.im_files = list(cache.keys()) # update
474
  self.label_files = img2label_paths(cache.keys()) # update
475
  n = len(shapes) # number of images
@@ -671,8 +680,7 @@ class LoadImagesAndLabels(Dataset):
671
  interp = cv2.INTER_LINEAR if (self.augment or r > 1) else cv2.INTER_AREA
672
  im = cv2.resize(im, (int(w0 * r), int(h0 * r)), interpolation=interp)
673
  return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized
674
- else:
675
- return self.ims[i], self.im_hw0[i], self.im_hw[i] # im, hw_original, hw_resized
676
 
677
  def cache_images_to_disk(self, i):
678
  # Saves an image as an *.npy file for faster loading
@@ -849,18 +857,13 @@ class LoadImagesAndLabels(Dataset):
849
 
850
 
851
  # Ancillary functions --------------------------------------------------------------------------------------------------
852
- def create_folder(path='./new'):
853
- # Create folder
854
- if os.path.exists(path):
855
- shutil.rmtree(path) # delete output folder
856
- os.makedirs(path) # make new output folder
857
-
858
-
859
  def flatten_recursive(path=DATASETS_DIR / 'coco128'):
860
  # Flatten a recursive directory by bringing all files to top level
861
- new_path = Path(str(path) + '_flat')
862
- create_folder(new_path)
863
- for file in tqdm(glob.glob(str(Path(path)) + '/**/*.*', recursive=True)):
 
 
864
  shutil.copyfile(file, new_path / Path(file).name)
865
 
866
 
@@ -919,7 +922,7 @@ def autosplit(path=DATASETS_DIR / 'coco128/images', weights=(0.9, 0.1, 0.0), ann
919
  for i, img in tqdm(zip(indices, files), total=n):
920
  if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check label
921
  with open(path.parent / txt[i], 'a') as f:
922
- f.write('./' + img.relative_to(path.parent).as_posix() + '\n') # add image to txt file
923
 
924
 
925
  def verify_image_label(args):
@@ -974,21 +977,35 @@ def verify_image_label(args):
974
  return [None, None, None, None, nm, nf, ne, nc, msg]
975
 
976
 
977
- def dataset_stats(path='coco128.yaml', autodownload=False, verbose=False, profile=False, hub=False):
978
  """ Return dataset statistics dictionary with images and instances counts per split per class
979
  To run in parent directory: export PYTHONPATH="$PWD/yolov5"
980
- Usage1: from utils.dataloaders import *; dataset_stats('coco128.yaml', autodownload=True)
981
- Usage2: from utils.dataloaders import *; dataset_stats('path/to/coco128_with_yaml.zip')
982
  Arguments
983
  path: Path to data.yaml or data.zip (with data.yaml inside data.zip)
984
  autodownload: Attempt to download dataset if not found locally
985
- verbose: Print stats dictionary
986
  """
987
 
988
- def _round_labels(labels):
989
- # Update labels to integer class and 6 decimal place floats
990
- return [[int(c), *(round(x, 4) for x in points)] for c, *points in labels]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
991
 
 
992
  def _find_yaml(dir):
993
  # Return data.yaml file
994
  files = list(dir.glob('*.yaml')) or list(dir.rglob('*.yaml')) # try root level first and then recursive
@@ -999,26 +1016,25 @@ def dataset_stats(path='coco128.yaml', autodownload=False, verbose=False, profil
999
  assert len(files) == 1, f'Multiple *.yaml files found: {files}, only 1 *.yaml file allowed in {dir}'
1000
  return files[0]
1001
 
1002
- def _unzip(path):
1003
  # Unzip data.zip
1004
- if str(path).endswith('.zip'): # path is data.zip
1005
- assert Path(path).is_file(), f'Error unzipping {path}, file not found'
1006
- ZipFile(path).extractall(path=path.parent) # unzip
1007
- dir = path.with_suffix('') # dataset directory == zip name
1008
- assert dir.is_dir(), f'Error unzipping {path}, {dir} not found. path/to/abc.zip MUST unzip to path/to/abc/'
1009
- return True, str(dir), _find_yaml(dir) # zipped, data_dir, yaml_path
1010
- else: # path is data.yaml
1011
  return False, None, path
 
 
 
 
 
1012
 
1013
- def _hub_ops(f, max_dim=1920):
1014
  # HUB ops for 1 image 'f': resize and save at reduced quality in /dataset-hub for web/app viewing
1015
- f_new = im_dir / Path(f).name # dataset-hub image filename
1016
  try: # use PIL
1017
  im = Image.open(f)
1018
  r = max_dim / max(im.height, im.width) # ratio
1019
  if r < 1.0: # image too large
1020
  im = im.resize((int(im.width * r), int(im.height * r)))
1021
- im.save(f_new, 'JPEG', quality=75, optimize=True) # save
1022
  except Exception as e: # use OpenCV
1023
  print(f'WARNING: HUB ops PIL failure {f}: {e}')
1024
  im = cv2.imread(f)
@@ -1028,69 +1044,49 @@ def dataset_stats(path='coco128.yaml', autodownload=False, verbose=False, profil
1028
  im = cv2.resize(im, (int(im_width * r), int(im_height * r)), interpolation=cv2.INTER_AREA)
1029
  cv2.imwrite(str(f_new), im)
1030
 
1031
- zipped, data_dir, yaml_path = _unzip(Path(path))
1032
- try:
1033
- with open(check_yaml(yaml_path), errors='ignore') as f:
1034
- data = yaml.safe_load(f) # data dict
1035
- if zipped:
1036
- data['path'] = data_dir # TODO: should this be dir.resolve()?`
1037
- except Exception:
1038
- raise Exception("error/HUB/dataset_stats/yaml_load")
1039
-
1040
- check_dataset(data, autodownload) # download dataset if missing
1041
- hub_dir = Path(data['path'] + ('-hub' if hub else ''))
1042
- stats = {'nc': data['nc'], 'names': data['names']} # statistics dictionary
1043
- for split in 'train', 'val', 'test':
1044
- if data.get(split) is None:
1045
- stats[split] = None # i.e. no test set
1046
- continue
1047
- x = []
1048
- dataset = LoadImagesAndLabels(data[split]) # load dataset
1049
- for label in tqdm(dataset.labels, total=dataset.n, desc='Statistics'):
1050
- x.append(np.bincount(label[:, 0].astype(int), minlength=data['nc']))
1051
- x = np.array(x) # shape(128x80)
1052
- stats[split] = {
1053
- 'instance_stats': {
1054
- 'total': int(x.sum()),
1055
- 'per_class': x.sum(0).tolist()},
1056
- 'image_stats': {
1057
- 'total': dataset.n,
1058
- 'unlabelled': int(np.all(x == 0, 1).sum()),
1059
- 'per_class': (x > 0).sum(0).tolist()},
1060
- 'labels': [{
1061
- str(Path(k).name): _round_labels(v.tolist())} for k, v in zip(dataset.im_files, dataset.labels)]}
1062
-
1063
- if hub:
1064
- im_dir = hub_dir / 'images'
1065
- im_dir.mkdir(parents=True, exist_ok=True)
1066
- for _ in tqdm(ThreadPool(NUM_THREADS).imap(_hub_ops, dataset.im_files), total=dataset.n, desc='HUB Ops'):
 
 
 
 
 
 
 
1067
  pass
1068
-
1069
- # Profile
1070
- stats_path = hub_dir / 'stats.json'
1071
- if profile:
1072
- for _ in range(1):
1073
- file = stats_path.with_suffix('.npy')
1074
- t1 = time.time()
1075
- np.save(file, stats)
1076
- t2 = time.time()
1077
- x = np.load(file, allow_pickle=True)
1078
- print(f'stats.npy times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write')
1079
-
1080
- file = stats_path.with_suffix('.json')
1081
- t1 = time.time()
1082
- with open(file, 'w') as f:
1083
- json.dump(stats, f) # save stats *.json
1084
- t2 = time.time()
1085
- with open(file) as f:
1086
- x = json.load(f) # load hyps dict
1087
- print(f'stats.json times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write')
1088
-
1089
- # Save, print and return
1090
- if hub:
1091
- print(f'Saving {stats_path.resolve()}...')
1092
- with open(stats_path, 'w') as f:
1093
- json.dump(stats, f) # save stats.json
1094
- if verbose:
1095
- print(json.dumps(stats, indent=2, sort_keys=False))
1096
- return stats
 
3
  Dataloaders and dataset utils
4
  """
5
 
6
+ import contextlib
7
  import glob
8
  import hashlib
9
  import json
 
56
  def exif_size(img):
57
  # Returns exif-corrected PIL size
58
  s = img.size # (width, height)
59
+ with contextlib.suppress(Exception):
60
  rotation = dict(img._getexif().items())[orientation]
61
  if rotation in [6, 8]: # rotation 270 or 90
62
  s = (s[1], s[0])
 
 
 
63
  return s
64
 
65
 
 
89
  return image
90
 
91
 
92
+ def seed_worker(worker_id):
93
+ # Set dataloader worker seed https://pytorch.org/docs/stable/notes/randomness.html#dataloader
94
+ worker_seed = torch.initial_seed() % 2 ** 32
95
+ np.random.seed(worker_seed)
96
+ random.seed(worker_seed)
97
+
98
+
99
  def create_dataloader(path,
100
  imgsz,
101
  batch_size,
 
135
  nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers]) # number of workers
136
  sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)
137
  loader = DataLoader if image_weights else InfiniteDataLoader # only DataLoader allows for attribute updates
138
+ generator = torch.Generator()
139
+ generator.manual_seed(0)
140
  return loader(dataset,
141
  batch_size=batch_size,
142
  shuffle=shuffle and sampler is None,
143
  num_workers=nw,
144
  sampler=sampler,
145
  pin_memory=True,
146
+ collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn,
147
+ worker_init_fn=seed_worker,
148
+ generator=generator), dataset
149
 
150
 
151
  class InfiniteDataLoader(dataloader.DataLoader):
 
478
  [cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove items
479
  labels, shapes, self.segments = zip(*cache.values())
480
  self.labels = list(labels)
481
+ self.shapes = np.array(shapes)
482
  self.im_files = list(cache.keys()) # update
483
  self.label_files = img2label_paths(cache.keys()) # update
484
  n = len(shapes) # number of images
 
680
  interp = cv2.INTER_LINEAR if (self.augment or r > 1) else cv2.INTER_AREA
681
  im = cv2.resize(im, (int(w0 * r), int(h0 * r)), interpolation=interp)
682
  return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized
683
+ return self.ims[i], self.im_hw0[i], self.im_hw[i] # im, hw_original, hw_resized
 
684
 
685
  def cache_images_to_disk(self, i):
686
  # Saves an image as an *.npy file for faster loading
 
857
 
858
 
859
  # Ancillary functions --------------------------------------------------------------------------------------------------
 
 
 
 
 
 
 
860
  def flatten_recursive(path=DATASETS_DIR / 'coco128'):
861
  # Flatten a recursive directory by bringing all files to top level
862
+ new_path = Path(f'{str(path)}_flat')
863
+ if os.path.exists(new_path):
864
+ shutil.rmtree(new_path) # delete output folder
865
+ os.makedirs(new_path) # make new output folder
866
+ for file in tqdm(glob.glob(f'{str(Path(path))}/**/*.*', recursive=True)):
867
  shutil.copyfile(file, new_path / Path(file).name)
868
 
869
 
 
922
  for i, img in tqdm(zip(indices, files), total=n):
923
  if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check label
924
  with open(path.parent / txt[i], 'a') as f:
925
+ f.write(f'./{img.relative_to(path.parent).as_posix()}' + '\n') # add image to txt file
926
 
927
 
928
  def verify_image_label(args):
 
977
  return [None, None, None, None, nm, nf, ne, nc, msg]
978
 
979
 
980
+ class HUBDatasetStats():
981
  """ Return dataset statistics dictionary with images and instances counts per split per class
982
  To run in parent directory: export PYTHONPATH="$PWD/yolov5"
983
+ Usage1: from utils.dataloaders import *; HUBDatasetStats('coco128.yaml', autodownload=True)
984
+ Usage2: from utils.dataloaders import *; HUBDatasetStats('path/to/coco128_with_yaml.zip')
985
  Arguments
986
  path: Path to data.yaml or data.zip (with data.yaml inside data.zip)
987
  autodownload: Attempt to download dataset if not found locally
 
988
  """
989
 
990
+ def __init__(self, path='coco128.yaml', autodownload=False):
991
+ # Initialize class
992
+ zipped, data_dir, yaml_path = self._unzip(Path(path))
993
+ try:
994
+ with open(check_yaml(yaml_path), errors='ignore') as f:
995
+ data = yaml.safe_load(f) # data dict
996
+ if zipped:
997
+ data['path'] = data_dir
998
+ except Exception as e:
999
+ raise Exception("error/HUB/dataset_stats/yaml_load") from e
1000
+
1001
+ check_dataset(data, autodownload) # download dataset if missing
1002
+ self.hub_dir = Path(data['path'] + '-hub')
1003
+ self.im_dir = self.hub_dir / 'images'
1004
+ self.im_dir.mkdir(parents=True, exist_ok=True) # makes /images
1005
+ self.stats = {'nc': data['nc'], 'names': data['names']} # statistics dictionary
1006
+ self.data = data
1007
 
1008
+ @staticmethod
1009
  def _find_yaml(dir):
1010
  # Return data.yaml file
1011
  files = list(dir.glob('*.yaml')) or list(dir.rglob('*.yaml')) # try root level first and then recursive
 
1016
  assert len(files) == 1, f'Multiple *.yaml files found: {files}, only 1 *.yaml file allowed in {dir}'
1017
  return files[0]
1018
 
1019
+ def _unzip(self, path):
1020
  # Unzip data.zip
1021
+ if not str(path).endswith('.zip'): # path is data.yaml
 
 
 
 
 
 
1022
  return False, None, path
1023
+ assert Path(path).is_file(), f'Error unzipping {path}, file not found'
1024
+ ZipFile(path).extractall(path=path.parent) # unzip
1025
+ dir = path.with_suffix('') # dataset directory == zip name
1026
+ assert dir.is_dir(), f'Error unzipping {path}, {dir} not found. path/to/abc.zip MUST unzip to path/to/abc/'
1027
+ return True, str(dir), self._find_yaml(dir) # zipped, data_dir, yaml_path
1028
 
1029
+ def _hub_ops(self, f, max_dim=1920):
1030
  # HUB ops for 1 image 'f': resize and save at reduced quality in /dataset-hub for web/app viewing
1031
+ f_new = self.im_dir / Path(f).name # dataset-hub image filename
1032
  try: # use PIL
1033
  im = Image.open(f)
1034
  r = max_dim / max(im.height, im.width) # ratio
1035
  if r < 1.0: # image too large
1036
  im = im.resize((int(im.width * r), int(im.height * r)))
1037
+ im.save(f_new, 'JPEG', quality=50, optimize=True) # save
1038
  except Exception as e: # use OpenCV
1039
  print(f'WARNING: HUB ops PIL failure {f}: {e}')
1040
  im = cv2.imread(f)
 
1044
  im = cv2.resize(im, (int(im_width * r), int(im_height * r)), interpolation=cv2.INTER_AREA)
1045
  cv2.imwrite(str(f_new), im)
1046
 
1047
+ def get_json(self, save=False, verbose=False):
1048
+ # Return dataset JSON for Ultralytics HUB
1049
+ def _round(labels):
1050
+ # Update labels to integer class and 6 decimal place floats
1051
+ return [[int(c), *(round(x, 4) for x in points)] for c, *points in labels]
1052
+
1053
+ for split in 'train', 'val', 'test':
1054
+ if self.data.get(split) is None:
1055
+ self.stats[split] = None # i.e. no test set
1056
+ continue
1057
+ dataset = LoadImagesAndLabels(self.data[split]) # load dataset
1058
+ x = np.array([
1059
+ np.bincount(label[:, 0].astype(int), minlength=self.data['nc'])
1060
+ for label in tqdm(dataset.labels, total=dataset.n, desc='Statistics')]) # shape(128x80)
1061
+ self.stats[split] = {
1062
+ 'instance_stats': {
1063
+ 'total': int(x.sum()),
1064
+ 'per_class': x.sum(0).tolist()},
1065
+ 'image_stats': {
1066
+ 'total': dataset.n,
1067
+ 'unlabelled': int(np.all(x == 0, 1).sum()),
1068
+ 'per_class': (x > 0).sum(0).tolist()},
1069
+ 'labels': [{
1070
+ str(Path(k).name): _round(v.tolist())} for k, v in zip(dataset.im_files, dataset.labels)]}
1071
+
1072
+ # Save, print and return
1073
+ if save:
1074
+ stats_path = self.hub_dir / 'stats.json'
1075
+ print(f'Saving {stats_path.resolve()}...')
1076
+ with open(stats_path, 'w') as f:
1077
+ json.dump(self.stats, f) # save stats.json
1078
+ if verbose:
1079
+ print(json.dumps(self.stats, indent=2, sort_keys=False))
1080
+ return self.stats
1081
+
1082
+ def process_images(self):
1083
+ # Compress images for Ultralytics HUB
1084
+ for split in 'train', 'val', 'test':
1085
+ if self.data.get(split) is None:
1086
+ continue
1087
+ dataset = LoadImagesAndLabels(self.data[split]) # load dataset
1088
+ desc = f'{split} images'
1089
+ for _ in tqdm(ThreadPool(NUM_THREADS).imap(self._hub_ops, dataset.im_files), total=dataset.n, desc=desc):
1090
  pass
1091
+ print(f'Done. All images saved to {self.im_dir}')
1092
+ return self.im_dir
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
utils/downloads.py CHANGED
@@ -16,12 +16,14 @@ import requests
16
  import torch
17
 
18
 
19
- def is_url(url):
20
  # Check if online file exists
21
  try:
22
- r = urllib.request.urlopen(url) # response
23
- return r.getcode() == 200
24
- except urllib.request.HTTPError:
 
 
25
  return False
26
 
27
 
 
16
  import torch
17
 
18
 
19
+ def is_url(url, check_online=True):
20
  # Check if online file exists
21
  try:
22
+ url = str(url)
23
+ result = urllib.parse.urlparse(url)
24
+ assert all([result.scheme, result.netloc, result.path]) # check if is url
25
+ return (urllib.request.urlopen(url).getcode() == 200) if check_online else True # check if exists online
26
+ except (AssertionError, urllib.request.HTTPError):
27
  return False
28
 
29
 
utils/general.py CHANGED
@@ -14,6 +14,7 @@ import random
14
  import re
15
  import shutil
16
  import signal
 
17
  import threading
18
  import time
19
  import urllib
@@ -52,7 +53,7 @@ np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format})
52
  pd.options.display.max_columns = 10
53
  cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
54
  os.environ['NUMEXPR_MAX_THREADS'] = str(NUM_THREADS) # NumExpr max threads
55
- os.environ['OMP_NUM_THREADS'] = str(NUM_THREADS) # OpenMP max threads (PyTorch and SciPy)
56
 
57
 
58
  def is_kaggle():
@@ -68,7 +69,7 @@ def is_kaggle():
68
  def is_writeable(dir, test=False):
69
  # Return True if directory has write permissions, test opening a file with write permissions if test=True
70
  if not test:
71
- return os.access(dir, os.R_OK) # possible issues on Windows
72
  file = Path(dir) / 'tmp.txt'
73
  try:
74
  with open(file, 'w'): # open file with write permissions
@@ -96,6 +97,9 @@ def set_logging(name=None, verbose=VERBOSE):
96
 
97
  set_logging() # run before defining LOGGER
98
  LOGGER = logging.getLogger("yolov5") # define globally (used in train.py, val.py, detect.py, etc.)
 
 
 
99
 
100
 
101
  def user_config_dir(dir='Ultralytics', env_var='YOLOV5_CONFIG_DIR'):
@@ -203,14 +207,14 @@ def init_seeds(seed=0, deterministic=False):
203
  if deterministic and check_version(torch.__version__, '1.12.0'): # https://github.com/ultralytics/yolov5/pull/8213
204
  torch.use_deterministic_algorithms(True)
205
  os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
206
- # os.environ['PYTHONHASHSEED'] = str(seed)
207
 
208
  random.seed(seed)
209
  np.random.seed(seed)
210
  torch.manual_seed(seed)
211
  cudnn.benchmark, cudnn.deterministic = (False, True) if seed == 0 else (True, False)
212
- # torch.cuda.manual_seed(seed)
213
- # torch.cuda.manual_seed_all(seed) # for multi GPU, exception safe
214
 
215
 
216
  def intersect_dicts(da, db, exclude=()):
@@ -224,9 +228,15 @@ def get_latest_run(search_dir='.'):
224
  return max(last_list, key=os.path.getctime) if last_list else ''
225
 
226
 
227
- def is_docker():
228
- # Is environment a Docker container?
229
- return Path('/workspace').exists() # or Path('/.dockerenv').exists()
 
 
 
 
 
 
230
 
231
 
232
  def is_colab():
@@ -304,23 +314,30 @@ def git_describe(path=ROOT): # path must be a directory
304
 
305
  @try_except
306
  @WorkingDirectory(ROOT)
307
- def check_git_status():
308
- # Recommend 'git pull' if code is out of date
309
- msg = ', for updates see https://github.com/ultralytics/yolov5'
 
310
  s = colorstr('github: ') # string
311
  assert Path('.git').exists(), s + 'skipping check (not a git repository)' + msg
312
- assert not is_docker(), s + 'skipping check (Docker image)' + msg
313
  assert check_online(), s + 'skipping check (offline)' + msg
314
 
315
- cmd = 'git fetch && git config --get remote.origin.url'
316
- url = check_output(cmd, shell=True, timeout=5).decode().strip().rstrip('.git') # git fetch
 
 
 
 
 
 
317
  branch = check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out
318
- n = int(check_output(f'git rev-list {branch}..origin/master --count', shell=True)) # commits behind
319
  if n > 0:
320
- s += f"⚠️ YOLOv5 is out of date by {n} commit{'s' * (n > 1)}. Use `git pull` or `git clone {url}` to update."
 
321
  else:
322
  s += f'up to date with {url} βœ…'
323
- LOGGER.info(emojis(s)) # emoji-safe
324
 
325
 
326
  def check_python(minimum='3.7.0'):
@@ -374,7 +391,7 @@ def check_requirements(requirements=ROOT / 'requirements.txt', exclude=(), insta
374
  source = file.resolve() if 'file' in locals() else requirements
375
  s = f"{prefix} {n} package{'s' * (n > 1)} updated per {source}\n" \
376
  f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n"
377
- LOGGER.info(emojis(s))
378
 
379
 
380
  def check_img_size(imgsz, s=32, floor=0):
@@ -436,6 +453,9 @@ def check_file(file, suffix=''):
436
  torch.hub.download_url_to_file(url, file)
437
  assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}' # check
438
  return file
 
 
 
439
  else: # search
440
  files = []
441
  for d in 'data', 'models', 'utils': # search directories
@@ -461,7 +481,7 @@ def check_dataset(data, autodownload=True):
461
  # Download (optional)
462
  extract_dir = ''
463
  if isinstance(data, (str, Path)) and str(data).endswith('.zip'): # i.e. gs://bucket/dir/coco128.zip
464
- download(data, dir=DATASETS_DIR, unzip=True, delete=False, curl=False, threads=1)
465
  data = next((DATASETS_DIR / Path(data).stem).rglob('*.yaml'))
466
  extract_dir, autodownload = data.parent, False
467
 
@@ -472,9 +492,9 @@ def check_dataset(data, autodownload=True):
472
 
473
  # Checks
474
  for k in 'train', 'val', 'nc':
475
- assert k in data, emojis(f"data.yaml '{k}:' field missing ❌")
476
  if 'names' not in data:
477
- LOGGER.warning(emojis("data.yaml 'names:' field missing ⚠, assigning default names 'class0', 'class1', etc."))
478
  data['names'] = [f'class{i}' for i in range(data['nc'])] # default names
479
 
480
  # Resolve paths
@@ -490,9 +510,9 @@ def check_dataset(data, autodownload=True):
490
  if val:
491
  val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path
492
  if not all(x.exists() for x in val):
493
- LOGGER.info(emojis('\nDataset not found ⚠, missing paths %s' % [str(x) for x in val if not x.exists()]))
494
  if not s or not autodownload:
495
- raise Exception(emojis('Dataset not found ❌'))
496
  t = time.time()
497
  root = path.parent if 'path' in data else '..' # unzip directory i.e. '../'
498
  if s.startswith('http') and s.endswith('.zip'): # URL
@@ -510,7 +530,7 @@ def check_dataset(data, autodownload=True):
510
  r = exec(s, {'yaml': data}) # return None
511
  dt = f'({round(time.time() - t, 1)}s)'
512
  s = f"success βœ… {dt}, saved to {colorstr('bold', root)}" if r in (0, None) else f"failure {dt} ❌"
513
- LOGGER.info(emojis(f"Dataset download {s}"))
514
  check_font('Arial.ttf' if is_ascii(data['names']) else 'Arial.Unicode.ttf', progress=True) # download fonts
515
  return data # dictionary
516
 
@@ -535,11 +555,11 @@ def check_amp(model):
535
  im = f if f.exists() else 'https://ultralytics.com/images/bus.jpg' if check_online() else np.ones((640, 640, 3))
536
  try:
537
  assert amp_allclose(model, im) or amp_allclose(DetectMultiBackend('yolov5n.pt', device), im)
538
- LOGGER.info(emojis(f'{prefix}checks passed βœ…'))
539
  return True
540
  except Exception:
541
  help_url = 'https://github.com/ultralytics/yolov5/issues/7908'
542
- LOGGER.warning(emojis(f'{prefix}checks failed ❌, disabling Automatic Mixed Precision. See {help_url}'))
543
  return False
544
 
545
 
 
14
  import re
15
  import shutil
16
  import signal
17
+ import sys
18
  import threading
19
  import time
20
  import urllib
 
53
  pd.options.display.max_columns = 10
54
  cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
55
  os.environ['NUMEXPR_MAX_THREADS'] = str(NUM_THREADS) # NumExpr max threads
56
+ os.environ['OMP_NUM_THREADS'] = '1' if platform.system() == 'darwin' else str(NUM_THREADS) # OpenMP (PyTorch and SciPy)
57
 
58
 
59
  def is_kaggle():
 
69
  def is_writeable(dir, test=False):
70
  # Return True if directory has write permissions, test opening a file with write permissions if test=True
71
  if not test:
72
+ return os.access(dir, os.W_OK) # possible issues on Windows
73
  file = Path(dir) / 'tmp.txt'
74
  try:
75
  with open(file, 'w'): # open file with write permissions
 
97
 
98
  set_logging() # run before defining LOGGER
99
  LOGGER = logging.getLogger("yolov5") # define globally (used in train.py, val.py, detect.py, etc.)
100
+ if platform.system() == 'Windows':
101
+ for fn in LOGGER.info, LOGGER.warning:
102
+ setattr(LOGGER, fn.__name__, lambda x: fn(emojis(x))) # emoji safe logging
103
 
104
 
105
  def user_config_dir(dir='Ultralytics', env_var='YOLOV5_CONFIG_DIR'):
 
207
  if deterministic and check_version(torch.__version__, '1.12.0'): # https://github.com/ultralytics/yolov5/pull/8213
208
  torch.use_deterministic_algorithms(True)
209
  os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
210
+ os.environ['PYTHONHASHSEED'] = str(seed)
211
 
212
  random.seed(seed)
213
  np.random.seed(seed)
214
  torch.manual_seed(seed)
215
  cudnn.benchmark, cudnn.deterministic = (False, True) if seed == 0 else (True, False)
216
+ torch.cuda.manual_seed(seed)
217
+ torch.cuda.manual_seed_all(seed) # for Multi-GPU, exception safe
218
 
219
 
220
  def intersect_dicts(da, db, exclude=()):
 
228
  return max(last_list, key=os.path.getctime) if last_list else ''
229
 
230
 
231
+ def is_docker() -> bool:
232
+ """Check if the process runs inside a docker container."""
233
+ if Path("/.dockerenv").exists():
234
+ return True
235
+ try: # check if docker is in control groups
236
+ with open("/proc/self/cgroup") as file:
237
+ return any("docker" in line for line in file)
238
+ except OSError:
239
+ return False
240
 
241
 
242
  def is_colab():
 
314
 
315
  @try_except
316
  @WorkingDirectory(ROOT)
317
+ def check_git_status(repo='ultralytics/yolov5'):
318
+ # YOLOv5 status check, recommend 'git pull' if code is out of date
319
+ url = f'https://github.com/{repo}'
320
+ msg = f', for updates see {url}'
321
  s = colorstr('github: ') # string
322
  assert Path('.git').exists(), s + 'skipping check (not a git repository)' + msg
 
323
  assert check_online(), s + 'skipping check (offline)' + msg
324
 
325
+ splits = re.split(pattern=r'\s', string=check_output('git remote -v', shell=True).decode())
326
+ matches = [repo in s for s in splits]
327
+ if any(matches):
328
+ remote = splits[matches.index(True) - 1]
329
+ else:
330
+ remote = 'ultralytics'
331
+ check_output(f'git remote add {remote} {url}', shell=True)
332
+ check_output(f'git fetch {remote}', shell=True, timeout=5) # git fetch
333
  branch = check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out
334
+ n = int(check_output(f'git rev-list {branch}..{remote}/master --count', shell=True)) # commits behind
335
  if n > 0:
336
+ pull = 'git pull' if remote == 'origin' else f'git pull {remote} master'
337
+ s += f"⚠️ YOLOv5 is out of date by {n} commit{'s' * (n > 1)}. Use `{pull}` or `git clone {url}` to update."
338
  else:
339
  s += f'up to date with {url} βœ…'
340
+ LOGGER.info(s)
341
 
342
 
343
  def check_python(minimum='3.7.0'):
 
391
  source = file.resolve() if 'file' in locals() else requirements
392
  s = f"{prefix} {n} package{'s' * (n > 1)} updated per {source}\n" \
393
  f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n"
394
+ LOGGER.info(s)
395
 
396
 
397
  def check_img_size(imgsz, s=32, floor=0):
 
453
  torch.hub.download_url_to_file(url, file)
454
  assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}' # check
455
  return file
456
+ elif file.startswith('clearml://'): # ClearML Dataset ID
457
+ assert 'clearml' in sys.modules, "ClearML is not installed, so cannot use ClearML dataset. Try running 'pip install clearml'."
458
+ return file
459
  else: # search
460
  files = []
461
  for d in 'data', 'models', 'utils': # search directories
 
481
  # Download (optional)
482
  extract_dir = ''
483
  if isinstance(data, (str, Path)) and str(data).endswith('.zip'): # i.e. gs://bucket/dir/coco128.zip
484
+ download(data, dir=f'{DATASETS_DIR}/{Path(data).stem}', unzip=True, delete=False, curl=False, threads=1)
485
  data = next((DATASETS_DIR / Path(data).stem).rglob('*.yaml'))
486
  extract_dir, autodownload = data.parent, False
487
 
 
492
 
493
  # Checks
494
  for k in 'train', 'val', 'nc':
495
+ assert k in data, f"data.yaml '{k}:' field missing ❌"
496
  if 'names' not in data:
497
+ LOGGER.warning("data.yaml 'names:' field missing ⚠️, assigning default names 'class0', 'class1', etc.")
498
  data['names'] = [f'class{i}' for i in range(data['nc'])] # default names
499
 
500
  # Resolve paths
 
510
  if val:
511
  val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path
512
  if not all(x.exists() for x in val):
513
+ LOGGER.info('\nDataset not found ⚠️, missing paths %s' % [str(x) for x in val if not x.exists()])
514
  if not s or not autodownload:
515
+ raise Exception('Dataset not found ❌')
516
  t = time.time()
517
  root = path.parent if 'path' in data else '..' # unzip directory i.e. '../'
518
  if s.startswith('http') and s.endswith('.zip'): # URL
 
530
  r = exec(s, {'yaml': data}) # return None
531
  dt = f'({round(time.time() - t, 1)}s)'
532
  s = f"success βœ… {dt}, saved to {colorstr('bold', root)}" if r in (0, None) else f"failure {dt} ❌"
533
+ LOGGER.info(f"Dataset download {s}")
534
  check_font('Arial.ttf' if is_ascii(data['names']) else 'Arial.Unicode.ttf', progress=True) # download fonts
535
  return data # dictionary
536
 
 
555
  im = f if f.exists() else 'https://ultralytics.com/images/bus.jpg' if check_online() else np.ones((640, 640, 3))
556
  try:
557
  assert amp_allclose(model, im) or amp_allclose(DetectMultiBackend('yolov5n.pt', device), im)
558
+ LOGGER.info(f'{prefix}checks passed βœ…')
559
  return True
560
  except Exception:
561
  help_url = 'https://github.com/ultralytics/yolov5/issues/7908'
562
+ LOGGER.warning(f'{prefix}checks failed ❌, disabling Automatic Mixed Precision. See {help_url}')
563
  return False
564
 
565
 
utils/metrics.py CHANGED
@@ -139,6 +139,12 @@ class ConfusionMatrix:
139
  Returns:
140
  None, updates confusion matrix accordingly
141
  """
 
 
 
 
 
 
142
  detections = detections[detections[:, 4] > self.conf]
143
  gt_classes = labels[:, 0].int()
144
  detection_classes = detections[:, 5].int()
@@ -203,6 +209,7 @@ class ConfusionMatrix:
203
  yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1))
204
  fig.axes[0].set_xlabel('True')
205
  fig.axes[0].set_ylabel('Predicted')
 
206
  fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
207
  plt.close()
208
  except Exception as e:
@@ -330,6 +337,7 @@ def plot_pr_curve(px, py, ap, save_dir=Path('pr_curve.png'), names=()):
330
  ax.set_xlim(0, 1)
331
  ax.set_ylim(0, 1)
332
  plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
 
333
  fig.savefig(save_dir, dpi=250)
334
  plt.close()
335
 
@@ -351,5 +359,6 @@ def plot_mc_curve(px, py, save_dir=Path('mc_curve.png'), names=(), xlabel='Confi
351
  ax.set_xlim(0, 1)
352
  ax.set_ylim(0, 1)
353
  plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
 
354
  fig.savefig(save_dir, dpi=250)
355
  plt.close()
 
139
  Returns:
140
  None, updates confusion matrix accordingly
141
  """
142
+ if detections is None:
143
+ gt_classes = labels.int()
144
+ for i, gc in enumerate(gt_classes):
145
+ self.matrix[self.nc, gc] += 1 # background FN
146
+ return
147
+
148
  detections = detections[detections[:, 4] > self.conf]
149
  gt_classes = labels[:, 0].int()
150
  detection_classes = detections[:, 5].int()
 
209
  yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1))
210
  fig.axes[0].set_xlabel('True')
211
  fig.axes[0].set_ylabel('Predicted')
212
+ plt.title('Confusion Matrix')
213
  fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
214
  plt.close()
215
  except Exception as e:
 
337
  ax.set_xlim(0, 1)
338
  ax.set_ylim(0, 1)
339
  plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
340
+ plt.title('Precision-Recall Curve')
341
  fig.savefig(save_dir, dpi=250)
342
  plt.close()
343
 
 
359
  ax.set_xlim(0, 1)
360
  ax.set_ylim(0, 1)
361
  plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
362
+ plt.title(f'{ylabel}-Confidence Curve')
363
  fig.savefig(save_dir, dpi=250)
364
  plt.close()
utils/plots.py CHANGED
@@ -148,6 +148,7 @@ def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detec
148
  ax[i].axis('off')
149
 
150
  LOGGER.info(f'Saving {f}... ({n}/{channels})')
 
151
  plt.savefig(f, dpi=300, bbox_inches='tight')
152
  plt.close()
153
  np.save(str(f.with_suffix('.npy')), x[0].cpu().numpy()) # npy save
@@ -484,6 +485,6 @@ def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False,
484
  if save:
485
  file.parent.mkdir(parents=True, exist_ok=True) # make directory
486
  f = str(increment_path(file).with_suffix('.jpg'))
487
- # cv2.imwrite(f, crop) # https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue
488
- Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)).save(f, quality=95, subsampling=0)
489
  return crop
 
148
  ax[i].axis('off')
149
 
150
  LOGGER.info(f'Saving {f}... ({n}/{channels})')
151
+ plt.title('Features')
152
  plt.savefig(f, dpi=300, bbox_inches='tight')
153
  plt.close()
154
  np.save(str(f.with_suffix('.npy')), x[0].cpu().numpy()) # npy save
 
485
  if save:
486
  file.parent.mkdir(parents=True, exist_ok=True) # make directory
487
  f = str(increment_path(file).with_suffix('.jpg'))
488
+ # cv2.imwrite(f, crop) # save BGR, https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue
489
+ Image.fromarray(crop[..., ::-1]).save(f, quality=95, subsampling=0) # save RGB
490
  return crop
utils/torch_utils.py CHANGED
@@ -17,8 +17,13 @@ import torch
17
  import torch.distributed as dist
18
  import torch.nn as nn
19
  import torch.nn.functional as F
 
20
 
21
- from utils.general import LOGGER, file_date, git_describe
 
 
 
 
22
 
23
  try:
24
  import thop # for FLOPs computation
@@ -29,6 +34,25 @@ except ImportError:
29
  warnings.filterwarnings('ignore', message='User provided device_type of \'cuda\', but CUDA is not available. Disabling')
30
 
31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
32
  @contextmanager
33
  def torch_distributed_zero_first(local_rank: int):
34
  # Decorator to make all processes in distributed training wait for each local_master to do something
@@ -81,7 +105,7 @@ def select_device(device='', batch_size=0, newline=True):
81
 
82
  if not newline:
83
  s = s.rstrip()
84
- LOGGER.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe
85
  return torch.device(arg)
86
 
87
 
@@ -183,12 +207,11 @@ def sparsity(model):
183
  def prune(model, amount=0.3):
184
  # Prune model to requested global sparsity
185
  import torch.nn.utils.prune as prune
186
- print('Pruning model... ', end='')
187
  for name, m in model.named_modules():
188
  if isinstance(m, nn.Conv2d):
189
  prune.l1_unstructured(m, name='weight', amount=amount) # prune
190
  prune.remove(m, 'weight') # make permanent
191
- print(' %.3g global sparsity' % sparsity(model))
192
 
193
 
194
  def fuse_conv_and_bn(conv, bn):
@@ -214,7 +237,7 @@ def fuse_conv_and_bn(conv, bn):
214
  return fusedconv
215
 
216
 
217
- def model_info(model, verbose=False, img_size=640):
218
  # Model information. img_size may be int or list, i.e. img_size=640 or img_size=[640, 320]
219
  n_p = sum(x.numel() for x in model.parameters()) # number parameters
220
  n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients
@@ -226,12 +249,12 @@ def model_info(model, verbose=False, img_size=640):
226
  (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
227
 
228
  try: # FLOPs
229
- from thop import profile
230
- stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32
231
- img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input
232
- flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPs
233
- img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float
234
- fs = ', %.1f GFLOPs' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPs
235
  except Exception:
236
  fs = ''
237
 
@@ -260,6 +283,56 @@ def copy_attr(a, b, include=(), exclude=()):
260
  setattr(a, k, v)
261
 
262
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
263
  class EarlyStopping:
264
  # YOLOv5 simple early stopper
265
  def __init__(self, patience=30):
@@ -299,17 +372,17 @@ class ModelEMA:
299
  for p in self.ema.parameters():
300
  p.requires_grad_(False)
301
 
 
302
  def update(self, model):
303
  # Update EMA parameters
304
- with torch.no_grad():
305
- self.updates += 1
306
- d = self.decay(self.updates)
307
-
308
- msd = de_parallel(model).state_dict() # model state_dict
309
- for k, v in self.ema.state_dict().items():
310
- if v.dtype.is_floating_point:
311
- v *= d
312
- v += (1 - d) * msd[k].detach()
313
 
314
  def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
315
  # Update EMA attributes
 
17
  import torch.distributed as dist
18
  import torch.nn as nn
19
  import torch.nn.functional as F
20
+ from torch.nn.parallel import DistributedDataParallel as DDP
21
 
22
+ from utils.general import LOGGER, check_version, colorstr, file_date, git_describe
23
+
24
+ LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html
25
+ RANK = int(os.getenv('RANK', -1))
26
+ WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
27
 
28
  try:
29
  import thop # for FLOPs computation
 
34
  warnings.filterwarnings('ignore', message='User provided device_type of \'cuda\', but CUDA is not available. Disabling')
35
 
36
 
37
+ def smart_inference_mode(torch_1_9=check_version(torch.__version__, '1.9.0')):
38
+ # Applies torch.inference_mode() decorator if torch>=1.9.0 else torch.no_grad() decorator
39
+ def decorate(fn):
40
+ return (torch.inference_mode if torch_1_9 else torch.no_grad)()(fn)
41
+
42
+ return decorate
43
+
44
+
45
+ def smart_DDP(model):
46
+ # Model DDP creation with checks
47
+ assert not check_version(torch.__version__, '1.12.0', pinned=True), \
48
+ 'torch==1.12.0 torchvision==0.13.0 DDP training is not supported due to a known issue. ' \
49
+ 'Please upgrade or downgrade torch to use DDP. See https://github.com/ultralytics/yolov5/issues/8395'
50
+ if check_version(torch.__version__, '1.11.0'):
51
+ return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK, static_graph=True)
52
+ else:
53
+ return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
54
+
55
+
56
  @contextmanager
57
  def torch_distributed_zero_first(local_rank: int):
58
  # Decorator to make all processes in distributed training wait for each local_master to do something
 
105
 
106
  if not newline:
107
  s = s.rstrip()
108
+ LOGGER.info(s)
109
  return torch.device(arg)
110
 
111
 
 
207
  def prune(model, amount=0.3):
208
  # Prune model to requested global sparsity
209
  import torch.nn.utils.prune as prune
 
210
  for name, m in model.named_modules():
211
  if isinstance(m, nn.Conv2d):
212
  prune.l1_unstructured(m, name='weight', amount=amount) # prune
213
  prune.remove(m, 'weight') # make permanent
214
+ LOGGER.info(f'Model pruned to {sparsity(model):.3g} global sparsity')
215
 
216
 
217
  def fuse_conv_and_bn(conv, bn):
 
237
  return fusedconv
238
 
239
 
240
+ def model_info(model, verbose=False, imgsz=640):
241
  # Model information. img_size may be int or list, i.e. img_size=640 or img_size=[640, 320]
242
  n_p = sum(x.numel() for x in model.parameters()) # number parameters
243
  n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients
 
249
  (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
250
 
251
  try: # FLOPs
252
+ p = next(model.parameters())
253
+ stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32 # max stride
254
+ im = torch.zeros((1, p.shape[1], stride, stride), device=p.device) # input image in BCHW format
255
+ flops = thop.profile(deepcopy(model), inputs=(im,), verbose=False)[0] / 1E9 * 2 # stride GFLOPs
256
+ imgsz = imgsz if isinstance(imgsz, list) else [imgsz, imgsz] # expand if int/float
257
+ fs = f', {flops * imgsz[0] / stride * imgsz[1] / stride:.1f} GFLOPs' # 640x640 GFLOPs
258
  except Exception:
259
  fs = ''
260
 
 
283
  setattr(a, k, v)
284
 
285
 
286
+ def smart_optimizer(model, name='Adam', lr=0.001, momentum=0.9, decay=1e-5):
287
+ # YOLOv5 3-param group optimizer: 0) weights with decay, 1) weights no decay, 2) biases no decay
288
+ g = [], [], [] # optimizer parameter groups
289
+ bn = tuple(v for k, v in nn.__dict__.items() if 'Norm' in k) # normalization layers, i.e. BatchNorm2d()
290
+ for v in model.modules():
291
+ if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): # bias (no decay)
292
+ g[2].append(v.bias)
293
+ if isinstance(v, bn): # weight (no decay)
294
+ g[1].append(v.weight)
295
+ elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): # weight (with decay)
296
+ g[0].append(v.weight)
297
+
298
+ if name == 'Adam':
299
+ optimizer = torch.optim.Adam(g[2], lr=lr, betas=(momentum, 0.999)) # adjust beta1 to momentum
300
+ elif name == 'AdamW':
301
+ optimizer = torch.optim.AdamW(g[2], lr=lr, betas=(momentum, 0.999), weight_decay=0.0)
302
+ elif name == 'RMSProp':
303
+ optimizer = torch.optim.RMSprop(g[2], lr=lr, momentum=momentum)
304
+ elif name == 'SGD':
305
+ optimizer = torch.optim.SGD(g[2], lr=lr, momentum=momentum, nesterov=True)
306
+ else:
307
+ raise NotImplementedError(f'Optimizer {name} not implemented.')
308
+
309
+ optimizer.add_param_group({'params': g[0], 'weight_decay': decay}) # add g0 with weight_decay
310
+ optimizer.add_param_group({'params': g[1], 'weight_decay': 0.0}) # add g1 (BatchNorm2d weights)
311
+ LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}) with parameter groups "
312
+ f"{len(g[1])} weight(decay=0.0), {len(g[0])} weight(decay={decay}), {len(g[2])} bias")
313
+ return optimizer
314
+
315
+
316
+ def smart_resume(ckpt, optimizer, ema=None, weights='yolov5s.pt', epochs=300, resume=True):
317
+ # Resume training from a partially trained checkpoint
318
+ best_fitness = 0.0
319
+ start_epoch = ckpt['epoch'] + 1
320
+ if ckpt['optimizer'] is not None:
321
+ optimizer.load_state_dict(ckpt['optimizer']) # optimizer
322
+ best_fitness = ckpt['best_fitness']
323
+ if ema and ckpt.get('ema'):
324
+ ema.ema.load_state_dict(ckpt['ema'].float().state_dict()) # EMA
325
+ ema.updates = ckpt['updates']
326
+ if resume:
327
+ assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.\n' \
328
+ f"Start a new training without --resume, i.e. 'python train.py --weights {weights}'"
329
+ LOGGER.info(f'Resuming training from {weights} from epoch {start_epoch} to {epochs} total epochs')
330
+ if epochs < start_epoch:
331
+ LOGGER.info(f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs.")
332
+ epochs += ckpt['epoch'] # finetune additional epochs
333
+ return best_fitness, start_epoch, epochs
334
+
335
+
336
  class EarlyStopping:
337
  # YOLOv5 simple early stopper
338
  def __init__(self, patience=30):
 
372
  for p in self.ema.parameters():
373
  p.requires_grad_(False)
374
 
375
+ @smart_inference_mode()
376
  def update(self, model):
377
  # Update EMA parameters
378
+ self.updates += 1
379
+ d = self.decay(self.updates)
380
+
381
+ msd = de_parallel(model).state_dict() # model state_dict
382
+ for k, v in self.ema.state_dict().items():
383
+ if v.dtype.is_floating_point:
384
+ v *= d
385
+ v += (1 - d) * msd[k].detach()
 
386
 
387
  def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
388
  # Update EMA attributes
val.py CHANGED
@@ -38,11 +38,11 @@ from models.common import DetectMultiBackend
38
  from utils.callbacks import Callbacks
39
  from utils.dataloaders import create_dataloader
40
  from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, check_yaml,
41
- coco80_to_coco91_class, colorstr, emojis, increment_path, non_max_suppression, print_args,
42
  scale_coords, xywh2xyxy, xyxy2xywh)
43
  from utils.metrics import ConfusionMatrix, ap_per_class, box_iou
44
  from utils.plots import output_to_target, plot_images, plot_val_study
45
- from utils.torch_utils import select_device, time_sync
46
 
47
 
48
  def save_one_txt(predn, save_conf, shape, file):
@@ -93,7 +93,7 @@ def process_batch(detections, labels, iouv):
93
  return torch.tensor(correct, dtype=torch.bool, device=iouv.device)
94
 
95
 
96
- @torch.no_grad()
97
  def run(
98
  data,
99
  weights=None, # model.pt path(s)
@@ -182,7 +182,7 @@ def run(
182
 
183
  seen = 0
184
  confusion_matrix = ConfusionMatrix(nc=nc)
185
- names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
186
  class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
187
  s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', '[email protected]', '[email protected]:.95')
188
  dt, p, r, f1, mp, mr, map50, map = [0.0, 0.0, 0.0], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
@@ -228,6 +228,8 @@ def run(
228
  if npr == 0:
229
  if nl:
230
  stats.append((correct, *torch.zeros((2, 0), device=device), labels[:, 0]))
 
 
231
  continue
232
 
233
  # Predictions
@@ -248,7 +250,7 @@ def run(
248
 
249
  # Save/log
250
  if save_txt:
251
- save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / (path.stem + '.txt'))
252
  if save_json:
253
  save_one_json(predn, jdict, path, class_map) # append to COCO-JSON dictionary
254
  callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
@@ -266,13 +268,13 @@ def run(
266
  tp, fp, p, r, f1, ap, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
267
  ap50, ap = ap[:, 0], ap.mean(1) # [email protected], [email protected]:0.95
268
  mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
269
- nt = np.bincount(stats[3].astype(int), minlength=nc) # number of targets per class
270
- else:
271
- nt = torch.zeros(1)
272
 
273
  # Print results
274
  pf = '%20s' + '%11i' * 2 + '%11.3g' * 4 # print format
275
  LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
 
 
276
 
277
  # Print results per class
278
  if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
@@ -363,7 +365,7 @@ def main(opt):
363
 
364
  if opt.task in ('train', 'val', 'test'): # run normally
365
  if opt.conf_thres > 0.001: # https://github.com/ultralytics/yolov5/issues/1466
366
- LOGGER.info(emojis(f'WARNING: confidence threshold {opt.conf_thres} > 0.001 produces invalid results ⚠️'))
367
  run(**vars(opt))
368
 
369
  else:
 
38
  from utils.callbacks import Callbacks
39
  from utils.dataloaders import create_dataloader
40
  from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, check_yaml,
41
+ coco80_to_coco91_class, colorstr, increment_path, non_max_suppression, print_args,
42
  scale_coords, xywh2xyxy, xyxy2xywh)
43
  from utils.metrics import ConfusionMatrix, ap_per_class, box_iou
44
  from utils.plots import output_to_target, plot_images, plot_val_study
45
+ from utils.torch_utils import select_device, smart_inference_mode, time_sync
46
 
47
 
48
  def save_one_txt(predn, save_conf, shape, file):
 
93
  return torch.tensor(correct, dtype=torch.bool, device=iouv.device)
94
 
95
 
96
+ @smart_inference_mode()
97
  def run(
98
  data,
99
  weights=None, # model.pt path(s)
 
182
 
183
  seen = 0
184
  confusion_matrix = ConfusionMatrix(nc=nc)
185
+ names = dict(enumerate(model.names if hasattr(model, 'names') else model.module.names))
186
  class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
187
  s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', '[email protected]', '[email protected]:.95')
188
  dt, p, r, f1, mp, mr, map50, map = [0.0, 0.0, 0.0], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
 
228
  if npr == 0:
229
  if nl:
230
  stats.append((correct, *torch.zeros((2, 0), device=device), labels[:, 0]))
231
+ if plots:
232
+ confusion_matrix.process_batch(detections=None, labels=labels[:, 0])
233
  continue
234
 
235
  # Predictions
 
250
 
251
  # Save/log
252
  if save_txt:
253
+ save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt')
254
  if save_json:
255
  save_one_json(predn, jdict, path, class_map) # append to COCO-JSON dictionary
256
  callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
 
268
  tp, fp, p, r, f1, ap, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
269
  ap50, ap = ap[:, 0], ap.mean(1) # [email protected], [email protected]:0.95
270
  mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
271
+ nt = np.bincount(stats[3].astype(int), minlength=nc) # number of targets per class
 
 
272
 
273
  # Print results
274
  pf = '%20s' + '%11i' * 2 + '%11.3g' * 4 # print format
275
  LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
276
+ if nt.sum() == 0:
277
+ LOGGER.warning(f'WARNING: no labels found in {task} set, can not compute metrics without labels ⚠️')
278
 
279
  # Print results per class
280
  if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
 
365
 
366
  if opt.task in ('train', 'val', 'test'): # run normally
367
  if opt.conf_thres > 0.001: # https://github.com/ultralytics/yolov5/issues/1466
368
+ LOGGER.info(f'WARNING: confidence threshold {opt.conf_thres} > 0.001 produces invalid results ⚠️')
369
  run(**vars(opt))
370
 
371
  else: