Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

1221 line
50KB

  1. import glob
  2. import math
  3. import os
  4. import random
  5. import shutil
  6. import subprocess
  7. import time
  8. from copy import copy
  9. from pathlib import Path
  10. from sys import platform
  11. import cv2
  12. import matplotlib
  13. import matplotlib.pyplot as plt
  14. import numpy as np
  15. import torch
  16. import torch.nn as nn
  17. import torchvision
  18. import yaml
  19. from scipy.signal import butter, filtfilt
  20. from tqdm import tqdm
  21. from . import torch_utils #  torch_utils, google_utils
  22. # Set printoptions
  23. torch.set_printoptions(linewidth=320, precision=5, profile='long')
  24. np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5
  25. matplotlib.rc('font', **{'size': 11})
  26. # Prevent OpenCV from multithreading (to use PyTorch DataLoader)
  27. cv2.setNumThreads(0)
  28. def init_seeds(seed=0):
  29. random.seed(seed)
  30. np.random.seed(seed)
  31. torch_utils.init_seeds(seed=seed)
  32. def get_latest_run(search_dir='./runs'):
  33. # Return path to most recent 'last.pt' in /runs (i.e. to --resume from)
  34. last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True)
  35. return max(last_list, key=os.path.getctime)
  36. def check_git_status():
  37. # Suggest 'git pull' if repo is out of date
  38. if platform in ['linux', 'darwin'] and not os.path.isfile('/.dockerenv'):
  39. s = subprocess.check_output('if [ -d .git ]; then git fetch && git status -uno; fi', shell=True).decode('utf-8')
  40. if 'Your branch is behind' in s:
  41. print(s[s.find('Your branch is behind'):s.find('\n\n')] + '\n')
  42. def check_img_size(img_size, s=32):
  43. # Verify img_size is a multiple of stride s
  44. new_size = make_divisible(img_size, int(s)) # ceil gs-multiple
  45. if new_size != img_size:
  46. print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size))
  47. return new_size
  48. def check_anchors(dataset, model, thr=4.0, imgsz=640):
  49. # Check anchor fit to data, recompute if necessary
  50. print('\nAnalyzing anchors... ', end='')
  51. m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect()
  52. shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True)
  53. scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale
  54. wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh
  55. def metric(k): # compute metric
  56. r = wh[:, None] / k[None]
  57. x = torch.min(r, 1. / r).min(2)[0] # ratio metric
  58. best = x.max(1)[0] # best_x
  59. return (best > 1. / thr).float().mean() #  best possible recall
  60. bpr = metric(m.anchor_grid.clone().cpu().view(-1, 2))
  61. print('Best Possible Recall (BPR) = %.4f' % bpr, end='')
  62. if bpr < 0.99: # threshold to recompute
  63. print('. Attempting to generate improved anchors, please wait...' % bpr)
  64. na = m.anchor_grid.numel() // 2 # number of anchors
  65. new_anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
  66. new_bpr = metric(new_anchors.reshape(-1, 2))
  67. if new_bpr > bpr: # replace anchors
  68. new_anchors = torch.tensor(new_anchors, device=m.anchors.device).type_as(m.anchors)
  69. m.anchor_grid[:] = new_anchors.clone().view_as(m.anchor_grid) # for inference
  70. m.anchors[:] = new_anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss
  71. check_anchor_order(m)
  72. print('New anchors saved to model. Update model *.yaml to use these anchors in the future.')
  73. else:
  74. print('Original anchors better than new anchors. Proceeding with original anchors.')
  75. print('') # newline
  76. def check_anchor_order(m):
  77. # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary
  78. a = m.anchor_grid.prod(-1).view(-1) # anchor area
  79. da = a[-1] - a[0] # delta a
  80. ds = m.stride[-1] - m.stride[0] # delta s
  81. if da.sign() != ds.sign(): # same order
  82. m.anchors[:] = m.anchors.flip(0)
  83. m.anchor_grid[:] = m.anchor_grid.flip(0)
  84. def check_file(file):
  85. # Searches for file if not found locally
  86. if os.path.isfile(file):
  87. return file
  88. else:
  89. files = glob.glob('./**/' + file, recursive=True) # find file
  90. assert len(files), 'File Not Found: %s' % file # assert file was found
  91. return files[0] # return first file if multiple found
  92. def make_divisible(x, divisor):
  93. # Returns x evenly divisble by divisor
  94. return math.ceil(x / divisor) * divisor
  95. def labels_to_class_weights(labels, nc=80):
  96. # Get class weights (inverse frequency) from training labels
  97. if labels[0] is None: # no labels loaded
  98. return torch.Tensor()
  99. labels = np.concatenate(labels, 0) # labels.shape = (866643, 5) for COCO
  100. classes = labels[:, 0].astype(np.int) # labels = [class xywh]
  101. weights = np.bincount(classes, minlength=nc) # occurences per class
  102. # Prepend gridpoint count (for uCE trianing)
  103. # gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum() # gridpoints per image
  104. # weights = np.hstack([gpi * len(labels) - weights.sum() * 9, weights * 9]) ** 0.5 # prepend gridpoints to start
  105. weights[weights == 0] = 1 # replace empty bins with 1
  106. weights = 1 / weights # number of targets per class
  107. weights /= weights.sum() # normalize
  108. return torch.from_numpy(weights)
  109. def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):
  110. # Produces image weights based on class mAPs
  111. n = len(labels)
  112. class_counts = np.array([np.bincount(labels[i][:, 0].astype(np.int), minlength=nc) for i in range(n)])
  113. image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1)
  114. # index = random.choices(range(n), weights=image_weights, k=1) # weight image sample
  115. return image_weights
  116. def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper)
  117. # https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
  118. # a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n')
  119. # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
  120. # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco
  121. # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet
  122. x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
  123. 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
  124. 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
  125. return x
  126. def xyxy2xywh(x):
  127. # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
  128. y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)
  129. y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center
  130. y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center
  131. y[:, 2] = x[:, 2] - x[:, 0] # width
  132. y[:, 3] = x[:, 3] - x[:, 1] # height
  133. return y
  134. def xywh2xyxy(x):
  135. # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
  136. y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)
  137. y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
  138. y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y
  139. y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x
  140. y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y
  141. return y
  142. def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
  143. # Rescale coords (xyxy) from img1_shape to img0_shape
  144. if ratio_pad is None: # calculate from img0_shape
  145. gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
  146. pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
  147. else:
  148. gain = ratio_pad[0][0]
  149. pad = ratio_pad[1]
  150. coords[:, [0, 2]] -= pad[0] # x padding
  151. coords[:, [1, 3]] -= pad[1] # y padding
  152. coords[:, :4] /= gain
  153. clip_coords(coords, img0_shape)
  154. return coords
  155. def clip_coords(boxes, img_shape):
  156. # Clip bounding xyxy bounding boxes to image shape (height, width)
  157. boxes[:, 0].clamp_(0, img_shape[1]) # x1
  158. boxes[:, 1].clamp_(0, img_shape[0]) # y1
  159. boxes[:, 2].clamp_(0, img_shape[1]) # x2
  160. boxes[:, 3].clamp_(0, img_shape[0]) # y2
  161. def ap_per_class(tp, conf, pred_cls, target_cls):
  162. """ Compute the average precision, given the recall and precision curves.
  163. Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
  164. # Arguments
  165. tp: True positives (nparray, nx1 or nx10).
  166. conf: Objectness value from 0-1 (nparray).
  167. pred_cls: Predicted object classes (nparray).
  168. target_cls: True object classes (nparray).
  169. # Returns
  170. The average precision as computed in py-faster-rcnn.
  171. """
  172. # Sort by objectness
  173. i = np.argsort(-conf)
  174. tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]
  175. # Find unique classes
  176. unique_classes = np.unique(target_cls)
  177. # Create Precision-Recall curve and compute AP for each class
  178. pr_score = 0.1 # score to evaluate P and R https://github.com/ultralytics/yolov3/issues/898
  179. s = [unique_classes.shape[0], tp.shape[1]] # number class, number iou thresholds (i.e. 10 for mAP0.5...0.95)
  180. ap, p, r = np.zeros(s), np.zeros(s), np.zeros(s)
  181. for ci, c in enumerate(unique_classes):
  182. i = pred_cls == c
  183. n_gt = (target_cls == c).sum() # Number of ground truth objects
  184. n_p = i.sum() # Number of predicted objects
  185. if n_p == 0 or n_gt == 0:
  186. continue
  187. else:
  188. # Accumulate FPs and TPs
  189. fpc = (1 - tp[i]).cumsum(0)
  190. tpc = tp[i].cumsum(0)
  191. # Recall
  192. recall = tpc / (n_gt + 1e-16) # recall curve
  193. r[ci] = np.interp(-pr_score, -conf[i], recall[:, 0]) # r at pr_score, negative x, xp because xp decreases
  194. # Precision
  195. precision = tpc / (tpc + fpc) # precision curve
  196. p[ci] = np.interp(-pr_score, -conf[i], precision[:, 0]) # p at pr_score
  197. # AP from recall-precision curve
  198. for j in range(tp.shape[1]):
  199. ap[ci, j] = compute_ap(recall[:, j], precision[:, j])
  200. # Plot
  201. # fig, ax = plt.subplots(1, 1, figsize=(5, 5))
  202. # ax.plot(recall, precision)
  203. # ax.set_xlabel('Recall')
  204. # ax.set_ylabel('Precision')
  205. # ax.set_xlim(0, 1.01)
  206. # ax.set_ylim(0, 1.01)
  207. # fig.tight_layout()
  208. # fig.savefig('PR_curve.png', dpi=300)
  209. # Compute F1 score (harmonic mean of precision and recall)
  210. f1 = 2 * p * r / (p + r + 1e-16)
  211. return p, r, ap, f1, unique_classes.astype('int32')
  212. def compute_ap(recall, precision):
  213. """ Compute the average precision, given the recall and precision curves.
  214. Source: https://github.com/rbgirshick/py-faster-rcnn.
  215. # Arguments
  216. recall: The recall curve (list).
  217. precision: The precision curve (list).
  218. # Returns
  219. The average precision as computed in py-faster-rcnn.
  220. """
  221. # Append sentinel values to beginning and end
  222. mrec = np.concatenate(([0.], recall, [min(recall[-1] + 1E-3, 1.)]))
  223. mpre = np.concatenate(([0.], precision, [0.]))
  224. # Compute the precision envelope
  225. mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
  226. # Integrate area under curve
  227. method = 'interp' # methods: 'continuous', 'interp'
  228. if method == 'interp':
  229. x = np.linspace(0, 1, 101) # 101-point interp (COCO)
  230. ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate
  231. else: # 'continuous'
  232. i = np.where(mrec[1:] != mrec[:-1])[0] # points where x axis (recall) changes
  233. ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) # area under curve
  234. return ap
  235. def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False):
  236. # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
  237. box2 = box2.t()
  238. # Get the coordinates of bounding boxes
  239. if x1y1x2y2: # x1, y1, x2, y2 = box1
  240. b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
  241. b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
  242. else: # transform from xywh to xyxy
  243. b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
  244. b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
  245. b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
  246. b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2
  247. # Intersection area
  248. inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
  249. (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
  250. # Union Area
  251. w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1
  252. w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1
  253. union = (w1 * h1 + 1e-16) + w2 * h2 - inter
  254. iou = inter / union # iou
  255. if GIoU or DIoU or CIoU:
  256. cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width
  257. ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
  258. if GIoU: # Generalized IoU https://arxiv.org/pdf/1902.09630.pdf
  259. c_area = cw * ch + 1e-16 # convex area
  260. return iou - (c_area - union) / c_area # GIoU
  261. if DIoU or CIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
  262. # convex diagonal squared
  263. c2 = cw ** 2 + ch ** 2 + 1e-16
  264. # centerpoint distance squared
  265. rho2 = ((b2_x1 + b2_x2) - (b1_x1 + b1_x2)) ** 2 / 4 + ((b2_y1 + b2_y2) - (b1_y1 + b1_y2)) ** 2 / 4
  266. if DIoU:
  267. return iou - rho2 / c2 # DIoU
  268. elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
  269. v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
  270. with torch.no_grad():
  271. alpha = v / (1 - iou + v)
  272. return iou - (rho2 / c2 + v * alpha) # CIoU
  273. return iou
  274. def box_iou(box1, box2):
  275. # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
  276. """
  277. Return intersection-over-union (Jaccard index) of boxes.
  278. Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
  279. Arguments:
  280. box1 (Tensor[N, 4])
  281. box2 (Tensor[M, 4])
  282. Returns:
  283. iou (Tensor[N, M]): the NxM matrix containing the pairwise
  284. IoU values for every element in boxes1 and boxes2
  285. """
  286. def box_area(box):
  287. # box = 4xn
  288. return (box[2] - box[0]) * (box[3] - box[1])
  289. area1 = box_area(box1.t())
  290. area2 = box_area(box2.t())
  291. # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
  292. inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
  293. return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)
  294. def wh_iou(wh1, wh2):
  295. # Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2
  296. wh1 = wh1[:, None] # [N,1,2]
  297. wh2 = wh2[None] # [1,M,2]
  298. inter = torch.min(wh1, wh2).prod(2) # [N,M]
  299. return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter)
  300. class FocalLoss(nn.Module):
  301. # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
  302. def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
  303. super(FocalLoss, self).__init__()
  304. self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
  305. self.gamma = gamma
  306. self.alpha = alpha
  307. self.reduction = loss_fcn.reduction
  308. self.loss_fcn.reduction = 'none' # required to apply FL to each element
  309. def forward(self, pred, true):
  310. loss = self.loss_fcn(pred, true)
  311. # p_t = torch.exp(-loss)
  312. # loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability
  313. # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
  314. pred_prob = torch.sigmoid(pred) # prob from logits
  315. p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
  316. alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
  317. modulating_factor = (1.0 - p_t) ** self.gamma
  318. loss *= alpha_factor * modulating_factor
  319. if self.reduction == 'mean':
  320. return loss.mean()
  321. elif self.reduction == 'sum':
  322. return loss.sum()
  323. else: # 'none'
  324. return loss
  325. def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441
  326. # return positive, negative label smoothing BCE targets
  327. return 1.0 - 0.5 * eps, 0.5 * eps
  328. class BCEBlurWithLogitsLoss(nn.Module):
  329. # BCEwithLogitLoss() with reduced missing label effects.
  330. def __init__(self, alpha=0.05):
  331. super(BCEBlurWithLogitsLoss, self).__init__()
  332. self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none') # must be nn.BCEWithLogitsLoss()
  333. self.alpha = alpha
  334. def forward(self, pred, true):
  335. loss = self.loss_fcn(pred, true)
  336. pred = torch.sigmoid(pred) # prob from logits
  337. dx = pred - true # reduce only missing label effects
  338. # dx = (pred - true).abs() # reduce missing label and false label effects
  339. alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4))
  340. loss *= alpha_factor
  341. return loss.mean()
  342. def compute_loss(p, targets, model): # predictions, targets, model
  343. ft = torch.cuda.FloatTensor if p[0].is_cuda else torch.Tensor
  344. lcls, lbox, lobj = ft([0]), ft([0]), ft([0])
  345. tcls, tbox, indices, anchors = build_targets(p, targets, model) # targets
  346. h = model.hyp # hyperparameters
  347. red = 'mean' # Loss reduction (sum or mean)
  348. # Define criteria
  349. BCEcls = nn.BCEWithLogitsLoss(pos_weight=ft([h['cls_pw']]), reduction=red)
  350. BCEobj = nn.BCEWithLogitsLoss(pos_weight=ft([h['obj_pw']]), reduction=red)
  351. # class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
  352. cp, cn = smooth_BCE(eps=0.0)
  353. # focal loss
  354. g = h['fl_gamma'] # focal loss gamma
  355. if g > 0:
  356. BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
  357. # per output
  358. nt = 0 # number of targets
  359. np = len(p) # number of outputs
  360. balance = [1.0, 1.0, 1.0]
  361. for i, pi in enumerate(p): # layer index, layer predictions
  362. b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
  363. tobj = torch.zeros_like(pi[..., 0]) # target obj
  364. nb = b.shape[0] # number of targets
  365. if nb:
  366. nt += nb # cumulative targets
  367. ps = pi[b, a, gj, gi] # prediction subset corresponding to targets
  368. # GIoU
  369. pxy = ps[:, :2].sigmoid() * 2. - 0.5
  370. pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
  371. pbox = torch.cat((pxy, pwh), 1) # predicted box
  372. giou = bbox_iou(pbox.t(), tbox[i], x1y1x2y2=False, GIoU=True) # giou(prediction, target)
  373. lbox += (1.0 - giou).sum() if red == 'sum' else (1.0 - giou).mean() # giou loss
  374. # Obj
  375. tobj[b, a, gj, gi] = (1.0 - model.gr) + model.gr * giou.detach().clamp(0).type(tobj.dtype) # giou ratio
  376. # Class
  377. if model.nc > 1: # cls loss (only if multiple classes)
  378. t = torch.full_like(ps[:, 5:], cn) # targets
  379. t[range(nb), tcls[i]] = cp
  380. lcls += BCEcls(ps[:, 5:], t) # BCE
  381. # Append targets to text file
  382. # with open('targets.txt', 'a') as file:
  383. # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
  384. lobj += BCEobj(pi[..., 4], tobj) * balance[i] # obj loss
  385. s = 3 / np # output count scaling
  386. lbox *= h['giou'] * s
  387. lobj *= h['obj'] * s
  388. lcls *= h['cls'] * s
  389. bs = tobj.shape[0] # batch size
  390. if red == 'sum':
  391. g = 3.0 # loss gain
  392. lobj *= g / bs
  393. if nt:
  394. lcls *= g / nt / model.nc
  395. lbox *= g / nt
  396. loss = lbox + lobj + lcls
  397. return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()
  398. def build_targets(p, targets, model):
  399. # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
  400. det = model.module.model[-1] if type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel) \
  401. else model.model[-1] # Detect() module
  402. na, nt = det.na, targets.shape[0] # number of anchors, targets
  403. tcls, tbox, indices, anch = [], [], [], []
  404. gain = torch.ones(6, device=targets.device) # normalized to gridspace gain
  405. off = torch.tensor([[1, 0], [0, 1], [-1, 0], [0, -1]], device=targets.device).float() # overlap offsets
  406. at = torch.arange(na).view(na, 1).repeat(1, nt) # anchor tensor, same as .repeat_interleave(nt)
  407. g = 0.5 # offset
  408. style = 'rect4'
  409. for i in range(det.nl):
  410. anchors = det.anchors[i]
  411. gain[2:] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain
  412. # Match targets to anchors
  413. a, t, offsets = [], targets * gain, 0
  414. if nt:
  415. r = t[None, :, 4:6] / anchors[:, None] # wh ratio
  416. j = torch.max(r, 1. / r).max(2)[0] < model.hyp['anchor_t'] # compare
  417. # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n) = wh_iou(anchors(3,2), gwh(n,2))
  418. a, t = at[j], t.repeat(na, 1, 1)[j] # filter
  419. # overlaps
  420. gxy = t[:, 2:4] # grid xy
  421. z = torch.zeros_like(gxy)
  422. if style == 'rect2':
  423. j, k = ((gxy % 1. < g) & (gxy > 1.)).T
  424. a, t = torch.cat((a, a[j], a[k]), 0), torch.cat((t, t[j], t[k]), 0)
  425. offsets = torch.cat((z, z[j] + off[0], z[k] + off[1]), 0) * g
  426. elif style == 'rect4':
  427. j, k = ((gxy % 1. < g) & (gxy > 1.)).T
  428. l, m = ((gxy % 1. > (1 - g)) & (gxy < (gain[[2, 3]] - 1.))).T
  429. a, t = torch.cat((a, a[j], a[k], a[l], a[m]), 0), torch.cat((t, t[j], t[k], t[l], t[m]), 0)
  430. offsets = torch.cat((z, z[j] + off[0], z[k] + off[1], z[l] + off[2], z[m] + off[3]), 0) * g
  431. # Define
  432. b, c = t[:, :2].long().T # image, class
  433. gxy = t[:, 2:4] # grid xy
  434. gwh = t[:, 4:6] # grid wh
  435. gij = (gxy - offsets).long()
  436. gi, gj = gij.T # grid xy indices
  437. # Append
  438. indices.append((b, a, gj, gi)) # image, anchor, grid indices
  439. tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
  440. anch.append(anchors[a]) # anchors
  441. tcls.append(c) # class
  442. return tcls, tbox, indices, anch
  443. def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False):
  444. """Performs Non-Maximum Suppression (NMS) on inference results
  445. Returns:
  446. detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
  447. """
  448. if prediction.dtype is torch.float16:
  449. prediction = prediction.float() # to FP32
  450. nc = prediction[0].shape[1] - 5 # number of classes
  451. xc = prediction[..., 4] > conf_thres # candidates
  452. # Settings
  453. min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
  454. max_det = 300 # maximum number of detections per image
  455. time_limit = 10.0 # seconds to quit after
  456. redundant = True # require redundant detections
  457. multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img)
  458. t = time.time()
  459. output = [None] * prediction.shape[0]
  460. for xi, x in enumerate(prediction): # image index, image inference
  461. # Apply constraints
  462. # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
  463. x = x[xc[xi]] # confidence
  464. # If none remain process next image
  465. if not x.shape[0]:
  466. continue
  467. # Compute conf
  468. x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf
  469. # Box (center x, center y, width, height) to (x1, y1, x2, y2)
  470. box = xywh2xyxy(x[:, :4])
  471. # Detections matrix nx6 (xyxy, conf, cls)
  472. if multi_label:
  473. i, j = (x[:, 5:] > conf_thres).nonzero().t()
  474. x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
  475. else: # best class only
  476. conf, j = x[:, 5:].max(1, keepdim=True)
  477. x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
  478. # Filter by class
  479. if classes:
  480. x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
  481. # Apply finite constraint
  482. # if not torch.isfinite(x).all():
  483. # x = x[torch.isfinite(x).all(1)]
  484. # If none remain process next image
  485. n = x.shape[0] # number of boxes
  486. if not n:
  487. continue
  488. # Sort by confidence
  489. # x = x[x[:, 4].argsort(descending=True)]
  490. # Batched NMS
  491. c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
  492. boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
  493. i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
  494. if i.shape[0] > max_det: # limit detections
  495. i = i[:max_det]
  496. if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
  497. try: # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
  498. iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
  499. weights = iou * scores[None] # box weights
  500. x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
  501. if redundant:
  502. i = i[iou.sum(1) > 1] # require redundancy
  503. except: # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139
  504. print(x, i, x.shape, i.shape)
  505. pass
  506. output[xi] = x[i]
  507. if (time.time() - t) > time_limit:
  508. break # time limit exceeded
  509. return output
  510. def strip_optimizer(f='weights/best.pt'): # from utils.utils import *; strip_optimizer()
  511. # Strip optimizer from *.pt files for lighter files (reduced by 1/2 size)
  512. x = torch.load(f, map_location=torch.device('cpu'))
  513. x['optimizer'] = None
  514. x['model'].half() # to FP16
  515. torch.save(x, f)
  516. print('Optimizer stripped from %s, %.1fMB' % (f, os.path.getsize(f) / 1E6))
  517. def create_pretrained(f='weights/best.pt', s='weights/pretrained.pt'): # from utils.utils import *; create_pretrained()
  518. # create pretrained checkpoint 's' from 'f' (create_pretrained(x, x) for x in glob.glob('./*.pt'))
  519. x = torch.load(f, map_location=torch.device('cpu'))
  520. x['optimizer'] = None
  521. x['training_results'] = None
  522. x['epoch'] = -1
  523. x['model'].half() # to FP16
  524. for p in x['model'].parameters():
  525. p.requires_grad = True
  526. torch.save(x, s)
  527. print('%s saved as pretrained checkpoint %s, %.1fMB' % (f, s, os.path.getsize(s) / 1E6))
  528. def coco_class_count(path='../coco/labels/train2014/'):
  529. # Histogram of occurrences per class
  530. nc = 80 # number classes
  531. x = np.zeros(nc, dtype='int32')
  532. files = sorted(glob.glob('%s/*.*' % path))
  533. for i, file in enumerate(files):
  534. labels = np.loadtxt(file, dtype=np.float32).reshape(-1, 5)
  535. x += np.bincount(labels[:, 0].astype('int32'), minlength=nc)
  536. print(i, len(files))
  537. def coco_only_people(path='../coco/labels/train2017/'): # from utils.utils import *; coco_only_people()
  538. # Find images with only people
  539. files = sorted(glob.glob('%s/*.*' % path))
  540. for i, file in enumerate(files):
  541. labels = np.loadtxt(file, dtype=np.float32).reshape(-1, 5)
  542. if all(labels[:, 0] == 0):
  543. print(labels.shape[0], file)
  544. def crop_images_random(path='../images/', scale=0.50): # from utils.utils import *; crop_images_random()
  545. # crops images into random squares up to scale fraction
  546. # WARNING: overwrites images!
  547. for file in tqdm(sorted(glob.glob('%s/*.*' % path))):
  548. img = cv2.imread(file) # BGR
  549. if img is not None:
  550. h, w = img.shape[:2]
  551. # create random mask
  552. a = 30 # minimum size (pixels)
  553. mask_h = random.randint(a, int(max(a, h * scale))) # mask height
  554. mask_w = mask_h # mask width
  555. # box
  556. xmin = max(0, random.randint(0, w) - mask_w // 2)
  557. ymin = max(0, random.randint(0, h) - mask_h // 2)
  558. xmax = min(w, xmin + mask_w)
  559. ymax = min(h, ymin + mask_h)
  560. # apply random color mask
  561. cv2.imwrite(file, img[ymin:ymax, xmin:xmax])
  562. def coco_single_class_labels(path='../coco/labels/train2014/', label_class=43):
  563. # Makes single-class coco datasets. from utils.utils import *; coco_single_class_labels()
  564. if os.path.exists('new/'):
  565. shutil.rmtree('new/') # delete output folder
  566. os.makedirs('new/') # make new output folder
  567. os.makedirs('new/labels/')
  568. os.makedirs('new/images/')
  569. for file in tqdm(sorted(glob.glob('%s/*.*' % path))):
  570. with open(file, 'r') as f:
  571. labels = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32)
  572. i = labels[:, 0] == label_class
  573. if any(i):
  574. img_file = file.replace('labels', 'images').replace('txt', 'jpg')
  575. labels[:, 0] = 0 # reset class to 0
  576. with open('new/images.txt', 'a') as f: # add image to dataset list
  577. f.write(img_file + '\n')
  578. with open('new/labels/' + Path(file).name, 'a') as f: # write label
  579. for l in labels[i]:
  580. f.write('%g %.6f %.6f %.6f %.6f\n' % tuple(l))
  581. shutil.copyfile(src=img_file, dst='new/images/' + Path(file).name.replace('txt', 'jpg')) # copy images
  582. def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
  583. """ Creates kmeans-evolved anchors from training dataset
  584. Arguments:
  585. path: path to dataset *.yaml, or a loaded dataset
  586. n: number of anchors
  587. img_size: image size used for training
  588. thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0
  589. gen: generations to evolve anchors using genetic algorithm
  590. Return:
  591. k: kmeans evolved anchors
  592. Usage:
  593. from utils.utils import *; _ = kmean_anchors()
  594. """
  595. thr = 1. / thr
  596. def metric(k, wh): # compute metrics
  597. r = wh[:, None] / k[None]
  598. x = torch.min(r, 1. / r).min(2)[0] # ratio metric
  599. # x = wh_iou(wh, torch.tensor(k)) # iou metric
  600. return x, x.max(1)[0] # x, best_x
  601. def fitness(k): # mutation fitness
  602. _, best = metric(torch.tensor(k, dtype=torch.float32), wh)
  603. return (best * (best > thr).float()).mean() # fitness
  604. def print_results(k):
  605. k = k[np.argsort(k.prod(1))] # sort small to large
  606. x, best = metric(k, wh0)
  607. bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr
  608. print('thr=%.2f: %.4f best possible recall, %.2f anchors past thr' % (thr, bpr, aat))
  609. print('n=%g, img_size=%s, metric_all=%.3f/%.3f-mean/best, past_thr=%.3f-mean: ' %
  610. (n, img_size, x.mean(), best.mean(), x[x > thr].mean()), end='')
  611. for i, x in enumerate(k):
  612. print('%i,%i' % (round(x[0]), round(x[1])), end=', ' if i < len(k) - 1 else '\n') # use in *.cfg
  613. return k
  614. if isinstance(path, str): # *.yaml file
  615. with open(path) as f:
  616. data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict
  617. from utils.datasets import LoadImagesAndLabels
  618. dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
  619. else:
  620. dataset = path # dataset
  621. # Get label wh
  622. shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True)
  623. wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh
  624. # Filter
  625. i = (wh0 < 3.0).any(1).sum()
  626. if i:
  627. print('WARNING: Extremely small objects found. '
  628. '%g of %g labels are < 3 pixels in width or height.' % (i, len(wh0)))
  629. wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
  630. # Kmeans calculation
  631. from scipy.cluster.vq import kmeans
  632. print('Running kmeans for %g anchors on %g points...' % (n, len(wh)))
  633. s = wh.std(0) # sigmas for whitening
  634. k, dist = kmeans(wh / s, n, iter=30) # points, mean distance
  635. k *= s
  636. wh = torch.tensor(wh, dtype=torch.float32) # filtered
  637. wh0 = torch.tensor(wh0, dtype=torch.float32) # unflitered
  638. k = print_results(k)
  639. # Plot
  640. # k, d = [None] * 20, [None] * 20
  641. # for i in tqdm(range(1, 21)):
  642. # k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance
  643. # fig, ax = plt.subplots(1, 2, figsize=(14, 7))
  644. # ax = ax.ravel()
  645. # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.')
  646. # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh
  647. # ax[0].hist(wh[wh[:, 0]<100, 0],400)
  648. # ax[1].hist(wh[wh[:, 1]<100, 1],400)
  649. # fig.tight_layout()
  650. # fig.savefig('wh.png', dpi=200)
  651. # Evolve
  652. npr = np.random
  653. f, sh, mp, s = fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma
  654. pbar = tqdm(range(gen), desc='Evolving anchors with Genetic Algorithm') # progress bar
  655. for _ in pbar:
  656. v = np.ones(sh)
  657. while (v == 1).all(): # mutate until a change occurs (prevent duplicates)
  658. v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0)
  659. kg = (k.copy() * v).clip(min=2.0)
  660. fg = fitness(kg)
  661. if fg > f:
  662. f, k = fg, kg.copy()
  663. pbar.desc = 'Evolving anchors with Genetic Algorithm: fitness = %.4f' % f
  664. if verbose:
  665. print_results(k)
  666. return print_results(k)
  667. def print_mutation(hyp, results, bucket=''):
  668. # Print mutation results to evolve.txt (for use with train.py --evolve)
  669. a = '%10s' * len(hyp) % tuple(hyp.keys()) # hyperparam keys
  670. b = '%10.3g' * len(hyp) % tuple(hyp.values()) # hyperparam values
  671. c = '%10.4g' * len(results) % results # results (P, R, mAP, F1, test_loss)
  672. print('\n%s\n%s\nEvolved fitness: %s\n' % (a, b, c))
  673. if bucket:
  674. os.system('gsutil cp gs://%s/evolve.txt .' % bucket) # download evolve.txt
  675. with open('evolve.txt', 'a') as f: # append result
  676. f.write(c + b + '\n')
  677. x = np.unique(np.loadtxt('evolve.txt', ndmin=2), axis=0) # load unique rows
  678. np.savetxt('evolve.txt', x[np.argsort(-fitness(x))], '%10.3g') # save sort by fitness
  679. if bucket:
  680. os.system('gsutil cp evolve.txt gs://%s' % bucket) # upload evolve.txt
  681. def apply_classifier(x, model, img, im0):
  682. # applies a second stage classifier to yolo outputs
  683. im0 = [im0] if isinstance(im0, np.ndarray) else im0
  684. for i, d in enumerate(x): # per image
  685. if d is not None and len(d):
  686. d = d.clone()
  687. # Reshape and pad cutouts
  688. b = xyxy2xywh(d[:, :4]) # boxes
  689. b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # rectangle to square
  690. b[:, 2:] = b[:, 2:] * 1.3 + 30 # pad
  691. d[:, :4] = xywh2xyxy(b).long()
  692. # Rescale boxes from img_size to im0 size
  693. scale_coords(img.shape[2:], d[:, :4], im0[i].shape)
  694. # Classes
  695. pred_cls1 = d[:, 5].long()
  696. ims = []
  697. for j, a in enumerate(d): # per item
  698. cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])]
  699. im = cv2.resize(cutout, (224, 224)) # BGR
  700. # cv2.imwrite('test%i.jpg' % j, cutout)
  701. im = im[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
  702. im = np.ascontiguousarray(im, dtype=np.float32) # uint8 to float32
  703. im /= 255.0 # 0 - 255 to 0.0 - 1.0
  704. ims.append(im)
  705. pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1) # classifier prediction
  706. x[i] = x[i][pred_cls1 == pred_cls2] # retain matching class detections
  707. return x
  708. def fitness(x):
  709. # Returns fitness (for use with results.txt or evolve.txt)
  710. w = [0.0, 0.0, 0.1, 0.9] # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
  711. return (x[:, :4] * w).sum(1)
  712. def output_to_target(output, width, height):
  713. # Convert model output to target format [batch_id, class_id, x, y, w, h, conf]
  714. if isinstance(output, torch.Tensor):
  715. output = output.cpu().numpy()
  716. targets = []
  717. for i, o in enumerate(output):
  718. if o is not None:
  719. for pred in o:
  720. box = pred[:4]
  721. w = (box[2] - box[0]) / width
  722. h = (box[3] - box[1]) / height
  723. x = box[0] / width + w / 2
  724. y = box[1] / height + h / 2
  725. conf = pred[4]
  726. cls = int(pred[5])
  727. targets.append([i, cls, x, y, w, h, conf])
  728. return np.array(targets)
  729. def increment_dir(dir, comment=''):
  730. # Increments a directory runs/exp1 --> runs/exp2_comment
  731. n = 0 # number
  732. dir = str(Path(dir)) # os-agnostic
  733. d = sorted(glob.glob(dir + '*')) # directories
  734. if len(d):
  735. n = max([int(x[len(dir):x.find('_') if '_' in x else None]) for x in d]) + 1 # increment
  736. return dir + str(n) + ('_' + comment if comment else '')
  737. # Plotting functions ---------------------------------------------------------------------------------------------------
  738. def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
  739. # https://stackoverflow.com/questions/28536191/how-to-filter-smooth-with-scipy-numpy
  740. def butter_lowpass(cutoff, fs, order):
  741. nyq = 0.5 * fs
  742. normal_cutoff = cutoff / nyq
  743. b, a = butter(order, normal_cutoff, btype='low', analog=False)
  744. return b, a
  745. b, a = butter_lowpass(cutoff, fs, order=order)
  746. return filtfilt(b, a, data) # forward-backward filter
  747. def plot_one_box(x, img, color=None, label=None, line_thickness=None):
  748. # Plots one bounding box on image img
  749. tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness
  750. color = color or [random.randint(0, 255) for _ in range(3)]
  751. c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
  752. cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
  753. if label:
  754. tf = max(tl - 1, 1) # font thickness
  755. t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
  756. c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
  757. cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled
  758. cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
  759. def plot_wh_methods(): # from utils.utils import *; plot_wh_methods()
  760. # Compares the two methods for width-height anchor multiplication
  761. # https://github.com/ultralytics/yolov3/issues/168
  762. x = np.arange(-4.0, 4.0, .1)
  763. ya = np.exp(x)
  764. yb = torch.sigmoid(torch.from_numpy(x)).numpy() * 2
  765. fig = plt.figure(figsize=(6, 3), dpi=150)
  766. plt.plot(x, ya, '.-', label='yolo method')
  767. plt.plot(x, yb ** 2, '.-', label='^2 power method')
  768. plt.plot(x, yb ** 2.5, '.-', label='^2.5 power method')
  769. plt.xlim(left=-4, right=4)
  770. plt.ylim(bottom=0, top=6)
  771. plt.xlabel('input')
  772. plt.ylabel('output')
  773. plt.legend()
  774. fig.tight_layout()
  775. fig.savefig('comparison.png', dpi=200)
  776. def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max_size=640, max_subplots=16):
  777. tl = 3 # line thickness
  778. tf = max(tl - 1, 1) # font thickness
  779. if os.path.isfile(fname): # do not overwrite
  780. return None
  781. if isinstance(images, torch.Tensor):
  782. images = images.cpu().float().numpy()
  783. if isinstance(targets, torch.Tensor):
  784. targets = targets.cpu().numpy()
  785. # un-normalise
  786. if np.max(images[0]) <= 1:
  787. images *= 255
  788. bs, _, h, w = images.shape # batch size, _, height, width
  789. bs = min(bs, max_subplots) # limit plot images
  790. ns = np.ceil(bs ** 0.5) # number of subplots (square)
  791. # Check if we should resize
  792. scale_factor = max_size / max(h, w)
  793. if scale_factor < 1:
  794. h = math.ceil(scale_factor * h)
  795. w = math.ceil(scale_factor * w)
  796. # Empty array for output
  797. mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8)
  798. # Fix class - colour map
  799. prop_cycle = plt.rcParams['axes.prop_cycle']
  800. # https://stackoverflow.com/questions/51350872/python-from-color-name-to-rgb
  801. hex2rgb = lambda h: tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
  802. color_lut = [hex2rgb(h) for h in prop_cycle.by_key()['color']]
  803. for i, img in enumerate(images):
  804. if i == max_subplots: # if last batch has fewer images than we expect
  805. break
  806. block_x = int(w * (i // ns))
  807. block_y = int(h * (i % ns))
  808. img = img.transpose(1, 2, 0)
  809. if scale_factor < 1:
  810. img = cv2.resize(img, (w, h))
  811. mosaic[block_y:block_y + h, block_x:block_x + w, :] = img
  812. if len(targets) > 0:
  813. image_targets = targets[targets[:, 0] == i]
  814. boxes = xywh2xyxy(image_targets[:, 2:6]).T
  815. classes = image_targets[:, 1].astype('int')
  816. gt = image_targets.shape[1] == 6 # ground truth if no conf column
  817. conf = None if gt else image_targets[:, 6] # check for confidence presence (gt vs pred)
  818. boxes[[0, 2]] *= w
  819. boxes[[0, 2]] += block_x
  820. boxes[[1, 3]] *= h
  821. boxes[[1, 3]] += block_y
  822. for j, box in enumerate(boxes.T):
  823. cls = int(classes[j])
  824. color = color_lut[cls % len(color_lut)]
  825. cls = names[cls] if names else cls
  826. if gt or conf[j] > 0.3: # 0.3 conf thresh
  827. label = '%s' % cls if gt else '%s %.1f' % (cls, conf[j])
  828. plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl)
  829. # Draw image filename labels
  830. if paths is not None:
  831. label = os.path.basename(paths[i])[:40] # trim to 40 char
  832. t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
  833. cv2.putText(mosaic, label, (block_x + 5, block_y + t_size[1] + 5), 0, tl / 3, [220, 220, 220], thickness=tf,
  834. lineType=cv2.LINE_AA)
  835. # Image border
  836. cv2.rectangle(mosaic, (block_x, block_y), (block_x + w, block_y + h), (255, 255, 255), thickness=3)
  837. if fname is not None:
  838. mosaic = cv2.resize(mosaic, (int(ns * w * 0.5), int(ns * h * 0.5)), interpolation=cv2.INTER_AREA)
  839. cv2.imwrite(fname, cv2.cvtColor(mosaic, cv2.COLOR_BGR2RGB))
  840. return mosaic
  841. def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=''):
  842. # Plot LR simulating training for full epochs
  843. optimizer, scheduler = copy(optimizer), copy(scheduler) # do not modify originals
  844. y = []
  845. for _ in range(epochs):
  846. scheduler.step()
  847. y.append(optimizer.param_groups[0]['lr'])
  848. plt.plot(y, '.-', label='LR')
  849. plt.xlabel('epoch')
  850. plt.ylabel('LR')
  851. plt.grid()
  852. plt.xlim(0, epochs)
  853. plt.ylim(0)
  854. plt.tight_layout()
  855. plt.savefig(Path(save_dir) / 'LR.png', dpi=200)
  856. def plot_test_txt(): # from utils.utils import *; plot_test()
  857. # Plot test.txt histograms
  858. x = np.loadtxt('test.txt', dtype=np.float32)
  859. box = xyxy2xywh(x[:, :4])
  860. cx, cy = box[:, 0], box[:, 1]
  861. fig, ax = plt.subplots(1, 1, figsize=(6, 6), tight_layout=True)
  862. ax.hist2d(cx, cy, bins=600, cmax=10, cmin=0)
  863. ax.set_aspect('equal')
  864. plt.savefig('hist2d.png', dpi=300)
  865. fig, ax = plt.subplots(1, 2, figsize=(12, 6), tight_layout=True)
  866. ax[0].hist(cx, bins=600)
  867. ax[1].hist(cy, bins=600)
  868. plt.savefig('hist1d.png', dpi=200)
  869. def plot_targets_txt(): # from utils.utils import *; plot_targets_txt()
  870. # Plot targets.txt histograms
  871. x = np.loadtxt('targets.txt', dtype=np.float32).T
  872. s = ['x targets', 'y targets', 'width targets', 'height targets']
  873. fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)
  874. ax = ax.ravel()
  875. for i in range(4):
  876. ax[i].hist(x[i], bins=100, label='%.3g +/- %.3g' % (x[i].mean(), x[i].std()))
  877. ax[i].legend()
  878. ax[i].set_title(s[i])
  879. plt.savefig('targets.jpg', dpi=200)
  880. def plot_study_txt(f='study.txt', x=None): # from utils.utils import *; plot_study_txt()
  881. # Plot study.txt generated by test.py
  882. fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)
  883. ax = ax.ravel()
  884. fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
  885. for f in ['coco_study/study_coco_yolov5%s.txt' % x for x in ['s', 'm', 'l', 'x']]:
  886. y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
  887. x = np.arange(y.shape[1]) if x is None else np.array(x)
  888. s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_inference (ms/img)', 't_NMS (ms/img)', 't_total (ms/img)']
  889. for i in range(7):
  890. ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8)
  891. ax[i].set_title(s[i])
  892. j = y[3].argmax() + 1
  893. ax2.plot(y[6, :j], y[3, :j] * 1E2, '.-', linewidth=2, markersize=8,
  894. label=Path(f).stem.replace('study_coco_', '').replace('yolo', 'YOLO'))
  895. ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [33.5, 39.1, 42.5, 45.9, 49., 50.5],
  896. 'k.-', linewidth=2, markersize=8, alpha=.25, label='EfficientDet')
  897. ax2.grid()
  898. ax2.set_xlim(0, 30)
  899. ax2.set_ylim(28, 50)
  900. ax2.set_yticks(np.arange(30, 55, 5))
  901. ax2.set_xlabel('GPU Speed (ms/img)')
  902. ax2.set_ylabel('COCO AP val')
  903. ax2.legend(loc='lower right')
  904. plt.savefig('study_mAP_latency.png', dpi=300)
  905. plt.savefig(f.replace('.txt', '.png'), dpi=200)
  906. def plot_labels(labels, save_dir=''):
  907. # plot dataset labels
  908. c, b = labels[:, 0], labels[:, 1:].transpose() # classees, boxes
  909. def hist2d(x, y, n=100):
  910. xedges, yedges = np.linspace(x.min(), x.max(), n), np.linspace(y.min(), y.max(), n)
  911. hist, xedges, yedges = np.histogram2d(x, y, (xedges, yedges))
  912. xidx = np.clip(np.digitize(x, xedges) - 1, 0, hist.shape[0] - 1)
  913. yidx = np.clip(np.digitize(y, yedges) - 1, 0, hist.shape[1] - 1)
  914. return np.log(hist[xidx, yidx])
  915. fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)
  916. ax = ax.ravel()
  917. ax[0].hist(c, bins=int(c.max() + 1))
  918. ax[0].set_xlabel('classes')
  919. ax[1].scatter(b[0], b[1], c=hist2d(b[0], b[1], 90), cmap='jet')
  920. ax[1].set_xlabel('x')
  921. ax[1].set_ylabel('y')
  922. ax[2].scatter(b[2], b[3], c=hist2d(b[2], b[3], 90), cmap='jet')
  923. ax[2].set_xlabel('width')
  924. ax[2].set_ylabel('height')
  925. plt.savefig(Path(save_dir) / 'labels.png', dpi=200)
  926. plt.close()
  927. def plot_evolution_results(hyp): # from utils.utils import *; plot_evolution_results(hyp)
  928. # Plot hyperparameter evolution results in evolve.txt
  929. x = np.loadtxt('evolve.txt', ndmin=2)
  930. f = fitness(x)
  931. # weights = (f - f.min()) ** 2 # for weighted results
  932. plt.figure(figsize=(12, 10), tight_layout=True)
  933. matplotlib.rc('font', **{'size': 8})
  934. for i, (k, v) in enumerate(hyp.items()):
  935. y = x[:, i + 7]
  936. # mu = (y * weights).sum() / weights.sum() # best weighted result
  937. mu = y[f.argmax()] # best single result
  938. plt.subplot(4, 5, i + 1)
  939. plt.plot(mu, f.max(), 'o', markersize=10)
  940. plt.plot(y, f, '.')
  941. plt.title('%s = %.3g' % (k, mu), fontdict={'size': 9}) # limit to 40 characters
  942. print('%15s: %.3g' % (k, mu))
  943. plt.savefig('evolve.png', dpi=200)
  944. def plot_results_overlay(start=0, stop=0): # from utils.utils import *; plot_results_overlay()
  945. # Plot training 'results*.txt', overlaying train and val losses
  946. s = ['train', 'train', 'train', 'Precision', 'mAP@0.5', 'val', 'val', 'val', 'Recall', 'mAP@0.5:0.95'] # legends
  947. t = ['GIoU', 'Objectness', 'Classification', 'P-R', 'mAP-F1'] # titles
  948. for f in sorted(glob.glob('results*.txt') + glob.glob('../../Downloads/results*.txt')):
  949. results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T
  950. n = results.shape[1] # number of rows
  951. x = range(start, min(stop, n) if stop else n)
  952. fig, ax = plt.subplots(1, 5, figsize=(14, 3.5), tight_layout=True)
  953. ax = ax.ravel()
  954. for i in range(5):
  955. for j in [i, i + 5]:
  956. y = results[j, x]
  957. ax[i].plot(x, y, marker='.', label=s[j])
  958. # y_smooth = butter_lowpass_filtfilt(y)
  959. # ax[i].plot(x, np.gradient(y_smooth), marker='.', label=s[j])
  960. ax[i].set_title(t[i])
  961. ax[i].legend()
  962. ax[i].set_ylabel(f) if i == 0 else None # add filename
  963. fig.savefig(f.replace('.txt', '.png'), dpi=200)
  964. def plot_results(start=0, stop=0, bucket='', id=(), labels=(),
  965. save_dir=''): # from utils.utils import *; plot_results()
  966. # Plot training 'results*.txt' as seen in https://github.com/ultralytics/yolov5#reproduce-our-training
  967. fig, ax = plt.subplots(2, 5, figsize=(12, 6))
  968. ax = ax.ravel()
  969. s = ['GIoU', 'Objectness', 'Classification', 'Precision', 'Recall',
  970. 'val GIoU', 'val Objectness', 'val Classification', 'mAP@0.5', 'mAP@0.5:0.95']
  971. if bucket:
  972. os.system('rm -rf storage.googleapis.com')
  973. files = ['https://storage.googleapis.com/%s/results%g.txt' % (bucket, x) for x in id]
  974. else:
  975. files = glob.glob(str(Path(save_dir) / 'results*.txt')) + glob.glob('../../Downloads/results*.txt')
  976. for fi, f in enumerate(files):
  977. try:
  978. results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T
  979. n = results.shape[1] # number of rows
  980. x = range(start, min(stop, n) if stop else n)
  981. for i in range(10):
  982. y = results[i, x]
  983. if i in [0, 1, 2, 5, 6, 7]:
  984. y[y == 0] = np.nan # dont show zero loss values
  985. # y /= y[0] # normalize
  986. label = labels[fi] if len(labels) else Path(f).stem
  987. ax[i].plot(x, y, marker='.', label=label, linewidth=2, markersize=8)
  988. ax[i].set_title(s[i])
  989. # if i in [5, 6, 7]: # share train and val loss y axes
  990. # ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
  991. except:
  992. print('Warning: Plotting error for %s, skipping file' % f)
  993. fig.tight_layout()
  994. ax[1].legend()
  995. fig.savefig(Path(save_dir) / 'results.png', dpi=200)