Add notebook_init() to utils/__init__.py (#5488)

* Update __init__.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * notebook_init * notebook_init * notebook_init * notebook_init * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * notebook_init * Created using Colaboratory * [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>
ultralytics · Nov 9, 2021 · 7ebb5e5 · 7ebb5e5
1 parent 7207fe9
commit 7ebb5e5
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 17 deletions.
diff --git a/tutorial.ipynb b/tutorial.ipynb
@@ -368,7 +368,7 @@
         "colab_type": "text"
       },
       "source": [
-        "<a href=\"https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+        "<a href=\"https://colab.research.google.com/github/ultralytics/yolov5/blob/update%2Fnotebook/tutorial.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
       ]
     },
     {
@@ -402,26 +402,24 @@
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
-        "outputId": "e2e839d5-d6fc-409c-e44c-0b0b6aa9319d"
+        "outputId": "3809e5a9-dd41-4577-fe62-5531abf7cca2"
       },
       "source": [
-        "!git clone https://github.com/ultralytics/yolov5  # clone repo\n",
+        "!git clone https://github.com/ultralytics/yolov5  # clone\n",
         "%cd yolov5\n",
-        "%pip install -qr requirements.txt  # install dependencies\n",
+        "%pip install -qr requirements.txt  # install\n",
         "\n",
-        "import torch\n",
-        "from IPython.display import Image, clear_output  # to display images\n",
-        "\n",
-        "clear_output()\n",
-        "print(f\"Setup complete. Using torch {torch.__version__} ({torch.cuda.get_device_properties(0).name if torch.cuda.is_available() else 'CPU'})\")"
+        "from yolov5 import utils\n",
+        "display = utils.notebook_init()  # checks"
       ],
-      "execution_count": 11,
+      "execution_count": 2,
       "outputs": [
         {
           "output_type": "stream",
           "name": "stdout",
           "text": [
-            "Setup complete. Using torch 1.10.0+cu102 (Tesla V100-SXM2-16GB)\n"
+            "YOLOv5 🚀 v6.0-48-g84a8099 torch 1.10.0+cu102 CUDA:0 (Tesla V100-SXM2-16GB, 16160MiB)\n",
+            "Setup complete ✅\n"
           ]
         }
       ]
@@ -458,9 +456,9 @@
       },
       "source": [
         "!python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source data/images\n",
-        "Image(filename='runs/detect/exp/zidane.jpg', width=600)"
+        "display.Image(filename='runs/detect/exp/zidane.jpg', width=600)"
       ],
-      "execution_count": 17,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "stream",
@@ -537,7 +535,7 @@
         "torch.hub.download_url_to_file('https://ultralytics.com/assets/coco2017val.zip', 'tmp.zip')\n",
         "!unzip -q tmp.zip -d ../datasets && rm tmp.zip"
       ],
-      "execution_count": 18,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "display_data",
@@ -568,7 +566,7 @@
         "# Run YOLOv5x on COCO val\n",
         "!python val.py --weights yolov5x.pt --data coco.yaml --img 640 --iou 0.65 --half"
       ],
-      "execution_count": 19,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "stream",
@@ -726,7 +724,7 @@
         "# Train YOLOv5s on COCO128 for 3 epochs\n",
         "!python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt --cache"
       ],
-      "execution_count": 24,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "stream",

diff --git a/utils/__init__.py b/utils/__init__.py
@@ -0,0 +1,18 @@
+# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
+"""
+utils/initialization
+"""
+
+
+def notebook_init():
+    # For YOLOv5 notebooks
+    print('Checking setup...')
+    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
+
+    display.clear_output()
+    select_device(newline=False)
+    print(emojis('Setup complete ✅'))
+    return display
diff --git a/utils/torch_utils.py b/utils/torch_utils.py
@@ -53,7 +53,7 @@ def git_describe(path=Path(__file__).parent):  # path must be a directory
         return ''  # not a git repository
 
 
-def select_device(device='', batch_size=None):
+def select_device(device='', batch_size=None, newline=True):
     # device = 'cpu' or '0' or '0,1,2,3'
     s = f'YOLOv5 🚀 {git_describe() or date_modified()} torch {torch.__version__} '  # string
     device = str(device).strip().lower().replace('cuda:', '')  # to string, 'cuda:0' to '0'
@@ -77,6 +77,8 @@ def select_device(device='', batch_size=None):
     else:
         s += 'CPU\n'
 
+    if not newline:
+        s = s.rstrip()
     LOGGER.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s)  # emoji-safe
     return torch.device('cuda:0' if cuda else 'cpu')