ulralytics 8.0.199 *.npy image loading exception handling (#5683)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: snyk-bot <snyk-bot@snyk.io>
Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

tests/test_integrations.py

@@ -1,12 +1,20 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
+import contextlib
+from pathlib import Path
 
 import pytest
 
-from ultralytics import YOLO
-from ultralytics.utils import SETTINGS, checks
+from ultralytics import YOLO, download
+from ultralytics.utils import ASSETS, DATASETS_DIR, ROOT, SETTINGS, WEIGHTS_DIR
+from ultralytics.utils.checks import check_requirements
+
+MODEL = WEIGHTS_DIR / 'path with spaces' / 'yolov8n.pt'  # test spaces in path
+CFG = 'yolov8n.yaml'
+SOURCE = ASSETS / 'bus.jpg'
+TMP = (ROOT / '../tests/tmp').resolve()  # temp directory for test files
 
 
-@pytest.mark.skipif(not checks.check_requirements('ray', install=False), reason='RayTune not installed')
+@pytest.mark.skipif(not check_requirements('ray', install=False), reason='ray[tune] not installed')
 def test_model_ray_tune():
     """Tune YOLO model with Ray optimization library."""
     YOLO('yolov8n-cls.yaml').tune(use_ray=True,
@@ -19,8 +27,90 @@ def test_model_ray_tune():
                                   device='cpu')
 
 
-@pytest.mark.skipif(not checks.check_requirements('mlflow', install=False), reason='MLflow not installed')
+@pytest.mark.skipif(not check_requirements('mlflow', install=False), reason='mlflow not installed')
 def test_mlflow():
     """Test training with MLflow tracking enabled."""
     SETTINGS['mlflow'] = True
     YOLO('yolov8n-cls.yaml').train(data='imagenet10', imgsz=32, epochs=3, plots=False, device='cpu')
+
+
+@pytest.mark.skipif(not check_requirements('tritonclient', install=False), reason='tritonclient[all] not installed')
+def test_triton():
+    """Test NVIDIA Triton Server functionalities."""
+    check_requirements('tritonclient[all]')
+    import subprocess
+    import time
+
+    from tritonclient.http import InferenceServerClient  # noqa
+
+    # Create variables
+    model_name = 'yolo'
+    triton_repo_path = TMP / 'triton_repo'
+    triton_model_path = triton_repo_path / model_name
+
+    # Export model to ONNX
+    f = YOLO(MODEL).export(format='onnx', dynamic=True)
+
+    # Prepare Triton repo
+    (triton_model_path / '1').mkdir(parents=True, exist_ok=True)
+    Path(f).rename(triton_model_path / '1' / 'model.onnx')
+    (triton_model_path / 'config.pdtxt').touch()
+
+    # Define image https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tritonserver
+    tag = 'nvcr.io/nvidia/tritonserver:23.09-py3'  # 6.4 GB
+
+    # Pull the image
+    subprocess.call(f'docker pull {tag}', shell=True)
+
+    # Run the Triton server and capture the container ID
+    container_id = subprocess.check_output(
+        f'docker run -d --rm -v {triton_repo_path}:/models -p 8000:8000 {tag} tritonserver --model-repository=/models',
+        shell=True).decode('utf-8').strip()
+
+    # Wait for the Triton server to start
+    triton_client = InferenceServerClient(url='localhost:8000', verbose=False, ssl=False)
+
+    # Wait until model is ready
+    for _ in range(10):
+        with contextlib.suppress(Exception):
+            assert triton_client.is_model_ready(model_name)
+            break
+        time.sleep(1)
+
+    # Check Triton inference
+    YOLO(f'http://localhost:8000/{model_name}', 'detect')(SOURCE)  # exported model inference
+
+    # Kill and remove the container at the end of the test
+    subprocess.call(f'docker kill {container_id}', shell=True)
+
+
+@pytest.mark.skipif(not check_requirements('pycocotools', install=False), reason='pycocotools not installed')
+def test_pycocotools():
+    """Validate model predictions using pycocotools."""
+    from ultralytics.models.yolo.detect import DetectionValidator
+    from ultralytics.models.yolo.pose import PoseValidator
+    from ultralytics.models.yolo.segment import SegmentationValidator
+
+    # Download annotations after each dataset downloads first
+    url = 'https://github.com/ultralytics/assets/releases/download/v0.0.0/'
+
+    args = {'model': 'yolov8n.pt', 'data': 'coco8.yaml', 'save_json': True, 'imgsz': 64}
+    validator = DetectionValidator(args=args)
+    validator()
+    validator.is_coco = True
+    download(f'{url}instances_val2017.json', dir=DATASETS_DIR / 'coco8/annotations')
+    _ = validator.eval_json(validator.stats)
+
+    args = {'model': 'yolov8n-seg.pt', 'data': 'coco8-seg.yaml', 'save_json': True, 'imgsz': 64}
+    validator = SegmentationValidator(args=args)
+    validator()
+    validator.is_coco = True
+    download(f'{url}instances_val2017.json', dir=DATASETS_DIR / 'coco8-seg/annotations')
+    _ = validator.eval_json(validator.stats)
+
+    args = {'model': 'yolov8n-pose.pt', 'data': 'coco8-pose.yaml', 'save_json': True, 'imgsz': 64}
+    validator = PoseValidator(args=args)
+    validator()
+    validator.is_coco = True
+    download(f'{url}person_keypoints_val2017.json', dir=DATASETS_DIR / 'coco8-pose/annotations')
+    _ = validator.eval_json(validator.stats)