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.

701 satır
33KB

  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), 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) # optional 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)), 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(
  185. GhostConv(c1, c_, 1, 1), # pw
  186. DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
  187. GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
  188. self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1,
  189. act=False)) if s == 2 else nn.Identity()
  190. def forward(self, x):
  191. return self.conv(x) + self.shortcut(x)
  192. class Contract(nn.Module):
  193. # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
  194. def __init__(self, gain=2):
  195. super().__init__()
  196. self.gain = gain
  197. def forward(self, x):
  198. b, c, h, w = x.size() # assert (h / s == 0) and (W / s == 0), 'Indivisible gain'
  199. s = self.gain
  200. x = x.view(b, c, h // s, s, w // s, s) # x(1,64,40,2,40,2)
  201. x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)
  202. return x.view(b, c * s * s, h // s, w // s) # x(1,256,40,40)
  203. class Expand(nn.Module):
  204. # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)
  205. def __init__(self, gain=2):
  206. super().__init__()
  207. self.gain = gain
  208. def forward(self, x):
  209. b, c, h, w = x.size() # assert C / s ** 2 == 0, 'Indivisible gain'
  210. s = self.gain
  211. x = x.view(b, s, s, c // s ** 2, h, w) # x(1,2,2,16,80,80)
  212. x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2)
  213. return x.view(b, c // s ** 2, h * s, w * s) # x(1,16,160,160)
  214. class Concat(nn.Module):
  215. # Concatenate a list of tensors along dimension
  216. def __init__(self, dimension=1):
  217. super().__init__()
  218. self.d = dimension
  219. def forward(self, x):
  220. return torch.cat(x, self.d)
  221. class DetectMultiBackend(nn.Module):
  222. # YOLOv5 MultiBackend class for python inference on various backends
  223. def __init__(self, weights='yolov5s.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False):
  224. # Usage:
  225. # PyTorch: weights = *.pt
  226. # TorchScript: *.torchscript
  227. # ONNX Runtime: *.onnx
  228. # ONNX OpenCV DNN: *.onnx with --dnn
  229. # OpenVINO: *.xml
  230. # CoreML: *.mlmodel
  231. # TensorRT: *.engine
  232. # TensorFlow SavedModel: *_saved_model
  233. # TensorFlow GraphDef: *.pb
  234. # TensorFlow Lite: *.tflite
  235. # TensorFlow Edge TPU: *_edgetpu.tflite
  236. from models.experimental import attempt_download, attempt_load # scoped to avoid circular import
  237. super().__init__()
  238. w = str(weights[0] if isinstance(weights, list) else weights)
  239. pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = self.model_type(w) # get backend
  240. stride, names = 32, [f'class{i}' for i in range(1000)] # assign defaults
  241. w = attempt_download(w) # download if not local
  242. fp16 &= (pt or jit or onnx or engine) and device.type != 'cpu' # FP16
  243. if data: # data.yaml path (optional)
  244. with open(data, errors='ignore') as f:
  245. names = yaml.safe_load(f)['names'] # class names
  246. if pt: # PyTorch
  247. model = attempt_load(weights if isinstance(weights, list) else w, map_location=device)
  248. stride = max(int(model.stride.max()), 32) # model stride
  249. names = model.module.names if hasattr(model, 'module') else model.names # get class names
  250. model.half() if fp16 else model.float()
  251. self.model = model # explicitly assign for to(), cpu(), cuda(), half()
  252. elif jit: # TorchScript
  253. LOGGER.info(f'Loading {w} for TorchScript inference...')
  254. extra_files = {'config.txt': ''} # model metadata
  255. model = torch.jit.load(w, _extra_files=extra_files)
  256. model.half() if fp16 else model.float()
  257. if extra_files['config.txt']:
  258. d = json.loads(extra_files['config.txt']) # extra_files dict
  259. stride, names = int(d['stride']), d['names']
  260. elif dnn: # ONNX OpenCV DNN
  261. LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')
  262. check_requirements(('opencv-python>=4.5.4',))
  263. net = cv2.dnn.readNetFromONNX(w)
  264. elif onnx: # ONNX Runtime
  265. LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
  266. cuda = torch.cuda.is_available()
  267. check_requirements(('onnx', 'onnxruntime-gpu' if cuda else 'onnxruntime'))
  268. import onnxruntime
  269. providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
  270. session = onnxruntime.InferenceSession(w, providers=providers)
  271. meta = session.get_modelmeta().custom_metadata_map # metadata
  272. if 'stride' in meta:
  273. stride, names = int(meta['stride']), eval(meta['names'])
  274. elif xml: # OpenVINO
  275. LOGGER.info(f'Loading {w} for OpenVINO inference...')
  276. check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/
  277. import openvino.inference_engine as ie
  278. core = ie.IECore()
  279. if not Path(w).is_file(): # if not *.xml
  280. w = next(Path(w).glob('*.xml')) # get *.xml file from *_openvino_model dir
  281. network = core.read_network(model=w, weights=Path(w).with_suffix('.bin')) # *.xml, *.bin paths
  282. executable_network = core.load_network(network, device_name='CPU', num_requests=1)
  283. elif engine: # TensorRT
  284. LOGGER.info(f'Loading {w} for TensorRT inference...')
  285. import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download
  286. check_version(trt.__version__, '7.0.0', hard=True) # require tensorrt>=7.0.0
  287. Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
  288. logger = trt.Logger(trt.Logger.INFO)
  289. with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
  290. model = runtime.deserialize_cuda_engine(f.read())
  291. bindings = OrderedDict()
  292. fp16 = False # default updated below
  293. for index in range(model.num_bindings):
  294. name = model.get_binding_name(index)
  295. dtype = trt.nptype(model.get_binding_dtype(index))
  296. shape = tuple(model.get_binding_shape(index))
  297. data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(device)
  298. bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr()))
  299. if model.binding_is_input(index) and dtype == np.float16:
  300. fp16 = True
  301. binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
  302. context = model.create_execution_context()
  303. batch_size = bindings['images'].shape[0]
  304. elif coreml: # CoreML
  305. LOGGER.info(f'Loading {w} for CoreML inference...')
  306. import coremltools as ct
  307. model = ct.models.MLModel(w)
  308. else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
  309. if saved_model: # SavedModel
  310. LOGGER.info(f'Loading {w} for TensorFlow SavedModel inference...')
  311. import tensorflow as tf
  312. keras = False # assume TF1 saved_model
  313. model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w)
  314. elif pb: # GraphDef https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
  315. LOGGER.info(f'Loading {w} for TensorFlow GraphDef inference...')
  316. import tensorflow as tf
  317. def wrap_frozen_graph(gd, inputs, outputs):
  318. x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), []) # wrapped
  319. ge = x.graph.as_graph_element
  320. return x.prune(tf.nest.map_structure(ge, inputs), tf.nest.map_structure(ge, outputs))
  321. gd = tf.Graph().as_graph_def() # graph_def
  322. with open(w, 'rb') as f:
  323. gd.ParseFromString(f.read())
  324. frozen_func = wrap_frozen_graph(gd, inputs="x:0", outputs="Identity:0")
  325. elif tflite or edgetpu: # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
  326. try: # https://coral.ai/docs/edgetpu/tflite-python/#update-existing-tf-lite-code-for-the-edge-tpu
  327. from tflite_runtime.interpreter import Interpreter, load_delegate
  328. except ImportError:
  329. import tensorflow as tf
  330. Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate,
  331. if edgetpu: # Edge TPU https://coral.ai/software/#edgetpu-runtime
  332. LOGGER.info(f'Loading {w} for TensorFlow Lite Edge TPU inference...')
  333. delegate = {
  334. 'Linux': 'libedgetpu.so.1',
  335. 'Darwin': 'libedgetpu.1.dylib',
  336. 'Windows': 'edgetpu.dll'}[platform.system()]
  337. interpreter = Interpreter(model_path=w, experimental_delegates=[load_delegate(delegate)])
  338. else: # Lite
  339. LOGGER.info(f'Loading {w} for TensorFlow Lite inference...')
  340. interpreter = Interpreter(model_path=w) # load TFLite model
  341. interpreter.allocate_tensors() # allocate
  342. input_details = interpreter.get_input_details() # inputs
  343. output_details = interpreter.get_output_details() # outputs
  344. elif tfjs:
  345. raise Exception('ERROR: YOLOv5 TF.js inference is not supported')
  346. self.__dict__.update(locals()) # assign all variables to self
  347. def forward(self, im, augment=False, visualize=False, val=False):
  348. # YOLOv5 MultiBackend inference
  349. b, ch, h, w = im.shape # batch, channel, height, width
  350. if self.pt: # PyTorch
  351. y = self.model(im, augment=augment, visualize=visualize)[0]
  352. elif self.jit: # TorchScript
  353. y = self.model(im)[0]
  354. elif self.dnn: # ONNX OpenCV DNN
  355. im = im.cpu().numpy() # torch to numpy
  356. self.net.setInput(im)
  357. y = self.net.forward()
  358. elif self.onnx: # ONNX Runtime
  359. im = im.cpu().numpy() # torch to numpy
  360. y = self.session.run([self.session.get_outputs()[0].name], {self.session.get_inputs()[0].name: im})[0]
  361. elif self.xml: # OpenVINO
  362. im = im.cpu().numpy() # FP32
  363. desc = self.ie.TensorDesc(precision='FP32', dims=im.shape, layout='NCHW') # Tensor Description
  364. request = self.executable_network.requests[0] # inference request
  365. request.set_blob(blob_name='images', blob=self.ie.Blob(desc, im)) # name=next(iter(request.input_blobs))
  366. request.infer()
  367. y = request.output_blobs['output'].buffer # name=next(iter(request.output_blobs))
  368. elif self.engine: # TensorRT
  369. assert im.shape == self.bindings['images'].shape, (im.shape, self.bindings['images'].shape)
  370. self.binding_addrs['images'] = int(im.data_ptr())
  371. self.context.execute_v2(list(self.binding_addrs.values()))
  372. y = self.bindings['output'].data
  373. elif self.coreml: # CoreML
  374. im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
  375. im = Image.fromarray((im[0] * 255).astype('uint8'))
  376. # im = im.resize((192, 320), Image.ANTIALIAS)
  377. y = self.model.predict({'image': im}) # coordinates are xywh normalized
  378. if 'confidence' in y:
  379. box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels
  380. conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
  381. y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
  382. else:
  383. k = 'var_' + str(sorted(int(k.replace('var_', '')) for k in y)[-1]) # output key
  384. y = y[k] # output
  385. else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
  386. im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
  387. if self.saved_model: # SavedModel
  388. y = (self.model(im, training=False) if self.keras else self.model(im)).numpy()
  389. elif self.pb: # GraphDef
  390. y = self.frozen_func(x=self.tf.constant(im)).numpy()
  391. else: # Lite or Edge TPU
  392. input, output = self.input_details[0], self.output_details[0]
  393. int8 = input['dtype'] == np.uint8 # is TFLite quantized uint8 model
  394. if int8:
  395. scale, zero_point = input['quantization']
  396. im = (im / scale + zero_point).astype(np.uint8) # de-scale
  397. self.interpreter.set_tensor(input['index'], im)
  398. self.interpreter.invoke()
  399. y = self.interpreter.get_tensor(output['index'])
  400. if int8:
  401. scale, zero_point = output['quantization']
  402. y = (y.astype(np.float32) - zero_point) * scale # re-scale
  403. y[..., :4] *= [w, h, w, h] # xywh normalized to pixels
  404. if isinstance(y, np.ndarray):
  405. y = torch.tensor(y, device=self.device)
  406. return (y, []) if val else y
  407. def warmup(self, imgsz=(1, 3, 640, 640)):
  408. # Warmup model by running inference once
  409. if any((self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb)): # warmup types
  410. if self.device.type != 'cpu': # only warmup GPU models
  411. im = torch.zeros(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device) # input
  412. for _ in range(2 if self.jit else 1): #
  413. self.forward(im) # warmup
  414. @staticmethod
  415. def model_type(p='path/to/model.pt'):
  416. # Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx
  417. from export import export_formats
  418. suffixes = list(export_formats().Suffix) + ['.xml'] # export suffixes
  419. check_suffix(p, suffixes) # checks
  420. p = Path(p).name # eliminate trailing separators
  421. pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, xml2 = (s in p for s in suffixes)
  422. xml |= xml2 # *_openvino_model or *.xml
  423. tflite &= not edgetpu # *.tflite
  424. return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs
  425. class AutoShape(nn.Module):
  426. # YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
  427. conf = 0.25 # NMS confidence threshold
  428. iou = 0.45 # NMS IoU threshold
  429. agnostic = False # NMS class-agnostic
  430. multi_label = False # NMS multiple labels per box
  431. classes = None # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
  432. max_det = 1000 # maximum number of detections per image
  433. amp = False # Automatic Mixed Precision (AMP) inference
  434. def __init__(self, model):
  435. super().__init__()
  436. LOGGER.info('Adding AutoShape... ')
  437. copy_attr(self, model, include=('yaml', 'nc', 'hyp', 'names', 'stride', 'abc'), exclude=()) # copy attributes
  438. self.dmb = isinstance(model, DetectMultiBackend) # DetectMultiBackend() instance
  439. self.pt = not self.dmb or model.pt # PyTorch model
  440. self.model = model.eval()
  441. def _apply(self, fn):
  442. # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
  443. self = super()._apply(fn)
  444. if self.pt:
  445. m = self.model.model.model[-1] if self.dmb else self.model.model[-1] # Detect()
  446. m.stride = fn(m.stride)
  447. m.grid = list(map(fn, m.grid))
  448. if isinstance(m.anchor_grid, list):
  449. m.anchor_grid = list(map(fn, m.anchor_grid))
  450. return self
  451. @torch.no_grad()
  452. def forward(self, imgs, size=640, augment=False, profile=False):
  453. # Inference from various sources. For height=640, width=1280, RGB images example inputs are:
  454. # file: imgs = 'data/images/zidane.jpg' # str or PosixPath
  455. # URI: = 'https://ultralytics.com/images/zidane.jpg'
  456. # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
  457. # PIL: = Image.open('image.jpg') or ImageGrab.grab() # HWC x(640,1280,3)
  458. # numpy: = np.zeros((640,1280,3)) # HWC
  459. # torch: = torch.zeros(16,3,320,640) # BCHW (scaled to size=640, 0-1 values)
  460. # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
  461. t = [time_sync()]
  462. p = next(self.model.parameters()) if self.pt else torch.zeros(1) # for device and type
  463. autocast = self.amp and (p.device.type != 'cpu') # Automatic Mixed Precision (AMP) inference
  464. if isinstance(imgs, torch.Tensor): # torch
  465. with amp.autocast(autocast):
  466. return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference
  467. # Pre-process
  468. n, imgs = (len(imgs), list(imgs)) if isinstance(imgs, (list, tuple)) else (1, [imgs]) # number, list of images
  469. shape0, shape1, files = [], [], [] # image and inference shapes, filenames
  470. for i, im in enumerate(imgs):
  471. f = f'image{i}' # filename
  472. if isinstance(im, (str, Path)): # filename or uri
  473. im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith('http') else im), im
  474. im = np.asarray(exif_transpose(im))
  475. elif isinstance(im, Image.Image): # PIL Image
  476. im, f = np.asarray(exif_transpose(im)), getattr(im, 'filename', f) or f
  477. files.append(Path(f).with_suffix('.jpg').name)
  478. if im.shape[0] < 5: # image in CHW
  479. im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
  480. im = im[..., :3] if im.ndim == 3 else np.tile(im[..., None], 3) # enforce 3ch input
  481. s = im.shape[:2] # HWC
  482. shape0.append(s) # image shape
  483. g = (size / max(s)) # gain
  484. shape1.append([y * g for y in s])
  485. imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
  486. shape1 = [make_divisible(x, self.stride) if self.pt else size for x in np.array(shape1).max(0)] # inf shape
  487. x = [letterbox(im, shape1, auto=False)[0] for im in imgs] # pad
  488. x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2))) # stack and BHWC to BCHW
  489. x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32
  490. t.append(time_sync())
  491. with amp.autocast(autocast):
  492. # Inference
  493. y = self.model(x, augment, profile) # forward
  494. t.append(time_sync())
  495. # Post-process
  496. y = non_max_suppression(y if self.dmb else y[0],
  497. self.conf,
  498. self.iou,
  499. self.classes,
  500. self.agnostic,
  501. self.multi_label,
  502. max_det=self.max_det) # NMS
  503. for i in range(n):
  504. scale_coords(shape1, y[i][:, :4], shape0[i])
  505. t.append(time_sync())
  506. return Detections(imgs, y, files, t, self.names, x.shape)
  507. class Detections:
  508. # YOLOv5 detections class for inference results
  509. def __init__(self, imgs, pred, files, times=(0, 0, 0, 0), names=None, shape=None):
  510. super().__init__()
  511. d = pred[0].device # device
  512. gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1, 1], device=d) for im in imgs] # normalizations
  513. self.imgs = imgs # list of images as numpy arrays
  514. self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
  515. self.names = names # class names
  516. self.files = files # image filenames
  517. self.times = times # profiling times
  518. self.xyxy = pred # xyxy pixels
  519. self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
  520. self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
  521. self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
  522. self.n = len(self.pred) # number of images (batch size)
  523. self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms)
  524. self.s = shape # inference BCHW shape
  525. def display(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path('')):
  526. crops = []
  527. for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
  528. s = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' # string
  529. if pred.shape[0]:
  530. for c in pred[:, -1].unique():
  531. n = (pred[:, -1] == c).sum() # detections per class
  532. s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
  533. if show or save or render or crop:
  534. annotator = Annotator(im, example=str(self.names))
  535. for *box, conf, cls in reversed(pred): # xyxy, confidence, class
  536. label = f'{self.names[int(cls)]} {conf:.2f}'
  537. if crop:
  538. file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None
  539. crops.append({
  540. 'box': box,
  541. 'conf': conf,
  542. 'cls': cls,
  543. 'label': label,
  544. 'im': save_one_box(box, im, file=file, save=save)})
  545. else: # all others
  546. annotator.box_label(box, label if labels else '', color=colors(cls))
  547. im = annotator.im
  548. else:
  549. s += '(no detections)'
  550. im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np
  551. if pprint:
  552. LOGGER.info(s.rstrip(', '))
  553. if show:
  554. im.show(self.files[i]) # show
  555. if save:
  556. f = self.files[i]
  557. im.save(save_dir / f) # save
  558. if i == self.n - 1:
  559. LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
  560. if render:
  561. self.imgs[i] = np.asarray(im)
  562. if crop:
  563. if save:
  564. LOGGER.info(f'Saved results to {save_dir}\n')
  565. return crops
  566. def print(self):
  567. self.display(pprint=True) # print results
  568. LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' %
  569. self.t)
  570. def show(self, labels=True):
  571. self.display(show=True, labels=labels) # show results
  572. def save(self, labels=True, save_dir='runs/detect/exp'):
  573. save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir
  574. self.display(save=True, labels=labels, save_dir=save_dir) # save results
  575. def crop(self, save=True, save_dir='runs/detect/exp'):
  576. save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) if save else None
  577. return self.display(crop=True, save=save, save_dir=save_dir) # crop results
  578. def render(self, labels=True):
  579. self.display(render=True, labels=labels) # render results
  580. return self.imgs
  581. def pandas(self):
  582. # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])
  583. new = copy(self) # return copy
  584. ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name' # xyxy columns
  585. cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns
  586. for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]):
  587. a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # update
  588. setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])
  589. return new
  590. def tolist(self):
  591. # return a list of Detections objects, i.e. 'for result in results.tolist():'
  592. r = range(self.n) # iterable
  593. x = [Detections([self.imgs[i]], [self.pred[i]], [self.files[i]], self.times, self.names, self.s) for i in r]
  594. # for d in x:
  595. # for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
  596. # setattr(d, k, getattr(d, k)[0]) # pop out of list
  597. return x
  598. def __len__(self):
  599. return self.n
  600. class Classify(nn.Module):
  601. # Classification head, i.e. x(b,c1,20,20) to x(b,c2)
  602. def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups
  603. super().__init__()
  604. self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1)
  605. self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g) # to x(b,c2,1,1)
  606. self.flat = nn.Flatten()
  607. def forward(self, x):
  608. z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list
  609. return self.flat(self.conv(z)) # flatten to x(b,c2)