zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

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>
This commit is contained in:
Jirka Borovec 2022-03-31 23:52:34 +09:00 committed by GitHub
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
.pre-commit-config.yaml
View File

@ -36,12 +36,11 @@ repos:
- id: isort - id: isort
name: Sort imports name: Sort imports
# TODO - repo: https://github.com/pre-commit/mirrors-yapf
#- repo: https://github.com/pre-commit/mirrors-yapf rev: v0.31.0
# rev: v0.31.0 hooks:
# hooks: - id: yapf
# - id: yapf name: formatting
# name: formatting
# TODO # TODO
#- repo: https://github.com/executablebooks/mdformat #- repo: https://github.com/executablebooks/mdformat

3
detect.py
View File

@ -47,7 +47,8 @@ from utils.torch_utils import select_device, time_sync
@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)

90
export.py
View File

@ -76,16 +76,11 @@ from utils.torch_utils import select_device
def export_formats(): def export_formats():
# YOLOv5 export formats # YOLOv5 export formats
x = [['PyTorch', '-', '.pt', True], x = [['PyTorch', '-', '.pt', True], ['TorchScript', 'torchscript', '.torchscript', True],
['TorchScript', 'torchscript', '.torchscript', True], ['ONNX', 'onnx', '.onnx', True], ['OpenVINO', 'openvino', '_openvino_model', False],
['ONNX', 'onnx', '.onnx', True], ['TensorRT', 'engine', '.engine', True], ['CoreML', 'coreml', '.mlmodel', False],
['OpenVINO', 'openvino', '_openvino_model', False], ['TensorFlow SavedModel', 'saved_model', '_saved_model', True], ['TensorFlow GraphDef', 'pb', '.pb', True],
['TensorRT', 'engine', '.engine', True], ['TensorFlow Lite', 'tflite', '.tflite', False], ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False],
['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'])
@ -119,13 +114,24 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
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, torch.onnx.export(
model,
im,
f,
verbose=False,
opset_version=opset,
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL, training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
do_constant_folding=not train, do_constant_folding=not train,
input_names=['images'], input_names=['images'],
output_names=['output'], output_names=['output'],
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640) dynamic_axes={
'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85) '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) } if dynamic else None)
# Checks # Checks
@ -140,8 +146,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
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, check = onnxsim.simplify(model_onnx,
model_onnx,
dynamic_input_shape=dynamic, dynamic_input_shape=dynamic,
input_shapes={'images': list(im.shape)} if dynamic else None) input_shapes={'images': list(im.shape)} if dynamic else None)
assert check, 'assert check failed' assert check, 'assert check failed'
@ -246,9 +251,18 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
LOGGER.info(f'\n{prefix} export failure: {e}') LOGGER.info(f'\n{prefix} export failure: {e}')
def export_saved_model(model, im, file, dynamic, def export_saved_model(model,
tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, im,
conf_thres=0.25, keras=False, prefix=colorstr('TensorFlow SavedModel:')): 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
@ -278,11 +292,10 @@ def export_saved_model(model, im, file, dynamic,
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( tf.saved_model.save(tfm,
tfm,
f, f,
options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) if options=tf.saved_model.SaveOptions(experimental_custom_gradients=False)
check_version(tf.__version__, '2.6') else tf.saved_model.SaveOptions()) 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:
@ -352,10 +365,10 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')):
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 -', 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', '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 update', 'sudo apt-get install edgetpu-compiler'):
'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]
@ -395,12 +408,10 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
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'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"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"}}}', r'"Identity_3": {"name": "Identity_3"}}}', json)
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)')
@ -410,7 +421,8 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
@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
@ -431,7 +443,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
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
@ -495,9 +507,16 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
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, model, f[5] = export_saved_model(model.cpu(),
agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class, im,
topk_all=topk_all, conf_thres=conf_thres, iou_thres=iou_thres) # keras model 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:
@ -542,7 +561,8 @@ def parse_opt():
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()

3
hubconf.py
View File

@ -132,7 +132,8 @@ if __name__ == '__main__':
from utils.general import cv2 from utils.general import cv2
imgs = ['data/images/zidane.jpg', # filename imgs = [
'data/images/zidane.jpg', # filename
Path('data/images/zidane.jpg'), # Path Path('data/images/zidane.jpg'), # Path
'https://ultralytics.com/images/zidane.jpg', # URI 'https://ultralytics.com/images/zidane.jpg', # URI
cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV

27
models/common.py
View File

@ -227,11 +227,12 @@ class GhostBottleneck(nn.Module):
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 self.conv = nn.Sequential(
GhostConv(c1, c_, 1, 1), # pw
DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
GhostConv(c_, c2, 1, 1, act=False)) # pw-linear GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1,
Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity() 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)
@ -387,7 +388,8 @@ class DetectMultiBackend(nn.Module):
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', delegate = {
'Linux': 'libedgetpu.so.1',
'Darwin': 'libedgetpu.1.dylib', 'Darwin': 'libedgetpu.1.dylib',
'Windows': 'edgetpu.dll'}[platform.system()] '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)])
@ -531,7 +533,7 @@ class AutoShape(nn.Module):
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
@ -561,8 +563,13 @@ class AutoShape(nn.Module):
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, y = non_max_suppression(y if self.dmb else y[0],
self.multi_label, max_det=self.max_det) # NMS 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])
@ -603,7 +610,11 @@ class Detections:
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, crops.append({
'box': box,
'conf': conf,
'cls': cls,
'label': label,
'im': save_one_box(box, im, file=file, save=save)}) '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))

4
models/experimental.py
View File

@ -63,8 +63,8 @@ class MixConv2d(nn.Module):
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( 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_)]) 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()

65
models/tf.py
View File

@ -69,7 +69,11 @@ class TFConv(keras.layers.Layer):
# 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])
@ -98,10 +102,10 @@ class TFFocus(keras.layers.Layer):
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, :], return self.conv(
inputs[:, 1::2, ::2, :], tf.concat(
inputs[:, ::2, 1::2, :], [inputs[:, ::2, ::2, :], inputs[:, 1::2, ::2, :], inputs[:, ::2, 1::2, :], inputs[:, 1::2, 1::2, :]],
inputs[:, 1::2, 1::2, :]], 3)) 3))
class TFBottleneck(keras.layers.Layer): class TFBottleneck(keras.layers.Layer):
@ -123,9 +127,14 @@ class TFConv2d(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)
@ -206,8 +215,7 @@ class TFDetect(keras.layers.Layer):
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.anchor_grid = tf.reshape(self.anchors * tf.reshape(self.stride, [self.nl, 1, 1]), [self.nl, 1, -1, 1, 2])
[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
@ -339,7 +347,13 @@ class TFModel:
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
@ -361,8 +375,13 @@ class TFModel:
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( nms = tf.image.combined_non_max_suppression(boxes,
boxes, scores, topk_per_class, topk_all, iou_thres, conf_thres, clip_boxes=False) 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, ...]
@ -383,7 +402,8 @@ class AgnosticNMS(keras.layers.Layer):
# 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')
@ -392,20 +412,26 @@ class AgnosticNMS(keras.layers.Layer):
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( selected_inds = tf.image.non_max_suppression(boxes,
boxes, scores_inp, max_output_size=topk_all, iou_threshold=iou_thres, score_threshold=conf_thres) 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
@ -421,7 +447,8 @@ def representative_dataset_gen(dataset, ncalib=100):
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

4
models/yolo.py
View File

@ -260,8 +260,8 @@ def parse_model(d, ch): # model_dict, input_channels(3)
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, if m in (Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]: 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
setup.cfg
View File

@ -1,5 +1,6 @@
# 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
@ -42,4 +43,17 @@ ignore =
[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

53
train.py
View File

@ -62,11 +62,7 @@ RANK = int(os.getenv('RANK', -1))
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 def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary
opt,
device,
callbacks
):
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
@ -220,20 +216,38 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary
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, train_loader, dataset = create_dataloader(train_path,
hyp=hyp, augment=True, cache=None if opt.cache == 'val' else opt.cache, imgsz,
rect=opt.rect, rank=LOCAL_RANK, workers=workers, batch_size // WORLD_SIZE,
image_weights=opt.image_weights, quad=opt.quad, gs,
prefix=colorstr('train: '), shuffle=True) 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, val_loader = create_dataloader(val_path,
hyp=hyp, cache=None if noval else opt.cache, imgsz,
rect=True, rank=-1, workers=workers * 2, pad=0.5, 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:
@ -350,8 +364,8 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary
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) % ( pbar.set_description(('%10s' * 2 + '%10.4g' * 5) %
f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1])) (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
@ -387,7 +401,8 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary
# 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, ckpt = {
'epoch': epoch,
'best_fitness': best_fitness, 'best_fitness': best_fitness,
'model': deepcopy(de_parallel(model)).half(), 'model': deepcopy(de_parallel(model)).half(),
'ema': deepcopy(ema.ema).half(), 'ema': deepcopy(ema.ema).half(),
@ -428,7 +443,8 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary
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, results, _, _ = val.run(
data_dict,
batch_size=batch_size // WORLD_SIZE * 2, batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz, imgsz=imgsz,
model=attempt_load(f, device).half(), model=attempt_load(f, device).half(),
@ -546,7 +562,8 @@ def main(opt, callbacks=Callbacks()):
# 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) 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) 'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)
'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1 'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1
'weight_decay': (1, 0.0, 0.001), # optimizer weight decay 'weight_decay': (1, 0.0, 0.001), # optimizer weight decay

2
utils/activations.py
View File

@ -64,7 +64,6 @@ class AconC(nn.Module):
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))
@ -81,7 +80,6 @@ class MetaAconC(nn.Module):
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)

15
utils/augmentations.py
View File

@ -21,15 +21,15 @@ class Albumentations:
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)], A.ImageCompression(quality_lower=75, p=0.0)] # transforms
bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'])) 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
@ -121,7 +121,14 @@ def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleF
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
utils/benchmarks.py
View File

@ -45,7 +45,8 @@ from utils.general import LOGGER, print_args
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

7
utils/callbacks.py
View File

@ -8,13 +8,11 @@ class Callbacks:
"""" """"
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': [],
@ -22,19 +20,16 @@ class Callbacks:
'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):

82
utils/datasets.py
View File

@ -77,14 +77,14 @@ def exif_transpose(image):
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, method = {
2: Image.FLIP_LEFT_RIGHT,
3: Image.ROTATE_180, 3: Image.ROTATE_180,
4: Image.FLIP_TOP_BOTTOM, 4: Image.FLIP_TOP_BOTTOM,
5: Image.TRANSPOSE, 5: Image.TRANSPOSE,
6: Image.ROTATE_270, 6: Image.ROTATE_270,
7: Image.TRANSVERSE, 7: Image.TRANSVERSE,
8: Image.ROTATE_90, 8: Image.ROTATE_90,}.get(orientation)
}.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]
@ -92,13 +92,30 @@ def exif_transpose(image):
return image return image
def create_dataloader(path, imgsz, batch_size, stride, single_cls=False, hyp=None, augment=False, cache=False, pad=0.0, def create_dataloader(path,
rect=False, rank=-1, workers=8, image_weights=False, quad=False, prefix='', shuffle=False): 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, dataset = LoadImagesAndLabels(
path,
imgsz,
batch_size,
augment=augment, # augmentation augment=augment, # augmentation
hyp=hyp, # hyperparameters hyp=hyp, # hyperparameters
rect=rect, # rectangular batches rect=rect, # rectangular batches
@ -128,7 +145,6 @@ class InfiniteDataLoader(dataloader.DataLoader):
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))
@ -148,7 +164,6 @@ class _RepeatSampler:
Args: Args:
sampler (Sampler) sampler (Sampler)
""" """
def __init__(self, sampler): def __init__(self, sampler):
self.sampler = sampler self.sampler = sampler
@ -380,8 +395,19 @@ class LoadImagesAndLabels(Dataset):
# 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, def __init__(self,
cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''): 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
@ -510,7 +536,9 @@ class LoadImagesAndLabels(Dataset):
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
@ -576,7 +604,8 @@ class LoadImagesAndLabels(Dataset):
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'],
@ -633,8 +662,7 @@ class LoadImagesAndLabels(Dataset):
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, im = cv2.resize(im, (int(w0 * r), int(h0 * r)),
(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:
@ -692,7 +720,9 @@ class LoadImagesAndLabels(Dataset):
# 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'],
@ -766,7 +796,9 @@ class LoadImagesAndLabels(Dataset):
# 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'],
@ -795,8 +827,8 @@ class LoadImagesAndLabels(Dataset):
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)[ im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2.0, mode='bilinear',
0].type(img[i].type()) 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)
@ -946,7 +978,6 @@ def dataset_stats(path='coco128.yaml', autodownload=False, verbose=False, profil
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]
@ -996,11 +1027,16 @@ def dataset_stats(path='coco128.yaml', autodownload=False, verbose=False, profil
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()}, stats[split] = {
'image_stats': {'total': dataset.n, 'unlabelled': int(np.all(x == 0, 1).sum()), '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()}, 'per_class': (x > 0).sum(0).tolist()},
'labels': [{str(Path(k).name): round_labels(v.tolist())} for k, v in 'labels': [{
zip(dataset.im_files, dataset.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'

9
utils/downloads.py
View File

@ -63,15 +63,17 @@ def attempt_download(file, repo='ultralytics/yolov5'): # from utils.downloads i
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', assets = [
'yolov5n6.pt', 'yolov5s6.pt', 'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt'] '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, safe_download(
file,
url=f'https://github.com/{repo}/releases/download/{tag}/{name}', url=f'https://github.com/{repo}/releases/download/{tag}/{name}',
# url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional) # url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional)
min_bytes=1E5, min_bytes=1E5,
@ -122,6 +124,7 @@ def get_token(cookie="./cookie"):
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 ----------------------------------------------
# #
# #

34
utils/general.py
View File

@ -536,7 +536,8 @@ def one_cycle(y1=0.0, y2=1.0, steps=100):
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 colors = {
'black': '\033[30m', # basic colors
'red': '\033[31m', 'red': '\033[31m',
'green': '\033[32m', 'green': '\033[32m',
'yellow': '\033[33m', 'yellow': '\033[33m',
@ -591,7 +592,8 @@ def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper)
# 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, 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, 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] 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
@ -701,8 +703,14 @@ def clip_coords(boxes, shape):
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, def non_max_suppression(prediction,
labels=(), max_det=300): 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:
@ -816,8 +824,8 @@ def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_op
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', keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 'val/box_loss',
'val/box_loss', 'val/obj_loss', 'val/cls_loss') + tuple(hyp.keys()) # [results + hyps] '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)
@ -839,17 +847,15 @@ def print_mutation(results, hyp, save_dir, bucket, prefix=colorstr('evolve: ')):
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.write('# YOLOv5 Hyperparameter Evolution Results\n' + f'# Best generation: {i}\n' +
f'# Best generation: {i}\n' + f'# Last generation: {generations - 1}\n' + '# ' + ', '.join(f'{x.strip():>20s}' for x in keys[:7]) +
f'# Last generation: {generations - 1}\n' + '\n' + '# ' + ', '.join(f'{x:>20.5g}' for x in data.values[i, :7]) + '\n\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' + LOGGER.info(prefix + f'{generations} generations finished, current result:\n' + prefix +
prefix + ', '.join(f'{x.strip():>20s}' for x in keys) + '\n' + ', '.join(f'{x.strip():>20s}' for x in keys) + '\n' + prefix + ', '.join(f'{x:20.5g}'
prefix + ', '.join(f'{x:20.5g}' for x in vals) + '\n\n') 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

21
utils/loggers/__init__.py
View File

@ -43,10 +43,20 @@ class Loggers():
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 self.keys = [
'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', # metrics 'train/box_loss',
'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss 'train/obj_loss',
'x/lr0', 'x/lr1', 'x/lr2'] # params '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
@ -155,7 +165,8 @@ class Loggers():
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()

84
utils/loggers/wandb/wandb_utils.py
View File

@ -46,10 +46,10 @@ def check_wandb_dataset(data_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 is_trainset_wandb_artifact = isinstance(data_dict['train'],
data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX)) str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX)
is_valset_wandb_artifact = (isinstance(data_dict['val'], str) and is_valset_wandb_artifact = isinstance(data_dict['val'],
data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX)) 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:
@ -116,7 +116,6 @@ class WandbLogger():
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
@ -181,8 +180,7 @@ class WandbLogger():
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}, self.wandb_run.config.update({'data_dict': self.wandb_artifact_data_dict}, allow_val_change=True)
allow_val_change=True)
self.setup_training(opt) self.setup_training(opt)
if self.job_type == 'Dataset Creation': if self.job_type == 'Dataset Creation':
@ -200,8 +198,7 @@ class WandbLogger():
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, config_path = self.log_dataset_artifact(opt.data, opt.single_cls,
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)
@ -230,10 +227,10 @@ class WandbLogger():
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'), self.train_artifact_path, self.train_artifact = self.download_dataset_artifact(
opt.artifact_alias) data_dict.get('train'), opt.artifact_alias)
self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(data_dict.get('val'), self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(
opt.artifact_alias) 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/'
@ -308,14 +305,15 @@ class WandbLogger():
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={ model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model',
type='model',
metadata={
'original_url': str(path), 'original_url': str(path),
'epochs_trained': epoch + 1, 'epochs_trained': epoch + 1,
'save period': opt.save_period, 'save period': opt.save_period,
'project': opt.project, 'project': opt.project,
'total_epochs': opt.epochs, 'total_epochs': opt.epochs,
'fitness_score': fitness_score '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 ''])
@ -344,13 +342,14 @@ class WandbLogger():
# log train set # log train set
if not log_val_only: if not log_val_only:
self.train_artifact = self.create_dataset_table(LoadImagesAndLabels( self.train_artifact = self.create_dataset_table(LoadImagesAndLabels(data['train'], rect=True, batch_size=1),
data['train'], rect=True, batch_size=1), names, name='train') if data.get('train') else None 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( self.val_artifact = self.create_dataset_table(
data['val'], rect=True, batch_size=1), names, name='val') if data.get('val') else None 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')
@ -412,15 +411,19 @@ class WandbLogger():
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), artifact.add_file(str(label_file), name='data/labels/' +
name='data/labels/' + label_file.name) if label_file.exists() else None 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]}, box_data.append({
"position": {
"middle": [xywh[0], xywh[1]],
"width": xywh[2],
"height": xywh[3]},
"class_id": cls, "class_id": cls,
"box_caption": "%s" % (class_to_id[cls])}) "box_caption": "%s" % (class_to_id[cls])})
img_classes[cls] = class_to_id[cls] img_classes[cls] = class_to_id[cls]
@ -446,11 +449,16 @@ class WandbLogger():
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( box_data.append({
{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]}, "position": {
"minX": xyxy[0],
"minY": xyxy[1],
"maxX": xyxy[2],
"maxY": xyxy[3]},
"class_id": cls, "class_id": cls,
"box_caption": f"{names[cls]} {conf:.3f}", "box_caption": f"{names[cls]} {conf:.3f}",
"scores": {"class_score": conf}, "scores": {
"class_score": conf},
"domain": "pixel"}) "domain": "pixel"})
avg_conf_per_class[cls] += conf avg_conf_per_class[cls] += conf
@ -464,12 +472,9 @@ class WandbLogger():
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, self.result_table.add_data(self.current_epoch, id, self.val_table.data[id][1],
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):
""" """
@ -485,10 +490,16 @@ class WandbLogger():
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]}, box_data = [{
"position": {
"minX": xyxy[0],
"minY": xyxy[1],
"maxX": xyxy[2],
"maxY": xyxy[3]},
"class_id": int(cls), "class_id": int(cls),
"box_caption": f"{names[int(cls)]} {conf:.3f}", "box_caption": f"{names[int(cls)]} {conf:.3f}",
"scores": {"class_score": conf}, "scores": {
"class_score": conf},
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()] "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))
@ -519,7 +530,8 @@ class WandbLogger():
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
@ -527,7 +539,9 @@ class WandbLogger():
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), wandb.log_artifact(self.result_artifact,
aliases=[
'latest', 'last', 'epoch ' + str(self.current_epoch),
('best' if best_result else '')]) ('best' if best_result else '')])
wandb.log({"evaluation": self.result_table}) wandb.log({"evaluation": self.result_table})

12
utils/loss.py
View File

@ -183,10 +183,16 @@ class ComputeLoss:
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], off = torch.tensor(
[1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m [
[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 # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
], device=self.device).float() * g # offsets ],
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]

11
utils/metrics.py
View File

@ -184,7 +184,14 @@ class ConfusionMatrix:
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')
@ -253,7 +260,6 @@ def box_iou(box1, box2):
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])
@ -300,6 +306,7 @@ def wh_iou(wh1, wh2):
# 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)

30
utils/plots.py
View File

@ -89,10 +89,11 @@ class Annotator:
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], self.draw.rectangle(
box[1] - h if outside else box[1], (box[0], box[1] - h if outside else box[1], box[0] + w + 1,
box[0] + w + 1, box[1] + 1 if outside else box[1] + h + 1),
box[1] + 1 if outside else box[1] + h + 1), fill=color) 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
@ -104,8 +105,13 @@ class Annotator:
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, cv2.putText(self.im,
thickness=tf, lineType=cv2.LINE_AA) 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)
@ -307,11 +313,19 @@ def plot_val_study(file='', dir='', x=None): # from utils.plots import *; plot_
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))

1
utils/torch_utils.py
View File

@ -284,7 +284,6 @@ class ModelEMA:
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
val.py
View File

@ -62,7 +62,8 @@ def save_one_json(predn, jdict, path, class_map):
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, jdict.append({
'image_id': image_id,
'category_id': class_map[int(p[5])], 'category_id': class_map[int(p[5])],
'bbox': [round(x, 3) for x in b], 'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)}) 'score': round(p[4], 5)})
@ -93,7 +94,8 @@ 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)
@ -164,8 +166,15 @@ def run(data,
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, dataloader = create_dataloader(data[task],
workers=workers, prefix=colorstr(f'{task}: '))[0] 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)