zjc
/
yolov5
Watch 1
Star 0
Fork 0
Code Issues 0 Pull Requests 0 Releases 0 Wiki Activity
Browse Source

precommit: yapf (#5494) 

* precommit: yapf

* align isort

* fix

# Conflicts:
#	utils/plots.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update setup.cfg

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update setup.cfg

* Update setup.cfg

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update wandb_utils.py

* Update augmentations.py

* Update setup.cfg

* Update yolo.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update val.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* simplify colorstr

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* val run fix

* export.py last comma

* Update export.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update hubconf.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* PyTorch Hub tuple fix

* PyTorch Hub tuple fix2

* PyTorch Hub tuple fix3

* Update setup

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
modifyDataloader
Jirka Borovec GitHub 2 years ago
parent
df9008ee69
commit
c3d5ac151e
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
24 changed files with 527 additions and 331 deletions
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. +5
    -6
      .pre-commit-config.yaml
  2. +3
    -2
      detect.py
  3. +65
    -45
      export.py
  4. +7
    -6
      hubconf.py
  5. +24
    -13
      models/common.py
  6. +2
    -2
      models/experimental.py
  7. +47
    -20
      models/tf.py
  8. +2
    -2
      models/yolo.py
  9. +14
    -0
      setup.cfg
  10. +82
    -65
      train.py
  11. +0
    -2
      utils/activations.py
  12. +11
    -4
      utils/augmentations.py
  13. +3
    -2
      utils/benchmarks.py
  14. +1
    -6
      utils/callbacks.py
  15. +74
    -38
      utils/datasets.py
  16. +10
    -7
      utils/downloads.py
  17. +40
    -34
      utils/general.py
  18. +16
    -5
      utils/loggers/__init__.py
  19. +63
    -49
      utils/loggers/wandb/wandb_utils.py
  20. +10
    -4
      utils/loss.py
  21. +9
    -2
      utils/metrics.py
  22. +22
    -8
      utils/plots.py
  23. +0
    -1
      utils/torch_utils.py
  24. +17
    -8
      val.py

+ 5
- 6
.pre-commit-config.yaml View File

- id: isort - id: isort
name: Sort imports name: Sort imports


# TODO
#- repo: https://github.com/pre-commit/mirrors-yapf
# rev: v0.31.0
# hooks:
# - id: yapf
# name: formatting
- repo: https://github.com/pre-commit/mirrors-yapf
rev: v0.31.0
hooks:
- id: yapf
name: formatting


# TODO # TODO
#- repo: https://github.com/executablebooks/mdformat #- repo: https://github.com/executablebooks/mdformat

+ 3
- 2
detect.py View File





@torch.no_grad() @torch.no_grad()
def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
def run(
weights=ROOT / 'yolov5s.pt', # model.pt path(s)
source=ROOT / 'data/images', # file/dir/URL/glob, 0 for webcam source=ROOT / 'data/images', # file/dir/URL/glob, 0 for webcam
data=ROOT / 'data/coco128.yaml', # dataset.yaml path data=ROOT / 'data/coco128.yaml', # dataset.yaml path
imgsz=(640, 640), # inference size (height, width) imgsz=(640, 640), # inference size (height, width)
hide_conf=False, # hide confidences hide_conf=False, # hide confidences
half=False, # use FP16 half-precision inference half=False, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference dnn=False, # use OpenCV DNN for ONNX inference
):
):
source = str(source) source = str(source)
save_img = not nosave and not source.endswith('.txt') # save inference images save_img = not nosave and not source.endswith('.txt') # save inference images
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)

+ 65
- 45
export.py View File



def export_formats(): def export_formats():
# YOLOv5 export formats # YOLOv5 export formats
x = [['PyTorch', '-', '.pt', True],
['TorchScript', 'torchscript', '.torchscript', True],
['ONNX', 'onnx', '.onnx', True],
['OpenVINO', 'openvino', '_openvino_model', False],
['TensorRT', 'engine', '.engine', True],
['CoreML', 'coreml', '.mlmodel', False],
['TensorFlow SavedModel', 'saved_model', '_saved_model', True],
['TensorFlow GraphDef', 'pb', '.pb', True],
['TensorFlow Lite', 'tflite', '.tflite', False],
['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False],
x = [['PyTorch', '-', '.pt', True], ['TorchScript', 'torchscript', '.torchscript', True],
['ONNX', 'onnx', '.onnx', True], ['OpenVINO', 'openvino', '_openvino_model', False],
['TensorRT', 'engine', '.engine', True], ['CoreML', 'coreml', '.mlmodel', False],
['TensorFlow SavedModel', 'saved_model', '_saved_model', True], ['TensorFlow GraphDef', 'pb', '.pb', True],
['TensorFlow Lite', 'tflite', '.tflite', False], ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False],
['TensorFlow.js', 'tfjs', '_web_model', False]] ['TensorFlow.js', 'tfjs', '_web_model', False]]
return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'GPU']) return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'GPU'])


LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...') LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
f = file.with_suffix('.onnx') f = file.with_suffix('.onnx')


torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
do_constant_folding=not train,
input_names=['images'],
output_names=['output'],
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640)
'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
} if dynamic else None)
torch.onnx.export(
model,
im,
f,
verbose=False,
opset_version=opset,
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
do_constant_folding=not train,
input_names=['images'],
output_names=['output'],
dynamic_axes={
'images': {
0: 'batch',
2: 'height',
3: 'width'}, # shape(1,3,640,640)
'output': {
0: 'batch',
1: 'anchors'} # shape(1,25200,85)
} if dynamic else None)


# Checks # Checks
model_onnx = onnx.load(f) # load onnx model model_onnx = onnx.load(f) # load onnx model
import onnxsim import onnxsim


LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...') LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
model_onnx, check = onnxsim.simplify(
model_onnx,
dynamic_input_shape=dynamic,
input_shapes={'images': list(im.shape)} if dynamic else None)
model_onnx, check = onnxsim.simplify(model_onnx,
dynamic_input_shape=dynamic,
input_shapes={'images': list(im.shape)} if dynamic else None)
assert check, 'assert check failed' assert check, 'assert check failed'
onnx.save(model_onnx, f) onnx.save(model_onnx, f)
except Exception as e: except Exception as e:
LOGGER.info(f'\n{prefix} export failure: {e}') LOGGER.info(f'\n{prefix} export failure: {e}')




def export_saved_model(model, im, file, dynamic,
tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
conf_thres=0.25, keras=False, prefix=colorstr('TensorFlow SavedModel:')):
def export_saved_model(model,
im,
file,
dynamic,
tf_nms=False,
agnostic_nms=False,
topk_per_class=100,
topk_all=100,
iou_thres=0.45,
conf_thres=0.25,
keras=False,
prefix=colorstr('TensorFlow SavedModel:')):
# YOLOv5 TensorFlow SavedModel export # YOLOv5 TensorFlow SavedModel export
try: try:
import tensorflow as tf import tensorflow as tf
tfm = tf.Module() tfm = tf.Module()
tfm.__call__ = tf.function(lambda x: frozen_func(x)[0], [spec]) tfm.__call__ = tf.function(lambda x: frozen_func(x)[0], [spec])
tfm.__call__(im) tfm.__call__(im)
tf.saved_model.save(
tfm,
f,
options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) if
check_version(tf.__version__, '2.6') else tf.saved_model.SaveOptions())
tf.saved_model.save(tfm,
f,
options=tf.saved_model.SaveOptions(experimental_custom_gradients=False)
if check_version(tf.__version__, '2.6') else tf.saved_model.SaveOptions())
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
return keras_model, f return keras_model, f
except Exception as e: except Exception as e:
if subprocess.run(cmd + ' >/dev/null', shell=True).returncode != 0: if subprocess.run(cmd + ' >/dev/null', shell=True).returncode != 0:
LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}') LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}')
sudo = subprocess.run('sudo --version >/dev/null', shell=True).returncode == 0 # sudo installed on system sudo = subprocess.run('sudo --version >/dev/null', shell=True).returncode == 0 # sudo installed on system
for c in ['curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -',
'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',
'sudo apt-get update',
'sudo apt-get install edgetpu-compiler']:
for c in (
'curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -',
'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',
'sudo apt-get update', 'sudo apt-get install edgetpu-compiler'):
subprocess.run(c if sudo else c.replace('sudo ', ''), shell=True, check=True) subprocess.run(c if sudo else c.replace('sudo ', ''), shell=True, check=True)
ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1] ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]


r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, ' r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
r'"Identity.?.?": {"name": "Identity.?.?"}, ' r'"Identity.?.?": {"name": "Identity.?.?"}, '
r'"Identity.?.?": {"name": "Identity.?.?"}, ' r'"Identity.?.?": {"name": "Identity.?.?"}, '
r'"Identity.?.?": {"name": "Identity.?.?"}}}',
r'{"outputs": {"Identity": {"name": "Identity"}, '
r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, '
r'"Identity_1": {"name": "Identity_1"}, ' r'"Identity_1": {"name": "Identity_1"}, '
r'"Identity_2": {"name": "Identity_2"}, ' r'"Identity_2": {"name": "Identity_2"}, '
r'"Identity_3": {"name": "Identity_3"}}}',
json)
r'"Identity_3": {"name": "Identity_3"}}}', json)
j.write(subst) j.write(subst)


LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')




@torch.no_grad() @torch.no_grad()
def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
def run(
data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
weights=ROOT / 'yolov5s.pt', # weights path weights=ROOT / 'yolov5s.pt', # weights path
imgsz=(640, 640), # image (height, width) imgsz=(640, 640), # image (height, width)
batch_size=1, # batch size batch_size=1, # batch size
topk_per_class=100, # TF.js NMS: topk per class to keep topk_per_class=100, # TF.js NMS: topk per class to keep
topk_all=100, # TF.js NMS: topk for all classes to keep topk_all=100, # TF.js NMS: topk for all classes to keep
iou_thres=0.45, # TF.js NMS: IoU threshold iou_thres=0.45, # TF.js NMS: IoU threshold
conf_thres=0.25 # TF.js NMS: confidence threshold
):
conf_thres=0.25, # TF.js NMS: confidence threshold
):
t = time.time() t = time.time()
include = [x.lower() for x in include] # to lowercase include = [x.lower() for x in include] # to lowercase
formats = tuple(export_formats()['Argument'][1:]) # --include arguments formats = tuple(export_formats()['Argument'][1:]) # --include arguments
if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707 if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow` check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.' assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
model, f[5] = export_saved_model(model.cpu(), im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class,
topk_all=topk_all, conf_thres=conf_thres, iou_thres=iou_thres) # keras model
model, f[5] = export_saved_model(model.cpu(),
im,
file,
dynamic,
tf_nms=nms or agnostic_nms or tfjs,
agnostic_nms=agnostic_nms or tfjs,
topk_per_class=topk_per_class,
topk_all=topk_all,
conf_thres=conf_thres,
iou_thres=iou_thres) # keras model
if pb or tfjs: # pb prerequisite to tfjs if pb or tfjs: # pb prerequisite to tfjs
f[6] = export_pb(model, im, file) f[6] = export_pb(model, im, file)
if tflite or edgetpu: if tflite or edgetpu:
parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep') parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep')
parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold') parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold')
parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold') parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
parser.add_argument('--include', nargs='+',
parser.add_argument('--include',
nargs='+',
default=['torchscript', 'onnx'], default=['torchscript', 'onnx'],
help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs') help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs')
opt = parser.parse_args() opt = parser.parse_args()

+ 7
- 6
hubconf.py View File



from utils.general import cv2 from utils.general import cv2


imgs = ['data/images/zidane.jpg', # filename
Path('data/images/zidane.jpg'), # Path
'https://ultralytics.com/images/zidane.jpg', # URI
cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV
Image.open('data/images/bus.jpg'), # PIL
np.zeros((320, 640, 3))] # numpy
imgs = [
'data/images/zidane.jpg', # filename
Path('data/images/zidane.jpg'), # Path
'https://ultralytics.com/images/zidane.jpg', # URI
cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV
Image.open('data/images/bus.jpg'), # PIL
np.zeros((320, 640, 3))] # numpy


results = model(imgs, size=320) # batched inference results = model(imgs, size=320) # batched inference
results.print() results.print()

+ 24
- 13
models/common.py View File

def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride
super().__init__() super().__init__()
c_ = c2 // 2 c_ = c2 // 2
self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw
DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False),
Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
self.conv = nn.Sequential(
GhostConv(c1, c_, 1, 1), # pw
DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1,
act=False)) if s == 2 else nn.Identity()


def forward(self, x): def forward(self, x):
return self.conv(x) + self.shortcut(x) return self.conv(x) + self.shortcut(x)
Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate, Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate,
if edgetpu: # Edge TPU https://coral.ai/software/#edgetpu-runtime if edgetpu: # Edge TPU https://coral.ai/software/#edgetpu-runtime
LOGGER.info(f'Loading {w} for TensorFlow Lite Edge TPU inference...') LOGGER.info(f'Loading {w} for TensorFlow Lite Edge TPU inference...')
delegate = {'Linux': 'libedgetpu.so.1',
'Darwin': 'libedgetpu.1.dylib',
'Windows': 'edgetpu.dll'}[platform.system()]
delegate = {
'Linux': 'libedgetpu.so.1',
'Darwin': 'libedgetpu.1.dylib',
'Windows': 'edgetpu.dll'}[platform.system()]
interpreter = Interpreter(model_path=w, experimental_delegates=[load_delegate(delegate)]) interpreter = Interpreter(model_path=w, experimental_delegates=[load_delegate(delegate)])
else: # Lite else: # Lite
LOGGER.info(f'Loading {w} for TensorFlow Lite inference...') LOGGER.info(f'Loading {w} for TensorFlow Lite inference...')
return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference


# Pre-process # Pre-process
n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images
n, imgs = (len(imgs), list(imgs)) if isinstance(imgs, (list, tuple)) else (1, [imgs]) # number, list of images
shape0, shape1, files = [], [], [] # image and inference shapes, filenames shape0, shape1, files = [], [], [] # image and inference shapes, filenames
for i, im in enumerate(imgs): for i, im in enumerate(imgs):
f = f'image{i}' # filename f = f'image{i}' # filename
t.append(time_sync()) t.append(time_sync())


# Post-process # Post-process
y = non_max_suppression(y if self.dmb else y[0], self.conf, self.iou, self.classes, self.agnostic,
self.multi_label, max_det=self.max_det) # NMS
y = non_max_suppression(y if self.dmb else y[0],
self.conf,
self.iou,
self.classes,
self.agnostic,
self.multi_label,
max_det=self.max_det) # NMS
for i in range(n): for i in range(n):
scale_coords(shape1, y[i][:, :4], shape0[i]) scale_coords(shape1, y[i][:, :4], shape0[i])


label = f'{self.names[int(cls)]} {conf:.2f}' label = f'{self.names[int(cls)]} {conf:.2f}'
if crop: if crop:
file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None
crops.append({'box': box, 'conf': conf, 'cls': cls, 'label': label,
'im': save_one_box(box, im, file=file, save=save)})
crops.append({
'box': box,
'conf': conf,
'cls': cls,
'label': label,
'im': save_one_box(box, im, file=file, save=save)})
else: # all others else: # all others
annotator.box_label(box, label if labels else '', color=colors(cls)) annotator.box_label(box, label if labels else '', color=colors(cls))
im = annotator.im im = annotator.im

+ 2
- 2
models/experimental.py View File

a[0] = 1 a[0] = 1
c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b


self.m = nn.ModuleList(
[nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)])
self.m = nn.ModuleList([
nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)])
self.bn = nn.BatchNorm2d(c2) self.bn = nn.BatchNorm2d(c2)
self.act = nn.SiLU() self.act = nn.SiLU()



+ 47
- 20
models/tf.py View File

# see https://stackoverflow.com/questions/52975843/comparing-conv2d-with-padding-between-tensorflow-and-pytorch # see https://stackoverflow.com/questions/52975843/comparing-conv2d-with-padding-between-tensorflow-and-pytorch


conv = keras.layers.Conv2D( conv = keras.layers.Conv2D(
c2, k, s, 'SAME' if s == 1 else 'VALID', use_bias=False if hasattr(w, 'bn') else True,
c2,
k,
s,
'SAME' if s == 1 else 'VALID',
use_bias=False if hasattr(w, 'bn') else True,
kernel_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy()), kernel_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy()),
bias_initializer='zeros' if hasattr(w, 'bn') else keras.initializers.Constant(w.conv.bias.numpy())) bias_initializer='zeros' if hasattr(w, 'bn') else keras.initializers.Constant(w.conv.bias.numpy()))
self.conv = conv if s == 1 else keras.Sequential([TFPad(autopad(k, p)), conv]) self.conv = conv if s == 1 else keras.Sequential([TFPad(autopad(k, p)), conv])


def call(self, inputs): # x(b,w,h,c) -> y(b,w/2,h/2,4c) def call(self, inputs): # x(b,w,h,c) -> y(b,w/2,h/2,4c)
# inputs = inputs / 255 # normalize 0-255 to 0-1 # inputs = inputs / 255 # normalize 0-255 to 0-1
return self.conv(tf.concat([inputs[:, ::2, ::2, :],
inputs[:, 1::2, ::2, :],
inputs[:, ::2, 1::2, :],
inputs[:, 1::2, 1::2, :]], 3))
return self.conv(
tf.concat(
[inputs[:, ::2, ::2, :], inputs[:, 1::2, ::2, :], inputs[:, ::2, 1::2, :], inputs[:, 1::2, 1::2, :]],
3))




class TFBottleneck(keras.layers.Layer): class TFBottleneck(keras.layers.Layer):
super().__init__() super().__init__()
assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument" assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument"
self.conv = keras.layers.Conv2D( self.conv = keras.layers.Conv2D(
c2, k, s, 'VALID', use_bias=bias,
c2,
k,
s,
'VALID',
use_bias=bias,
kernel_initializer=keras.initializers.Constant(w.weight.permute(2, 3, 1, 0).numpy()), kernel_initializer=keras.initializers.Constant(w.weight.permute(2, 3, 1, 0).numpy()),
bias_initializer=keras.initializers.Constant(w.bias.numpy()) if bias else None, )
bias_initializer=keras.initializers.Constant(w.bias.numpy()) if bias else None,
)


def call(self, inputs): def call(self, inputs):
return self.conv(inputs) return self.conv(inputs)
self.na = len(anchors[0]) // 2 # number of anchors self.na = len(anchors[0]) // 2 # number of anchors
self.grid = [tf.zeros(1)] * self.nl # init grid self.grid = [tf.zeros(1)] * self.nl # init grid
self.anchors = tf.convert_to_tensor(w.anchors.numpy(), dtype=tf.float32) self.anchors = tf.convert_to_tensor(w.anchors.numpy(), dtype=tf.float32)
self.anchor_grid = tf.reshape(self.anchors * tf.reshape(self.stride, [self.nl, 1, 1]),
[self.nl, 1, -1, 1, 2])
self.anchor_grid = tf.reshape(self.anchors * tf.reshape(self.stride, [self.nl, 1, 1]), [self.nl, 1, -1, 1, 2])
self.m = [TFConv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in enumerate(ch)] self.m = [TFConv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in enumerate(ch)]
self.training = False # set to False after building model self.training = False # set to False after building model
self.imgsz = imgsz self.imgsz = imgsz
self.yaml['nc'] = nc # override yaml value self.yaml['nc'] = nc # override yaml value
self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch], model=model, imgsz=imgsz) self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch], model=model, imgsz=imgsz)


def predict(self, inputs, tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
def predict(self,
inputs,
tf_nms=False,
agnostic_nms=False,
topk_per_class=100,
topk_all=100,
iou_thres=0.45,
conf_thres=0.25): conf_thres=0.25):
y = [] # outputs y = [] # outputs
x = inputs x = inputs
return nms, x[1] return nms, x[1]
else: else:
boxes = tf.expand_dims(boxes, 2) boxes = tf.expand_dims(boxes, 2)
nms = tf.image.combined_non_max_suppression(
boxes, scores, topk_per_class, topk_all, iou_thres, conf_thres, clip_boxes=False)
nms = tf.image.combined_non_max_suppression(boxes,
scores,
topk_per_class,
topk_all,
iou_thres,
conf_thres,
clip_boxes=False)
return nms, x[1] return nms, x[1]


return x[0] # output only first tensor [1,6300,85] = [xywh, conf, class0, class1, ...] return x[0] # output only first tensor [1,6300,85] = [xywh, conf, class0, class1, ...]
# TF Agnostic NMS # TF Agnostic NMS
def call(self, input, topk_all, iou_thres, conf_thres): def call(self, input, topk_all, iou_thres, conf_thres):
# wrap map_fn to avoid TypeSpec related error https://stackoverflow.com/a/65809989/3036450 # wrap map_fn to avoid TypeSpec related error https://stackoverflow.com/a/65809989/3036450
return tf.map_fn(lambda x: self._nms(x, topk_all, iou_thres, conf_thres), input,
return tf.map_fn(lambda x: self._nms(x, topk_all, iou_thres, conf_thres),
input,
fn_output_signature=(tf.float32, tf.float32, tf.float32, tf.int32), fn_output_signature=(tf.float32, tf.float32, tf.float32, tf.int32),
name='agnostic_nms') name='agnostic_nms')


boxes, classes, scores = x boxes, classes, scores = x
class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32) class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32)
scores_inp = tf.reduce_max(scores, -1) scores_inp = tf.reduce_max(scores, -1)
selected_inds = tf.image.non_max_suppression(
boxes, scores_inp, max_output_size=topk_all, iou_threshold=iou_thres, score_threshold=conf_thres)
selected_inds = tf.image.non_max_suppression(boxes,
scores_inp,
max_output_size=topk_all,
iou_threshold=iou_thres,
score_threshold=conf_thres)
selected_boxes = tf.gather(boxes, selected_inds) selected_boxes = tf.gather(boxes, selected_inds)
padded_boxes = tf.pad(selected_boxes, padded_boxes = tf.pad(selected_boxes,
paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]], paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]],
mode="CONSTANT", constant_values=0.0)
mode="CONSTANT",
constant_values=0.0)
selected_scores = tf.gather(scores_inp, selected_inds) selected_scores = tf.gather(scores_inp, selected_inds)
padded_scores = tf.pad(selected_scores, padded_scores = tf.pad(selected_scores,
paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
mode="CONSTANT", constant_values=-1.0)
mode="CONSTANT",
constant_values=-1.0)
selected_classes = tf.gather(class_inds, selected_inds) selected_classes = tf.gather(class_inds, selected_inds)
padded_classes = tf.pad(selected_classes, padded_classes = tf.pad(selected_classes,
paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
mode="CONSTANT", constant_values=-1.0)
mode="CONSTANT",
constant_values=-1.0)
valid_detections = tf.shape(selected_inds)[0] valid_detections = tf.shape(selected_inds)[0]
return padded_boxes, padded_scores, padded_classes, valid_detections return padded_boxes, padded_scores, padded_classes, valid_detections


break break




def run(weights=ROOT / 'yolov5s.pt', # weights path
def run(
weights=ROOT / 'yolov5s.pt', # weights path
imgsz=(640, 640), # inference size h,w imgsz=(640, 640), # inference size h,w
batch_size=1, # batch size batch_size=1, # batch size
dynamic=False, # dynamic batch size dynamic=False, # dynamic batch size
):
):
# PyTorch model # PyTorch model
im = torch.zeros((batch_size, 3, *imgsz)) # BCHW image im = torch.zeros((batch_size, 3, *imgsz)) # BCHW image
model = attempt_load(weights, map_location=torch.device('cpu'), inplace=True, fuse=False) model = attempt_load(weights, map_location=torch.device('cpu'), inplace=True, fuse=False)

+ 2
- 2
models/yolo.py View File

pass pass


n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]:
if m in (Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
BottleneckCSP, C3, C3TR, C3SPP, C3Ghost):
c1, c2 = ch[f], args[0] c1, c2 = ch[f], args[0]
if c2 != no: # if not output if c2 != no: # if not output
c2 = make_divisible(c2 * gw, 8) c2 = make_divisible(c2 * gw, 8)

+ 14
- 0
setup.cfg View File

# Project-wide configuration file, can be used for package metadata and other toll configurations # Project-wide configuration file, can be used for package metadata and other toll configurations
# Example usage: global configuration for PEP8 (via flake8) setting or default pytest arguments # Example usage: global configuration for PEP8 (via flake8) setting or default pytest arguments
# Local usage: pip install pre-commit, pre-commit run --all-files


[metadata] [metadata]
license_file = LICENSE license_file = LICENSE
[isort] [isort]
# https://pycqa.github.io/isort/docs/configuration/options.html # https://pycqa.github.io/isort/docs/configuration/options.html
line_length = 120 line_length = 120
# see: https://pycqa.github.io/isort/docs/configuration/multi_line_output_modes.html
multi_line_output = 0 multi_line_output = 0


[yapf]
based_on_style = pep8
spaces_before_comment = 2
COLUMN_LIMIT = 120
COALESCE_BRACKETS = True
SPACES_AROUND_POWER_OPERATOR = True
SPACE_BETWEEN_ENDING_COMMA_AND_CLOSING_BRACKET = False
SPLIT_BEFORE_CLOSING_BRACKET = False
SPLIT_BEFORE_FIRST_ARGUMENT = False
# EACH_DICT_ENTRY_ON_SEPARATE_LINE = False

+ 82
- 65
train.py View File

WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1)) WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))




def train(hyp, # path/to/hyp.yaml or hyp dictionary
opt,
device,
callbacks
):
def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary
save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = \ save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = \
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \ Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze
LOGGER.info('Using SyncBatchNorm()') LOGGER.info('Using SyncBatchNorm()')


# Trainloader # Trainloader
train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls,
hyp=hyp, augment=True, cache=None if opt.cache == 'val' else opt.cache,
rect=opt.rect, rank=LOCAL_RANK, workers=workers,
image_weights=opt.image_weights, quad=opt.quad,
prefix=colorstr('train: '), shuffle=True)
train_loader, dataset = create_dataloader(train_path,
imgsz,
batch_size // WORLD_SIZE,
gs,
single_cls,
hyp=hyp,
augment=True,
cache=None if opt.cache == 'val' else opt.cache,
rect=opt.rect,
rank=LOCAL_RANK,
workers=workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '),
shuffle=True)
mlc = int(np.concatenate(dataset.labels, 0)[:, 0].max()) # max label class mlc = int(np.concatenate(dataset.labels, 0)[:, 0].max()) # max label class
nb = len(train_loader) # number of batches nb = len(train_loader) # number of batches
assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}' assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'


# Process 0 # Process 0
if RANK in [-1, 0]: if RANK in [-1, 0]:
val_loader = create_dataloader(val_path, imgsz, batch_size // WORLD_SIZE * 2, gs, single_cls,
hyp=hyp, cache=None if noval else opt.cache,
rect=True, rank=-1, workers=workers * 2, pad=0.5,
val_loader = create_dataloader(val_path,
imgsz,
batch_size // WORLD_SIZE * 2,
gs,
single_cls,
hyp=hyp,
cache=None if noval else opt.cache,
rect=True,
rank=-1,
workers=workers * 2,
pad=0.5,
prefix=colorstr('val: '))[0] prefix=colorstr('val: '))[0]


if not resume: if not resume:
if RANK in [-1, 0]: if RANK in [-1, 0]:
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB) mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)
pbar.set_description(('%10s' * 2 + '%10.4g' * 5) % (
f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
pbar.set_description(('%10s' * 2 + '%10.4g' * 5) %
(f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots, opt.sync_bn) callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots, opt.sync_bn)
if callbacks.stop_training: if callbacks.stop_training:
return return


# Save model # Save model
if (not nosave) or (final_epoch and not evolve): # if save if (not nosave) or (final_epoch and not evolve): # if save
ckpt = {'epoch': epoch,
'best_fitness': best_fitness,
'model': deepcopy(de_parallel(model)).half(),
'ema': deepcopy(ema.ema).half(),
'updates': ema.updates,
'optimizer': optimizer.state_dict(),
'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None,
'date': datetime.now().isoformat()}
ckpt = {
'epoch': epoch,
'best_fitness': best_fitness,
'model': deepcopy(de_parallel(model)).half(),
'ema': deepcopy(ema.ema).half(),
'updates': ema.updates,
'optimizer': optimizer.state_dict(),
'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None,
'date': datetime.now().isoformat()}


# Save last, best and delete # Save last, best and delete
torch.save(ckpt, last) torch.save(ckpt, last)
strip_optimizer(f) # strip optimizers strip_optimizer(f) # strip optimizers
if f is best: if f is best:
LOGGER.info(f'\nValidating {f}...') LOGGER.info(f'\nValidating {f}...')
results, _, _ = val.run(data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
model=attempt_load(f, device).half(),
iou_thres=0.65 if is_coco else 0.60, # best pycocotools results at 0.65
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
save_json=is_coco,
verbose=True,
plots=True,
callbacks=callbacks,
compute_loss=compute_loss) # val best model with plots
results, _, _ = val.run(
data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
model=attempt_load(f, device).half(),
iou_thres=0.65 if is_coco else 0.60, # best pycocotools results at 0.65
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
save_json=is_coco,
verbose=True,
plots=True,
callbacks=callbacks,
compute_loss=compute_loss) # val best model with plots
if is_coco: if is_coco:
callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi) callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)


# Evolve hyperparameters (optional) # Evolve hyperparameters (optional)
else: else:
# Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit) # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
meta = {'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)
'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)
'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1
'weight_decay': (1, 0.0, 0.001), # optimizer weight decay
'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok)
'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum
'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr
'box': (1, 0.02, 0.2), # box loss gain
'cls': (1, 0.2, 4.0), # cls loss gain
'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight
'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels)
'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight
'iou_t': (0, 0.1, 0.7), # IoU training threshold
'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold
'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore)
'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)
'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction)
'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction)
'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction)
'degrees': (1, 0.0, 45.0), # image rotation (+/- deg)
'translate': (1, 0.0, 0.9), # image translation (+/- fraction)
'scale': (1, 0.0, 0.9), # image scale (+/- gain)
'shear': (1, 0.0, 10.0), # image shear (+/- deg)
'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001
'flipud': (1, 0.0, 1.0), # image flip up-down (probability)
'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)
'mosaic': (1, 0.0, 1.0), # image mixup (probability)
'mixup': (1, 0.0, 1.0), # image mixup (probability)
'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability)
meta = {
'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)
'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)
'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1
'weight_decay': (1, 0.0, 0.001), # optimizer weight decay
'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok)
'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum
'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr
'box': (1, 0.02, 0.2), # box loss gain
'cls': (1, 0.2, 4.0), # cls loss gain
'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight
'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels)
'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight
'iou_t': (0, 0.1, 0.7), # IoU training threshold
'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold
'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore)
'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)
'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction)
'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction)
'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction)
'degrees': (1, 0.0, 45.0), # image rotation (+/- deg)
'translate': (1, 0.0, 0.9), # image translation (+/- fraction)
'scale': (1, 0.0, 0.9), # image scale (+/- gain)
'shear': (1, 0.0, 10.0), # image shear (+/- deg)
'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001
'flipud': (1, 0.0, 1.0), # image flip up-down (probability)
'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)
'mosaic': (1, 0.0, 1.0), # image mixup (probability)
'mixup': (1, 0.0, 1.0), # image mixup (probability)
'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability)


with open(opt.hyp, errors='ignore') as f: with open(opt.hyp, errors='ignore') as f:
hyp = yaml.safe_load(f) # load hyps dict hyp = yaml.safe_load(f) # load hyps dict

+ 0
- 2
utils/activations.py View File

AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter
according to "Activate or Not: Learning Customized Activation" <https://arxiv.org/pdf/2009.04759.pdf>. according to "Activate or Not: Learning Customized Activation" <https://arxiv.org/pdf/2009.04759.pdf>.
""" """

def __init__(self, c1): def __init__(self, c1):
super().__init__() super().__init__()
self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1)) self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1))
MetaAconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is generated by a small network MetaAconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is generated by a small network
according to "Activate or Not: Learning Customized Activation" <https://arxiv.org/pdf/2009.04759.pdf>. according to "Activate or Not: Learning Customized Activation" <https://arxiv.org/pdf/2009.04759.pdf>.
""" """

def __init__(self, c1, k=1, s=1, r=16): # ch_in, kernel, stride, r def __init__(self, c1, k=1, s=1, r=16): # ch_in, kernel, stride, r
super().__init__() super().__init__()
c2 = max(r, c1 // r) c2 = max(r, c1 // r)

+ 11
- 4
utils/augmentations.py View File

import albumentations as A import albumentations as A
check_version(A.__version__, '1.0.3', hard=True) # version requirement check_version(A.__version__, '1.0.3', hard=True) # version requirement


self.transform = A.Compose([
T = [
A.Blur(p=0.01), A.Blur(p=0.01),
A.MedianBlur(p=0.01), A.MedianBlur(p=0.01),
A.ToGray(p=0.01), A.ToGray(p=0.01),
A.CLAHE(p=0.01), A.CLAHE(p=0.01),
A.RandomBrightnessContrast(p=0.0), A.RandomBrightnessContrast(p=0.0),
A.RandomGamma(p=0.0), A.RandomGamma(p=0.0),
A.ImageCompression(quality_lower=75, p=0.0)],
bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
A.ImageCompression(quality_lower=75, p=0.0)] # transforms
self.transform = A.Compose(T, bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))


LOGGER.info(colorstr('albumentations: ') + ', '.join(f'{x}' for x in self.transform.transforms if x.p)) LOGGER.info(colorstr('albumentations: ') + ', '.join(f'{x}' for x in self.transform.transforms if x.p))
except ImportError: # package not installed, skip except ImportError: # package not installed, skip
return im, ratio, (dw, dh) return im, ratio, (dw, dh)




def random_perspective(im, targets=(), segments=(), degrees=10, translate=.1, scale=.1, shear=10, perspective=0.0,
def random_perspective(im,
targets=(),
segments=(),
degrees=10,
translate=.1,
scale=.1,
shear=10,
perspective=0.0,
border=(0, 0)): border=(0, 0)):
# torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10)) # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10))
# targets = [cls, xyxy] # targets = [cls, xyxy]

+ 3
- 2
utils/benchmarks.py View File

from utils.torch_utils import select_device from utils.torch_utils import select_device




def run(weights=ROOT / 'yolov5s.pt', # weights path
def run(
weights=ROOT / 'yolov5s.pt', # weights path
imgsz=640, # inference size (pixels) imgsz=640, # inference size (pixels)
batch_size=1, # batch size batch_size=1, # batch size
data=ROOT / 'data/coco128.yaml', # dataset.yaml path data=ROOT / 'data/coco128.yaml', # dataset.yaml path
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
half=False, # use FP16 half-precision inference half=False, # use FP16 half-precision inference
):
):
y, t = [], time.time() y, t = [], time.time()
formats = export.export_formats() formats = export.export_formats()
device = select_device(device) device = select_device(device)

+ 1
- 6
utils/callbacks.py View File

"""" """"
Handles all registered callbacks for YOLOv5 Hooks Handles all registered callbacks for YOLOv5 Hooks
""" """

def __init__(self): def __init__(self):
# Define the available callbacks # Define the available callbacks
self._callbacks = { self._callbacks = {
'on_pretrain_routine_start': [], 'on_pretrain_routine_start': [],
'on_pretrain_routine_end': [], 'on_pretrain_routine_end': [],

'on_train_start': [], 'on_train_start': [],
'on_train_epoch_start': [], 'on_train_epoch_start': [],
'on_train_batch_start': [], 'on_train_batch_start': [],
'on_before_zero_grad': [], 'on_before_zero_grad': [],
'on_train_batch_end': [], 'on_train_batch_end': [],
'on_train_epoch_end': [], 'on_train_epoch_end': [],

'on_val_start': [], 'on_val_start': [],
'on_val_batch_start': [], 'on_val_batch_start': [],
'on_val_image_end': [], 'on_val_image_end': [],
'on_val_batch_end': [], 'on_val_batch_end': [],
'on_val_end': [], 'on_val_end': [],

'on_fit_epoch_end': [], # fit = train + val 'on_fit_epoch_end': [], # fit = train + val
'on_model_save': [], 'on_model_save': [],
'on_train_end': [], 'on_train_end': [],
'on_params_update': [], 'on_params_update': [],
'teardown': [],
}
'teardown': [],}
self.stop_training = False # set True to interrupt training self.stop_training = False # set True to interrupt training


def register_action(self, hook, name='', callback=None): def register_action(self, hook, name='', callback=None):

+ 74
- 38
utils/datasets.py View File

exif = image.getexif() exif = image.getexif()
orientation = exif.get(0x0112, 1) # default 1 orientation = exif.get(0x0112, 1) # default 1
if orientation > 1: if orientation > 1:
method = {2: Image.FLIP_LEFT_RIGHT,
3: Image.ROTATE_180,
4: Image.FLIP_TOP_BOTTOM,
5: Image.TRANSPOSE,
6: Image.ROTATE_270,
7: Image.TRANSVERSE,
8: Image.ROTATE_90,
}.get(orientation)
method = {
2: Image.FLIP_LEFT_RIGHT,
3: Image.ROTATE_180,
4: Image.FLIP_TOP_BOTTOM,
5: Image.TRANSPOSE,
6: Image.ROTATE_270,
7: Image.TRANSVERSE,
8: Image.ROTATE_90,}.get(orientation)
if method is not None: if method is not None:
image = image.transpose(method) image = image.transpose(method)
del exif[0x0112] del exif[0x0112]
return image return image




def create_dataloader(path, imgsz, batch_size, stride, single_cls=False, hyp=None, augment=False, cache=False, pad=0.0,
rect=False, rank=-1, workers=8, image_weights=False, quad=False, prefix='', shuffle=False):
def create_dataloader(path,
imgsz,
batch_size,
stride,
single_cls=False,
hyp=None,
augment=False,
cache=False,
pad=0.0,
rect=False,
rank=-1,
workers=8,
image_weights=False,
quad=False,
prefix='',
shuffle=False):
if rect and shuffle: if rect and shuffle:
LOGGER.warning('WARNING: --rect is incompatible with DataLoader shuffle, setting shuffle=False') LOGGER.warning('WARNING: --rect is incompatible with DataLoader shuffle, setting shuffle=False')
shuffle = False shuffle = False
with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP
dataset = LoadImagesAndLabels(path, imgsz, batch_size,
augment=augment, # augmentation
hyp=hyp, # hyperparameters
rect=rect, # rectangular batches
cache_images=cache,
single_cls=single_cls,
stride=int(stride),
pad=pad,
image_weights=image_weights,
prefix=prefix)
dataset = LoadImagesAndLabels(
path,
imgsz,
batch_size,
augment=augment, # augmentation
hyp=hyp, # hyperparameters
rect=rect, # rectangular batches
cache_images=cache,
single_cls=single_cls,
stride=int(stride),
pad=pad,
image_weights=image_weights,
prefix=prefix)


batch_size = min(batch_size, len(dataset)) batch_size = min(batch_size, len(dataset))
nd = torch.cuda.device_count() # number of CUDA devices nd = torch.cuda.device_count() # number of CUDA devices


Uses same syntax as vanilla DataLoader Uses same syntax as vanilla DataLoader
""" """

def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler)) object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))
Args: Args:
sampler (Sampler) sampler (Sampler)
""" """

def __init__(self, sampler): def __init__(self, sampler):
self.sampler = sampler self.sampler = sampler


# YOLOv5 train_loader/val_loader, loads images and labels for training and validation # YOLOv5 train_loader/val_loader, loads images and labels for training and validation
cache_version = 0.6 # dataset labels *.cache version cache_version = 0.6 # dataset labels *.cache version


def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''):
def __init__(self,
path,
img_size=640,
batch_size=16,
augment=False,
hyp=None,
rect=False,
image_weights=False,
cache_images=False,
single_cls=False,
stride=32,
pad=0.0,
prefix=''):
self.img_size = img_size self.img_size = img_size
self.augment = augment self.augment = augment
self.hyp = hyp self.hyp = hyp
desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels..." desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels..."
with Pool(NUM_THREADS) as pool: with Pool(NUM_THREADS) as pool:
pbar = tqdm(pool.imap(verify_image_label, zip(self.im_files, self.label_files, repeat(prefix))), pbar = tqdm(pool.imap(verify_image_label, zip(self.im_files, self.label_files, repeat(prefix))),
desc=desc, total=len(self.im_files), bar_format=BAR_FORMAT)
desc=desc,
total=len(self.im_files),
bar_format=BAR_FORMAT)
for im_file, lb, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar: for im_file, lb, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar:
nm += nm_f nm += nm_f
nf += nf_f nf += nf_f
labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1]) labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])


if self.augment: if self.augment:
img, labels = random_perspective(img, labels,
img, labels = random_perspective(img,
labels,
degrees=hyp['degrees'], degrees=hyp['degrees'],
translate=hyp['translate'], translate=hyp['translate'],
scale=hyp['scale'], scale=hyp['scale'],
h0, w0 = im.shape[:2] # orig hw h0, w0 = im.shape[:2] # orig hw
r = self.img_size / max(h0, w0) # ratio r = self.img_size / max(h0, w0) # ratio
if r != 1: # if sizes are not equal if r != 1: # if sizes are not equal
im = cv2.resize(im,
(int(w0 * r), int(h0 * r)),
im = cv2.resize(im, (int(w0 * r), int(h0 * r)),
interpolation=cv2.INTER_LINEAR if (self.augment or r > 1) else cv2.INTER_AREA) interpolation=cv2.INTER_LINEAR if (self.augment or r > 1) else cv2.INTER_AREA)
return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized
else: else:


# Augment # Augment
img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste']) img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste'])
img4, labels4 = random_perspective(img4, labels4, segments4,
img4, labels4 = random_perspective(img4,
labels4,
segments4,
degrees=self.hyp['degrees'], degrees=self.hyp['degrees'],
translate=self.hyp['translate'], translate=self.hyp['translate'],
scale=self.hyp['scale'], scale=self.hyp['scale'],
# img9, labels9 = replicate(img9, labels9) # replicate # img9, labels9 = replicate(img9, labels9) # replicate


# Augment # Augment
img9, labels9 = random_perspective(img9, labels9, segments9,
img9, labels9 = random_perspective(img9,
labels9,
segments9,
degrees=self.hyp['degrees'], degrees=self.hyp['degrees'],
translate=self.hyp['translate'], translate=self.hyp['translate'],
scale=self.hyp['scale'], scale=self.hyp['scale'],
for i in range(n): # zidane torch.zeros(16,3,720,1280) # BCHW for i in range(n): # zidane torch.zeros(16,3,720,1280) # BCHW
i *= 4 i *= 4
if random.random() < 0.5: if random.random() < 0.5:
im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2.0, mode='bilinear', align_corners=False)[
0].type(img[i].type())
im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2.0, mode='bilinear',
align_corners=False)[0].type(img[i].type())
lb = label[i] lb = label[i]
else: else:
im = torch.cat((torch.cat((img[i], img[i + 1]), 1), torch.cat((img[i + 2], img[i + 3]), 1)), 2) im = torch.cat((torch.cat((img[i], img[i + 1]), 1), torch.cat((img[i + 2], img[i + 3]), 1)), 2)
autodownload: Attempt to download dataset if not found locally autodownload: Attempt to download dataset if not found locally
verbose: Print stats dictionary verbose: Print stats dictionary
""" """

def round_labels(labels): def round_labels(labels):
# Update labels to integer class and 6 decimal place floats # Update labels to integer class and 6 decimal place floats
return [[int(c), *(round(x, 4) for x in points)] for c, *points in labels] return [[int(c), *(round(x, 4) for x in points)] for c, *points in labels]
for label in tqdm(dataset.labels, total=dataset.n, desc='Statistics'): for label in tqdm(dataset.labels, total=dataset.n, desc='Statistics'):
x.append(np.bincount(label[:, 0].astype(int), minlength=data['nc'])) x.append(np.bincount(label[:, 0].astype(int), minlength=data['nc']))
x = np.array(x) # shape(128x80) x = np.array(x) # shape(128x80)
stats[split] = {'instance_stats': {'total': int(x.sum()), 'per_class': x.sum(0).tolist()},
'image_stats': {'total': dataset.n, 'unlabelled': int(np.all(x == 0, 1).sum()),
'per_class': (x > 0).sum(0).tolist()},
'labels': [{str(Path(k).name): round_labels(v.tolist())} for k, v in
zip(dataset.im_files, dataset.labels)]}
stats[split] = {
'instance_stats': {
'total': int(x.sum()),
'per_class': x.sum(0).tolist()},
'image_stats': {
'total': dataset.n,
'unlabelled': int(np.all(x == 0, 1).sum()),
'per_class': (x > 0).sum(0).tolist()},
'labels': [{
str(Path(k).name): round_labels(v.tolist())} for k, v in zip(dataset.im_files, dataset.labels)]}


if hub: if hub:
im_dir = hub_dir / 'images' im_dir = hub_dir / 'images'

+ 10
- 7
utils/downloads.py View File

assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...] assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...]
tag = response['tag_name'] # i.e. 'v1.0' tag = response['tag_name'] # i.e. 'v1.0'
except Exception: # fallback plan except Exception: # fallback plan
assets = ['yolov5n.pt', 'yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt',
'yolov5n6.pt', 'yolov5s6.pt', 'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt']
assets = [
'yolov5n.pt', 'yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 'yolov5n6.pt', 'yolov5s6.pt',
'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt']
try: try:
tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1] tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
except Exception: except Exception:
tag = 'v6.0' # current release tag = 'v6.0' # current release


if name in assets: if name in assets:
safe_download(file,
url=f'https://github.com/{repo}/releases/download/{tag}/{name}',
# url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional)
min_bytes=1E5,
error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/')
safe_download(
file,
url=f'https://github.com/{repo}/releases/download/{tag}/{name}',
# url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional)
min_bytes=1E5,
error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/')


return str(file) return str(file)


return line.split()[-1] return line.split()[-1]
return "" return ""



# Google utils: https://cloud.google.com/storage/docs/reference/libraries ---------------------------------------------- # Google utils: https://cloud.google.com/storage/docs/reference/libraries ----------------------------------------------
# #
# #

+ 40
- 34
utils/general.py View File

def colorstr(*input): def colorstr(*input):
# Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e. colorstr('blue', 'hello world') # Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e. colorstr('blue', 'hello world')
*args, string = input if len(input) > 1 else ('blue', 'bold', input[0]) # color arguments, string *args, string = input if len(input) > 1 else ('blue', 'bold', input[0]) # color arguments, string
colors = {'black': '\033[30m', # basic colors
'red': '\033[31m',
'green': '\033[32m',
'yellow': '\033[33m',
'blue': '\033[34m',
'magenta': '\033[35m',
'cyan': '\033[36m',
'white': '\033[37m',
'bright_black': '\033[90m', # bright colors
'bright_red': '\033[91m',
'bright_green': '\033[92m',
'bright_yellow': '\033[93m',
'bright_blue': '\033[94m',
'bright_magenta': '\033[95m',
'bright_cyan': '\033[96m',
'bright_white': '\033[97m',
'end': '\033[0m', # misc
'bold': '\033[1m',
'underline': '\033[4m'}
colors = {
'black': '\033[30m', # basic colors
'red': '\033[31m',
'green': '\033[32m',
'yellow': '\033[33m',
'blue': '\033[34m',
'magenta': '\033[35m',
'cyan': '\033[36m',
'white': '\033[37m',
'bright_black': '\033[90m', # bright colors
'bright_red': '\033[91m',
'bright_green': '\033[92m',
'bright_yellow': '\033[93m',
'bright_blue': '\033[94m',
'bright_magenta': '\033[95m',
'bright_cyan': '\033[96m',
'bright_white': '\033[97m',
'end': '\033[0m', # misc
'bold': '\033[1m',
'underline': '\033[4m'}
return ''.join(colors[x] for x in args) + f'{string}' + colors['end'] return ''.join(colors[x] for x in args) + f'{string}' + colors['end']




# b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n') # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
# x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco
# x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet
x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
x = [
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
return x return x




boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0]) # y1, y2 boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0]) # y1, y2




def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
labels=(), max_det=300):
def non_max_suppression(prediction,
conf_thres=0.25,
iou_thres=0.45,
classes=None,
agnostic=False,
multi_label=False,
labels=(),
max_det=300):
"""Non-Maximum Suppression (NMS) on inference results to reject overlapping bounding boxes """Non-Maximum Suppression (NMS) on inference results to reject overlapping bounding boxes


Returns: Returns:
def print_mutation(results, hyp, save_dir, bucket, prefix=colorstr('evolve: ')): def print_mutation(results, hyp, save_dir, bucket, prefix=colorstr('evolve: ')):
evolve_csv = save_dir / 'evolve.csv' evolve_csv = save_dir / 'evolve.csv'
evolve_yaml = save_dir / 'hyp_evolve.yaml' evolve_yaml = save_dir / 'hyp_evolve.yaml'
keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
'val/box_loss', 'val/obj_loss', 'val/cls_loss') + tuple(hyp.keys()) # [results + hyps]
keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 'val/box_loss',
'val/obj_loss', 'val/cls_loss') + tuple(hyp.keys()) # [results + hyps]
keys = tuple(x.strip() for x in keys) keys = tuple(x.strip() for x in keys)
vals = results + tuple(hyp.values()) vals = results + tuple(hyp.values())
n = len(keys) n = len(keys)
data = data.rename(columns=lambda x: x.strip()) # strip keys data = data.rename(columns=lambda x: x.strip()) # strip keys
i = np.argmax(fitness(data.values[:, :4])) # i = np.argmax(fitness(data.values[:, :4])) #
generations = len(data) generations = len(data)
f.write('# YOLOv5 Hyperparameter Evolution Results\n' +
f'# Best generation: {i}\n' +
f'# Last generation: {generations - 1}\n' +
'# ' + ', '.join(f'{x.strip():>20s}' for x in keys[:7]) + '\n' +
'# ' + ', '.join(f'{x:>20.5g}' for x in data.values[i, :7]) + '\n\n')
f.write('# YOLOv5 Hyperparameter Evolution Results\n' + f'# Best generation: {i}\n' +
f'# Last generation: {generations - 1}\n' + '# ' + ', '.join(f'{x.strip():>20s}' for x in keys[:7]) +
'\n' + '# ' + ', '.join(f'{x:>20.5g}' for x in data.values[i, :7]) + '\n\n')
yaml.safe_dump(data.loc[i][7:].to_dict(), f, sort_keys=False) yaml.safe_dump(data.loc[i][7:].to_dict(), f, sort_keys=False)


# Print to screen # Print to screen
LOGGER.info(prefix + f'{generations} generations finished, current result:\n' +
prefix + ', '.join(f'{x.strip():>20s}' for x in keys) + '\n' +
prefix + ', '.join(f'{x:20.5g}' for x in vals) + '\n\n')
LOGGER.info(prefix + f'{generations} generations finished, current result:\n' + prefix +
', '.join(f'{x.strip():>20s}' for x in keys) + '\n' + prefix + ', '.join(f'{x:20.5g}'
for x in vals) + '\n\n')


if bucket: if bucket:
os.system(f'gsutil cp {evolve_csv} {evolve_yaml} gs://{bucket}') # upload os.system(f'gsutil cp {evolve_csv} {evolve_yaml} gs://{bucket}') # upload

+ 16
- 5
utils/loggers/__init__.py View File

self.hyp = hyp self.hyp = hyp
self.logger = logger # for printing results to console self.logger = logger # for printing results to console
self.include = include self.include = include
self.keys = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss
'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', # metrics
'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss
'x/lr0', 'x/lr1', 'x/lr2'] # params
self.keys = [
'train/box_loss',
'train/obj_loss',
'train/cls_loss', # train loss
'metrics/precision',
'metrics/recall',
'metrics/mAP_0.5',
'metrics/mAP_0.5:0.95', # metrics
'val/box_loss',
'val/obj_loss',
'val/cls_loss', # val loss
'x/lr0',
'x/lr1',
'x/lr2'] # params
self.best_keys = ['best/epoch', 'best/precision', 'best/recall', 'best/mAP_0.5', 'best/mAP_0.5:0.95'] self.best_keys = ['best/epoch', 'best/precision', 'best/recall', 'best/mAP_0.5', 'best/mAP_0.5:0.95']
for k in LOGGERS: for k in LOGGERS:
setattr(self, k, None) # init empty logger dictionary setattr(self, k, None) # init empty logger dictionary
self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]}) self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]})
# Calling wandb.log. TODO: Refactor this into WandbLogger.log_model # Calling wandb.log. TODO: Refactor this into WandbLogger.log_model
if not self.opt.evolve: if not self.opt.evolve:
wandb.log_artifact(str(best if best.exists() else last), type='model',
wandb.log_artifact(str(best if best.exists() else last),
type='model',
name='run_' + self.wandb.wandb_run.id + '_model', name='run_' + self.wandb.wandb_run.id + '_model',
aliases=['latest', 'best', 'stripped']) aliases=['latest', 'best', 'stripped'])
self.wandb.finish_run() self.wandb.finish_run()

+ 63
- 49
utils/loggers/wandb/wandb_utils.py View File

if check_file(data_file) and data_file.endswith('.yaml'): if check_file(data_file) and data_file.endswith('.yaml'):
with open(data_file, errors='ignore') as f: with open(data_file, errors='ignore') as f:
data_dict = yaml.safe_load(f) data_dict = yaml.safe_load(f)
is_trainset_wandb_artifact = (isinstance(data_dict['train'], str) and
data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX))
is_valset_wandb_artifact = (isinstance(data_dict['val'], str) and
data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX))
is_trainset_wandb_artifact = isinstance(data_dict['train'],
str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX)
is_valset_wandb_artifact = isinstance(data_dict['val'],
str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX)
if is_trainset_wandb_artifact or is_valset_wandb_artifact: if is_trainset_wandb_artifact or is_valset_wandb_artifact:
return data_dict return data_dict
else: else:
For more on how this logger is used, see the Weights & Biases documentation: For more on how this logger is used, see the Weights & Biases documentation:
https://docs.wandb.com/guides/integrations/yolov5 https://docs.wandb.com/guides/integrations/yolov5
""" """

def __init__(self, opt, run_id=None, job_type='Training'): def __init__(self, opt, run_id=None, job_type='Training'):
""" """
- Initialize WandbLogger instance - Initialize WandbLogger instance
self.wandb_artifact_data_dict = self.wandb_artifact_data_dict or self.data_dict self.wandb_artifact_data_dict = self.wandb_artifact_data_dict or self.data_dict


# write data_dict to config. useful for resuming from artifacts. Do this only when not resuming. # write data_dict to config. useful for resuming from artifacts. Do this only when not resuming.
self.wandb_run.config.update({'data_dict': self.wandb_artifact_data_dict},
allow_val_change=True)
self.wandb_run.config.update({'data_dict': self.wandb_artifact_data_dict}, allow_val_change=True)
self.setup_training(opt) self.setup_training(opt)


if self.job_type == 'Dataset Creation': if self.job_type == 'Dataset Creation':
Updated dataset info dictionary where local dataset paths are replaced by WAND_ARFACT_PREFIX links. Updated dataset info dictionary where local dataset paths are replaced by WAND_ARFACT_PREFIX links.
""" """
assert wandb, 'Install wandb to upload dataset' assert wandb, 'Install wandb to upload dataset'
config_path = self.log_dataset_artifact(opt.data,
opt.single_cls,
config_path = self.log_dataset_artifact(opt.data, opt.single_cls,
'YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem) 'YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem)
with open(config_path, errors='ignore') as f: with open(config_path, errors='ignore') as f:
wandb_data_dict = yaml.safe_load(f) wandb_data_dict = yaml.safe_load(f)
config.hyp, config.imgsz config.hyp, config.imgsz
data_dict = self.data_dict data_dict = self.data_dict
if self.val_artifact is None: # If --upload_dataset is set, use the existing artifact, don't download if self.val_artifact is None: # If --upload_dataset is set, use the existing artifact, don't download
self.train_artifact_path, self.train_artifact = self.download_dataset_artifact(data_dict.get('train'),
opt.artifact_alias)
self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(data_dict.get('val'),
opt.artifact_alias)
self.train_artifact_path, self.train_artifact = self.download_dataset_artifact(
data_dict.get('train'), opt.artifact_alias)
self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(
data_dict.get('val'), opt.artifact_alias)


if self.train_artifact_path is not None: if self.train_artifact_path is not None:
train_path = Path(self.train_artifact_path) / 'data/images/' train_path = Path(self.train_artifact_path) / 'data/images/'
fitness_score (float) -- fitness score for current epoch fitness_score (float) -- fitness score for current epoch
best_model (boolean) -- Boolean representing if the current checkpoint is the best yet. best_model (boolean) -- Boolean representing if the current checkpoint is the best yet.
""" """
model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={
'original_url': str(path),
'epochs_trained': epoch + 1,
'save period': opt.save_period,
'project': opt.project,
'total_epochs': opt.epochs,
'fitness_score': fitness_score
})
model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model',
type='model',
metadata={
'original_url': str(path),
'epochs_trained': epoch + 1,
'save period': opt.save_period,
'project': opt.project,
'total_epochs': opt.epochs,
'fitness_score': fitness_score})
model_artifact.add_file(str(path / 'last.pt'), name='last.pt') model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
wandb.log_artifact(model_artifact, wandb.log_artifact(model_artifact,
aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), 'best' if best_model else '']) aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), 'best' if best_model else ''])


# log train set # log train set
if not log_val_only: if not log_val_only:
self.train_artifact = self.create_dataset_table(LoadImagesAndLabels(
data['train'], rect=True, batch_size=1), names, name='train') if data.get('train') else None
self.train_artifact = self.create_dataset_table(LoadImagesAndLabels(data['train'], rect=True, batch_size=1),
names,
name='train') if data.get('train') else None
if data.get('train'): if data.get('train'):
data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'train') data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'train')


self.val_artifact = self.create_dataset_table(LoadImagesAndLabels(
data['val'], rect=True, batch_size=1), names, name='val') if data.get('val') else None
self.val_artifact = self.create_dataset_table(
LoadImagesAndLabels(data['val'], rect=True, batch_size=1), names, name='val') if data.get('val') else None
if data.get('val'): if data.get('val'):
data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'val') data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'val')


else: else:
artifact.add_file(img_file, name='data/images/' + Path(img_file).name) artifact.add_file(img_file, name='data/images/' + Path(img_file).name)
label_file = Path(img2label_paths([img_file])[0]) label_file = Path(img2label_paths([img_file])[0])
artifact.add_file(str(label_file),
name='data/labels/' + label_file.name) if label_file.exists() else None
artifact.add_file(str(label_file), name='data/labels/' +
label_file.name) if label_file.exists() else None
table = wandb.Table(columns=["id", "train_image", "Classes", "name"]) table = wandb.Table(columns=["id", "train_image", "Classes", "name"])
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()]) class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)): for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)):
box_data, img_classes = [], {} box_data, img_classes = [], {}
for cls, *xywh in labels[:, 1:].tolist(): for cls, *xywh in labels[:, 1:].tolist():
cls = int(cls) cls = int(cls)
box_data.append({"position": {"middle": [xywh[0], xywh[1]], "width": xywh[2], "height": xywh[3]},
"class_id": cls,
"box_caption": "%s" % (class_to_id[cls])})
box_data.append({
"position": {
"middle": [xywh[0], xywh[1]],
"width": xywh[2],
"height": xywh[3]},
"class_id": cls,
"box_caption": "%s" % (class_to_id[cls])})
img_classes[cls] = class_to_id[cls] img_classes[cls] = class_to_id[cls]
boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space
table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), list(img_classes.values()), table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), list(img_classes.values()),
for *xyxy, conf, cls in predn.tolist(): for *xyxy, conf, cls in predn.tolist():
if conf >= 0.25: if conf >= 0.25:
cls = int(cls) cls = int(cls)
box_data.append(
{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
"class_id": cls,
"box_caption": f"{names[cls]} {conf:.3f}",
"scores": {"class_score": conf},
"domain": "pixel"})
box_data.append({
"position": {
"minX": xyxy[0],
"minY": xyxy[1],
"maxX": xyxy[2],
"maxY": xyxy[3]},
"class_id": cls,
"box_caption": f"{names[cls]} {conf:.3f}",
"scores": {
"class_score": conf},
"domain": "pixel"})
avg_conf_per_class[cls] += conf avg_conf_per_class[cls] += conf


if cls in pred_class_count: if cls in pred_class_count:


boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
id = self.val_table_path_map[Path(path).name] id = self.val_table_path_map[Path(path).name]
self.result_table.add_data(self.current_epoch,
id,
self.val_table.data[id][1],
self.result_table.add_data(self.current_epoch, id, self.val_table.data[id][1],
wandb.Image(self.val_table.data[id][1], boxes=boxes, classes=class_set), wandb.Image(self.val_table.data[id][1], boxes=boxes, classes=class_set),
*avg_conf_per_class
)
*avg_conf_per_class)


def val_one_image(self, pred, predn, path, names, im): def val_one_image(self, pred, predn, path, names, im):
""" """


if len(self.bbox_media_panel_images) < self.max_imgs_to_log and self.current_epoch > 0: if len(self.bbox_media_panel_images) < self.max_imgs_to_log and self.current_epoch > 0:
if self.current_epoch % self.bbox_interval == 0: if self.current_epoch % self.bbox_interval == 0:
box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
"class_id": int(cls),
"box_caption": f"{names[int(cls)]} {conf:.3f}",
"scores": {"class_score": conf},
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
box_data = [{
"position": {
"minX": xyxy[0],
"minY": xyxy[1],
"maxX": xyxy[2],
"maxY": xyxy[3]},
"class_id": int(cls),
"box_caption": f"{names[int(cls)]} {conf:.3f}",
"scores": {
"class_score": conf},
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
self.bbox_media_panel_images.append(wandb.Image(im, boxes=boxes, caption=path.name)) self.bbox_media_panel_images.append(wandb.Image(im, boxes=boxes, caption=path.name))


wandb.log(self.log_dict) wandb.log(self.log_dict)
except BaseException as e: except BaseException as e:
LOGGER.info( LOGGER.info(
f"An error occurred in wandb logger. The training will proceed without interruption. More info\n{e}")
f"An error occurred in wandb logger. The training will proceed without interruption. More info\n{e}"
)
self.wandb_run.finish() self.wandb_run.finish()
self.wandb_run = None self.wandb_run = None


self.bbox_media_panel_images = [] self.bbox_media_panel_images = []
if self.result_artifact: if self.result_artifact:
self.result_artifact.add(self.result_table, 'result') self.result_artifact.add(self.result_table, 'result')
wandb.log_artifact(self.result_artifact, aliases=['latest', 'last', 'epoch ' + str(self.current_epoch),
('best' if best_result else '')])
wandb.log_artifact(self.result_artifact,
aliases=[
'latest', 'last', 'epoch ' + str(self.current_epoch),
('best' if best_result else '')])


wandb.log({"evaluation": self.result_table}) wandb.log({"evaluation": self.result_table})
columns = ["epoch", "id", "ground truth", "prediction"] columns = ["epoch", "id", "ground truth", "prediction"]

+ 10
- 4
utils/loss.py View File

targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices


g = 0.5 # bias g = 0.5 # bias
off = torch.tensor([[0, 0],
[1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
], device=self.device).float() * g # offsets
off = torch.tensor(
[
[0, 0],
[1, 0],
[0, 1],
[-1, 0],
[0, -1], # j,k,l,m
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
],
device=self.device).float() * g # offsets


for i in range(self.nl): for i in range(self.nl):
anchors = self.anchors[i] anchors = self.anchors[i]

+ 9
- 2
utils/metrics.py View File

labels = (0 < nn < 99) and (nn == nc) # apply names to ticklabels labels = (0 < nn < 99) and (nn == nc) # apply names to ticklabels
with warnings.catch_warnings(): with warnings.catch_warnings():
warnings.simplefilter('ignore') # suppress empty matrix RuntimeWarning: All-NaN slice encountered warnings.simplefilter('ignore') # suppress empty matrix RuntimeWarning: All-NaN slice encountered
sn.heatmap(array, annot=nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True, vmin=0.0,
sn.heatmap(array,
annot=nc < 30,
annot_kws={
"size": 8},
cmap='Blues',
fmt='.2f',
square=True,
vmin=0.0,
xticklabels=names + ['background FP'] if labels else "auto", xticklabels=names + ['background FP'] if labels else "auto",
yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1)) yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1))
fig.axes[0].set_xlabel('True') fig.axes[0].set_xlabel('True')
iou (Tensor[N, M]): the NxM matrix containing the pairwise iou (Tensor[N, M]): the NxM matrix containing the pairwise
IoU values for every element in boxes1 and boxes2 IoU values for every element in boxes1 and boxes2
""" """

def box_area(box): def box_area(box):
# box = 4xn # box = 4xn
return (box[2] - box[0]) * (box[3] - box[1]) return (box[2] - box[0]) * (box[3] - box[1])


# Plots ---------------------------------------------------------------------------------------------------------------- # Plots ----------------------------------------------------------------------------------------------------------------



def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()): def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()):
# Precision-recall curve # Precision-recall curve
fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True) fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)

+ 22
- 8
utils/plots.py View File

if label: if label:
w, h = self.font.getsize(label) # text width, height w, h = self.font.getsize(label) # text width, height
outside = box[1] - h >= 0 # label fits outside box outside = box[1] - h >= 0 # label fits outside box
self.draw.rectangle((box[0],
box[1] - h if outside else box[1],
box[0] + w + 1,
box[1] + 1 if outside else box[1] + h + 1), fill=color)
self.draw.rectangle(
(box[0], box[1] - h if outside else box[1], box[0] + w + 1,
box[1] + 1 if outside else box[1] + h + 1),
fill=color,
)
# self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls') # for PIL>8.0 # self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls') # for PIL>8.0
self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font) self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font)
else: # cv2 else: # cv2
outside = p1[1] - h - 3 >= 0 # label fits outside box outside = p1[1] - h - 3 >= 0 # label fits outside box
p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3 p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled
cv2.putText(self.im, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, self.lw / 3, txt_color,
thickness=tf, lineType=cv2.LINE_AA)
cv2.putText(self.im,
label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
0,
self.lw / 3,
txt_color,
thickness=tf,
lineType=cv2.LINE_AA)


def rectangle(self, xy, fill=None, outline=None, width=1): def rectangle(self, xy, fill=None, outline=None, width=1):
# Add rectangle to image (PIL-only) # Add rectangle to image (PIL-only)
ax[i].set_title(s[i]) ax[i].set_title(s[i])


j = y[3].argmax() + 1 j = y[3].argmax() + 1
ax2.plot(y[5, 1:j], y[3, 1:j] * 1E2, '.-', linewidth=2, markersize=8,
ax2.plot(y[5, 1:j],
y[3, 1:j] * 1E2,
'.-',
linewidth=2,
markersize=8,
label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO')) label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO'))


ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5], ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
'k.-', linewidth=2, markersize=8, alpha=.25, label='EfficientDet')
'k.-',
linewidth=2,
markersize=8,
alpha=.25,
label='EfficientDet')


ax2.grid(alpha=0.2) ax2.grid(alpha=0.2)
ax2.set_yticks(np.arange(20, 60, 5)) ax2.set_yticks(np.arange(20, 60, 5))

+ 0
- 1
utils/torch_utils.py View File

Keeps a moving average of everything in the model state_dict (parameters and buffers) Keeps a moving average of everything in the model state_dict (parameters and buffers)
For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
""" """

def __init__(self, model, decay=0.9999, tau=2000, updates=0): def __init__(self, model, decay=0.9999, tau=2000, updates=0):
# Create EMA # Create EMA
self.ema = deepcopy(de_parallel(model)).eval() # FP32 EMA self.ema = deepcopy(de_parallel(model)).eval() # FP32 EMA

+ 17
- 8
val.py View File

box = xyxy2xywh(predn[:, :4]) # xywh box = xyxy2xywh(predn[:, :4]) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for p, b in zip(predn.tolist(), box.tolist()): for p, b in zip(predn.tolist(), box.tolist()):
jdict.append({'image_id': image_id,
'category_id': class_map[int(p[5])],
'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)})
jdict.append({
'image_id': image_id,
'category_id': class_map[int(p[5])],
'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)})




def process_batch(detections, labels, iouv): def process_batch(detections, labels, iouv):




@torch.no_grad() @torch.no_grad()
def run(data,
def run(
data,
weights=None, # model.pt path(s) weights=None, # model.pt path(s)
batch_size=32, # batch size batch_size=32, # batch size
imgsz=640, # inference size (pixels) imgsz=640, # inference size (pixels)
plots=True, plots=True,
callbacks=Callbacks(), callbacks=Callbacks(),
compute_loss=None, compute_loss=None,
):
):
# Initialize/load model and set device # Initialize/load model and set device
training = model is not None training = model is not None
if training: # called by train.py if training: # called by train.py
pad = 0.0 if task in ('speed', 'benchmark') else 0.5 pad = 0.0 if task in ('speed', 'benchmark') else 0.5
rect = False if task == 'benchmark' else pt # square inference for benchmarks rect = False if task == 'benchmark' else pt # square inference for benchmarks
task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images
dataloader = create_dataloader(data[task], imgsz, batch_size, stride, single_cls, pad=pad, rect=rect,
workers=workers, prefix=colorstr(f'{task}: '))[0]
dataloader = create_dataloader(data[task],
imgsz,
batch_size,
stride,
single_cls,
pad=pad,
rect=rect,
workers=workers,
prefix=colorstr(f'{task}: '))[0]


seen = 0 seen = 0
confusion_matrix = ConfusionMatrix(nc=nc) confusion_matrix = ConfusionMatrix(nc=nc)

Write Preview
Loading…
Cancel
Save