You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

678 lines
32KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Common modules
  4. """
  5. import json
  6. import math
  7. import platform
  8. import warnings
  9. from collections import OrderedDict, namedtuple
  10. from copy import copy
  11. from pathlib import Path
  12. import cv2
  13. import numpy as np
  14. import pandas as pd
  15. import requests
  16. import torch
  17. import torch.nn as nn
  18. import yaml
  19. from PIL import Image
  20. from torch.cuda import amp
  21. from utils.datasets import exif_transpose, letterbox
  22. from utils.general import (LOGGER, check_requirements, check_suffix, check_version, colorstr, increment_path,
  23. make_divisible, non_max_suppression, scale_coords, xywh2xyxy, xyxy2xywh)
  24. from utils.plots import Annotator, colors, save_one_box
  25. from utils.torch_utils import copy_attr, time_sync
  26. def autopad(k, p=None): # kernel, padding
  27. # Pad to 'same'
  28. if p is None:
  29. p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
  30. return p
  31. class Conv(nn.Module):
  32. # Standard convolution
  33. def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
  34. super().__init__()
  35. self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
  36. self.bn = nn.BatchNorm2d(c2)
  37. self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
  38. def forward(self, x):
  39. return self.act(self.bn(self.conv(x)))
  40. def forward_fuse(self, x):
  41. return self.act(self.conv(x))
  42. class DWConv(Conv):
  43. # Depth-wise convolution class
  44. def __init__(self, c1, c2, k=1, s=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
  45. super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
  46. class TransformerLayer(nn.Module):
  47. # Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
  48. def __init__(self, c, num_heads):
  49. super().__init__()
  50. self.q = nn.Linear(c, c, bias=False)
  51. self.k = nn.Linear(c, c, bias=False)
  52. self.v = nn.Linear(c, c, bias=False)
  53. self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
  54. self.fc1 = nn.Linear(c, c, bias=False)
  55. self.fc2 = nn.Linear(c, c, bias=False)
  56. def forward(self, x):
  57. x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
  58. x = self.fc2(self.fc1(x)) + x
  59. return x
  60. class TransformerBlock(nn.Module):
  61. # Vision Transformer https://arxiv.org/abs/2010.11929
  62. def __init__(self, c1, c2, num_heads, num_layers):
  63. super().__init__()
  64. self.conv = None
  65. if c1 != c2:
  66. self.conv = Conv(c1, c2)
  67. self.linear = nn.Linear(c2, c2) # learnable position embedding
  68. self.tr = nn.Sequential(*(TransformerLayer(c2, num_heads) for _ in range(num_layers)))
  69. self.c2 = c2
  70. def forward(self, x):
  71. if self.conv is not None:
  72. x = self.conv(x)
  73. b, _, w, h = x.shape
  74. p = x.flatten(2).permute(2, 0, 1)
  75. return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, w, h)
  76. class Bottleneck(nn.Module):
  77. # Standard bottleneck
  78. def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
  79. super().__init__()
  80. c_ = int(c2 * e) # hidden channels
  81. self.cv1 = Conv(c1, c_, 1, 1)
  82. self.cv2 = Conv(c_, c2, 3, 1, g=g)
  83. self.add = shortcut and c1 == c2
  84. def forward(self, x):
  85. return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
  86. class BottleneckCSP(nn.Module):
  87. # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
  88. def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
  89. super().__init__()
  90. c_ = int(c2 * e) # hidden channels
  91. self.cv1 = Conv(c1, c_, 1, 1)
  92. self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
  93. self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
  94. self.cv4 = Conv(2 * c_, c2, 1, 1)
  95. self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
  96. self.act = nn.SiLU()
  97. self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
  98. def forward(self, x):
  99. y1 = self.cv3(self.m(self.cv1(x)))
  100. y2 = self.cv2(x)
  101. return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
  102. class C3(nn.Module):
  103. # CSP Bottleneck with 3 convolutions
  104. def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
  105. super().__init__()
  106. c_ = int(c2 * e) # hidden channels
  107. self.cv1 = Conv(c1, c_, 1, 1)
  108. self.cv2 = Conv(c1, c_, 1, 1)
  109. self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
  110. self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
  111. # self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
  112. def forward(self, x):
  113. return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
  114. class C3TR(C3):
  115. # C3 module with TransformerBlock()
  116. def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
  117. super().__init__(c1, c2, n, shortcut, g, e)
  118. c_ = int(c2 * e)
  119. self.m = TransformerBlock(c_, c_, 4, n)
  120. class C3SPP(C3):
  121. # C3 module with SPP()
  122. def __init__(self, c1, c2, k=(5, 9, 13), n=1, shortcut=True, g=1, e=0.5):
  123. super().__init__(c1, c2, n, shortcut, g, e)
  124. c_ = int(c2 * e)
  125. self.m = SPP(c_, c_, k)
  126. class C3Ghost(C3):
  127. # C3 module with GhostBottleneck()
  128. def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
  129. super().__init__(c1, c2, n, shortcut, g, e)
  130. c_ = int(c2 * e) # hidden channels
  131. self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))
  132. class SPP(nn.Module):
  133. # Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729
  134. def __init__(self, c1, c2, k=(5, 9, 13)):
  135. super().__init__()
  136. c_ = c1 // 2 # hidden channels
  137. self.cv1 = Conv(c1, c_, 1, 1)
  138. self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
  139. self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
  140. def forward(self, x):
  141. x = self.cv1(x)
  142. with warnings.catch_warnings():
  143. warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
  144. return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
  145. class SPPF(nn.Module):
  146. # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
  147. def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
  148. super().__init__()
  149. c_ = c1 // 2 # hidden channels
  150. self.cv1 = Conv(c1, c_, 1, 1)
  151. self.cv2 = Conv(c_ * 4, c2, 1, 1)
  152. self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
  153. def forward(self, x):
  154. x = self.cv1(x)
  155. with warnings.catch_warnings():
  156. warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
  157. y1 = self.m(x)
  158. y2 = self.m(y1)
  159. return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
  160. class Focus(nn.Module):
  161. # Focus wh information into c-space
  162. def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
  163. super().__init__()
  164. self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
  165. # self.contract = Contract(gain=2)
  166. def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
  167. return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
  168. # return self.conv(self.contract(x))
  169. class GhostConv(nn.Module):
  170. # Ghost Convolution https://github.com/huawei-noah/ghostnet
  171. def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups
  172. super().__init__()
  173. c_ = c2 // 2 # hidden channels
  174. self.cv1 = Conv(c1, c_, k, s, None, g, act)
  175. self.cv2 = Conv(c_, c_, 5, 1, None, c_, act)
  176. def forward(self, x):
  177. y = self.cv1(x)
  178. return torch.cat([y, self.cv2(y)], 1)
  179. class GhostBottleneck(nn.Module):
  180. # Ghost Bottleneck https://github.com/huawei-noah/ghostnet
  181. def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride
  182. super().__init__()
  183. c_ = c2 // 2
  184. self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw
  185. DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
  186. GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
  187. self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False),
  188. Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
  189. def forward(self, x):
  190. return self.conv(x) + self.shortcut(x)
  191. class Contract(nn.Module):
  192. # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
  193. def __init__(self, gain=2):
  194. super().__init__()
  195. self.gain = gain
  196. def forward(self, x):
  197. b, c, h, w = x.size() # assert (h / s == 0) and (W / s == 0), 'Indivisible gain'
  198. s = self.gain
  199. x = x.view(b, c, h // s, s, w // s, s) # x(1,64,40,2,40,2)
  200. x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)
  201. return x.view(b, c * s * s, h // s, w // s) # x(1,256,40,40)
  202. class Expand(nn.Module):
  203. # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)
  204. def __init__(self, gain=2):
  205. super().__init__()
  206. self.gain = gain
  207. def forward(self, x):
  208. b, c, h, w = x.size() # assert C / s ** 2 == 0, 'Indivisible gain'
  209. s = self.gain
  210. x = x.view(b, s, s, c // s ** 2, h, w) # x(1,2,2,16,80,80)
  211. x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2)
  212. return x.view(b, c // s ** 2, h * s, w * s) # x(1,16,160,160)
  213. class Concat(nn.Module):
  214. # Concatenate a list of tensors along dimension
  215. def __init__(self, dimension=1):
  216. super().__init__()
  217. self.d = dimension
  218. def forward(self, x):
  219. return torch.cat(x, self.d)
  220. class DetectMultiBackend(nn.Module):
  221. # YOLOv5 MultiBackend class for python inference on various backends
  222. def __init__(self, weights='yolov5s.pt', device=None, dnn=False, data=None):
  223. # Usage:
  224. # PyTorch: weights = *.pt
  225. # TorchScript: *.torchscript
  226. # ONNX Runtime: *.onnx
  227. # ONNX OpenCV DNN: *.onnx with --dnn
  228. # OpenVINO: *.xml
  229. # CoreML: *.mlmodel
  230. # TensorRT: *.engine
  231. # TensorFlow SavedModel: *_saved_model
  232. # TensorFlow GraphDef: *.pb
  233. # TensorFlow Lite: *.tflite
  234. # TensorFlow Edge TPU: *_edgetpu.tflite
  235. from models.experimental import attempt_download, attempt_load # scoped to avoid circular import
  236. super().__init__()
  237. w = str(weights[0] if isinstance(weights, list) else weights)
  238. pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = self.model_type(w) # get backend
  239. stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
  240. w = attempt_download(w) # download if not local
  241. if data: # data.yaml path (optional)
  242. with open(data, errors='ignore') as f:
  243. names = yaml.safe_load(f)['names'] # class names
  244. if pt: # PyTorch
  245. model = attempt_load(weights if isinstance(weights, list) else w, map_location=device)
  246. stride = max(int(model.stride.max()), 32) # model stride
  247. names = model.module.names if hasattr(model, 'module') else model.names # get class names
  248. self.model = model # explicitly assign for to(), cpu(), cuda(), half()
  249. elif jit: # TorchScript
  250. LOGGER.info(f'Loading {w} for TorchScript inference...')
  251. extra_files = {'config.txt': ''} # model metadata
  252. model = torch.jit.load(w, _extra_files=extra_files)
  253. if extra_files['config.txt']:
  254. d = json.loads(extra_files['config.txt']) # extra_files dict
  255. stride, names = int(d['stride']), d['names']
  256. elif dnn: # ONNX OpenCV DNN
  257. LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')
  258. check_requirements(('opencv-python>=4.5.4',))
  259. net = cv2.dnn.readNetFromONNX(w)
  260. elif onnx: # ONNX Runtime
  261. LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
  262. cuda = torch.cuda.is_available()
  263. check_requirements(('onnx', 'onnxruntime-gpu' if cuda else 'onnxruntime'))
  264. import onnxruntime
  265. providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
  266. session = onnxruntime.InferenceSession(w, providers=providers)
  267. elif xml: # OpenVINO
  268. LOGGER.info(f'Loading {w} for OpenVINO inference...')
  269. check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/
  270. import openvino.inference_engine as ie
  271. core = ie.IECore()
  272. if not Path(w).is_file(): # if not *.xml
  273. w = next(Path(w).glob('*.xml')) # get *.xml file from *_openvino_model dir
  274. network = core.read_network(model=w, weights=Path(w).with_suffix('.bin')) # *.xml, *.bin paths
  275. executable_network = core.load_network(network, device_name='CPU', num_requests=1)
  276. elif engine: # TensorRT
  277. LOGGER.info(f'Loading {w} for TensorRT inference...')
  278. import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download
  279. check_version(trt.__version__, '7.0.0', hard=True) # require tensorrt>=7.0.0
  280. Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
  281. logger = trt.Logger(trt.Logger.INFO)
  282. with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
  283. model = runtime.deserialize_cuda_engine(f.read())
  284. bindings = OrderedDict()
  285. for index in range(model.num_bindings):
  286. name = model.get_binding_name(index)
  287. dtype = trt.nptype(model.get_binding_dtype(index))
  288. shape = tuple(model.get_binding_shape(index))
  289. data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(device)
  290. bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr()))
  291. binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
  292. context = model.create_execution_context()
  293. batch_size = bindings['images'].shape[0]
  294. elif coreml: # CoreML
  295. LOGGER.info(f'Loading {w} for CoreML inference...')
  296. import coremltools as ct
  297. model = ct.models.MLModel(w)
  298. else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
  299. if saved_model: # SavedModel
  300. LOGGER.info(f'Loading {w} for TensorFlow SavedModel inference...')
  301. import tensorflow as tf
  302. keras = False # assume TF1 saved_model
  303. model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w)
  304. elif pb: # GraphDef https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
  305. LOGGER.info(f'Loading {w} for TensorFlow GraphDef inference...')
  306. import tensorflow as tf
  307. def wrap_frozen_graph(gd, inputs, outputs):
  308. x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), []) # wrapped
  309. ge = x.graph.as_graph_element
  310. return x.prune(tf.nest.map_structure(ge, inputs), tf.nest.map_structure(ge, outputs))
  311. gd = tf.Graph().as_graph_def() # graph_def
  312. gd.ParseFromString(open(w, 'rb').read())
  313. frozen_func = wrap_frozen_graph(gd, inputs="x:0", outputs="Identity:0")
  314. elif tflite or edgetpu: # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
  315. try: # https://coral.ai/docs/edgetpu/tflite-python/#update-existing-tf-lite-code-for-the-edge-tpu
  316. from tflite_runtime.interpreter import Interpreter, load_delegate
  317. except ImportError:
  318. import tensorflow as tf
  319. Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate,
  320. if edgetpu: # Edge TPU https://coral.ai/software/#edgetpu-runtime
  321. LOGGER.info(f'Loading {w} for TensorFlow Lite Edge TPU inference...')
  322. delegate = {'Linux': 'libedgetpu.so.1',
  323. 'Darwin': 'libedgetpu.1.dylib',
  324. 'Windows': 'edgetpu.dll'}[platform.system()]
  325. interpreter = Interpreter(model_path=w, experimental_delegates=[load_delegate(delegate)])
  326. else: # Lite
  327. LOGGER.info(f'Loading {w} for TensorFlow Lite inference...')
  328. interpreter = Interpreter(model_path=w) # load TFLite model
  329. interpreter.allocate_tensors() # allocate
  330. input_details = interpreter.get_input_details() # inputs
  331. output_details = interpreter.get_output_details() # outputs
  332. elif tfjs:
  333. raise Exception('ERROR: YOLOv5 TF.js inference is not supported')
  334. self.__dict__.update(locals()) # assign all variables to self
  335. def forward(self, im, augment=False, visualize=False, val=False):
  336. # YOLOv5 MultiBackend inference
  337. b, ch, h, w = im.shape # batch, channel, height, width
  338. if self.pt or self.jit: # PyTorch
  339. y = self.model(im) if self.jit else self.model(im, augment=augment, visualize=visualize)
  340. return y if val else y[0]
  341. elif self.dnn: # ONNX OpenCV DNN
  342. im = im.cpu().numpy() # torch to numpy
  343. self.net.setInput(im)
  344. y = self.net.forward()
  345. elif self.onnx: # ONNX Runtime
  346. im = im.cpu().numpy() # torch to numpy
  347. y = self.session.run([self.session.get_outputs()[0].name], {self.session.get_inputs()[0].name: im})[0]
  348. elif self.xml: # OpenVINO
  349. im = im.cpu().numpy() # FP32
  350. desc = self.ie.TensorDesc(precision='FP32', dims=im.shape, layout='NCHW') # Tensor Description
  351. request = self.executable_network.requests[0] # inference request
  352. request.set_blob(blob_name='images', blob=self.ie.Blob(desc, im)) # name=next(iter(request.input_blobs))
  353. request.infer()
  354. y = request.output_blobs['output'].buffer # name=next(iter(request.output_blobs))
  355. elif self.engine: # TensorRT
  356. assert im.shape == self.bindings['images'].shape, (im.shape, self.bindings['images'].shape)
  357. self.binding_addrs['images'] = int(im.data_ptr())
  358. self.context.execute_v2(list(self.binding_addrs.values()))
  359. y = self.bindings['output'].data
  360. elif self.coreml: # CoreML
  361. im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
  362. im = Image.fromarray((im[0] * 255).astype('uint8'))
  363. # im = im.resize((192, 320), Image.ANTIALIAS)
  364. y = self.model.predict({'image': im}) # coordinates are xywh normalized
  365. if 'confidence' in y:
  366. box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels
  367. conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
  368. y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
  369. else:
  370. k = 'var_' + str(sorted(int(k.replace('var_', '')) for k in y)[-1]) # output key
  371. y = y[k] # output
  372. else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
  373. im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
  374. if self.saved_model: # SavedModel
  375. y = (self.model(im, training=False) if self.keras else self.model(im)[0]).numpy()
  376. elif self.pb: # GraphDef
  377. y = self.frozen_func(x=self.tf.constant(im)).numpy()
  378. else: # Lite or Edge TPU
  379. input, output = self.input_details[0], self.output_details[0]
  380. int8 = input['dtype'] == np.uint8 # is TFLite quantized uint8 model
  381. if int8:
  382. scale, zero_point = input['quantization']
  383. im = (im / scale + zero_point).astype(np.uint8) # de-scale
  384. self.interpreter.set_tensor(input['index'], im)
  385. self.interpreter.invoke()
  386. y = self.interpreter.get_tensor(output['index'])
  387. if int8:
  388. scale, zero_point = output['quantization']
  389. y = (y.astype(np.float32) - zero_point) * scale # re-scale
  390. y[..., :4] *= [w, h, w, h] # xywh normalized to pixels
  391. y = torch.tensor(y) if isinstance(y, np.ndarray) else y
  392. return (y, []) if val else y
  393. def warmup(self, imgsz=(1, 3, 640, 640), half=False):
  394. # Warmup model by running inference once
  395. if self.pt or self.jit or self.onnx or self.engine: # warmup types
  396. if isinstance(self.device, torch.device) and self.device.type != 'cpu': # only warmup GPU models
  397. im = torch.zeros(*imgsz).to(self.device).type(torch.half if half else torch.float) # input image
  398. self.forward(im) # warmup
  399. @staticmethod
  400. def model_type(p='path/to/model.pt'):
  401. # Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx
  402. from export import export_formats
  403. suffixes = list(export_formats().Suffix) + ['.xml'] # export suffixes
  404. check_suffix(p, suffixes) # checks
  405. p = Path(p).name # eliminate trailing separators
  406. pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, xml2 = (s in p for s in suffixes)
  407. xml |= xml2 # *_openvino_model or *.xml
  408. tflite &= not edgetpu # *.tflite
  409. return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs
  410. class AutoShape(nn.Module):
  411. # YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
  412. conf = 0.25 # NMS confidence threshold
  413. iou = 0.45 # NMS IoU threshold
  414. agnostic = False # NMS class-agnostic
  415. multi_label = False # NMS multiple labels per box
  416. classes = None # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
  417. max_det = 1000 # maximum number of detections per image
  418. amp = False # Automatic Mixed Precision (AMP) inference
  419. def __init__(self, model):
  420. super().__init__()
  421. LOGGER.info('Adding AutoShape... ')
  422. copy_attr(self, model, include=('yaml', 'nc', 'hyp', 'names', 'stride', 'abc'), exclude=()) # copy attributes
  423. self.dmb = isinstance(model, DetectMultiBackend) # DetectMultiBackend() instance
  424. self.pt = not self.dmb or model.pt # PyTorch model
  425. self.model = model.eval()
  426. def _apply(self, fn):
  427. # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
  428. self = super()._apply(fn)
  429. if self.pt:
  430. m = self.model.model.model[-1] if self.dmb else self.model.model[-1] # Detect()
  431. m.stride = fn(m.stride)
  432. m.grid = list(map(fn, m.grid))
  433. if isinstance(m.anchor_grid, list):
  434. m.anchor_grid = list(map(fn, m.anchor_grid))
  435. return self
  436. @torch.no_grad()
  437. def forward(self, imgs, size=640, augment=False, profile=False):
  438. # Inference from various sources. For height=640, width=1280, RGB images example inputs are:
  439. # file: imgs = 'data/images/zidane.jpg' # str or PosixPath
  440. # URI: = 'https://ultralytics.com/images/zidane.jpg'
  441. # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
  442. # PIL: = Image.open('image.jpg') or ImageGrab.grab() # HWC x(640,1280,3)
  443. # numpy: = np.zeros((640,1280,3)) # HWC
  444. # torch: = torch.zeros(16,3,320,640) # BCHW (scaled to size=640, 0-1 values)
  445. # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
  446. t = [time_sync()]
  447. p = next(self.model.parameters()) if self.pt else torch.zeros(1) # for device and type
  448. autocast = self.amp and (p.device.type != 'cpu') # Automatic Mixed Precision (AMP) inference
  449. if isinstance(imgs, torch.Tensor): # torch
  450. with amp.autocast(enabled=autocast):
  451. return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference
  452. # Pre-process
  453. n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images
  454. shape0, shape1, files = [], [], [] # image and inference shapes, filenames
  455. for i, im in enumerate(imgs):
  456. f = f'image{i}' # filename
  457. if isinstance(im, (str, Path)): # filename or uri
  458. im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith('http') else im), im
  459. im = np.asarray(exif_transpose(im))
  460. elif isinstance(im, Image.Image): # PIL Image
  461. im, f = np.asarray(exif_transpose(im)), getattr(im, 'filename', f) or f
  462. files.append(Path(f).with_suffix('.jpg').name)
  463. if im.shape[0] < 5: # image in CHW
  464. im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
  465. im = im[..., :3] if im.ndim == 3 else np.tile(im[..., None], 3) # enforce 3ch input
  466. s = im.shape[:2] # HWC
  467. shape0.append(s) # image shape
  468. g = (size / max(s)) # gain
  469. shape1.append([y * g for y in s])
  470. imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
  471. shape1 = [make_divisible(x, self.stride) for x in np.stack(shape1, 0).max(0)] # inference shape
  472. x = [letterbox(im, new_shape=shape1 if self.pt else size, auto=False)[0] for im in imgs] # pad
  473. x = np.stack(x, 0) if n > 1 else x[0][None] # stack
  474. x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
  475. x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32
  476. t.append(time_sync())
  477. with amp.autocast(enabled=autocast):
  478. # Inference
  479. y = self.model(x, augment, profile) # forward
  480. t.append(time_sync())
  481. # Post-process
  482. y = non_max_suppression(y if self.dmb else y[0], self.conf, iou_thres=self.iou, classes=self.classes,
  483. agnostic=self.agnostic, multi_label=self.multi_label, max_det=self.max_det) # NMS
  484. for i in range(n):
  485. scale_coords(shape1, y[i][:, :4], shape0[i])
  486. t.append(time_sync())
  487. return Detections(imgs, y, files, t, self.names, x.shape)
  488. class Detections:
  489. # YOLOv5 detections class for inference results
  490. def __init__(self, imgs, pred, files, times=(0, 0, 0, 0), names=None, shape=None):
  491. super().__init__()
  492. d = pred[0].device # device
  493. gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1, 1], device=d) for im in imgs] # normalizations
  494. self.imgs = imgs # list of images as numpy arrays
  495. self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
  496. self.names = names # class names
  497. self.files = files # image filenames
  498. self.times = times # profiling times
  499. self.xyxy = pred # xyxy pixels
  500. self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
  501. self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
  502. self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
  503. self.n = len(self.pred) # number of images (batch size)
  504. self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms)
  505. self.s = shape # inference BCHW shape
  506. def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')):
  507. crops = []
  508. for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
  509. s = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' # string
  510. if pred.shape[0]:
  511. for c in pred[:, -1].unique():
  512. n = (pred[:, -1] == c).sum() # detections per class
  513. s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
  514. if show or save or render or crop:
  515. annotator = Annotator(im, example=str(self.names))
  516. for *box, conf, cls in reversed(pred): # xyxy, confidence, class
  517. label = f'{self.names[int(cls)]} {conf:.2f}'
  518. if crop:
  519. file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None
  520. crops.append({'box': box, 'conf': conf, 'cls': cls, 'label': label,
  521. 'im': save_one_box(box, im, file=file, save=save)})
  522. else: # all others
  523. annotator.box_label(box, label, color=colors(cls))
  524. im = annotator.im
  525. else:
  526. s += '(no detections)'
  527. im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np
  528. if pprint:
  529. LOGGER.info(s.rstrip(', '))
  530. if show:
  531. im.show(self.files[i]) # show
  532. if save:
  533. f = self.files[i]
  534. im.save(save_dir / f) # save
  535. if i == self.n - 1:
  536. LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
  537. if render:
  538. self.imgs[i] = np.asarray(im)
  539. if crop:
  540. if save:
  541. LOGGER.info(f'Saved results to {save_dir}\n')
  542. return crops
  543. def print(self):
  544. self.display(pprint=True) # print results
  545. LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' %
  546. self.t)
  547. def show(self):
  548. self.display(show=True) # show results
  549. def save(self, save_dir='runs/detect/exp'):
  550. save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir
  551. self.display(save=True, save_dir=save_dir) # save results
  552. def crop(self, save=True, save_dir='runs/detect/exp'):
  553. save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) if save else None
  554. return self.display(crop=True, save=save, save_dir=save_dir) # crop results
  555. def render(self):
  556. self.display(render=True) # render results
  557. return self.imgs
  558. def pandas(self):
  559. # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])
  560. new = copy(self) # return copy
  561. ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name' # xyxy columns
  562. cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns
  563. for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]):
  564. a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # update
  565. setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])
  566. return new
  567. def tolist(self):
  568. # return a list of Detections objects, i.e. 'for result in results.tolist():'
  569. r = range(self.n) # iterable
  570. x = [Detections([self.imgs[i]], [self.pred[i]], [self.files[i]], self.times, self.names, self.s) for i in r]
  571. # for d in x:
  572. # for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
  573. # setattr(d, k, getattr(d, k)[0]) # pop out of list
  574. return x
  575. def __len__(self):
  576. return self.n
  577. class Classify(nn.Module):
  578. # Classification head, i.e. x(b,c1,20,20) to x(b,c2)
  579. def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups
  580. super().__init__()
  581. self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1)
  582. self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g) # to x(b,c2,1,1)
  583. self.flat = nn.Flatten()
  584. def forward(self, x):
  585. z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list
  586. return self.flat(self.conv(z)) # flatten to x(b,c2)