YOLOV5.4，可能与之前版本不一样，但大同小异
general.py
1 # YOLOv5 general utils 2 3 import glob 4 import logging 5 import math 6 import os 7 import platform 8 import random 9 import re 10 import subprocess 11 import time 12 from pathlib import Path 13 14 import cv2 15 import numpy as np 16 import torch 17 import torchvision 18 import yaml 19 20 from utils.google_utils import gsutil_getsize 21 from utils.metrics import fitness 22 from utils.torch_utils import init_torch_seeds 23 24 # Settings 25 torch.set_printoptions(linewidth=320, precision=5, profile='long') 26 np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5 27 cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader) 28 os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads 29 30 31 def set_logging(rank=-1): 32 logging.basicConfig( 33 format="%(message)s", 34 level=logging.INFO if rank in [-1, 0] else logging.WARN) 35 36 37 def init_seeds(seed=0): 38 # Initialize random number generator (RNG) seeds 39 random.seed(seed) 40 np.random.seed(seed) 41 init_torch_seeds(seed) 42 43 44 def get_latest_run(search_dir='.'): 45 # Return path to most recent 'last.pt' in /runs (i.e. to --resume from) 46 last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True) 47 return max(last_list, key=os.path.getctime) if last_list else '' 48 49 50 def isdocker(): 51 # Is environment a Docker container 52 return Path('/workspace').exists() # or Path('/.dockerenv').exists() 53 54 55 def emojis(str=''): 56 # Return platform-dependent emoji-safe version of string 57 return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str 58 59 60 def check_online(): 61 # Check internet connectivity 62 import socket 63 try: 64 socket.create_connection(("1.1.1.1", 443), 5) # check host accesability 65 return True 66 except OSError: 67 return False 68 69 70 def check_git_status(): 71 # Recommend 'git pull' if code is out of date 72 print(colorstr('github: '), end='') 73 try: 74 assert Path('.git').exists(), 'skipping check (not a git repository)' 75 assert not isdocker(), 'skipping check (Docker image)' 76 assert check_online(), 'skipping check (offline)' 77 78 cmd = 'git fetch && git config --get remote.origin.url' 79 url = subprocess.check_output(cmd, shell=True).decode().strip().rstrip('.git') # github repo url 80 branch = subprocess.check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out 81 n = int(subprocess.check_output(f'git rev-list {branch}..origin/master --count', shell=True)) # commits behind 82 if n > 0: 83 s = f"⚠️ WARNING: code is out of date by {n} commit{'s' * (n > 1)}. " \ 84 f"Use 'git pull' to update or 'git clone {url}' to download latest." 85 else: 86 s = f'up to date with {url} ✅' 87 print(emojis(s)) # emoji-safe 88 except Exception as e: 89 print(e) 90 91 92 def check_requirements(file='requirements.txt', exclude=()): 93 # Check installed dependencies meet requirements 94 import pkg_resources as pkg 95 prefix = colorstr('red', 'bold', 'requirements:') 96 file = Path(file) 97 if not file.exists(): 98 print(f"{prefix} {file.resolve()} not found, check failed.") 99 return 100 101 n = 0 # number of packages updates 102 requirements = [f'{x.name}{x.specifier}' for x in pkg.parse_requirements(file.open()) if x.name not in exclude] 103 for r in requirements: 104 try: 105 pkg.require(r) 106 except Exception as e: # DistributionNotFound or VersionConflict if requirements not met 107 n += 1 108 print(f"{prefix} {e.req} not found and is required by YOLOv5, attempting auto-update...") 109 print(subprocess.check_output(f"pip install '{e.req}'", shell=True).decode()) 110 111 if n: # if packages updated 112 s = f"{prefix} {n} package{'s' * (n > 1)} updated per {file.resolve()}\n" \ 113 f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n" 114 print(emojis(s)) # emoji-safe 115 116 117 def check_img_size(img_size, s=32): 118 # Verify img_size is a multiple of stride s 119 new_size = make_divisible(img_size, int(s)) # ceil gs-multiple 120 if new_size != img_size: 121 print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size)) 122 return new_size 123 124 125 def check_imshow(): 126 # Check if environment supports image displays 127 try: 128 assert not isdocker(), 'cv2.imshow() is disabled in Docker environments' 129 cv2.imshow('test', np.zeros((1, 1, 3))) 130 cv2.waitKey(1) 131 cv2.destroyAllWindows() 132 cv2.waitKey(1) 133 return True 134 except Exception as e: 135 print(f'WARNING: Environment does not support cv2.imshow() or PIL Image.show() image displays\n{e}') 136 return False 137 138 139 def check_file(file): 140 # Search for file if not found 141 if os.path.isfile(file) or file == '': 142 return file 143 else: 144 files = glob.glob('./**/' + file, recursive=True) # find file 145 assert len(files), 'File Not Found: %s' % file # assert file was found 146 assert len(files) == 1, "Multiple files match '%s', specify exact path: %s" % (file, files) # assert unique 147 return files[0] # return file 148 149 150 def check_dataset(dict): 151 # Download dataset if not found locally 152 val, s = dict.get('val'), dict.get('download') 153 if val and len(val): 154 val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path 155 if not all(x.exists() for x in val): 156 print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()]) 157 if s and len(s): # download script 158 print('Downloading %s ...' % s) 159 if s.startswith('http') and s.endswith('.zip'): # URL 160 f = Path(s).name # filename 161 torch.hub.download_url_to_file(s, f) 162 r = os.system('unzip -q %s -d ../ && rm %s' % (f, f)) # unzip 163 else: # bash script 164 r = os.system(s) 165 print('Dataset autodownload %s\n' % ('success' if r == 0 else 'failure')) # analyze return value 166 else: 167 raise Exception('Dataset not found.') 168 169 170 def make_divisible(x, divisor): 171 # Returns x evenly divisible by divisor 172 return math.ceil(x / divisor) * divisor 173 174 175 def clean_str(s): 176 # Cleans a string by replacing special characters with underscore _ 177 return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s) 178 179 180 def one_cycle(y1=0.0, y2=1.0, steps=100): 181 # lambda function for sinusoidal ramp from y1 to y2 182 return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1 183 184 185 def colorstr(*input): 186 # Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e. colorstr('blue', 'hello world') 187 *args, string = input if len(input) > 1 else ('blue', 'bold', input[0]) # color arguments, string 188 colors = {'black': '\033[30m', # basic colors 189 'red': '\033[31m', 190 'green': '\033[32m', 191 'yellow': '\033[33m', 192 'blue': '\033[34m', 193 'magenta': '\033[35m', 194 'cyan': '\033[36m', 195 'white': '\033[37m', 196 'bright_black': '\033[90m', # bright colors 197 'bright_red': '\033[91m', 198 'bright_green': '\033[92m', 199 'bright_yellow': '\033[93m', 200 'bright_blue': '\033[94m', 201 'bright_magenta': '\033[95m', 202 'bright_cyan': '\033[96m', 203 'bright_white': '\033[97m', 204 'end': '\033[0m', # misc 205 'bold': '\033[1m', 206 'underline': '\033[4m'} 207 return ''.join(colors[x] for x in args) + f'{string}' + colors['end'] 208 209 210 def labels_to_class_weights(labels, nc=80): 211 # Get class weights (inverse frequency) from training labels 212 if labels[0] is None: # no labels loaded 213 return torch.Tensor() 214 215 labels = np.concatenate(labels, 0) # labels.shape = (866643, 5) for COCO 216 classes = labels[:, 0].astype(np.int) # labels = [class xywh] 217 weights = np.bincount(classes, minlength=nc) # occurrences per class 218 219 # Prepend gridpoint count (for uCE training) 220 # gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum() # gridpoints per image 221 # weights = np.hstack([gpi * len(labels) - weights.sum() * 9, weights * 9]) ** 0.5 # prepend gridpoints to start 222 223 weights[weights == 0] = 1 # replace empty bins with 1 224 weights = 1 / weights # number of targets per class 225 weights /= weights.sum() # normalize 226 return torch.from_numpy(weights) 227 228 229 def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)): 230 # Produces image weights based on class_weights and image contents 231 class_counts = np.array([np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels]) 232 image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1) 233 # index = random.choices(range(n), weights=image_weights, k=1) # weight image sample 234 return image_weights 235 236 237 def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper) 238 # https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/ 239 # a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n') 240 # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n') 241 # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco 242 # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet 243 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, 244 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, 245 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90] 246 return x 247 248 249 def xyxy2xywh(x): 250 # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right 251 y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x) 252 y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center 253 y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center 254 y[:, 2] = x[:, 2] - x[:, 0] # width 255 y[:, 3] = x[:, 3] - x[:, 1] # height 256 return y 257 258 259 def xywh2xyxy(x): 260 # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right 261 y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x) 262 y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x 263 y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y 264 y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x 265 y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y 266 return y 267 268 269 def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0): 270 # Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right 271 y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x) 272 y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw # top left x 273 y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh # top left y 274 y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw # bottom right x 275 y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh # bottom right y 276 return y 277 278 279 def xyn2xy(x, w=640, h=640, padw=0, padh=0): 280 # Convert normalized segments into pixel segments, shape (n,2) 281 y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x) 282 y[:, 0] = w * x[:, 0] + padw # top left x 283 y[:, 1] = h * x[:, 1] + padh # top left y 284 return y 285 286 287 def segment2box(segment, width=640, height=640): 288 # Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy) 289 x, y = segment.T # segment xy 290 inside = (x >= 0) & (y >= 0) & (x <= width) & (y <= height) 291 x, y, = x[inside], y[inside] 292 return np.array([x.min(), y.min(), x.max(), y.max()]) if any(x) else np.zeros((1, 4)) # xyxy 293 294 295 def segments2boxes(segments): 296 # Convert segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh) 297 boxes = [] 298 for s in segments: 299 x, y = s.T # segment xy 300 boxes.append([x.min(), y.min(), x.max(), y.max()]) # cls, xyxy 301 return xyxy2xywh(np.array(boxes)) # cls, xywh 302 303 304 def resample_segments(segments, n=1000): 305 # Up-sample an (n,2) segment 306 for i, s in enumerate(segments): 307 x = np.linspace(0, len(s) - 1, n) 308 xp = np.arange(len(s)) 309 segments[i] = np.concatenate([np.interp(x, xp, s[:, i]) for i in range(2)]).reshape(2, -1).T # segment xy 310 return segments 311 312 313 def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None): 314 # Rescale coords (xyxy) from img1_shape to img0_shape 315 if ratio_pad is None: # calculate from img0_shape 316 gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new 317 pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding 318 else: 319 gain = ratio_pad[0][0] 320 pad = ratio_pad[1] 321 322 coords[:, [0, 2]] -= pad[0] # x padding 323 coords[:, [1, 3]] -= pad[1] # y padding 324 coords[:, :4] /= gain 325 clip_coords(coords, img0_shape) 326 return coords 327 328 329 def clip_coords(boxes, img_shape): 330 # Clip bounding xyxy bounding boxes to image shape (height, width) 331 boxes[:, 0].clamp_(0, img_shape[1]) # x1 332 boxes[:, 1].clamp_(0, img_shape[0]) # y1 333 boxes[:, 2].clamp_(0, img_shape[1]) # x2 334 boxes[:, 3].clamp_(0, img_shape[0]) # y2 335 336 337 def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7): 338 # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4 339 box2 = box2.T 340 341 # Get the coordinates of bounding boxes 342 if x1y1x2y2: # x1, y1, x2, y2 = box1 343 b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3] 344 b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3] 345 else: # transform from xywh to xyxy 346 b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2 347 b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2 348 b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2 349 b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2 350 351 # Intersection area 352 inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \ 353 (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0) 354 355 # Union Area 356 w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps 357 w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps 358 union = w1 * h1 + w2 * h2 - inter + eps 359 360 iou = inter / union 361 if GIoU or DIoU or CIoU: 362 cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width 363 ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height 364 if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1 365 c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared 366 rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + 367 (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared 368 if DIoU: 369 return iou - rho2 / c2 # DIoU 370 elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 371 v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2) 372 with torch.no_grad(): 373 alpha = v / (v - iou + (1 + eps)) 374 return iou - (rho2 / c2 + v * alpha) # CIoU 375 else: # GIoU https://arxiv.org/pdf/1902.09630.pdf 376 c_area = cw * ch + eps # convex area 377 return iou - (c_area - union) / c_area # GIoU 378 else: 379 return iou # IoU 380 381 382 def box_iou(box1, box2): 383 # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py 384 """ 385 Return intersection-over-union (Jaccard index) of boxes. 386 Both sets of boxes are expected to be in (x1, y1, x2, y2) format. 387 Arguments: 388 box1 (Tensor[N, 4]) 389 box2 (Tensor[M, 4]) 390 Returns: 391 iou (Tensor[N, M]): the NxM matrix containing the pairwise 392 IoU values for every element in boxes1 and boxes2 393 """ 394 395 def box_area(box): 396 # box = 4xn 397 return (box[2] - box[0]) * (box[3] - box[1]) 398 399 area1 = box_area(box1.T) 400 area2 = box_area(box2.T) 401 402 # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2) 403 inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2) 404 return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter) 405 406 407 def wh_iou(wh1, wh2): 408 # Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2 409 wh1 = wh1[:, None] # [N,1,2] 410 wh2 = wh2[None] # [1,M,2] 411 inter = torch.min(wh1, wh2).prod(2) # [N,M] 412 return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter) 413 414 415 def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False, 416 labels=()): 417 """Runs Non-Maximum Suppression (NMS) on inference results 418 419 Returns: 420 list of detections, on (n,6) tensor per image [xyxy, conf, cls] 421 """ 422 # prediction.shape 423 # torch.Size([1, 10080, 7]) 424 nc = prediction.shape[2] - 5 # number of classes 为什么是减去5?,因为5表示4个位置信息(xyxy)，加一个置信度得分，置信度计算公式见下 425 # prediction[..., 4]：取出所有通道中矩阵的第四列 426 # xc是一个bool类型的list，IOU大于阈值的为true 427 # xc.shape 428 # torch.Size([1, 10080]) 429 xc = prediction[..., 4] > conf_thres # candidates 430 431 # Settings 432 min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height 433 max_det = 300 # maximum number of detections per image 434 max_nms = 30000 # maximum number of boxes into torchvision.ops.nms() 435 time_limit = 10.0 # seconds to quit after 436 redundant = True # require redundant detections 437 multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img) 438 merge = False # use merge-NMS 439 440 t = time.time() 441 # output 442 # [tensor([], size=(0, 6))]:[xyxy confidence classification] 443 output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0] 444 for xi, x in enumerate(prediction): # image index, image inference 445 # Apply constraints 446 # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height 447 # x.shape 448 # torch.Size([2387, 7]) 449 x = x[xc[xi]] # confidence ??????????????????????????????????? 450 451 # Cat apriori labels if autolabelling 452 if labels and len(labels[xi]): 453 l = labels[xi] 454 v = torch.zeros((len(l), nc + 5), device=x.device) 455 v[:, :4] = l[:, 1:5] # box 456 v[:, 4] = 1.0 # conf 457 v[range(len(l)), l[:, 0].long() + 5] = 1.0 # cls 458 x = torch.cat((x, v), 0) 459 460 # If none remain process next image 461 if not x.shape[0]: 462 continue 463 # 逐行相乘 464 # x.shape 465 # torch.Size([2387, 7]),有2387个候选框 466 467 # Compute conf; 468 # conf = cls_conf * obj_conf,即：Pr(Classi) = Pr(Classi|Object) * Pr(Object)，注意置信度公式和YOLOV123区别 469 # 其中 obj_conf = 候选框(bounding box)存在对象的概率; 470 # cls_conf = 如果当前网格(cell)存在对象，是类别i(i=1 2 …N,一共N个类别)的概率; 471 472 # 第5、6列分别乘第4列 473 # 第5、6列对应：Pr(Classi|Object)，对应cell两个类别概率 474 # 第4列对应：Pr(Object)，对应候选框存在对象概率 475 # 所以第0、1、2、3列就是候选框几何信息 476 # 综上x中每一行为：[x y x y Pr(Object) Pr(Class1|Object) Pr(Class2|Object)] 477 x[:, 5:] *= x[:, 4:5] 478 479 # 取出预测框位置信息 480 # Box (center x, center y, width, height) to (x1, y1, x2, y2) 481 # box.shape 482 # torch.Size([2387, 4]) 483 box = xywh2xyxy(x[:, :4]) 484 485 # Detections matrix nx6 (xyxy, conf, cls) 486 if multi_label: 487 i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T 488 x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) 489 else: # best class only 490 # 对比x的第5、6列的置信度，note：上面已经算过了，这里不再是概率 491 # conf.shape 492 # torch.Size([2387, 1]) 493 # j.shape 494 # torch.Size([2387, 1]) 495 # j的每个元素取值为0 or 1, 表示大的值对应索引 496 conf, j = x[:, 5:].max(1, keepdim=True) 497 # 这里按照(box conf j)重组x 498 # x.shape 499 # torch.Size([2387, 7]) 500 x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] 501 # x.shape 502 # torch.Size([2387, 6]) 503 504 # Filter by class 505 if classes is not None: 506 x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] 507 508 # Apply finite constraint 509 # if not torch.isfinite(x).all(): 510 # x = x[torch.isfinite(x).all(1)] 511 512 # Check shape 513 n = x.shape[0] # number of boxes 514 if not n: # no boxes 515 continue 516 elif n > max_nms: # excess boxes 517 # 将张量x按照第4列(置信度)进行排序，从大到小，最多取max_nms个，所以，x的行数是max_num，这一样大多数情况不会执行 518 x = x[x[:, 4].argsort(descending=True)[:max_nms]] # sort by confidence 519 520 # Batched NMS 521 # ?????????????????????????? 这个预测框的偏移量是如何计算的？？？？？？ 522 c = x[:, 5:6] * (0 if agnostic else max_wh) # classes 523 # scores:也就是上述置信度 524 boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores 525 # i:NMS后得到框的索引数组 526 i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS 527 if i.shape[0] > max_det: # limit detections 528 i = i[:max_det] 529 if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) 530 # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4) 531 iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix 532 weights = iou * scores[None] # box weights 533 x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes 534 if redundant: 535 i = i[iou.sum(1) > 1] # require redundancy 536 # output[0].shape 537 # torch.Size([225, 6]) 538 output[xi] = x[i] 539 if (time.time() - t) > time_limit: 540 print(f'WARNING: NMS time limit {time_limit}s exceeded') 541 break # time limit exceeded 542 543 return output 544 545 546 def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_optimizer() 547 # Strip optimizer from 'f' to finalize training, optionally save as 's' 548 x = torch.load(f, map_location=torch.device('cpu')) 549 if x.get('ema'): 550 x['model'] = x['ema'] # replace model with ema 551 for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates': # keys 552 x[k] = None 553 x['epoch'] = -1 554 x['model'].half() # to FP16 555 for p in x['model'].parameters(): 556 p.requires_grad = False 557 torch.save(x, s or f) 558 mb = os.path.getsize(s or f) / 1E6 # filesize 559 print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB") 560 561 562 def print_mutation(hyp, results, yaml_file='hyp_evolved.yaml', bucket=''): 563 # Print mutation results to evolve.txt (for use with train.py --evolve) 564 a = '%10s' * len(hyp) % tuple(hyp.keys()) # hyperparam keys 565 b = '%10.3g' * len(hyp) % tuple(hyp.values()) # hyperparam values 566 c = '%10.4g' * len(results) % results # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3) 567 print('\n%s\n%s\nEvolved fitness: %s\n' % (a, b, c)) 568 569 if bucket: 570 url = 'gs://%s/evolve.txt' % bucket 571 if gsutil_getsize(url) > (os.path.getsize('evolve.txt') if os.path.exists('evolve.txt') else 0): 572 os.system('gsutil cp %s .' % url) # download evolve.txt if larger than local 573 574 with open('evolve.txt', 'a') as f: # append result 575 f.write(c + b + '\n') 576 x = np.unique(np.loadtxt('evolve.txt', ndmin=2), axis=0) # load unique rows 577 x = x[np.argsort(-fitness(x))] # sort 578 np.savetxt('evolve.txt', x, '%10.3g') # save sort by fitness 579 580 # Save yaml 581 for i, k in enumerate(hyp.keys()): 582 hyp[k] = float(x[0, i + 7]) 583 with open(yaml_file, 'w') as f: 584 results = tuple(x[0, :7]) 585 c = '%10.4g' * len(results) % results # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3) 586 f.write('# Hyperparameter Evolution Results\n# Generations: %g\n# Metrics: ' % len(x) + c + '\n\n') 587 yaml.dump(hyp, f, sort_keys=False) 588 589 if bucket: 590 os.system('gsutil cp evolve.txt %s gs://%s' % (yaml_file, bucket)) # upload 591 592 593 def apply_classifier(x, model, img, im0): 594 # applies a second stage classifier to yolo outputs 595 im0 = [im0] if isinstance(im0, np.ndarray) else im0 596 for i, d in enumerate(x): # per image 597 if d is not None and len(d): 598 d = d.clone() 599 600 # Reshape and pad cutouts 601 b = xyxy2xywh(d[:, :4]) # boxes 602 b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # rectangle to square 603 b[:, 2:] = b[:, 2:] * 1.3 + 30 # pad 604 d[:, :4] = xywh2xyxy(b).long() 605 606 # Rescale boxes from img_size to im0 size 607 scale_coords(img.shape[2:], d[:, :4], im0[i].shape) 608 609 # Classes 610 pred_cls1 = d[:, 5].long() 611 ims = [] 612 for j, a in enumerate(d): # per item 613 cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])] 614 im = cv2.resize(cutout, (224, 224)) # BGR 615 # cv2.imwrite('test%i.jpg' % j, cutout) 616 617 im = im[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 618 im = np.ascontiguousarray(im, dtype=np.float32) # uint8 to float32 619 im /= 255.0 # 0 - 255 to 0.0 - 1.0 620 ims.append(im) 621 622 pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1) # classifier prediction 623 x[i] = x[i][pred_cls1 == pred_cls2] # retain matching class detections 624 625 return x 626 627 628 def increment_path(path, exist_ok=True, sep=''): 629 # Increment path, i.e. runs/exp --> runs/exp{sep}0, runs/exp{sep}1 etc. 630 path = Path(path) # os-agnostic 631 if (path.exists() and exist_ok) or (not path.exists()): 632 return str(path) 633 else: 634 dirs = glob.glob(f"{path}{sep}*") # similar paths 635 matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs] 636 i = [int(m.groups()[0]) for m in matches if m] # indices 637 n = max(i) + 1 if i else 2 # increment number 638 return f"{path}{sep}{n}" # update path
detect.py
1 import argparse 2 import time 3 from pathlib import Path 4 5 import cv2 6 import torch 7 import torch.backends.cudnn as cudnn 8 from numpy import random 9 10 from models.experimental import attempt_load 11 from utils.datasets import LoadStreams, LoadImages 12 from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \ 13 scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path 14 from utils.plots import plot_one_box 15 from utils.torch_utils import select_device, load_classifier, time_synchronized 16 17 18 def detect(save_img=False): 19 source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size 20 save_img = not opt.nosave and not source.endswith('.txt') # save inference images 21 webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith( 22 ('rtsp://', 'rtmp://', 'http://')) 23 24 # Directories 25 save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run 26 (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir 27 28 # Initialize 29 set_logging() 30 device = select_device(opt.device) 31 half = device.type != 'cpu' # half precision only supported on CUDA 32 33 # Load model 34 model = attempt_load(weights, map_location=device) # load FP32 model 35 stride = int(model.stride.max()) # model stride 36 imgsz = check_img_size(imgsz, s=stride) # check img_size 37 if half: 38 model.half() # to FP16 39 40 # Second-stage classifier 41 classify = False 42 if classify: 43 modelc = load_classifier(name='resnet101', n=2) # initialize 44 modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval() 45 46 # Set Dataloader 47 vid_path, vid_writer = None, None 48 if webcam: 49 view_img = check_imshow() 50 cudnn.benchmark = True # set True to speed up constant image size inference 51 dataset = LoadStreams(source, img_size=imgsz, stride=stride) 52 else: 53 # 这里底层使用opencv读取图片，接着使用letterbox函数进行padded resize图片 54 dataset = LoadImages(source, img_size=imgsz, stride=stride) 55 56 # Get names and colors 57 names = model.module.names if hasattr(model, 'module') else model.names 58 colors = [[random.randint(0, 255) for _ in range(3)] for _ in names] 59 60 # Run inference 61 if device.type != 'cpu': 62 model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once 63 t0 = time.time() 64 # resize后分辨率 实际分辨率 65 # img.shape (3, 256, 640) im0s.shape (360, 900, 3) 66 # img.shape (3, 480, 640) im0s.shape (375, 500, 3) 67 for path, img, im0s, vid_cap in dataset: 68 69 #img_ = cv2.cvtColor(img , cv2.COLOR_RGB2BGR) 70 #cv2.imshow("img", img_) 71 #cv2.imshow("im0s", im0s) 72 #cv2.waitKey(0) 73 74 img = torch.from_numpy(img).to(device) 75 img = img.half() if half else img.float() # uint8 to fp16/32 76 img /= 255.0 # 0 - 255 to 0.0 - 1.0 77 if img.ndimension() == 3: 78 # img.shape torch.Size([3, 256, 640]) 79 img = img.unsqueeze(0) # 增加一个维度 80 # img.shape torch.Size([1, 3, 256, 640]) 81 # Inference 82 t1 = time_synchronized() 83 """ 84 前向传播 返回pred的shape是(1, num_boxes, 5+num_class) 85 h,w为传入网络图片的长和宽，注意dataset在检测时使用了矩形推理，所以这里h不一定等于w 86 # 矩形推理：https://blog.csdn.net/songwsx/article/details/102639770 87 num_boxes = h/32 * w/32 + h/16 * w/16 + h/8 * w/8 88 pred[..., 0:4]为预测框坐标 89 预测框坐标为xywh(中心点+宽长)格式 90 pred[..., 4]为objectness置信度 91 pred[..., 5:-1]为分类结果 92 """ 93 pred = model(img, augment=opt.augment)[0] 94 95 # Apply NMS 96 pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms) 97 # pred[0].shape 98 # torch.Size([225, 6]) 99 t2 = time_synchronized() 100 101 # Apply Classifier 102 if classify: 103 pred = apply_classifier(pred, modelc, img, im0s) 104 105 # Process detections 106 # det: detction 107 for i, det in enumerate(pred): # detections per image 108 if webcam: # batch_size >= 1 109 p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count 110 else: 111 p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0) 112 113 p = Path(p) # to Path 114 save_path = str(save_dir / p.name) # img.jpg 115 txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt 116 s += '%gx%g ' % img.shape[2:] # print string 117 # im0(source image) (360, 900, 3) 118 # gn tensor([900, 360, 900, 360]) 119 gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh 120 if len(det): 121 # Rescale boxes from img_size to im0 size 122 # 调整框大小，缩放到实际原图中去 123 det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() 124 125 # Print results 126 for c in det[:, -1].unique(): 127 n = (det[:, -1] == c).sum() # detections per class 128 s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string 129 130 # Write results 131 # conf:置信度得分 cls：类别 132 for *xyxy, conf, cls in reversed(det): 133 if save_txt: # Write to file 134 # 框的几何信息转换：xyxy -> xywh 135 xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh 136 line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format 137 with open(txt_path + '.txt', 'a') as f: 138 f.write(('%g ' * len(line)).rstrip() % line + '\n') 139 140 if save_img or view_img: # Add bbox to image 141 label = f'{names[int(cls)]} {conf:.2f}' 142 #plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) 143 plot_one_box(xyxy, im0, label='', color=colors[int(cls)], line_thickness=2) 144 145 # Print time (inference + NMS) 146 print(f'{s}Done. ({t2 - t1:.3f}s)') 147 148 # Stream results 149 if view_img: 150 cv2.imshow(str(p), im0) 151 cv2.waitKey(1) # 1 millisecond 152 153 # Save results (image with detections) 154 if save_img: 155 if dataset.mode == 'image': 156 cv2.imwrite(save_path, im0) 157 else: # 'video' or 'stream' 158 if vid_path != save_path: # new video 159 vid_path = save_path 160 if isinstance(vid_writer, cv2.VideoWriter): 161 vid_writer.release() # release previous video writer 162 if vid_cap: # video 163 fps = vid_cap.get(cv2.CAP_PROP_FPS) 164 w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) 165 h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) 166 else: # stream 167 fps, w, h = 30, im0.shape[1], im0.shape[0] 168 save_path += '.mp4' 169 vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) 170 vid_writer.write(im0) 171 172 if save_txt or save_img: 173 s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else '' 174 print(f"Results saved to {save_dir}{s}") 175 176 print(f'Done. ({time.time() - t0:.3f}s)') 177 178 # --source D:/Data/yolo --weights weights/best.pt --device cpu 179 if __name__ == '__main__': 180 parser = argparse.ArgumentParser() 181 parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') 182 parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam 183 parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') 184 parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold') 185 parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS') 186 parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 187 parser.add_argument('--view-img', action='store_true', help='display results') 188 parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') 189 parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels') 190 parser.add_argument('--nosave', action='store_true', help='do not save images/videos') 191 parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3') 192 parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') 193 parser.add_argument('--augment', action='store_true', help='augmented inference') 194 parser.add_argument('--update', action='store_true', help='update all models') 195 parser.add_argument('--project', default='runs/detect', help='save results to project/name') 196 parser.add_argument('--name', default='exp', help='save results to project/name') 197 parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') 198 opt = parser.parse_args() 199 print(opt) 200 check_requirements(exclude=('pycocotools', 'thop')) 201 202 with torch.no_grad(): 203 if opt.update: # update all models (to fix SourceChangeWarning) 204 for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: 205 detect() 206 strip_optimizer(opt.weights) 207 else: 208 detect()