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.

396 lines
16KB

  1. # YOLOv5 common modules
  2. import math
  3. from copy import copy
  4. from pathlib import Path
  5. import numpy as np
  6. import pandas as pd
  7. import requests
  8. import torch
  9. import torch.nn as nn
  10. from PIL import Image
  11. from torch.cuda import amp
  12. from utils.datasets import letterbox
  13. from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh, save_one_box
  14. from utils.plots import colors, plot_one_box
  15. from utils.torch_utils import time_synchronized
  16. def autopad(k, p=None): # kernel, padding
  17. # Pad to 'same'
  18. if p is None:
  19. p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
  20. return p
  21. def DWConv(c1, c2, k=1, s=1, act=True):
  22. # Depthwise convolution
  23. return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
  24. class Conv(nn.Module):
  25. # Standard convolution
  26. def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
  27. super(Conv, self).__init__()
  28. self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
  29. self.bn = nn.BatchNorm2d(c2)
  30. self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
  31. def forward(self, x):
  32. return self.act(self.bn(self.conv(x)))
  33. def fuseforward(self, x):
  34. return self.act(self.conv(x))
  35. class TransformerLayer(nn.Module):
  36. # Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
  37. def __init__(self, c, num_heads):
  38. super().__init__()
  39. self.q = nn.Linear(c, c, bias=False)
  40. self.k = nn.Linear(c, c, bias=False)
  41. self.v = nn.Linear(c, c, bias=False)
  42. self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
  43. self.fc1 = nn.Linear(c, c, bias=False)
  44. self.fc2 = nn.Linear(c, c, bias=False)
  45. def forward(self, x):
  46. x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
  47. x = self.fc2(self.fc1(x)) + x
  48. return x
  49. class TransformerBlock(nn.Module):
  50. # Vision Transformer https://arxiv.org/abs/2010.11929
  51. def __init__(self, c1, c2, num_heads, num_layers):
  52. super().__init__()
  53. self.conv = None
  54. if c1 != c2:
  55. self.conv = Conv(c1, c2)
  56. self.linear = nn.Linear(c2, c2) # learnable position embedding
  57. self.tr = nn.Sequential(*[TransformerLayer(c2, num_heads) for _ in range(num_layers)])
  58. self.c2 = c2
  59. def forward(self, x):
  60. if self.conv is not None:
  61. x = self.conv(x)
  62. b, _, w, h = x.shape
  63. p = x.flatten(2)
  64. p = p.unsqueeze(0)
  65. p = p.transpose(0, 3)
  66. p = p.squeeze(3)
  67. e = self.linear(p)
  68. x = p + e
  69. x = self.tr(x)
  70. x = x.unsqueeze(3)
  71. x = x.transpose(0, 3)
  72. x = x.reshape(b, self.c2, w, h)
  73. return x
  74. class Bottleneck(nn.Module):
  75. # Standard bottleneck
  76. def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
  77. super(Bottleneck, self).__init__()
  78. c_ = int(c2 * e) # hidden channels
  79. self.cv1 = Conv(c1, c_, 1, 1)
  80. self.cv2 = Conv(c_, c2, 3, 1, g=g)
  81. self.add = shortcut and c1 == c2
  82. def forward(self, x):
  83. return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
  84. class BottleneckCSP(nn.Module):
  85. # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
  86. def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
  87. super(BottleneckCSP, self).__init__()
  88. c_ = int(c2 * e) # hidden channels
  89. self.cv1 = Conv(c1, c_, 1, 1)
  90. self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
  91. self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
  92. self.cv4 = Conv(2 * c_, c2, 1, 1)
  93. self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
  94. self.act = nn.LeakyReLU(0.1, inplace=True)
  95. self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
  96. def forward(self, x):
  97. y1 = self.cv3(self.m(self.cv1(x)))
  98. y2 = self.cv2(x)
  99. return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
  100. class C3(nn.Module):
  101. # CSP Bottleneck with 3 convolutions
  102. def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
  103. super(C3, self).__init__()
  104. c_ = int(c2 * e) # hidden channels
  105. self.cv1 = Conv(c1, c_, 1, 1)
  106. self.cv2 = Conv(c1, c_, 1, 1)
  107. self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
  108. self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
  109. # self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
  110. def forward(self, x):
  111. return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
  112. class C3TR(C3):
  113. # C3 module with TransformerBlock()
  114. def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
  115. super().__init__(c1, c2, n, shortcut, g, e)
  116. c_ = int(c2 * e)
  117. self.m = TransformerBlock(c_, c_, 4, n)
  118. class SPP(nn.Module):
  119. # Spatial pyramid pooling layer used in YOLOv3-SPP
  120. def __init__(self, c1, c2, k=(5, 9, 13)):
  121. super(SPP, self).__init__()
  122. c_ = c1 // 2 # hidden channels
  123. self.cv1 = Conv(c1, c_, 1, 1)
  124. self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
  125. self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
  126. def forward(self, x):
  127. x = self.cv1(x)
  128. return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
  129. class Focus(nn.Module):
  130. # Focus wh information into c-space
  131. def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
  132. super(Focus, self).__init__()
  133. self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
  134. # self.contract = Contract(gain=2)
  135. def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
  136. return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
  137. # return self.conv(self.contract(x))
  138. class Contract(nn.Module):
  139. # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
  140. def __init__(self, gain=2):
  141. super().__init__()
  142. self.gain = gain
  143. def forward(self, x):
  144. N, C, H, W = x.size() # assert (H / s == 0) and (W / s == 0), 'Indivisible gain'
  145. s = self.gain
  146. x = x.view(N, C, H // s, s, W // s, s) # x(1,64,40,2,40,2)
  147. x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)
  148. return x.view(N, C * s * s, H // s, W // s) # x(1,256,40,40)
  149. class Expand(nn.Module):
  150. # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)
  151. def __init__(self, gain=2):
  152. super().__init__()
  153. self.gain = gain
  154. def forward(self, x):
  155. N, C, H, W = x.size() # assert C / s ** 2 == 0, 'Indivisible gain'
  156. s = self.gain
  157. x = x.view(N, s, s, C // s ** 2, H, W) # x(1,2,2,16,80,80)
  158. x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2)
  159. return x.view(N, C // s ** 2, H * s, W * s) # x(1,16,160,160)
  160. class Concat(nn.Module):
  161. # Concatenate a list of tensors along dimension
  162. def __init__(self, dimension=1):
  163. super(Concat, self).__init__()
  164. self.d = dimension
  165. def forward(self, x):
  166. return torch.cat(x, self.d)
  167. class NMS(nn.Module):
  168. # Non-Maximum Suppression (NMS) module
  169. conf = 0.25 # confidence threshold
  170. iou = 0.45 # IoU threshold
  171. classes = None # (optional list) filter by class
  172. max_det = 1000 # maximum number of detections per image
  173. def __init__(self):
  174. super(NMS, self).__init__()
  175. def forward(self, x):
  176. return non_max_suppression(x[0], self.conf, iou_thres=self.iou, classes=self.classes, max_det=self.max_det)
  177. class AutoShape(nn.Module):
  178. # input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
  179. conf = 0.25 # NMS confidence threshold
  180. iou = 0.45 # NMS IoU threshold
  181. classes = None # (optional list) filter by class
  182. max_det = 1000 # maximum number of detections per image
  183. def __init__(self, model):
  184. super(AutoShape, self).__init__()
  185. self.model = model.eval()
  186. def autoshape(self):
  187. print('AutoShape already enabled, skipping... ') # model already converted to model.autoshape()
  188. return self
  189. @torch.no_grad()
  190. def forward(self, imgs, size=640, augment=False, profile=False):
  191. # Inference from various sources. For height=640, width=1280, RGB images example inputs are:
  192. # filename: imgs = 'data/images/zidane.jpg'
  193. # URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
  194. # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
  195. # PIL: = Image.open('image.jpg') # HWC x(640,1280,3)
  196. # numpy: = np.zeros((640,1280,3)) # HWC
  197. # torch: = torch.zeros(16,3,320,640) # BCHW (scaled to size=640, 0-1 values)
  198. # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
  199. t = [time_synchronized()]
  200. p = next(self.model.parameters()) # for device and type
  201. if isinstance(imgs, torch.Tensor): # torch
  202. with amp.autocast(enabled=p.device.type != 'cpu'):
  203. return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference
  204. # Pre-process
  205. n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images
  206. shape0, shape1, files = [], [], [] # image and inference shapes, filenames
  207. for i, im in enumerate(imgs):
  208. f = f'image{i}' # filename
  209. if isinstance(im, str): # filename or uri
  210. im, f = np.asarray(Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im)), im
  211. elif isinstance(im, Image.Image): # PIL Image
  212. im, f = np.asarray(im), getattr(im, 'filename', f) or f
  213. files.append(Path(f).with_suffix('.jpg').name)
  214. if im.shape[0] < 5: # image in CHW
  215. im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
  216. im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input
  217. s = im.shape[:2] # HWC
  218. shape0.append(s) # image shape
  219. g = (size / max(s)) # gain
  220. shape1.append([y * g for y in s])
  221. imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
  222. shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape
  223. x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
  224. x = np.stack(x, 0) if n > 1 else x[0][None] # stack
  225. x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
  226. x = torch.from_numpy(x).to(p.device).type_as(p) / 255. # uint8 to fp16/32
  227. t.append(time_synchronized())
  228. with amp.autocast(enabled=p.device.type != 'cpu'):
  229. # Inference
  230. y = self.model(x, augment, profile)[0] # forward
  231. t.append(time_synchronized())
  232. # Post-process
  233. y = non_max_suppression(y, self.conf, iou_thres=self.iou, classes=self.classes, max_det=self.max_det) # NMS
  234. for i in range(n):
  235. scale_coords(shape1, y[i][:, :4], shape0[i])
  236. t.append(time_synchronized())
  237. return Detections(imgs, y, files, t, self.names, x.shape)
  238. class Detections:
  239. # detections class for YOLOv5 inference results
  240. def __init__(self, imgs, pred, files, times=None, names=None, shape=None):
  241. super(Detections, self).__init__()
  242. d = pred[0].device # device
  243. gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations
  244. self.imgs = imgs # list of images as numpy arrays
  245. self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
  246. self.names = names # class names
  247. self.files = files # image filenames
  248. self.xyxy = pred # xyxy pixels
  249. self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
  250. self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
  251. self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
  252. self.n = len(self.pred) # number of images (batch size)
  253. self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms)
  254. self.s = shape # inference BCHW shape
  255. def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')):
  256. for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
  257. str = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} '
  258. if pred is not None:
  259. for c in pred[:, -1].unique():
  260. n = (pred[:, -1] == c).sum() # detections per class
  261. str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
  262. if show or save or render or crop:
  263. for *box, conf, cls in pred: # xyxy, confidence, class
  264. label = f'{self.names[int(cls)]} {conf:.2f}'
  265. if crop:
  266. save_one_box(box, im, file=save_dir / 'crops' / self.names[int(cls)] / self.files[i])
  267. else: # all others
  268. plot_one_box(box, im, label=label, color=colors(cls))
  269. im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np
  270. if pprint:
  271. print(str.rstrip(', '))
  272. if show:
  273. im.show(self.files[i]) # show
  274. if save:
  275. f = self.files[i]
  276. im.save(save_dir / f) # save
  277. print(f"{'Saved' * (i == 0)} {f}", end=',' if i < self.n - 1 else f' to {save_dir}\n')
  278. if render:
  279. self.imgs[i] = np.asarray(im)
  280. def print(self):
  281. self.display(pprint=True) # print results
  282. print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % self.t)
  283. def show(self):
  284. self.display(show=True) # show results
  285. def save(self, save_dir='runs/hub/exp'):
  286. save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp', mkdir=True) # increment save_dir
  287. self.display(save=True, save_dir=save_dir) # save results
  288. def crop(self, save_dir='runs/hub/exp'):
  289. save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp', mkdir=True) # increment save_dir
  290. self.display(crop=True, save_dir=save_dir) # crop results
  291. print(f'Saved results to {save_dir}\n')
  292. def render(self):
  293. self.display(render=True) # render results
  294. return self.imgs
  295. def pandas(self):
  296. # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])
  297. new = copy(self) # return copy
  298. ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name' # xyxy columns
  299. cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns
  300. for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]):
  301. a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # update
  302. setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])
  303. return new
  304. def tolist(self):
  305. # return a list of Detections objects, i.e. 'for result in results.tolist():'
  306. x = [Detections([self.imgs[i]], [self.pred[i]], self.names, self.s) for i in range(self.n)]
  307. for d in x:
  308. for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
  309. setattr(d, k, getattr(d, k)[0]) # pop out of list
  310. return x
  311. def __len__(self):
  312. return self.n
  313. class Classify(nn.Module):
  314. # Classification head, i.e. x(b,c1,20,20) to x(b,c2)
  315. def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups
  316. super(Classify, self).__init__()
  317. self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1)
  318. self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g) # to x(b,c2,1,1)
  319. self.flat = nn.Flatten()
  320. def forward(self, x):
  321. z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list
  322. return self.flat(self.conv(z)) # flatten to x(b,c2)