Add hardware checks to notebook_init() (#5919)

* Update notebook

* Update notebook

* update string

* update string

* Updates

* Updates

* Updates

* check both ipython and psutil

* remove sample_data if is_colab

* cleanup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

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

tutorial.ipynb

@@ -409,6 +409,7 @@
         "%cd yolov5\n",
         "%pip install -qr requirements.txt  # install\n",
         "\n",
+        "import torch\n",
         "from yolov5 import utils\n",
         "display = utils.notebook_init()  # checks"
       ],
@@ -983,7 +984,7 @@
       "source": [
         "# Reproduce\n",
         "for x in 'yolov5s', 'yolov5m', 'yolov5l', 'yolov5x':\n",
-        "  !python val.py --weights {x}.pt --data coco.yaml --img 640 --conf 0.25 --iou 0.45  # speed\n",
+        "  !python val.py --weights {x}.pt --data coco.yaml --img 640 --task speed  # speed\n",
         "  !python val.py --weights {x}.pt --data coco.yaml --img 640 --conf 0.001 --iou 0.65  # mAP"
       ],
       "execution_count": null,

utils/__init__.py

@@ -4,15 +4,34 @@
 """
 
 
-def notebook_init():
-    # For YOLOv5 notebooks
+def notebook_init(verbose=True):
+    # Check system software and hardware
     print('Checking setup...')
+
+    import os
+    import shutil
+
+    from utils.general import check_requirements, emojis, is_colab
+    from utils.torch_utils import select_device  # imports
+
+    check_requirements(('psutil', 'IPython'))
+    import psutil
     from IPython import display  # to display images and clear console output
 
-    from utils.general import emojis
-    from utils.torch_utils import select_device  # YOLOv5 imports
+    if is_colab():
+        shutil.rmtree('sample_data', ignore_errors=True)  # remove colab /sample_data directory
+
+    if verbose:
+        # System info
+        # gb = 1 / 1000 ** 3  # bytes to GB
+        gib = 1 / 1024 ** 3  # bytes to GiB
+        ram = psutil.virtual_memory().total
+        total, used, free = shutil.disk_usage("/")
+        display.clear_output()
+        s = f'({os.cpu_count()} CPUs, {ram * gib:.1f} GB RAM, {(total - free) * gib:.1f}/{total * gib:.1f} GB disk)'
+    else:
+        s = ''
 
-    display.clear_output()
     select_device(newline=False)
-    print(emojis('Setup complete ✅'))
+    print(emojis(f'Setup complete ✅ {s}'))
     return display