|
|
|
""" |
|
TensorFlow, Keras and TFLite versions of YOLOv5 |
|
Authored by https://github.com/zldrobit in PR https://github.com/ultralytics/yolov5/pull/1127 |
|
|
|
Usage: |
|
$ python models/tf.py --weights yolov5s.pt |
|
|
|
Export: |
|
$ python path/to/export.py --weights yolov5s.pt --include saved_model pb tflite tfjs |
|
""" |
|
|
|
import argparse |
|
import logging |
|
import sys |
|
from copy import deepcopy |
|
from pathlib import Path |
|
|
|
FILE = Path(__file__).resolve() |
|
ROOT = FILE.parents[1] |
|
if str(ROOT) not in sys.path: |
|
sys.path.append(str(ROOT)) |
|
|
|
|
|
import numpy as np |
|
import tensorflow as tf |
|
import torch |
|
import torch.nn as nn |
|
from tensorflow import keras |
|
|
|
from models.common import C3, SPP, SPPF, Bottleneck, BottleneckCSP, Concat, Conv, DWConv, Focus, autopad |
|
from models.experimental import CrossConv, MixConv2d, attempt_load |
|
from models.yolo import Detect |
|
from utils.activations import SiLU |
|
from utils.general import LOGGER, make_divisible, print_args |
|
|
|
|
|
class TFBN(keras.layers.Layer): |
|
|
|
def __init__(self, w=None): |
|
super().__init__() |
|
self.bn = keras.layers.BatchNormalization( |
|
beta_initializer=keras.initializers.Constant(w.bias.numpy()), |
|
gamma_initializer=keras.initializers.Constant(w.weight.numpy()), |
|
moving_mean_initializer=keras.initializers.Constant(w.running_mean.numpy()), |
|
moving_variance_initializer=keras.initializers.Constant(w.running_var.numpy()), |
|
epsilon=w.eps) |
|
|
|
def call(self, inputs): |
|
return self.bn(inputs) |
|
|
|
|
|
class TFPad(keras.layers.Layer): |
|
def __init__(self, pad): |
|
super().__init__() |
|
self.pad = tf.constant([[0, 0], [pad, pad], [pad, pad], [0, 0]]) |
|
|
|
def call(self, inputs): |
|
return tf.pad(inputs, self.pad, mode='constant', constant_values=0) |
|
|
|
|
|
class TFConv(keras.layers.Layer): |
|
|
|
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None): |
|
|
|
super().__init__() |
|
assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument" |
|
assert isinstance(k, int), "Convolution with multiple kernels are not allowed." |
|
|
|
|
|
|
|
conv = keras.layers.Conv2D( |
|
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()), |
|
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.bn = TFBN(w.bn) if hasattr(w, 'bn') else tf.identity |
|
|
|
|
|
if isinstance(w.act, nn.LeakyReLU): |
|
self.act = (lambda x: keras.activations.relu(x, alpha=0.1)) if act else tf.identity |
|
elif isinstance(w.act, nn.Hardswish): |
|
self.act = (lambda x: x * tf.nn.relu6(x + 3) * 0.166666667) if act else tf.identity |
|
elif isinstance(w.act, (nn.SiLU, SiLU)): |
|
self.act = (lambda x: keras.activations.swish(x)) if act else tf.identity |
|
else: |
|
raise Exception(f'no matching TensorFlow activation found for {w.act}') |
|
|
|
def call(self, inputs): |
|
return self.act(self.bn(self.conv(inputs))) |
|
|
|
|
|
class TFFocus(keras.layers.Layer): |
|
|
|
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None): |
|
|
|
super().__init__() |
|
self.conv = TFConv(c1 * 4, c2, k, s, p, g, act, w.conv) |
|
|
|
def call(self, inputs): |
|
|
|
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): |
|
|
|
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5, w=None): |
|
super().__init__() |
|
c_ = int(c2 * e) |
|
self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1) |
|
self.cv2 = TFConv(c_, c2, 3, 1, g=g, w=w.cv2) |
|
self.add = shortcut and c1 == c2 |
|
|
|
def call(self, inputs): |
|
return inputs + self.cv2(self.cv1(inputs)) if self.add else self.cv2(self.cv1(inputs)) |
|
|
|
|
|
class TFConv2d(keras.layers.Layer): |
|
|
|
def __init__(self, c1, c2, k, s=1, g=1, bias=True, w=None): |
|
super().__init__() |
|
assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument" |
|
self.conv = keras.layers.Conv2D( |
|
c2, k, s, 'VALID', use_bias=bias, |
|
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, ) |
|
|
|
def call(self, inputs): |
|
return self.conv(inputs) |
|
|
|
|
|
class TFBottleneckCSP(keras.layers.Layer): |
|
|
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None): |
|
|
|
super().__init__() |
|
c_ = int(c2 * e) |
|
self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1) |
|
self.cv2 = TFConv2d(c1, c_, 1, 1, bias=False, w=w.cv2) |
|
self.cv3 = TFConv2d(c_, c_, 1, 1, bias=False, w=w.cv3) |
|
self.cv4 = TFConv(2 * c_, c2, 1, 1, w=w.cv4) |
|
self.bn = TFBN(w.bn) |
|
self.act = lambda x: keras.activations.relu(x, alpha=0.1) |
|
self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)]) |
|
|
|
def call(self, inputs): |
|
y1 = self.cv3(self.m(self.cv1(inputs))) |
|
y2 = self.cv2(inputs) |
|
return self.cv4(self.act(self.bn(tf.concat((y1, y2), axis=3)))) |
|
|
|
|
|
class TFC3(keras.layers.Layer): |
|
|
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None): |
|
|
|
super().__init__() |
|
c_ = int(c2 * e) |
|
self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1) |
|
self.cv2 = TFConv(c1, c_, 1, 1, w=w.cv2) |
|
self.cv3 = TFConv(2 * c_, c2, 1, 1, w=w.cv3) |
|
self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)]) |
|
|
|
def call(self, inputs): |
|
return self.cv3(tf.concat((self.m(self.cv1(inputs)), self.cv2(inputs)), axis=3)) |
|
|
|
|
|
class TFSPP(keras.layers.Layer): |
|
|
|
def __init__(self, c1, c2, k=(5, 9, 13), w=None): |
|
super().__init__() |
|
c_ = c1 // 2 |
|
self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1) |
|
self.cv2 = TFConv(c_ * (len(k) + 1), c2, 1, 1, w=w.cv2) |
|
self.m = [keras.layers.MaxPool2D(pool_size=x, strides=1, padding='SAME') for x in k] |
|
|
|
def call(self, inputs): |
|
x = self.cv1(inputs) |
|
return self.cv2(tf.concat([x] + [m(x) for m in self.m], 3)) |
|
|
|
|
|
class TFSPPF(keras.layers.Layer): |
|
|
|
def __init__(self, c1, c2, k=5, w=None): |
|
super().__init__() |
|
c_ = c1 // 2 |
|
self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1) |
|
self.cv2 = TFConv(c_ * 4, c2, 1, 1, w=w.cv2) |
|
self.m = keras.layers.MaxPool2D(pool_size=k, strides=1, padding='SAME') |
|
|
|
def call(self, inputs): |
|
x = self.cv1(inputs) |
|
y1 = self.m(x) |
|
y2 = self.m(y1) |
|
return self.cv2(tf.concat([x, y1, y2, self.m(y2)], 3)) |
|
|
|
|
|
class TFDetect(keras.layers.Layer): |
|
def __init__(self, nc=80, anchors=(), ch=(), imgsz=(640, 640), w=None): |
|
super().__init__() |
|
self.stride = tf.convert_to_tensor(w.stride.numpy(), dtype=tf.float32) |
|
self.nc = nc |
|
self.no = nc + 5 |
|
self.nl = len(anchors) |
|
self.na = len(anchors[0]) // 2 |
|
self.grid = [tf.zeros(1)] * self.nl |
|
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.m = [TFConv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in enumerate(ch)] |
|
self.training = False |
|
self.imgsz = imgsz |
|
for i in range(self.nl): |
|
ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1] // self.stride[i] |
|
self.grid[i] = self._make_grid(nx, ny) |
|
|
|
def call(self, inputs): |
|
z = [] |
|
x = [] |
|
for i in range(self.nl): |
|
x.append(self.m[i](inputs[i])) |
|
|
|
ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1] // self.stride[i] |
|
x[i] = tf.transpose(tf.reshape(x[i], [-1, ny * nx, self.na, self.no]), [0, 2, 1, 3]) |
|
|
|
if not self.training: |
|
y = tf.sigmoid(x[i]) |
|
xy = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i] |
|
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] |
|
|
|
xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32) |
|
wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32) |
|
y = tf.concat([xy, wh, y[..., 4:]], -1) |
|
z.append(tf.reshape(y, [-1, 3 * ny * nx, self.no])) |
|
|
|
return x if self.training else (tf.concat(z, 1), x) |
|
|
|
@staticmethod |
|
def _make_grid(nx=20, ny=20): |
|
|
|
|
|
xv, yv = tf.meshgrid(tf.range(nx), tf.range(ny)) |
|
return tf.cast(tf.reshape(tf.stack([xv, yv], 2), [1, 1, ny * nx, 2]), dtype=tf.float32) |
|
|
|
|
|
class TFUpsample(keras.layers.Layer): |
|
def __init__(self, size, scale_factor, mode, w=None): |
|
super().__init__() |
|
assert scale_factor == 2, "scale_factor must be 2" |
|
self.upsample = lambda x: tf.image.resize(x, (x.shape[1] * 2, x.shape[2] * 2), method=mode) |
|
|
|
|
|
|
|
|
|
|
|
def call(self, inputs): |
|
return self.upsample(inputs) |
|
|
|
|
|
class TFConcat(keras.layers.Layer): |
|
def __init__(self, dimension=1, w=None): |
|
super().__init__() |
|
assert dimension == 1, "convert only NCHW to NHWC concat" |
|
self.d = 3 |
|
|
|
def call(self, inputs): |
|
return tf.concat(inputs, self.d) |
|
|
|
|
|
def parse_model(d, ch, model, imgsz): |
|
LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}") |
|
anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] |
|
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors |
|
no = na * (nc + 5) |
|
|
|
layers, save, c2 = [], [], ch[-1] |
|
for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): |
|
m_str = m |
|
m = eval(m) if isinstance(m, str) else m |
|
for j, a in enumerate(args): |
|
try: |
|
args[j] = eval(a) if isinstance(a, str) else a |
|
except NameError: |
|
pass |
|
|
|
n = max(round(n * gd), 1) if n > 1 else n |
|
if m in [nn.Conv2d, Conv, Bottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]: |
|
c1, c2 = ch[f], args[0] |
|
c2 = make_divisible(c2 * gw, 8) if c2 != no else c2 |
|
|
|
args = [c1, c2, *args[1:]] |
|
if m in [BottleneckCSP, C3]: |
|
args.insert(2, n) |
|
n = 1 |
|
elif m is nn.BatchNorm2d: |
|
args = [ch[f]] |
|
elif m is Concat: |
|
c2 = sum(ch[-1 if x == -1 else x + 1] for x in f) |
|
elif m is Detect: |
|
args.append([ch[x + 1] for x in f]) |
|
if isinstance(args[1], int): |
|
args[1] = [list(range(args[1] * 2))] * len(f) |
|
args.append(imgsz) |
|
else: |
|
c2 = ch[f] |
|
|
|
tf_m = eval('TF' + m_str.replace('nn.', '')) |
|
m_ = keras.Sequential([tf_m(*args, w=model.model[i][j]) for j in range(n)]) if n > 1 \ |
|
else tf_m(*args, w=model.model[i]) |
|
|
|
torch_m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) |
|
t = str(m)[8:-2].replace('__main__.', '') |
|
np = sum(x.numel() for x in torch_m_.parameters()) |
|
m_.i, m_.f, m_.type, m_.np = i, f, t, np |
|
LOGGER.info(f'{i:>3}{str(f):>18}{str(n):>3}{np:>10} {t:<40}{str(args):<30}') |
|
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) |
|
layers.append(m_) |
|
ch.append(c2) |
|
return keras.Sequential(layers), sorted(save) |
|
|
|
|
|
class TFModel: |
|
def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, model=None, imgsz=(640, 640)): |
|
super().__init__() |
|
if isinstance(cfg, dict): |
|
self.yaml = cfg |
|
else: |
|
import yaml |
|
self.yaml_file = Path(cfg).name |
|
with open(cfg) as f: |
|
self.yaml = yaml.load(f, Loader=yaml.FullLoader) |
|
|
|
|
|
if nc and nc != self.yaml['nc']: |
|
LOGGER.info(f"Overriding {cfg} nc={self.yaml['nc']} with nc={nc}") |
|
self.yaml['nc'] = nc |
|
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, |
|
conf_thres=0.25): |
|
y = [] |
|
x = inputs |
|
for i, m in enumerate(self.model.layers): |
|
if m.f != -1: |
|
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] |
|
|
|
x = m(x) |
|
y.append(x if m.i in self.savelist else None) |
|
|
|
|
|
if tf_nms: |
|
boxes = self._xywh2xyxy(x[0][..., :4]) |
|
probs = x[0][:, :, 4:5] |
|
classes = x[0][:, :, 5:] |
|
scores = probs * classes |
|
if agnostic_nms: |
|
nms = AgnosticNMS()((boxes, classes, scores), topk_all, iou_thres, conf_thres) |
|
return nms, x[1] |
|
else: |
|
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) |
|
return nms, x[1] |
|
|
|
return x[0] |
|
|
|
|
|
|
|
|
|
|
|
|
|
@staticmethod |
|
def _xywh2xyxy(xywh): |
|
|
|
x, y, w, h = tf.split(xywh, num_or_size_splits=4, axis=-1) |
|
return tf.concat([x - w / 2, y - h / 2, x + w / 2, y + h / 2], axis=-1) |
|
|
|
|
|
class AgnosticNMS(keras.layers.Layer): |
|
|
|
def call(self, input, topk_all, iou_thres, conf_thres): |
|
|
|
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), |
|
name='agnostic_nms') |
|
|
|
@staticmethod |
|
def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25): |
|
boxes, classes, scores = x |
|
class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32) |
|
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_boxes = tf.gather(boxes, selected_inds) |
|
padded_boxes = tf.pad(selected_boxes, |
|
paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]], |
|
mode="CONSTANT", constant_values=0.0) |
|
selected_scores = tf.gather(scores_inp, selected_inds) |
|
padded_scores = tf.pad(selected_scores, |
|
paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], |
|
mode="CONSTANT", constant_values=-1.0) |
|
selected_classes = tf.gather(class_inds, selected_inds) |
|
padded_classes = tf.pad(selected_classes, |
|
paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], |
|
mode="CONSTANT", constant_values=-1.0) |
|
valid_detections = tf.shape(selected_inds)[0] |
|
return padded_boxes, padded_scores, padded_classes, valid_detections |
|
|
|
|
|
def representative_dataset_gen(dataset, ncalib=100): |
|
|
|
for n, (path, img, im0s, vid_cap, string) in enumerate(dataset): |
|
input = np.transpose(img, [1, 2, 0]) |
|
input = np.expand_dims(input, axis=0).astype(np.float32) |
|
input /= 255 |
|
yield [input] |
|
if n >= ncalib: |
|
break |
|
|
|
|
|
def run(weights=ROOT / 'yolov5s.pt', |
|
imgsz=(640, 640), |
|
batch_size=1, |
|
dynamic=False, |
|
): |
|
|
|
im = torch.zeros((batch_size, 3, *imgsz)) |
|
model = attempt_load(weights, map_location=torch.device('cpu'), inplace=True, fuse=False) |
|
y = model(im) |
|
model.info() |
|
|
|
|
|
im = tf.zeros((batch_size, *imgsz, 3)) |
|
tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz) |
|
y = tf_model.predict(im) |
|
|
|
|
|
im = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size) |
|
keras_model = keras.Model(inputs=im, outputs=tf_model.predict(im)) |
|
keras_model.summary() |
|
|
|
LOGGER.info('PyTorch, TensorFlow and Keras models successfully verified.\nUse export.py for TF model export.') |
|
|
|
|
|
def parse_opt(): |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path') |
|
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w') |
|
parser.add_argument('--batch-size', type=int, default=1, help='batch size') |
|
parser.add_argument('--dynamic', action='store_true', help='dynamic batch size') |
|
opt = parser.parse_args() |
|
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 |
|
print_args(FILE.stem, opt) |
|
return opt |
|
|
|
|
|
def main(opt): |
|
run(**vars(opt)) |
|
|
|
|
|
if __name__ == "__main__": |
|
opt = parse_opt() |
|
main(opt) |
|
|