Add is_colab() and is_kaggle() (#41)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/tests/data/dataloader/yolopose.py

@@ -1,6 +1,5 @@
 import cv2
 import numpy as np
-import torch
 from omegaconf import OmegaConf
 
 from ultralytics.yolo.data import build_dataloader

ultralytics/tests/functional/test_loaders.py

@@ -1,5 +1,3 @@
-import torch
-
 from ultralytics.yolo.utils.checks import check_yaml
 from ultralytics.yolo.utils.modeling.tasks import DetectionModel
 

ultralytics/yolo/data/augment.py

@@ -8,10 +8,9 @@ import numpy as np
 import torch
 import torchvision.transforms as T
 
-from ..utils import LOGGER
+from ..utils import LOGGER, colorstr
 from ..utils.checks import check_version
 from ..utils.instance import Instances
-from ..utils.loggers import colorstr
 from ..utils.metrics import bbox_ioa
 from ..utils.ops import segment2box
 from .utils import IMAGENET_MEAN, IMAGENET_STD, polygons2masks, polygons2masks_overlap

ultralytics/yolo/engine/trainer.py

@@ -11,7 +11,6 @@ import time
 from collections import defaultdict
 from datetime import datetime
 from pathlib import Path
-from telnetlib import TLS
 from typing import Dict, Union
 
 import torch

ultralytics/yolo/engine/validator.py

@@ -1,7 +1,7 @@
 import logging
 
 import torch
-from omegaconf import DictConfig, OmegaConf
+from omegaconf import OmegaConf
 from tqdm import tqdm
 
 from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG

ultralytics/yolo/utils/__init__.py

@@ -2,13 +2,10 @@ import contextlib
 import logging
 import os
 import platform
+import sys
 from pathlib import Path
 
-from .files import user_config_dir
-from .loggers import emojis, set_logging
-
 # Constants
-
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[2]  # YOLOv5 root directory
 RANK = int(os.getenv('RANK', -1))
@@ -16,8 +13,94 @@ 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
+
+
+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=None, verbose=VERBOSE):
+    # Sets level and returns logger
+    if is_colab() or is_kaggle():
+        for h in logging.root.handlers:
+            logging.root.removeHandler(h)  # remove all handlers associated with the root logger object
+    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
+    log = logging.getLogger(name)
+    log.setLevel(level)
+    handler = logging.StreamHandler()
+    handler.setFormatter(logging.Formatter("%(message)s"))
+    handler.setLevel(level)
+    log.addHandler(handler)
 
-CONFIG_DIR = user_config_dir()  # Ultralytics settings dir
 
 set_logging()  # run before defining LOGGER
 LOGGER = logging.getLogger("yolov5")  # define globally

ultralytics/yolo/utils/anchors.py

@@ -11,9 +11,7 @@ import yaml
 from tqdm import tqdm
 
 from ultralytics.yolo.data import BaseDataset
-from ultralytics.yolo.utils import LOGGER, TryExcept
+from ultralytics.yolo.utils import LOGGER, TryExcept, colorstr
-
-from .loggers import colorstr
 
 PREFIX = colorstr('AutoAnchor: ')
 

ultralytics/yolo/utils/checks.py

@@ -1,5 +1,4 @@
 import glob
-import os
 import platform
 import sys
 import urllib
@@ -9,9 +8,13 @@ from subprocess import check_output
 import pkg_resources as pkg
 import torch
 
-from ultralytics.yolo.utils import AUTOINSTALL, CONFIG_DIR, FONT, LOGGER, ROOT, TryExcept
+from ultralytics.yolo.utils import AUTOINSTALL, FONT, LOGGER, ROOT, USER_CONFIG_DIR, TryExcept, colorstr, emojis
 
-from .loggers import colorstr, emojis
+
+def is_ascii(s=''):
+    # Is string composed of all ASCII (no UTF) characters? (note str().isascii() introduced in python 3.7)
+    s = str(s)  # convert list, tuple, None, etc. to str
+    return len(s.encode().decode('ascii', 'ignore')) == len(s)
 
 
 def check_version(current="0.0.0", minimum="0.0.0", name="version ", pinned=False, hard=False, verbose=False):
@@ -29,7 +32,7 @@ def check_version(current="0.0.0", minimum="0.0.0", name="version ", pinned=Fals
 def check_font(font=FONT, progress=False):
     # Download font to CONFIG_DIR if necessary
     font = Path(font)
-    file = CONFIG_DIR / font.name
+    file = USER_CONFIG_DIR / font.name
     if not font.exists() and not file.exists():
         url = f'https://ultralytics.com/assets/{font.name}'
         LOGGER.info(f'Downloading {url} to {file}...')
@@ -86,12 +89,6 @@ def check_requirements(requirements=ROOT / 'requirements.txt', exclude=(), insta
             LOGGER.warning(f'{prefix} ❌ {e}')
 
 
-def is_ascii(s=''):
-    # Is string composed of all ASCII (no UTF) characters? (note str().isascii() introduced in python 3.7)
-    s = str(s)  # convert list, tuple, None, etc. to str
-    return len(s.encode().decode('ascii', 'ignore')) == len(s)
-
-
 def check_suffix(file='yolov5s.pt', suffix=('.pt',), msg=''):
     # Check file(s) for acceptable suffix
     if file and suffix:
@@ -120,7 +117,7 @@ def check_file(file, suffix=''):
             assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}'  # check
         return file
     elif file.startswith('clearml://'):  # ClearML Dataset ID
-        assert 'clearml' in sys.modules, "ClearML is not installed, so cannot use ClearML dataset. Try running 'pip install clearml'."
+        assert 'clearml' in sys.modules, "Can not use ClearML dataset. Run 'pip install clearml' to install"
         return file
     else:  # search
         files = []

ultralytics/yolo/utils/files.py

@@ -1,6 +1,5 @@
 import contextlib
 import os
-import platform
 from pathlib import Path
 
 import yaml
@@ -19,33 +18,6 @@ class WorkingDirectory(contextlib.ContextDecorator):
         os.chdir(self.cwd)
 
 
-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
-
-
 def increment_path(path, exist_ok=False, sep='', mkdir=False):
     """
     Increment file or directory path, i.e. runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.

ultralytics/yolo/utils/loggers/__init__.py

@@ -1,57 +1 @@
-import logging
-import os
-import platform
-
 from .base import default_callbacks
-
-VERBOSE = str(os.getenv('YOLOv5_VERBOSE', True)).lower() == 'true'  # global verbose mode
-
-
-# console logging utils
-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=None, verbose=VERBOSE):
-    # Sets level and returns logger
-    is_kaggle = os.environ.get("PWD") == "/kaggle/working" and os.environ.get(
-        "KAGGLE_URL_BASE") == "https://www.kaggle.com"
-    is_colab = "COLAB_GPU" in os.environ
-    if is_colab or is_kaggle:
-        for h in logging.root.handlers:
-            logging.root.removeHandler(h)  # remove all handlers associated with the root logger object
-    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
-    log = logging.getLogger(name)
-    log.setLevel(level)
-    handler = logging.StreamHandler()
-    handler.setFormatter(logging.Formatter("%(message)s"))
-    handler.setLevel(level)
-    log.addHandler(handler)

ultralytics/yolo/utils/modeling/modules.py

@@ -3,7 +3,6 @@
 Common modules
 """
 
-import argparse
 import math
 import warnings
 from copy import copy
@@ -19,10 +18,9 @@ from PIL import Image, ImageOps
 from torch.cuda import amp
 
 from ultralytics.yolo.data.augment import LetterBox
-from ultralytics.yolo.utils import LOGGER
+from ultralytics.yolo.utils import LOGGER, colorstr
 from ultralytics.yolo.utils.checks import check_version
 from ultralytics.yolo.utils.files import increment_path
-from ultralytics.yolo.utils.loggers import colorstr
 from ultralytics.yolo.utils.ops import Profile, make_divisible, non_max_suppression, scale_boxes, xyxy2xywh
 from ultralytics.yolo.utils.plotting import Annotator, colors, save_one_box
 from ultralytics.yolo.utils.torch_utils import copy_attr, smart_inference_mode

ultralytics/yolo/utils/modeling/tasks.py

@@ -1,10 +1,7 @@
-import time
 from copy import deepcopy
 
 import thop
-import torch.nn as nn
 
-from ultralytics.yolo.utils import LOGGER
 from ultralytics.yolo.utils.anchors import check_anchor_order
 from ultralytics.yolo.utils.modeling import parse_model
 from ultralytics.yolo.utils.modeling.modules import *

ultralytics/yolo/utils/plotting.py

@@ -6,7 +6,7 @@ import numpy as np
 import torch
 from PIL import Image, ImageDraw, ImageFont
 
-from ultralytics.yolo.utils import CONFIG_DIR, FONT
+from ultralytics.yolo.utils import FONT, USER_CONFIG_DIR
 
 from .checks import check_font, check_requirements, is_ascii
 from .files import increment_path
@@ -150,7 +150,7 @@ class Annotator:
 def check_pil_font(font=FONT, size=10):
     # Return a PIL TrueType Font, downloading to CONFIG_DIR if necessary
     font = Path(font)
-    font = font if font.exists() else (CONFIG_DIR / font.name)
+    font = font if font.exists() else (USER_CONFIG_DIR / font.name)
     try:
         return ImageFont.truetype(str(font) if font.exists() else font.name, size)
     except Exception:  # download if missing

ultralytics/yolo/v8/classify/train.py

@@ -8,9 +8,9 @@ import torch
 from ultralytics.yolo import v8
 from ultralytics.yolo.data import build_classification_dataloader
 from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG, BaseTrainer
+from ultralytics.yolo.utils import colorstr
 from ultralytics.yolo.utils.downloads import download
 from ultralytics.yolo.utils.files import WorkingDirectory
-from ultralytics.yolo.utils.loggers import colorstr
 from ultralytics.yolo.utils.torch_utils import LOCAL_RANK, torch_distributed_zero_first