From 7f9d7142c2aa4b9ce73836a3d00a6e74d28fe206 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 11 Jan 2023 00:13:44 +0100
Subject: [PATCH] Minor updates and improvements (#216)

Co-authored-by: Hardik Dava <39372750+hardikdava@users.noreply.github.com>
Co-authored-by: Onuralp Sezer <thunderbirdtr@fedoraproject.org>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .github/workflows/ci.yaml           | 16 ++++++++--------
 docker/Dockerfile                   |  4 ++--
 docker/Dockerfile-cpu               |  2 +-
 examples/tutorial.ipynb             | 14 +++++++-------
 ultralytics/yolo/engine/exporter.py |  2 +-
 ultralytics/yolo/v8/detect/val.py   |  3 ++-
 6 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index acd21bf3..28db853a 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -84,17 +84,17 @@ jobs:
       - name: Test detection
         shell: bash  # for Windows compatibility
         run: |
-          yolo task=detect mode=train model=yolov8n.yaml data=coco128.yaml epochs=1 imgsz=64
-          yolo task=detect mode=val model=runs/detect/train/weights/last.pt imgsz=64
-          yolo task=detect mode=predict model=runs/detect/train/weights/last.pt imgsz=64 source=ultralytics/assets/bus.jpg
-          yolo mode=export model=runs/detect/train/weights/last.pt imgsz=64 format=torchscript
+          yolo task=detect mode=train model=yolov8n.yaml data=coco128.yaml epochs=1 imgsz=32
+          yolo task=detect mode=val model=runs/detect/train/weights/last.pt imgsz=32
+          yolo task=detect mode=predict model=runs/detect/train/weights/last.pt imgsz=32 source=ultralytics/assets/bus.jpg
+          yolo mode=export model=runs/detect/train/weights/last.pt imgsz=32 format=torchscript
       - name: Test segmentation
         shell: bash  # for Windows compatibility
         run: |
-          yolo task=segment mode=train model=yolov8n-seg.yaml data=coco128-seg.yaml epochs=1 imgsz=64
-          yolo task=segment mode=val model=runs/segment/train/weights/last.pt data=coco128-seg.yaml imgsz=64
-          yolo task=segment mode=predict model=runs/segment/train/weights/last.pt imgsz=64 source=ultralytics/assets/bus.jpg
-          yolo mode=export model=runs/segment/train/weights/last.pt imgsz=64 format=torchscript
+          yolo task=segment mode=train model=yolov8n-seg.yaml data=coco128-seg.yaml epochs=1 imgsz=32
+          yolo task=segment mode=val model=runs/segment/train/weights/last.pt data=coco128-seg.yaml imgsz=32
+          yolo task=segment mode=predict model=runs/segment/train/weights/last.pt imgsz=32 source=ultralytics/assets/bus.jpg
+          yolo mode=export model=runs/segment/train/weights/last.pt imgsz=32 format=torchscript
       - name: Test classification
         shell: bash  # for Windows compatibility
         run: |
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 146293f3..99f219dd 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, GPL-3.0 license
 # Builds ultralytics/ultralytics:latest image on DockerHub https://hub.docker.com/r/ultralytics/ultralytics
-# Image is CUDA-optimized for YOLOv5 single/multi-GPU training and inference
+# Image is CUDA-optimized for YOLOv8 single/multi-GPU training and inference
 
 # Start FROM NVIDIA PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
 FROM nvcr.io/nvidia/pytorch:22.12-py3
@@ -22,7 +22,7 @@ RUN git clone https://github.com/ultralytics/ultralytics /usr/src/ultralytics
 
 # Install pip packages
 RUN python -m pip install --upgrade pip wheel
-RUN pip uninstall -y Pillow torchtext  # torch torchvision
+RUN pip uninstall -y Pillow torchtext torch torchvision
 RUN pip install --no-cache ultralytics albumentations comet gsutil notebook Pillow>=9.1.0 \
     'opencv-python<4.6.0.66' \
     --extra-index-url https://download.pytorch.org/whl/cu113
diff --git a/docker/Dockerfile-cpu b/docker/Dockerfile-cpu
index dcd7b2d4..41b655c8 100644
--- a/docker/Dockerfile-cpu
+++ b/docker/Dockerfile-cpu
@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, GPL-3.0 license
 # Builds ultralytics/ultralytics:latest-cpu image on DockerHub https://hub.docker.com/r/ultralytics/ultralytics
-# Image is CPU-optimized for ONNX, OpenVINO and PyTorch YOLOv5 deployments
+# Image is CPU-optimized for ONNX, OpenVINO and PyTorch YOLOv8 deployments
 
 # Start FROM Ubuntu image https://hub.docker.com/_/ubuntu
 FROM ubuntu:20.04
diff --git a/examples/tutorial.ipynb b/examples/tutorial.ipynb
index d2be37c5..6a05261e 100644
--- a/examples/tutorial.ipynb
+++ b/examples/tutorial.ipynb
@@ -533,14 +533,14 @@
         "from ultralytics import YOLO\n",
         "\n",
         "# Load a model\n",
-        "model = YOLO('yolov8n.yaml')    # build a new model from scratch\n",
-        "model = YOLO('yolov8n.pt')      # load a pretrained model (recommended for training)\n",
+        "model = YOLO('yolov8n.yaml')  # build a new model from scratch\n",
+        "model = YOLO('yolov8n.pt')  # load a pretrained model (recommended for training)\n",
         "\n",
-        "# Model usage\n",
-        "results = model.train(data='coco128.yaml', epochs=3)                      # train the model\n",
-        "results = model.val(data='coco128.yaml')                                  # evaluate model performance on the validation set\n",
-        "results = model.predict(source='https://ultralytics.com/images/bus.jpg')  # predict on an image\n",
-        "success = YOLO('yolov8n.pt').export(format='onnx')                        # export a model to ONNX format"
+        "# Use the model\n",
+        "results = model.train(data='coco128.yaml', epochs=3)  # train the model\n",
+        "results = model.val()  # evaluate model performance on the validation set\n",
+        "results = model('https://ultralytics.com/images/bus.jpg')  # predict on an image\n",
+        "success = YOLO('yolov8n.pt').export(format='onnx')  # export a model to ONNX format"
       ],
       "metadata": {
         "id": "bpF9-vS_DAaf"
diff --git a/ultralytics/yolo/engine/exporter.py b/ultralytics/yolo/engine/exporter.py
index 8b58ce14..c3188594 100644
--- a/ultralytics/yolo/engine/exporter.py
+++ b/ultralytics/yolo/engine/exporter.py
@@ -378,7 +378,7 @@ class Exporter:
         LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
         f = self.file.with_suffix('.mlmodel')
 
-        model = iOSModel(self.model, self.im) if self.args.nms else self.model
+        model = iOSModel(self.model, self.im).eval() if self.args.nms else self.model
         ts = torch.jit.trace(model, self.im, strict=False)  # TorchScript model
         ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=self.im.shape, scale=1 / 255, bias=[0, 0, 0])])
         bits, mode = (8, 'kmeans_lut') if self.args.int8 else (16, 'linear') if self.args.half else (32, None)
diff --git a/ultralytics/yolo/v8/detect/val.py b/ultralytics/yolo/v8/detect/val.py
index de62002e..0fb4f903 100644
--- a/ultralytics/yolo/v8/detect/val.py
+++ b/ultralytics/yolo/v8/detect/val.py
@@ -43,7 +43,8 @@ class DetectionValidator(BaseValidator):
 
     def init_metrics(self, model):
         head = model.model[-1] if self.training else model.model.model[-1]
-        self.is_coco = self.data.get('val', '').endswith(f'coco{os.sep}val2017.txt')  # is COCO dataset
+        val = self.data.get('val', '')  # validation path
+        self.is_coco = isinstance(val, str) and val.endswith(f'coco{os.sep}val2017.txt')  # is COCO dataset
         self.class_map = ops.coco80_to_coco91_class() if self.is_coco else list(range(1000))
         self.args.save_json |= self.is_coco and not self.training  # run on final val if training COCO
         self.nc = head.nc