You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
132 lines
4.6 KiB
132 lines
4.6 KiB
import contextlib
|
|
import logging
|
|
import os
|
|
import platform
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Constants
|
|
FILE = Path(__file__).resolve()
|
|
ROOT = FILE.parents[2] # YOLO
|
|
RANK = int(os.getenv('RANK', -1))
|
|
DATASETS_DIR = ROOT.parent / 'datasets' # YOLOv5 datasets directory
|
|
NUM_THREADS = min(8, max(1, os.cpu_count() - 1)) # number of YOLOv5 multiprocessing threads
|
|
AUTOINSTALL = str(os.getenv('YOLOv5_AUTOINSTALL', True)).lower() == 'true' # global auto-install mode
|
|
FONT = 'Arial.ttf' # https://ultralytics.com/assets/Arial.ttf
|
|
VERBOSE = str(os.getenv('YOLOv5_VERBOSE', True)).lower() == 'true' # global verbose mode
|
|
LOGGING_NAME = 'yolov5'
|
|
|
|
|
|
def is_colab():
|
|
# Is environment a Google Colab instance?
|
|
return 'google.colab' in sys.modules
|
|
|
|
|
|
def is_kaggle():
|
|
# Is environment a Kaggle Notebook?
|
|
return os.environ.get('PWD') == '/kaggle/working' and os.environ.get('KAGGLE_URL_BASE') == 'https://www.kaggle.com'
|
|
|
|
|
|
def is_writeable(dir, test=False):
|
|
# Return True if directory has write permissions, test opening a file with write permissions if test=True
|
|
if not test:
|
|
return os.access(dir, os.W_OK) # possible issues on Windows
|
|
file = Path(dir) / 'tmp.txt'
|
|
try:
|
|
with open(file, 'w'): # open file with write permissions
|
|
pass
|
|
file.unlink() # remove file
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def user_config_dir(dir='Ultralytics', env_var='YOLOV5_CONFIG_DIR'):
|
|
# Return path of user configuration directory. Prefer environment variable if exists. Make dir if required.
|
|
env = os.getenv(env_var)
|
|
if env:
|
|
path = Path(env) # use environment variable
|
|
else:
|
|
cfg = {'Windows': 'AppData/Roaming', 'Linux': '.config', 'Darwin': 'Library/Application Support'} # 3 OS dirs
|
|
path = Path.home() / cfg.get(platform.system(), '') # OS-specific config dir
|
|
path = (path if is_writeable(path) else Path('/tmp')) / dir # GCP and AWS lambda fix, only /tmp is writeable
|
|
path.mkdir(exist_ok=True) # make if required
|
|
return path
|
|
|
|
|
|
USER_CONFIG_DIR = user_config_dir() # Ultralytics settings dir
|
|
|
|
|
|
def emojis(str=''):
|
|
# Return platform-dependent emoji-safe version of string
|
|
return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
|
|
|
|
|
|
def colorstr(*input):
|
|
# Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e. colorstr('blue', 'hello world')
|
|
*args, string = input if len(input) > 1 else ("blue", "bold", input[0]) # color arguments, string
|
|
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",
|
|
"underline": "\033[4m",}
|
|
return "".join(colors[x] for x in args) + f"{string}" + colors["end"]
|
|
|
|
|
|
def set_logging(name=LOGGING_NAME, verbose=True):
|
|
# sets up logging for the given name
|
|
rank = int(os.getenv('RANK', -1)) # rank in world for Multi-GPU trainings
|
|
level = logging.INFO if verbose and rank in {-1, 0} else logging.ERROR
|
|
logging.config.dictConfig({
|
|
"version": 1,
|
|
"disable_existing_loggers": False,
|
|
"formatters": {
|
|
name: {
|
|
"format": "%(message)s"}},
|
|
"handlers": {
|
|
name: {
|
|
"class": "logging.StreamHandler",
|
|
"formatter": name,
|
|
"level": level,}},
|
|
"loggers": {
|
|
name: {
|
|
"level": level,
|
|
"handlers": [name],
|
|
"propagate": False,}}})
|
|
|
|
|
|
set_logging(LOGGING_NAME) # run before defining LOGGER
|
|
LOGGER = logging.getLogger(LOGGING_NAME) # define globally (used in train.py, val.py, detect.py, etc.)
|
|
if platform.system() == 'Windows':
|
|
for fn in LOGGER.info, LOGGER.warning:
|
|
setattr(LOGGER, fn.__name__, lambda x: fn(emojis(x))) # emoji safe logging
|
|
|
|
|
|
class TryExcept(contextlib.ContextDecorator):
|
|
# YOLOv5 TryExcept class. Usage: @TryExcept() decorator or 'with TryExcept():' context manager
|
|
def __init__(self, msg=''):
|
|
self.msg = msg
|
|
|
|
def __enter__(self):
|
|
pass
|
|
|
|
def __exit__(self, exc_type, value, traceback):
|
|
if value:
|
|
print(emojis(f"{self.msg}{': ' if self.msg else ''}{value}"))
|
|
return True
|