Du kan inte välja fler än 25 ämnen Ämnen måste starta med en bokstav eller siffra, kan innehålla bindestreck ('-') och vara max 35 tecken långa.

812 lines
33KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. General utils
  4. """
  5. import contextlib
  6. import glob
  7. import logging
  8. import math
  9. import os
  10. import platform
  11. import random
  12. import re
  13. import signal
  14. import time
  15. import urllib
  16. from itertools import repeat
  17. from multiprocessing.pool import ThreadPool
  18. from pathlib import Path
  19. from subprocess import check_output
  20. import cv2
  21. import numpy as np
  22. import pandas as pd
  23. import pkg_resources as pkg
  24. import torch
  25. import torchvision
  26. import yaml
  27. from utils.downloads import gsutil_getsize
  28. from utils.metrics import box_iou, fitness
  29. # Settings
  30. torch.set_printoptions(linewidth=320, precision=5, profile='long')
  31. np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5
  32. pd.options.display.max_columns = 10
  33. cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
  34. os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads
  35. FILE = Path(__file__).resolve()
  36. ROOT = FILE.parents[1] # YOLOv5 root directory
  37. class Profile(contextlib.ContextDecorator):
  38. # Usage: @Profile() decorator or 'with Profile():' context manager
  39. def __enter__(self):
  40. self.start = time.time()
  41. def __exit__(self, type, value, traceback):
  42. print(f'Profile results: {time.time() - self.start:.5f}s')
  43. class Timeout(contextlib.ContextDecorator):
  44. # Usage: @Timeout(seconds) decorator or 'with Timeout(seconds):' context manager
  45. def __init__(self, seconds, *, timeout_msg='', suppress_timeout_errors=True):
  46. self.seconds = int(seconds)
  47. self.timeout_message = timeout_msg
  48. self.suppress = bool(suppress_timeout_errors)
  49. def _timeout_handler(self, signum, frame):
  50. raise TimeoutError(self.timeout_message)
  51. def __enter__(self):
  52. signal.signal(signal.SIGALRM, self._timeout_handler) # Set handler for SIGALRM
  53. signal.alarm(self.seconds) # start countdown for SIGALRM to be raised
  54. def __exit__(self, exc_type, exc_val, exc_tb):
  55. signal.alarm(0) # Cancel SIGALRM if it's scheduled
  56. if self.suppress and exc_type is TimeoutError: # Suppress TimeoutError
  57. return True
  58. def try_except(func):
  59. # try-except function. Usage: @try_except decorator
  60. def handler(*args, **kwargs):
  61. try:
  62. func(*args, **kwargs)
  63. except Exception as e:
  64. print(e)
  65. return handler
  66. def methods(instance):
  67. # Get class/instance methods
  68. return [f for f in dir(instance) if callable(getattr(instance, f)) and not f.startswith("__")]
  69. def set_logging(rank=-1, verbose=True):
  70. logging.basicConfig(
  71. format="%(message)s",
  72. level=logging.INFO if (verbose and rank in [-1, 0]) else logging.WARN)
  73. def print_args(name, opt):
  74. # Print argparser arguments
  75. print(colorstr(f'{name}: ') + ', '.join(f'{k}={v}' for k, v in vars(opt).items()))
  76. def init_seeds(seed=0):
  77. # Initialize random number generator (RNG) seeds https://pytorch.org/docs/stable/notes/randomness.html
  78. # cudnn seed 0 settings are slower and more reproducible, else faster and less reproducible
  79. import torch.backends.cudnn as cudnn
  80. random.seed(seed)
  81. np.random.seed(seed)
  82. torch.manual_seed(seed)
  83. cudnn.benchmark, cudnn.deterministic = (False, True) if seed == 0 else (True, False)
  84. def get_latest_run(search_dir='.'):
  85. # Return path to most recent 'last.pt' in /runs (i.e. to --resume from)
  86. last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True)
  87. return max(last_list, key=os.path.getctime) if last_list else ''
  88. def user_config_dir(dir='Ultralytics', env_var='YOLOV5_CONFIG_DIR'):
  89. # Return path of user configuration directory. Prefer environment variable if exists. Make dir if required.
  90. env = os.getenv(env_var)
  91. if env:
  92. path = Path(env) # use environment variable
  93. else:
  94. cfg = {'Windows': 'AppData/Roaming', 'Linux': '.config', 'Darwin': 'Library/Application Support'} # 3 OS dirs
  95. path = Path.home() / cfg.get(platform.system(), '') # OS-specific config dir
  96. path = (path if is_writeable(path) else Path('/tmp')) / dir # GCP and AWS lambda fix, only /tmp is writeable
  97. path.mkdir(exist_ok=True) # make if required
  98. return path
  99. def is_writeable(dir, test=False):
  100. # Return True if directory has write permissions, test opening a file with write permissions if test=True
  101. if test: # method 1
  102. file = Path(dir) / 'tmp.txt'
  103. try:
  104. with open(file, 'w'): # open file with write permissions
  105. pass
  106. file.unlink() # remove file
  107. return True
  108. except IOError:
  109. return False
  110. else: # method 2
  111. return os.access(dir, os.R_OK) # possible issues on Windows
  112. def is_docker():
  113. # Is environment a Docker container?
  114. return Path('/workspace').exists() # or Path('/.dockerenv').exists()
  115. def is_colab():
  116. # Is environment a Google Colab instance?
  117. try:
  118. import google.colab
  119. return True
  120. except Exception as e:
  121. return False
  122. def is_pip():
  123. # Is file in a pip package?
  124. return 'site-packages' in Path(__file__).resolve().parts
  125. def is_ascii(s=''):
  126. # Is string composed of all ASCII (no UTF) characters?
  127. s = str(s) # convert list, tuple, None, etc. to str
  128. return len(s.encode().decode('ascii', 'ignore')) == len(s)
  129. def emojis(str=''):
  130. # Return platform-dependent emoji-safe version of string
  131. return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
  132. def file_size(path):
  133. # Return file/dir size (MB)
  134. path = Path(path)
  135. if path.is_file():
  136. return path.stat().st_size / 1E6
  137. elif path.is_dir():
  138. return sum(f.stat().st_size for f in path.glob('**/*') if f.is_file()) / 1E6
  139. else:
  140. return 0.0
  141. def check_online():
  142. # Check internet connectivity
  143. import socket
  144. try:
  145. socket.create_connection(("1.1.1.1", 443), 5) # check host accessibility
  146. return True
  147. except OSError:
  148. return False
  149. @try_except
  150. def check_git_status():
  151. # Recommend 'git pull' if code is out of date
  152. msg = ', for updates see https://github.com/ultralytics/yolov5'
  153. print(colorstr('github: '), end='')
  154. assert Path('.git').exists(), 'skipping check (not a git repository)' + msg
  155. assert not is_docker(), 'skipping check (Docker image)' + msg
  156. assert check_online(), 'skipping check (offline)' + msg
  157. cmd = 'git fetch && git config --get remote.origin.url'
  158. url = check_output(cmd, shell=True, timeout=5).decode().strip().rstrip('.git') # git fetch
  159. branch = check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out
  160. n = int(check_output(f'git rev-list {branch}..origin/master --count', shell=True)) # commits behind
  161. if n > 0:
  162. s = f"⚠️ YOLOv5 is out of date by {n} commit{'s' * (n > 1)}. Use `git pull` or `git clone {url}` to update."
  163. else:
  164. s = f'up to date with {url} ✅'
  165. print(emojis(s)) # emoji-safe
  166. def check_python(minimum='3.6.2'):
  167. # Check current python version vs. required python version
  168. check_version(platform.python_version(), minimum, name='Python ')
  169. def check_version(current='0.0.0', minimum='0.0.0', name='version ', pinned=False):
  170. # Check version vs. required version
  171. current, minimum = (pkg.parse_version(x) for x in (current, minimum))
  172. result = (current == minimum) if pinned else (current >= minimum)
  173. assert result, f'{name}{minimum} required by YOLOv5, but {name}{current} is currently installed'
  174. @try_except
  175. def check_requirements(requirements=ROOT / 'requirements.txt', exclude=(), install=True):
  176. # Check installed dependencies meet requirements (pass *.txt file or list of packages)
  177. prefix = colorstr('red', 'bold', 'requirements:')
  178. check_python() # check python version
  179. if isinstance(requirements, (str, Path)): # requirements.txt file
  180. file = Path(requirements)
  181. assert file.exists(), f"{prefix} {file.resolve()} not found, check failed."
  182. requirements = [f'{x.name}{x.specifier}' for x in pkg.parse_requirements(file.open()) if x.name not in exclude]
  183. else: # list or tuple of packages
  184. requirements = [x for x in requirements if x not in exclude]
  185. n = 0 # number of packages updates
  186. for r in requirements:
  187. try:
  188. pkg.require(r)
  189. except Exception as e: # DistributionNotFound or VersionConflict if requirements not met
  190. s = f"{prefix} {r} not found and is required by YOLOv5"
  191. if install:
  192. print(f"{s}, attempting auto-update...")
  193. try:
  194. assert check_online(), f"'pip install {r}' skipped (offline)"
  195. print(check_output(f"pip install '{r}'", shell=True).decode())
  196. n += 1
  197. except Exception as e:
  198. print(f'{prefix} {e}')
  199. else:
  200. print(f'{s}. Please install and rerun your command.')
  201. if n: # if packages updated
  202. source = file.resolve() if 'file' in locals() else requirements
  203. s = f"{prefix} {n} package{'s' * (n > 1)} updated per {source}\n" \
  204. f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n"
  205. print(emojis(s))
  206. def check_img_size(imgsz, s=32, floor=0):
  207. # Verify image size is a multiple of stride s in each dimension
  208. if isinstance(imgsz, int): # integer i.e. img_size=640
  209. new_size = max(make_divisible(imgsz, int(s)), floor)
  210. else: # list i.e. img_size=[640, 480]
  211. new_size = [max(make_divisible(x, int(s)), floor) for x in imgsz]
  212. if new_size != imgsz:
  213. print(f'WARNING: --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}')
  214. return new_size
  215. def check_imshow():
  216. # Check if environment supports image displays
  217. try:
  218. assert not is_docker(), 'cv2.imshow() is disabled in Docker environments'
  219. assert not is_colab(), 'cv2.imshow() is disabled in Google Colab environments'
  220. cv2.imshow('test', np.zeros((1, 1, 3)))
  221. cv2.waitKey(1)
  222. cv2.destroyAllWindows()
  223. cv2.waitKey(1)
  224. return True
  225. except Exception as e:
  226. print(f'WARNING: Environment does not support cv2.imshow() or PIL Image.show() image displays\n{e}')
  227. return False
  228. def check_suffix(file='yolov5s.pt', suffix=('.pt',), msg=''):
  229. # Check file(s) for acceptable suffixes
  230. if file and suffix:
  231. if isinstance(suffix, str):
  232. suffix = [suffix]
  233. for f in file if isinstance(file, (list, tuple)) else [file]:
  234. assert Path(f).suffix.lower() in suffix, f"{msg}{f} acceptable suffix is {suffix}"
  235. def check_yaml(file, suffix=('.yaml', '.yml')):
  236. # Search/download YAML file (if necessary) and return path, checking suffix
  237. return check_file(file, suffix)
  238. def check_file(file, suffix=''):
  239. # Search/download file (if necessary) and return path
  240. check_suffix(file, suffix) # optional
  241. file = str(file) # convert to str()
  242. if Path(file).is_file() or file == '': # exists
  243. return file
  244. elif file.startswith(('http:/', 'https:/')): # download
  245. url = str(Path(file)).replace(':/', '://') # Pathlib turns :// -> :/
  246. file = Path(urllib.parse.unquote(file)).name.split('?')[0] # '%2F' to '/', split https://url.com/file.txt?auth
  247. print(f'Downloading {url} to {file}...')
  248. torch.hub.download_url_to_file(url, file)
  249. assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}' # check
  250. return file
  251. else: # search
  252. files = glob.glob('./**/' + file, recursive=True) # find file
  253. assert len(files), f'File not found: {file}' # assert file was found
  254. assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}" # assert unique
  255. return files[0] # return file
  256. def check_dataset(data, autodownload=True):
  257. # Download and/or unzip dataset if not found locally
  258. # Usage: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128_with_yaml.zip
  259. # Download (optional)
  260. extract_dir = ''
  261. if isinstance(data, (str, Path)) and str(data).endswith('.zip'): # i.e. gs://bucket/dir/coco128.zip
  262. download(data, dir='../datasets', unzip=True, delete=False, curl=False, threads=1)
  263. data = next((Path('../datasets') / Path(data).stem).rglob('*.yaml'))
  264. extract_dir, autodownload = data.parent, False
  265. # Read yaml (optional)
  266. if isinstance(data, (str, Path)):
  267. with open(data, errors='ignore') as f:
  268. data = yaml.safe_load(f) # dictionary
  269. # Parse yaml
  270. path = extract_dir or Path(data.get('path') or '') # optional 'path' default to '.'
  271. for k in 'train', 'val', 'test':
  272. if data.get(k): # prepend path
  273. data[k] = str(path / data[k]) if isinstance(data[k], str) else [str(path / x) for x in data[k]]
  274. assert 'nc' in data, "Dataset 'nc' key missing."
  275. if 'names' not in data:
  276. data['names'] = [f'class{i}' for i in range(data['nc'])] # assign class names if missing
  277. train, val, test, s = [data.get(x) for x in ('train', 'val', 'test', 'download')]
  278. if val:
  279. val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path
  280. if not all(x.exists() for x in val):
  281. print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()])
  282. if s and autodownload: # download script
  283. if s.startswith('http') and s.endswith('.zip'): # URL
  284. f = Path(s).name # filename
  285. print(f'Downloading {s} ...')
  286. torch.hub.download_url_to_file(s, f)
  287. root = path.parent if 'path' in data else '..' # unzip directory i.e. '../'
  288. Path(root).mkdir(parents=True, exist_ok=True) # create root
  289. r = os.system(f'unzip -q {f} -d {root} && rm {f}') # unzip
  290. elif s.startswith('bash '): # bash script
  291. print(f'Running {s} ...')
  292. r = os.system(s)
  293. else: # python script
  294. r = exec(s, {'yaml': data}) # return None
  295. print('Dataset autodownload %s\n' % ('success' if r in (0, None) else 'failure')) # print result
  296. else:
  297. raise Exception('Dataset not found.')
  298. return data # dictionary
  299. def url2file(url):
  300. # Convert URL to filename, i.e. https://url.com/file.txt?auth -> file.txt
  301. url = str(Path(url)).replace(':/', '://') # Pathlib turns :// -> :/
  302. file = Path(urllib.parse.unquote(url)).name.split('?')[0] # '%2F' to '/', split https://url.com/file.txt?auth
  303. return file
  304. def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1):
  305. # Multi-threaded file download and unzip function, used in data.yaml for autodownload
  306. def download_one(url, dir):
  307. # Download 1 file
  308. f = dir / Path(url).name # filename
  309. if Path(url).is_file(): # exists in current path
  310. Path(url).rename(f) # move to dir
  311. elif not f.exists():
  312. print(f'Downloading {url} to {f}...')
  313. if curl:
  314. os.system(f"curl -L '{url}' -o '{f}' --retry 9 -C -") # curl download, retry and resume on fail
  315. else:
  316. torch.hub.download_url_to_file(url, f, progress=True) # torch download
  317. if unzip and f.suffix in ('.zip', '.gz'):
  318. print(f'Unzipping {f}...')
  319. if f.suffix == '.zip':
  320. s = f'unzip -qo {f} -d {dir}' # unzip -quiet -overwrite
  321. elif f.suffix == '.gz':
  322. s = f'tar xfz {f} --directory {f.parent}' # unzip
  323. if delete: # delete zip file after unzip
  324. s += f' && rm {f}'
  325. os.system(s)
  326. dir = Path(dir)
  327. dir.mkdir(parents=True, exist_ok=True) # make directory
  328. if threads > 1:
  329. pool = ThreadPool(threads)
  330. pool.imap(lambda x: download_one(*x), zip(url, repeat(dir))) # multi-threaded
  331. pool.close()
  332. pool.join()
  333. else:
  334. for u in [url] if isinstance(url, (str, Path)) else url:
  335. download_one(u, dir)
  336. def make_divisible(x, divisor):
  337. # Returns x evenly divisible by divisor
  338. return math.ceil(x / divisor) * divisor
  339. def clean_str(s):
  340. # Cleans a string by replacing special characters with underscore _
  341. return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s)
  342. def one_cycle(y1=0.0, y2=1.0, steps=100):
  343. # lambda function for sinusoidal ramp from y1 to y2 https://arxiv.org/pdf/1812.01187.pdf
  344. return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1
  345. def colorstr(*input):
  346. # Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e. colorstr('blue', 'hello world')
  347. *args, string = input if len(input) > 1 else ('blue', 'bold', input[0]) # color arguments, string
  348. colors = {'black': '\033[30m', # basic colors
  349. 'red': '\033[31m',
  350. 'green': '\033[32m',
  351. 'yellow': '\033[33m',
  352. 'blue': '\033[34m',
  353. 'magenta': '\033[35m',
  354. 'cyan': '\033[36m',
  355. 'white': '\033[37m',
  356. 'bright_black': '\033[90m', # bright colors
  357. 'bright_red': '\033[91m',
  358. 'bright_green': '\033[92m',
  359. 'bright_yellow': '\033[93m',
  360. 'bright_blue': '\033[94m',
  361. 'bright_magenta': '\033[95m',
  362. 'bright_cyan': '\033[96m',
  363. 'bright_white': '\033[97m',
  364. 'end': '\033[0m', # misc
  365. 'bold': '\033[1m',
  366. 'underline': '\033[4m'}
  367. return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
  368. def labels_to_class_weights(labels, nc=80):
  369. # Get class weights (inverse frequency) from training labels
  370. if labels[0] is None: # no labels loaded
  371. return torch.Tensor()
  372. labels = np.concatenate(labels, 0) # labels.shape = (866643, 5) for COCO
  373. classes = labels[:, 0].astype(np.int) # labels = [class xywh]
  374. weights = np.bincount(classes, minlength=nc) # occurrences per class
  375. # Prepend gridpoint count (for uCE training)
  376. # gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum() # gridpoints per image
  377. # weights = np.hstack([gpi * len(labels) - weights.sum() * 9, weights * 9]) ** 0.5 # prepend gridpoints to start
  378. weights[weights == 0] = 1 # replace empty bins with 1
  379. weights = 1 / weights # number of targets per class
  380. weights /= weights.sum() # normalize
  381. return torch.from_numpy(weights)
  382. def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):
  383. # Produces image weights based on class_weights and image contents
  384. class_counts = np.array([np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels])
  385. image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1)
  386. # index = random.choices(range(n), weights=image_weights, k=1) # weight image sample
  387. return image_weights
  388. def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper)
  389. # https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
  390. # a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n')
  391. # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
  392. # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco
  393. # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet
  394. 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,
  395. 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,
  396. 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
  397. return x
  398. def xyxy2xywh(x):
  399. # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
  400. y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
  401. y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center
  402. y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center
  403. y[:, 2] = x[:, 2] - x[:, 0] # width
  404. y[:, 3] = x[:, 3] - x[:, 1] # height
  405. return y
  406. def xywh2xyxy(x):
  407. # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
  408. y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
  409. y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
  410. y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y
  411. y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x
  412. y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y
  413. return y
  414. def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
  415. # Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
  416. y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
  417. y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw # top left x
  418. y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh # top left y
  419. y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw # bottom right x
  420. y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh # bottom right y
  421. return y
  422. def xyxy2xywhn(x, w=640, h=640, clip=False, eps=0.0):
  423. # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] normalized where xy1=top-left, xy2=bottom-right
  424. if clip:
  425. clip_coords(x, (h - eps, w - eps)) # warning: inplace clip
  426. y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
  427. y[:, 0] = ((x[:, 0] + x[:, 2]) / 2) / w # x center
  428. y[:, 1] = ((x[:, 1] + x[:, 3]) / 2) / h # y center
  429. y[:, 2] = (x[:, 2] - x[:, 0]) / w # width
  430. y[:, 3] = (x[:, 3] - x[:, 1]) / h # height
  431. return y
  432. def xyn2xy(x, w=640, h=640, padw=0, padh=0):
  433. # Convert normalized segments into pixel segments, shape (n,2)
  434. y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
  435. y[:, 0] = w * x[:, 0] + padw # top left x
  436. y[:, 1] = h * x[:, 1] + padh # top left y
  437. return y
  438. def segment2box(segment, width=640, height=640):
  439. # Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy)
  440. x, y = segment.T # segment xy
  441. inside = (x >= 0) & (y >= 0) & (x <= width) & (y <= height)
  442. x, y, = x[inside], y[inside]
  443. return np.array([x.min(), y.min(), x.max(), y.max()]) if any(x) else np.zeros((1, 4)) # xyxy
  444. def segments2boxes(segments):
  445. # Convert segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh)
  446. boxes = []
  447. for s in segments:
  448. x, y = s.T # segment xy
  449. boxes.append([x.min(), y.min(), x.max(), y.max()]) # cls, xyxy
  450. return xyxy2xywh(np.array(boxes)) # cls, xywh
  451. def resample_segments(segments, n=1000):
  452. # Up-sample an (n,2) segment
  453. for i, s in enumerate(segments):
  454. x = np.linspace(0, len(s) - 1, n)
  455. xp = np.arange(len(s))
  456. segments[i] = np.concatenate([np.interp(x, xp, s[:, i]) for i in range(2)]).reshape(2, -1).T # segment xy
  457. return segments
  458. def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
  459. # Rescale coords (xyxy) from img1_shape to img0_shape
  460. if ratio_pad is None: # calculate from img0_shape
  461. gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
  462. pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
  463. else:
  464. gain = ratio_pad[0][0]
  465. pad = ratio_pad[1]
  466. coords[:, [0, 2]] -= pad[0] # x padding
  467. coords[:, [1, 3]] -= pad[1] # y padding
  468. coords[:, :4] /= gain
  469. clip_coords(coords, img0_shape)
  470. return coords
  471. def clip_coords(boxes, shape):
  472. # Clip bounding xyxy bounding boxes to image shape (height, width)
  473. if isinstance(boxes, torch.Tensor): # faster individually
  474. boxes[:, 0].clamp_(0, shape[1]) # x1
  475. boxes[:, 1].clamp_(0, shape[0]) # y1
  476. boxes[:, 2].clamp_(0, shape[1]) # x2
  477. boxes[:, 3].clamp_(0, shape[0]) # y2
  478. else: # np.array (faster grouped)
  479. boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, shape[1]) # x1, x2
  480. boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0]) # y1, y2
  481. def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
  482. labels=(), max_det=300):
  483. """Runs Non-Maximum Suppression (NMS) on inference results
  484. Returns:
  485. list of detections, on (n,6) tensor per image [xyxy, conf, cls]
  486. """
  487. nc = prediction.shape[2] - 5 # number of classes
  488. xc = prediction[..., 4] > conf_thres # candidates
  489. # Checks
  490. assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
  491. assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
  492. # Settings
  493. min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
  494. max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
  495. time_limit = 10.0 # seconds to quit after
  496. redundant = True # require redundant detections
  497. multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
  498. merge = False # use merge-NMS
  499. t = time.time()
  500. output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
  501. for xi, x in enumerate(prediction): # image index, image inference
  502. # Apply constraints
  503. # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
  504. x = x[xc[xi]] # confidence
  505. # Cat apriori labels if autolabelling
  506. if labels and len(labels[xi]):
  507. l = labels[xi]
  508. v = torch.zeros((len(l), nc + 5), device=x.device)
  509. v[:, :4] = l[:, 1:5] # box
  510. v[:, 4] = 1.0 # conf
  511. v[range(len(l)), l[:, 0].long() + 5] = 1.0 # cls
  512. x = torch.cat((x, v), 0)
  513. # If none remain process next image
  514. if not x.shape[0]:
  515. continue
  516. # Compute conf
  517. x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf
  518. # Box (center x, center y, width, height) to (x1, y1, x2, y2)
  519. box = xywh2xyxy(x[:, :4])
  520. # Detections matrix nx6 (xyxy, conf, cls)
  521. if multi_label:
  522. i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
  523. x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
  524. else: # best class only
  525. conf, j = x[:, 5:].max(1, keepdim=True)
  526. x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
  527. # Filter by class
  528. if classes is not None:
  529. x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
  530. # Apply finite constraint
  531. # if not torch.isfinite(x).all():
  532. # x = x[torch.isfinite(x).all(1)]
  533. # Check shape
  534. n = x.shape[0] # number of boxes
  535. if not n: # no boxes
  536. continue
  537. elif n > max_nms: # excess boxes
  538. x = x[x[:, 4].argsort(descending=True)[:max_nms]] # sort by confidence
  539. # Batched NMS
  540. c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
  541. boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
  542. i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
  543. if i.shape[0] > max_det: # limit detections
  544. i = i[:max_det]
  545. if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
  546. # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
  547. iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
  548. weights = iou * scores[None] # box weights
  549. x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
  550. if redundant:
  551. i = i[iou.sum(1) > 1] # require redundancy
  552. output[xi] = x[i]
  553. if (time.time() - t) > time_limit:
  554. print(f'WARNING: NMS time limit {time_limit}s exceeded')
  555. break # time limit exceeded
  556. return output
  557. def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_optimizer()
  558. # Strip optimizer from 'f' to finalize training, optionally save as 's'
  559. x = torch.load(f, map_location=torch.device('cpu'))
  560. if x.get('ema'):
  561. x['model'] = x['ema'] # replace model with ema
  562. for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates': # keys
  563. x[k] = None
  564. x['epoch'] = -1
  565. x['model'].half() # to FP16
  566. for p in x['model'].parameters():
  567. p.requires_grad = False
  568. torch.save(x, s or f)
  569. mb = os.path.getsize(s or f) / 1E6 # filesize
  570. print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")
  571. def print_mutation(results, hyp, save_dir, bucket):
  572. evolve_csv, results_csv, evolve_yaml = save_dir / 'evolve.csv', save_dir / 'results.csv', save_dir / 'hyp_evolve.yaml'
  573. keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
  574. 'val/box_loss', 'val/obj_loss', 'val/cls_loss') + tuple(hyp.keys()) # [results + hyps]
  575. keys = tuple(x.strip() for x in keys)
  576. vals = results + tuple(hyp.values())
  577. n = len(keys)
  578. # Download (optional)
  579. if bucket:
  580. url = f'gs://{bucket}/evolve.csv'
  581. if gsutil_getsize(url) > (os.path.getsize(evolve_csv) if os.path.exists(evolve_csv) else 0):
  582. os.system(f'gsutil cp {url} {save_dir}') # download evolve.csv if larger than local
  583. # Log to evolve.csv
  584. s = '' if evolve_csv.exists() else (('%20s,' * n % keys).rstrip(',') + '\n') # add header
  585. with open(evolve_csv, 'a') as f:
  586. f.write(s + ('%20.5g,' * n % vals).rstrip(',') + '\n')
  587. # Print to screen
  588. print(colorstr('evolve: ') + ', '.join(f'{x.strip():>20s}' for x in keys))
  589. print(colorstr('evolve: ') + ', '.join(f'{x:20.5g}' for x in vals), end='\n\n\n')
  590. # Save yaml
  591. with open(evolve_yaml, 'w') as f:
  592. data = pd.read_csv(evolve_csv)
  593. data = data.rename(columns=lambda x: x.strip()) # strip keys
  594. i = np.argmax(fitness(data.values[:, :7])) #
  595. f.write(f'# YOLOv5 Hyperparameter Evolution Results\n' +
  596. f'# Best generation: {i}\n' +
  597. f'# Last generation: {len(data)}\n' +
  598. f'# ' + ', '.join(f'{x.strip():>20s}' for x in keys[:7]) + '\n' +
  599. f'# ' + ', '.join(f'{x:>20.5g}' for x in data.values[i, :7]) + '\n\n')
  600. yaml.safe_dump(hyp, f, sort_keys=False)
  601. if bucket:
  602. os.system(f'gsutil cp {evolve_csv} {evolve_yaml} gs://{bucket}') # upload
  603. def apply_classifier(x, model, img, im0):
  604. # Apply a second stage classifier to yolo outputs
  605. im0 = [im0] if isinstance(im0, np.ndarray) else im0
  606. for i, d in enumerate(x): # per image
  607. if d is not None and len(d):
  608. d = d.clone()
  609. # Reshape and pad cutouts
  610. b = xyxy2xywh(d[:, :4]) # boxes
  611. b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # rectangle to square
  612. b[:, 2:] = b[:, 2:] * 1.3 + 30 # pad
  613. d[:, :4] = xywh2xyxy(b).long()
  614. # Rescale boxes from img_size to im0 size
  615. scale_coords(img.shape[2:], d[:, :4], im0[i].shape)
  616. # Classes
  617. pred_cls1 = d[:, 5].long()
  618. ims = []
  619. for j, a in enumerate(d): # per item
  620. cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])]
  621. im = cv2.resize(cutout, (224, 224)) # BGR
  622. # cv2.imwrite('example%i.jpg' % j, cutout)
  623. im = im[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
  624. im = np.ascontiguousarray(im, dtype=np.float32) # uint8 to float32
  625. im /= 255.0 # 0 - 255 to 0.0 - 1.0
  626. ims.append(im)
  627. pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1) # classifier prediction
  628. x[i] = x[i][pred_cls1 == pred_cls2] # retain matching class detections
  629. return x
  630. def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BGR=False, save=True):
  631. # Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
  632. xyxy = torch.tensor(xyxy).view(-1, 4)
  633. b = xyxy2xywh(xyxy) # boxes
  634. if square:
  635. b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square
  636. b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
  637. xyxy = xywh2xyxy(b).long()
  638. clip_coords(xyxy, im.shape)
  639. crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
  640. if save:
  641. cv2.imwrite(str(increment_path(file, mkdir=True).with_suffix('.jpg')), crop)
  642. return crop
  643. def increment_path(path, exist_ok=False, sep='', mkdir=False):
  644. # Increment file or directory path, i.e. runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
  645. path = Path(path) # os-agnostic
  646. if path.exists() and not exist_ok:
  647. suffix = path.suffix
  648. path = path.with_suffix('')
  649. dirs = glob.glob(f"{path}{sep}*") # similar paths
  650. matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
  651. i = [int(m.groups()[0]) for m in matches if m] # indices
  652. n = max(i) + 1 if i else 2 # increment number
  653. path = Path(f"{path}{sep}{n}{suffix}") # update path
  654. dir = path if path.suffix == '' else path.parent # directory
  655. if not dir.exists() and mkdir:
  656. dir.mkdir(parents=True, exist_ok=True) # make directory
  657. return path