__init__.py 2.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
  2. """
  3. utils/initialization
  4. """
  5. import contextlib
  6. import platform
  7. import threading
  8. def emojis(str=''):
  9. # Return platform-dependent emoji-safe version of string
  10. return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
  11. class TryExcept(contextlib.ContextDecorator):
  12. # YOLOv5 TryExcept class. Usage: @TryExcept() decorator or 'with TryExcept():' context manager
  13. def __init__(self, msg=''):
  14. self.msg = msg
  15. def __enter__(self):
  16. pass
  17. def __exit__(self, exc_type, value, traceback):
  18. if value:
  19. print(emojis(f"{self.msg}{': ' if self.msg else ''}{value}"))
  20. return True
  21. def threaded(func):
  22. # Multi-threads a target function and returns thread. Usage: @threaded decorator
  23. def wrapper(*args, **kwargs):
  24. thread = threading.Thread(target=func, args=args, kwargs=kwargs, daemon=True)
  25. thread.start()
  26. return thread
  27. return wrapper
  28. def join_threads(verbose=False):
  29. # Join all daemon threads, i.e. atexit.register(lambda: join_threads())
  30. main_thread = threading.current_thread()
  31. for t in threading.enumerate():
  32. if t is not main_thread:
  33. if verbose:
  34. print(f'Joining thread {t.name}')
  35. t.join()
  36. def notebook_init(verbose=True):
  37. # Check system software and hardware
  38. print('Checking setup...')
  39. import os
  40. import shutil
  41. from ultralytics.yolo.utils.checks import check_requirements
  42. from utils.general import check_font, is_colab
  43. from utils.torch_utils import select_device # imports
  44. check_font()
  45. import psutil
  46. if check_requirements('wandb', install=False):
  47. os.system('pip uninstall -y wandb') # eliminate unexpected account creation prompt with infinite hang
  48. if is_colab():
  49. shutil.rmtree('/content/sample_data', ignore_errors=True) # remove colab /sample_data directory
  50. # System info
  51. display = None
  52. if verbose:
  53. gb = 1 << 30 # bytes to GiB (1024 ** 3)
  54. ram = psutil.virtual_memory().total
  55. total, used, free = shutil.disk_usage('/')
  56. with contextlib.suppress(Exception): # clear display if ipython is installed
  57. from IPython import display
  58. display.clear_output()
  59. s = f'({os.cpu_count()} CPUs, {ram / gb:.1f} GB RAM, {(total - free) / gb:.1f}/{total / gb:.1f} GB disk)'
  60. else:
  61. s = ''
  62. select_device(newline=False)
  63. print(emojis(f'Setup complete ✅ {s}'))
  64. return display