diff --git a/tutorial.ipynb b/tutorial.ipynb
index b013fe694..7763a2606 100644
--- a/tutorial.ipynb
+++ b/tutorial.ipynb
@@ -368,7 +368,7 @@
"colab_type": "text"
},
"source": [
- "
"
+ "
"
]
},
{
@@ -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
index e69de29bb..2b0c89636 100644
--- 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
index b65b69fe1..16289104e 100644
--- 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')