general.py 30.2 KB
Newer Older
1
2
3
4
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
General utils
"""
Glenn Jocher's avatar
Glenn Jocher committed
5

6
import contextlib
Glenn Jocher's avatar
Glenn Jocher committed
7
import glob
Glenn Jocher's avatar
Glenn Jocher committed
8
import logging
9
import math
Glenn Jocher's avatar
Glenn Jocher committed
10
import os
11
import platform
Glenn Jocher's avatar
Glenn Jocher committed
12
import random
13
import re
14
import signal
Glenn Jocher's avatar
Glenn Jocher committed
15
import time
16
import urllib
17
18
from itertools import repeat
from multiprocessing.pool import ThreadPool
Glenn Jocher's avatar
Glenn Jocher committed
19
from pathlib import Path
20
from subprocess import check_output
Glenn Jocher's avatar
Glenn Jocher committed
21
22
23

import cv2
import numpy as np
24
import pandas as pd
Glenn Jocher's avatar
Glenn Jocher committed
25
import pkg_resources as pkg
Glenn Jocher's avatar
Glenn Jocher committed
26
import torch
Glenn Jocher's avatar
Glenn Jocher committed
27
import torchvision
Glenn Jocher's avatar
Glenn Jocher committed
28
import yaml
Glenn Jocher's avatar
Glenn Jocher committed
29

30
from utils.downloads import gsutil_getsize
31
from utils.metrics import box_iou, fitness
Glenn Jocher's avatar
Glenn Jocher committed
32
from utils.torch_utils import init_torch_seeds
Glenn Jocher's avatar
Glenn Jocher committed
33

34
# Settings
Glenn Jocher's avatar
Glenn Jocher committed
35
36
torch.set_printoptions(linewidth=320, precision=5, profile='long')
np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format})  # format short g, %precision=5
37
pd.options.display.max_columns = 10
38
cv2.setNumThreads(0)  # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
39
os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8))  # NumExpr max threads
Glenn Jocher's avatar
Glenn Jocher committed
40
41


42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
class timeout(contextlib.ContextDecorator):
    # Usage: @timeout(seconds) decorator or 'with timeout(seconds):' context manager
    def __init__(self, seconds, *, timeout_msg='', suppress_timeout_errors=True):
        self.seconds = int(seconds)
        self.timeout_message = timeout_msg
        self.suppress = bool(suppress_timeout_errors)

    def _timeout_handler(self, signum, frame):
        raise TimeoutError(self.timeout_message)

    def __enter__(self):
        signal.signal(signal.SIGALRM, self._timeout_handler)  # Set handler for SIGALRM
        signal.alarm(self.seconds)  # start countdown for SIGALRM to be raised

    def __exit__(self, exc_type, exc_val, exc_tb):
        signal.alarm(0)  # Cancel SIGALRM if it's scheduled
        if self.suppress and exc_type is TimeoutError:  # Suppress TimeoutError
            return True


62
63
64
65
66
67
68
69
70
71
72
def try_except(func):
    # try-except function. Usage: @try_except decorator
    def handler(*args, **kwargs):
        try:
            func(*args, **kwargs)
        except Exception as e:
            print(e)

    return handler


73
74
75
76
77
def methods(instance):
    # Get class/instance methods
    return [f for f in dir(instance) if callable(getattr(instance, f)) and not f.startswith("__")]


78
def set_logging(rank=-1, verbose=True):
79
80
    logging.basicConfig(
        format="%(message)s",
81
        level=logging.INFO if (verbose and rank in [-1, 0]) else logging.WARN)
82
83


Glenn Jocher's avatar
Glenn Jocher committed
84
def init_seeds(seed=0):
85
    # Initialize random number generator (RNG) seeds
Glenn Jocher's avatar
Glenn Jocher committed
86
87
    random.seed(seed)
    np.random.seed(seed)
88
    init_torch_seeds(seed)
Glenn Jocher's avatar
Glenn Jocher committed
89

Glenn Jocher's avatar
Glenn Jocher committed
90

91
def get_latest_run(search_dir='.'):
Glenn Jocher's avatar
Glenn Jocher committed
92
    # Return path to most recent 'last.pt' in /runs (i.e. to --resume from)
Glenn Jocher's avatar
Glenn Jocher committed
93
    last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True)
Glenn Jocher's avatar
Glenn Jocher committed
94
    return max(last_list, key=os.path.getctime) if last_list else ''
Glenn Jocher's avatar
Glenn Jocher committed
95

Glenn Jocher's avatar
Glenn Jocher committed
96

Glenn Jocher's avatar
Glenn Jocher committed
97
def is_docker():
98
    # Is environment a Docker container?
Glenn Jocher's avatar
Glenn Jocher committed
99
100
101
    return Path('/workspace').exists()  # or Path('/.dockerenv').exists()


Glenn Jocher's avatar
Glenn Jocher committed
102
def is_colab():
103
    # Is environment a Google Colab instance?
Glenn Jocher's avatar
Glenn Jocher committed
104
105
106
107
108
109
110
    try:
        import google.colab
        return True
    except Exception as e:
        return False


111
112
113
114
115
def is_pip():
    # Is file in a pip package?
    return 'site-packages' in Path(__file__).absolute().parts


116
117
118
119
120
def is_ascii(str=''):
    # Is string composed of all ASCII (no UTF) characters?
    return len(str.encode().decode('ascii', 'ignore')) == len(str)


121
122
def emojis(str=''):
    # Return platform-dependent emoji-safe version of string
123
    return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
124
125


126
127
128
129
130
def file_size(file):
    # Return file size in MB
    return Path(file).stat().st_size / 1e6


131
132
133
134
def check_online():
    # Check internet connectivity
    import socket
    try:
135
        socket.create_connection(("1.1.1.1", 443), 5)  # check host accessibility
136
137
138
139
140
        return True
    except OSError:
        return False


141
142
@try_except
def check_git_status():
143
    # Recommend 'git pull' if code is out of date
144
    msg = ', for updates see https://github.com/ultralytics/yolov5'
145
    print(colorstr('github: '), end='')
146
147
148
149
150
151
152
153
154
155
156
157
158
159
    assert Path('.git').exists(), 'skipping check (not a git repository)' + msg
    assert not is_docker(), 'skipping check (Docker image)' + msg
    assert check_online(), 'skipping check (offline)' + msg

    cmd = 'git fetch && git config --get remote.origin.url'
    url = check_output(cmd, shell=True, timeout=5).decode().strip().rstrip('.git')  # git fetch
    branch = check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip()  # checked out
    n = int(check_output(f'git rev-list {branch}..origin/master --count', shell=True))  # commits behind
    if n > 0:
        s = f"⚠️ WARNING: code is out of date by {n} commit{'s' * (n > 1)}. " \
            f"Use 'git pull' to update or 'git clone {url}' to download latest."
    else:
        s = f'up to date with {url} ✅'
    print(emojis(s))  # emoji-safe
Glenn Jocher's avatar
Glenn Jocher committed
160
161


162
def check_python(minimum='3.6.2'):
Glenn Jocher's avatar
Glenn Jocher committed
163
    # Check current python version vs. required python version
164
165
166
167
168
169
170
171
    check_version(platform.python_version(), minimum, name='Python ')


def check_version(current='0.0.0', minimum='0.0.0', name='version ', pinned=False):
    # Check version vs. required version
    current, minimum = (pkg.parse_version(x) for x in (current, minimum))
    result = (current == minimum) if pinned else (current >= minimum)
    assert result, f'{name}{minimum} required by YOLOv5, but {name}{current} is currently installed'
Glenn Jocher's avatar
Glenn Jocher committed
172
173


174
@try_except
175
176
def check_requirements(requirements='requirements.txt', exclude=()):
    # Check installed dependencies meet requirements (pass *.txt file or list of packages)
177
    prefix = colorstr('red', 'bold', 'requirements:')
Glenn Jocher's avatar
Glenn Jocher committed
178
    check_python()  # check python version
179
180
    if isinstance(requirements, (str, Path)):  # requirements.txt file
        file = Path(requirements)
181
        assert file.exists(), f"{prefix} {file.resolve()} not found, check failed."
182
183
184
        requirements = [f'{x.name}{x.specifier}' for x in pkg.parse_requirements(file.open()) if x.name not in exclude]
    else:  # list or tuple of packages
        requirements = [x for x in requirements if x not in exclude]
185

186
    n = 0  # number of packages updates
187
188
189
190
    for r in requirements:
        try:
            pkg.require(r)
        except Exception as e:  # DistributionNotFound or VersionConflict if requirements not met
191
            print(f"{prefix} {r} not found and is required by YOLOv5, attempting auto-update...")
192
            try:
193
194
195
                assert check_online(), f"'pip install {r}' skipped (offline)"
                print(check_output(f"pip install '{r}'", shell=True).decode())
                n += 1
196
197
            except Exception as e:
                print(f'{prefix} {e}')
198
199

    if n:  # if packages updated
200
201
        source = file.resolve() if 'file' in locals() else requirements
        s = f"{prefix} {n} package{'s' * (n > 1)} updated per {source}\n" \
202
            f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n"
203
        print(emojis(s))
204
205


Glenn Jocher's avatar
Glenn Jocher committed
206
207
208
209
210
211
212
213
def check_img_size(imgsz, s=32, floor=0):
    # Verify image size is a multiple of stride s in each dimension
    if isinstance(imgsz, int):  # integer i.e. img_size=640
        new_size = max(make_divisible(imgsz, int(s)), floor)
    else:  # list i.e. img_size=[640, 480]
        new_size = [max(make_divisible(x, int(s)), floor) for x in imgsz]
    if new_size != imgsz:
        print(f'WARNING: --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}')
Glenn Jocher's avatar
Glenn Jocher committed
214
    return new_size
215
216


Glenn Jocher's avatar
Glenn Jocher committed
217
218
219
def check_imshow():
    # Check if environment supports image displays
    try:
Glenn Jocher's avatar
Glenn Jocher committed
220
221
        assert not is_docker(), 'cv2.imshow() is disabled in Docker environments'
        assert not is_colab(), 'cv2.imshow() is disabled in Google Colab environments'
Glenn Jocher's avatar
Glenn Jocher committed
222
223
224
225
226
227
        cv2.imshow('test', np.zeros((1, 1, 3)))
        cv2.waitKey(1)
        cv2.destroyAllWindows()
        cv2.waitKey(1)
        return True
    except Exception as e:
Glenn Jocher's avatar
Glenn Jocher committed
228
        print(f'WARNING: Environment does not support cv2.imshow() or PIL Image.show() image displays\n{e}')
Glenn Jocher's avatar
Glenn Jocher committed
229
230
231
        return False


Glenn Jocher's avatar
Glenn Jocher committed
232
def check_file(file):
233
234
235
    # Search/download file (if necessary) and return path
    file = str(file)  # convert to str()
    if Path(file).is_file() or file == '':  # exists
Glenn Jocher's avatar
Glenn Jocher committed
236
        return file
Glenn Jocher's avatar
Glenn Jocher committed
237
238
239
    elif file.startswith(('http:/', 'https:/')):  # download
        url = str(Path(file)).replace(':/', '://')  # Pathlib turns :// -> :/
        file = Path(urllib.parse.unquote(file)).name.split('?')[0]  # '%2F' to '/', split https://url.com/file.txt?auth
240
241
242
243
244
        print(f'Downloading {url} to {file}...')
        torch.hub.download_url_to_file(url, file)
        assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}'  # check
        return file
    else:  # search
Glenn Jocher's avatar
Glenn Jocher committed
245
        files = glob.glob('./**/' + file, recursive=True)  # find file
246
        assert len(files), f'File not found: {file}'  # assert file was found
247
        assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}"  # assert unique
Glenn Jocher's avatar
Glenn Jocher committed
248
        return files[0]  # return file
Glenn Jocher's avatar
Glenn Jocher committed
249
250


251
def check_dataset(data, autodownload=True):
252
253
254
255
256
257
258
259
260
261
262
263
    # Download and/or unzip dataset if not found locally
    # Usage: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128_with_yaml.zip

    # Download (optional)
    extract_dir = ''
    if isinstance(data, (str, Path)) and str(data).endswith('.zip'):  # i.e. gs://bucket/dir/coco128.zip
        download(data, dir='../datasets', unzip=True, delete=False, curl=False, threads=1)
        data = next((Path('../datasets') / Path(data).stem).rglob('*.yaml'))
        extract_dir, autodownload = data.parent, False

    # Read yaml (optional)
    if isinstance(data, (str, Path)):
264
        with open(data, errors='ignore') as f:
265
266
267
268
269
270
271
            data = yaml.safe_load(f)  # dictionary

    # Parse yaml
    path = extract_dir or Path(data.get('path') or '')  # optional 'path' default to '.'
    for k in 'train', 'val', 'test':
        if data.get(k):  # prepend path
            data[k] = str(path / data[k]) if isinstance(data[k], str) else [str(path / x) for x in data[k]]
272

273
274
    assert 'nc' in data, "Dataset 'nc' key missing."
    if 'names' not in data:
275
        data['names'] = [f'class{i}' for i in range(data['nc'])]  # assign class names if missing
276
    train, val, test, s = [data.get(x) for x in ('train', 'val', 'test', 'download')]
277
    if val:
278
279
280
        val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])]  # val path
        if not all(x.exists() for x in val):
            print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()])
281
            if s and autodownload:  # download script
Hatovix's avatar
Hatovix committed
282
283
                if s.startswith('http') and s.endswith('.zip'):  # URL
                    f = Path(s).name  # filename
284
                    print(f'Downloading {s} ...')
285
                    torch.hub.download_url_to_file(s, f)
286
287
                    root = path.parent if 'path' in data else '..'  # unzip directory i.e. '../'
                    Path(root).mkdir(parents=True, exist_ok=True)  # create root
288
                    r = os.system(f'unzip -q {f} -d {root} && rm {f}')  # unzip
289
290
                elif s.startswith('bash '):  # bash script
                    print(f'Running {s} ...')
Hatovix's avatar
Hatovix committed
291
                    r = os.system(s)
292
                else:  # python script
293
                    r = exec(s, {'yaml': data})  # return None
294
                print('Dataset autodownload %s\n' % ('success' if r in (0, None) else 'failure'))  # print result
Hatovix's avatar
Hatovix committed
295
296
            else:
                raise Exception('Dataset not found.')
297

298
299
    return data  # dictionary

300

Glenn Jocher's avatar
Glenn Jocher committed
301
def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1):
302
    # Multi-threaded file download and unzip function, used in data.yaml for autodownload
303
304
305
    def download_one(url, dir):
        # Download 1 file
        f = dir / Path(url).name  # filename
306
307
308
        if Path(url).is_file():  # exists in current path
            Path(url).rename(f)  # move to dir
        elif not f.exists():
309
            print(f'Downloading {url} to {f}...')
Glenn Jocher's avatar
Glenn Jocher committed
310
311
312
313
            if curl:
                os.system(f"curl -L '{url}' -o '{f}' --retry 9 -C -")  # curl download, retry and resume on fail
            else:
                torch.hub.download_url_to_file(url, f, progress=True)  # torch download
314
        if unzip and f.suffix in ('.zip', '.gz'):
315
316
            print(f'Unzipping {f}...')
            if f.suffix == '.zip':
317
                s = f'unzip -qo {f} -d {dir}'  # unzip -quiet -overwrite
318
            elif f.suffix == '.gz':
Glenn Jocher's avatar
Glenn Jocher committed
319
320
321
322
                s = f'tar xfz {f} --directory {f.parent}'  # unzip
            if delete:  # delete zip file after unzip
                s += f' && rm {f}'
            os.system(s)
323
324
325

    dir = Path(dir)
    dir.mkdir(parents=True, exist_ok=True)  # make directory
Glenn Jocher's avatar
Glenn Jocher committed
326
    if threads > 1:
327
328
329
330
        pool = ThreadPool(threads)
        pool.imap(lambda x: download_one(*x), zip(url, repeat(dir)))  # multi-threaded
        pool.close()
        pool.join()
331
    else:
332
        for u in [url] if isinstance(url, (str, Path)) else url:
333
334
335
            download_one(u, dir)


Glenn Jocher's avatar
Glenn Jocher committed
336
def make_divisible(x, divisor):
337
    # Returns x evenly divisible by divisor
Glenn Jocher's avatar
Glenn Jocher committed
338
339
340
    return math.ceil(x / divisor) * divisor


341
342
343
344
345
def clean_str(s):
    # Cleans a string by replacing special characters with underscore _
    return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s)


346
def one_cycle(y1=0.0, y2=1.0, steps=100):
Glenn Jocher's avatar
Glenn Jocher committed
347
    # lambda function for sinusoidal ramp from y1 to y2 https://arxiv.org/pdf/1812.01187.pdf
348
349
350
    return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1


Glenn Jocher's avatar
Glenn Jocher committed
351
352
def colorstr(*input):
    # Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e.  colorstr('blue', 'hello world')
Glenn Jocher's avatar
Glenn Jocher committed
353
    *args, string = input if len(input) > 1 else ('blue', 'bold', input[0])  # color arguments, string
Glenn Jocher's avatar
Glenn Jocher committed
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
    colors = {'black': '\033[30m',  # basic colors
              'red': '\033[31m',
              'green': '\033[32m',
              'yellow': '\033[33m',
              'blue': '\033[34m',
              'magenta': '\033[35m',
              'cyan': '\033[36m',
              'white': '\033[37m',
              'bright_black': '\033[90m',  # bright colors
              'bright_red': '\033[91m',
              'bright_green': '\033[92m',
              'bright_yellow': '\033[93m',
              'bright_blue': '\033[94m',
              'bright_magenta': '\033[95m',
              'bright_cyan': '\033[96m',
              'bright_white': '\033[97m',
              'end': '\033[0m',  # misc
              'bold': '\033[1m',
372
              'underline': '\033[4m'}
Glenn Jocher's avatar
Glenn Jocher committed
373
    return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
Glenn Jocher's avatar
Glenn Jocher committed
374
375


Glenn Jocher's avatar
Glenn Jocher committed
376
377
378
379
380
381
382
def labels_to_class_weights(labels, nc=80):
    # Get class weights (inverse frequency) from training labels
    if labels[0] is None:  # no labels loaded
        return torch.Tensor()

    labels = np.concatenate(labels, 0)  # labels.shape = (866643, 5) for COCO
    classes = labels[:, 0].astype(np.int)  # labels = [class xywh]
383
    weights = np.bincount(classes, minlength=nc)  # occurrences per class
Glenn Jocher's avatar
Glenn Jocher committed
384

385
    # Prepend gridpoint count (for uCE training)
Glenn Jocher's avatar
Glenn Jocher committed
386
387
388
389
390
391
392
393
394
395
    # gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum()  # gridpoints per image
    # weights = np.hstack([gpi * len(labels)  - weights.sum() * 9, weights * 9]) ** 0.5  # prepend gridpoints to start

    weights[weights == 0] = 1  # replace empty bins with 1
    weights = 1 / weights  # number of targets per class
    weights /= weights.sum()  # normalize
    return torch.from_numpy(weights)


def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):
396
397
    # Produces image weights based on class_weights and image contents
    class_counts = np.array([np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels])
Glenn Jocher's avatar
Glenn Jocher committed
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
    image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1)
    # index = random.choices(range(n), weights=image_weights, k=1)  # weight image sample
    return image_weights


def coco80_to_coco91_class():  # converts 80-index (val2014) to 91-index (paper)
    # https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
    # a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n')
    # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
    # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)]  # darknet to coco
    # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)]  # coco to darknet
    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,
         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,
         64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
    return x


def xyxy2xywh(x):
    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
417
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
Glenn Jocher's avatar
Glenn Jocher committed
418
419
420
421
422
423
424
425
426
    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
    y[:, 2] = x[:, 2] - x[:, 0]  # width
    y[:, 3] = x[:, 3] - x[:, 1]  # height
    return y


def xywh2xyxy(x):
    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
427
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
Glenn Jocher's avatar
Glenn Jocher committed
428
429
430
431
432
433
434
    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
    return y


435
def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
Glenn Jocher's avatar
Glenn Jocher committed
436
437
438
439
440
441
442
443
444
    # Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw  # top left x
    y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh  # top left y
    y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw  # bottom right x
    y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh  # bottom right y
    return y


445
def xyxy2xywhn(x, w=640, h=640, clip=False, eps=0.0):
Yonghye Kwon's avatar
Yonghye Kwon committed
446
    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] normalized where xy1=top-left, xy2=bottom-right
447
    if clip:
448
        clip_coords(x, (h - eps, w - eps))  # warning: inplace clip
Yonghye Kwon's avatar
Yonghye Kwon committed
449
450
451
452
453
454
455
456
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = ((x[:, 0] + x[:, 2]) / 2) / w  # x center
    y[:, 1] = ((x[:, 1] + x[:, 3]) / 2) / h  # y center
    y[:, 2] = (x[:, 2] - x[:, 0]) / w  # width
    y[:, 3] = (x[:, 3] - x[:, 1]) / h  # height
    return y


457
458
459
460
461
462
463
464
465
466
467
468
469
def xyn2xy(x, w=640, h=640, padw=0, padh=0):
    # Convert normalized segments into pixel segments, shape (n,2)
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = w * x[:, 0] + padw  # top left x
    y[:, 1] = h * x[:, 1] + padh  # top left y
    return y


def segment2box(segment, width=640, height=640):
    # Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy)
    x, y = segment.T  # segment xy
    inside = (x >= 0) & (y >= 0) & (x <= width) & (y <= height)
    x, y, = x[inside], y[inside]
470
    return np.array([x.min(), y.min(), x.max(), y.max()]) if any(x) else np.zeros((1, 4))  # xyxy
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490


def segments2boxes(segments):
    # Convert segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh)
    boxes = []
    for s in segments:
        x, y = s.T  # segment xy
        boxes.append([x.min(), y.min(), x.max(), y.max()])  # cls, xyxy
    return xyxy2xywh(np.array(boxes))  # cls, xywh


def resample_segments(segments, n=1000):
    # Up-sample an (n,2) segment
    for i, s in enumerate(segments):
        x = np.linspace(0, len(s) - 1, n)
        xp = np.arange(len(s))
        segments[i] = np.concatenate([np.interp(x, xp, s[:, i]) for i in range(2)]).reshape(2, -1).T  # segment xy
    return segments


Glenn Jocher's avatar
Glenn Jocher committed
491
492
493
def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
    # Rescale coords (xyxy) from img1_shape to img0_shape
    if ratio_pad is None:  # calculate from img0_shape
494
        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
Glenn Jocher's avatar
Glenn Jocher committed
495
496
497
498
499
500
501
502
503
504
505
506
        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
    else:
        gain = ratio_pad[0][0]
        pad = ratio_pad[1]

    coords[:, [0, 2]] -= pad[0]  # x padding
    coords[:, [1, 3]] -= pad[1]  # y padding
    coords[:, :4] /= gain
    clip_coords(coords, img0_shape)
    return coords


507
def clip_coords(boxes, shape):
Glenn Jocher's avatar
Glenn Jocher committed
508
    # Clip bounding xyxy bounding boxes to image shape (height, width)
509
510
511
512
513
514
515
516
    if isinstance(boxes, torch.Tensor):  # faster individually
        boxes[:, 0].clamp_(0, shape[1])  # x1
        boxes[:, 1].clamp_(0, shape[0])  # y1
        boxes[:, 2].clamp_(0, shape[1])  # x2
        boxes[:, 3].clamp_(0, shape[0])  # y2
    else:  # np.array (faster grouped)
        boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, shape[1])  # x1, x2
        boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0])  # y1, y2
Glenn Jocher's avatar
Glenn Jocher committed
517
518


519
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
520
                        labels=(), max_det=300):
521
    """Runs Non-Maximum Suppression (NMS) on inference results
Glenn Jocher's avatar
Glenn Jocher committed
522
523

    Returns:
524
         list of detections, on (n,6) tensor per image [xyxy, conf, cls]
Glenn Jocher's avatar
Glenn Jocher committed
525
    """
Glenn Jocher's avatar
Glenn Jocher committed
526

527
    nc = prediction.shape[2] - 5  # number of classes
Glenn Jocher's avatar
Glenn Jocher committed
528
    xc = prediction[..., 4] > conf_thres  # candidates
Glenn Jocher's avatar
Glenn Jocher committed
529

530
531
532
533
    # Checks
    assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
    assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'

Glenn Jocher's avatar
Glenn Jocher committed
534
535
    # Settings
    min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
Glenn Jocher's avatar
Glenn Jocher committed
536
    max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
Glenn Jocher's avatar
Glenn Jocher committed
537
    time_limit = 10.0  # seconds to quit after
Glenn Jocher's avatar
Glenn Jocher committed
538
    redundant = True  # require redundant detections
539
    multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
540
    merge = False  # use merge-NMS
Glenn Jocher's avatar
Glenn Jocher committed
541
542

    t = time.time()
543
    output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
544
    mad_output = None
545
    logits = []
Glenn Jocher's avatar
Glenn Jocher committed
546
547
    for xi, x in enumerate(prediction):  # image index, image inference
        # Apply constraints
Glenn Jocher's avatar
Glenn Jocher committed
548
        # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
Glenn Jocher's avatar
Glenn Jocher committed
549
        x = x[xc[xi]]  # confidence
Glenn Jocher's avatar
Glenn Jocher committed
550

551
552
553
554
555
556
557
558
559
        # Cat apriori labels if autolabelling
        if labels and len(labels[xi]):
            l = labels[xi]
            v = torch.zeros((len(l), nc + 5), device=x.device)
            v[:, :4] = l[:, 1:5]  # box
            v[:, 4] = 1.0  # conf
            v[range(len(l)), l[:, 0].long() + 5] = 1.0  # cls
            x = torch.cat((x, v), 0)

Glenn Jocher's avatar
Glenn Jocher committed
560
561
562
563
564
        # If none remain process next image
        if not x.shape[0]:
            continue

        # Compute conf
Pavlo Beylin's avatar
Pavlo Beylin committed
565
566
567
        # x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
        x_hat = x.detach().clone()
        x[:, 5:] *= x_hat[:, 4:5]  # conf = obj_conf * cls_conf
Glenn Jocher's avatar
Glenn Jocher committed
568
569
570
571

        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
        box = xywh2xyxy(x[:, :4])

572
573
574
        mad_output = x.clone()
        mad_output[:, :4] = box

Glenn Jocher's avatar
Glenn Jocher committed
575
576
        # Detections matrix nx6 (xyxy, conf, cls)
        if multi_label:
Glenn Jocher's avatar
Glenn Jocher committed
577
            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
Glenn Jocher's avatar
Glenn Jocher committed
578
            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
Glenn Jocher's avatar
Glenn Jocher committed
579
        else:  # best class only
Glenn Jocher's avatar
Glenn Jocher committed
580
            conf, j = x[:, 5:].max(1, keepdim=True)
581
            logits.append(x[:, 5:][conf.view(-1) > conf_thres])
Glenn Jocher's avatar
Glenn Jocher committed
582
            x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
Glenn Jocher's avatar
Glenn Jocher committed
583
584

        # Filter by class
585
        if classes is not None:
Glenn Jocher's avatar
Glenn Jocher committed
586
            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
Glenn Jocher's avatar
Glenn Jocher committed
587
588
589
590
591

        # Apply finite constraint
        # if not torch.isfinite(x).all():
        #     x = x[torch.isfinite(x).all(1)]

Glenn Jocher's avatar
Glenn Jocher committed
592
        # Check shape
Glenn Jocher's avatar
Glenn Jocher committed
593
        n = x.shape[0]  # number of boxes
Glenn Jocher's avatar
Glenn Jocher committed
594
        if not n:  # no boxes
Glenn Jocher's avatar
Glenn Jocher committed
595
            continue
Glenn Jocher's avatar
Glenn Jocher committed
596
597
        elif n > max_nms:  # excess boxes
            x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence
Glenn Jocher's avatar
Glenn Jocher committed
598
599

        # Batched NMS
Glenn Jocher's avatar
Glenn Jocher committed
600
601
        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
Glenn Jocher's avatar
Glenn Jocher committed
602
        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
Glenn Jocher's avatar
updates    
Glenn Jocher committed
603
604
        if i.shape[0] > max_det:  # limit detections
            i = i[:max_det]
Glenn Jocher's avatar
Glenn Jocher committed
605
        if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
Glenn Jocher's avatar
Glenn Jocher committed
606
607
608
609
610
611
            # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
            weights = iou * scores[None]  # box weights
            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
            if redundant:
                i = i[iou.sum(1) > 1]  # require redundancy
Glenn Jocher's avatar
Glenn Jocher committed
612
613

        output[xi] = x[i]
614
        logits[xi] = [logits[0][x] for x in i]
Glenn Jocher's avatar
Glenn Jocher committed
615
        if (time.time() - t) > time_limit:
Glenn Jocher's avatar
Glenn Jocher committed
616
            print(f'WARNING: NMS time limit {time_limit}s exceeded')
Glenn Jocher's avatar
Glenn Jocher committed
617
618
            break  # time limit exceeded

619
    return output, mad_output, logits[0]
Glenn Jocher's avatar
Glenn Jocher committed
620
621


Glenn Jocher's avatar
Glenn Jocher committed
622
def strip_optimizer(f='best.pt', s=''):  # from utils.general import *; strip_optimizer()
623
    # Strip optimizer from 'f' to finalize training, optionally save as 's'
Glenn Jocher's avatar
Glenn Jocher committed
624
    x = torch.load(f, map_location=torch.device('cpu'))
Glenn Jocher's avatar
Glenn Jocher committed
625
626
627
    if x.get('ema'):
        x['model'] = x['ema']  # replace model with ema
    for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates':  # keys
628
        x[k] = None
Glenn Jocher's avatar
Glenn Jocher committed
629
    x['epoch'] = -1
630
    x['model'].half()  # to FP16
Glenn Jocher's avatar
Glenn Jocher committed
631
    for p in x['model'].parameters():
632
633
634
        p.requires_grad = False
    torch.save(x, s or f)
    mb = os.path.getsize(s or f) / 1E6  # filesize
Glenn Jocher's avatar
Glenn Jocher committed
635
    print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")
Glenn Jocher's avatar
Glenn Jocher committed
636
637


Glenn Jocher's avatar
Glenn Jocher committed
638
639
640
641
642
643
644
def print_mutation(results, hyp, save_dir, bucket):
    evolve_csv, results_csv, evolve_yaml = save_dir / 'evolve.csv', save_dir / 'results.csv', save_dir / 'hyp_evolve.yaml'
    keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
            'val/box_loss', 'val/obj_loss', 'val/cls_loss') + tuple(hyp.keys())  # [results + hyps]
    keys = tuple(x.strip() for x in keys)
    vals = results + tuple(hyp.values())
    n = len(keys)
Glenn Jocher's avatar
Glenn Jocher committed
645

Glenn Jocher's avatar
Glenn Jocher committed
646
    # Download (optional)
Glenn Jocher's avatar
Glenn Jocher committed
647
    if bucket:
Glenn Jocher's avatar
Glenn Jocher committed
648
649
650
651
652
653
654
655
        url = f'gs://{bucket}/evolve.csv'
        if gsutil_getsize(url) > (os.path.getsize(evolve_csv) if os.path.exists(evolve_csv) else 0):
            os.system(f'gsutil cp {url} {save_dir}')  # download evolve.csv if larger than local

    # Log to evolve.csv
    s = '' if evolve_csv.exists() else (('%20s,' * n % keys).rstrip(',') + '\n')  # add header
    with open(evolve_csv, 'a') as f:
        f.write(s + ('%20.5g,' * n % vals).rstrip(',') + '\n')
Glenn Jocher's avatar
Glenn Jocher committed
656

Glenn Jocher's avatar
Glenn Jocher committed
657
658
659
    # Print to screen
    print(colorstr('evolve: ') + ', '.join(f'{x.strip():>20s}' for x in keys))
    print(colorstr('evolve: ') + ', '.join(f'{x:20.5g}' for x in vals), end='\n\n\n')
Glenn Jocher's avatar
Glenn Jocher committed
660

661
    # Save yaml
Glenn Jocher's avatar
Glenn Jocher committed
662
663
664
665
666
667
668
669
670
    with open(evolve_yaml, 'w') as f:
        data = pd.read_csv(evolve_csv)
        data = data.rename(columns=lambda x: x.strip())  # strip keys
        i = np.argmax(fitness(data.values[:, :7]))  #
        f.write(f'# YOLOv5 Hyperparameter Evolution Results\n' +
                f'# Best generation: {i}\n' +
                f'# Last generation: {len(data)}\n' +
                f'# ' + ', '.join(f'{x.strip():>20s}' for x in keys[:7]) + '\n' +
                f'# ' + ', '.join(f'{x:>20.5g}' for x in data.values[i, :7]) + '\n\n')
671
        yaml.safe_dump(hyp, f, sort_keys=False)
672

673
    if bucket:
Glenn Jocher's avatar
Glenn Jocher committed
674
        os.system(f'gsutil cp {evolve_csv} {evolve_yaml} gs://{bucket}')  # upload
675

Glenn Jocher's avatar
Glenn Jocher committed
676
677

def apply_classifier(x, model, img, im0):
678
    # Apply a second stage classifier to yolo outputs
Glenn Jocher's avatar
Glenn Jocher committed
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
    im0 = [im0] if isinstance(im0, np.ndarray) else im0
    for i, d in enumerate(x):  # per image
        if d is not None and len(d):
            d = d.clone()

            # Reshape and pad cutouts
            b = xyxy2xywh(d[:, :4])  # boxes
            b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1)  # rectangle to square
            b[:, 2:] = b[:, 2:] * 1.3 + 30  # pad
            d[:, :4] = xywh2xyxy(b).long()

            # Rescale boxes from img_size to im0 size
            scale_coords(img.shape[2:], d[:, :4], im0[i].shape)

            # Classes
            pred_cls1 = d[:, 5].long()
            ims = []
            for j, a in enumerate(d):  # per item
                cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])]
                im = cv2.resize(cutout, (224, 224))  # BGR
699
                # cv2.imwrite('example%i.jpg' % j, cutout)
Glenn Jocher's avatar
Glenn Jocher committed
700
701
702
703
704
705
706
707
708
709
710
711

                im = im[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
                im = np.ascontiguousarray(im, dtype=np.float32)  # uint8 to float32
                im /= 255.0  # 0 - 255 to 0.0 - 1.0
                ims.append(im)

            pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1)  # classifier prediction
            x[i] = x[i][pred_cls1 == pred_cls2]  # retain matching class detections

    return x


712
713
def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BGR=False, save=True):
    # Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
714
715
716
717
718
719
720
    xyxy = torch.tensor(xyxy).view(-1, 4)
    b = xyxy2xywh(xyxy)  # boxes
    if square:
        b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1)  # attempt rectangle to square
    b[:, 2:] = b[:, 2:] * gain + pad  # box wh * gain + pad
    xyxy = xywh2xyxy(b).long()
    clip_coords(xyxy, im.shape)
721
722
723
724
    crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
    if save:
        cv2.imwrite(str(increment_path(file, mkdir=True).with_suffix('.jpg')), crop)
    return crop
725
726
727


def increment_path(path, exist_ok=False, sep='', mkdir=False):
728
    # Increment file or directory path, i.e. runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
729
    path = Path(path)  # os-agnostic
730
    if path.exists() and not exist_ok:
731
732
        suffix = path.suffix
        path = path.with_suffix('')
733
734
735
736
        dirs = glob.glob(f"{path}{sep}*")  # similar paths
        matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
        i = [int(m.groups()[0]) for m in matches if m]  # indices
        n = max(i) + 1 if i else 2  # increment number
737
738
739
740
741
        path = Path(f"{path}{sep}{n}{suffix}")  # update path
    dir = path if path.suffix == '' else path.parent  # directory
    if not dir.exists() and mkdir:
        dir.mkdir(parents=True, exist_ok=True)  # make directory
    return path