From 38737bdafa7253b5de640bb0a661704522d5ba69 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 8 Aug 2023 01:56:02 +0200 Subject: [PATCH] Fix missing `ultralytics` package on `torch.hub.load()` (#11950) * Fix missing `ultralytics` package on `torch.hub.load()` https://github.com/ultralytics/yolov5/issues/11945 Signed-off-by: Glenn Jocher * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update common.py Signed-off-by: Glenn Jocher * Update common.py Signed-off-by: Glenn Jocher * Update general.py Signed-off-by: Glenn Jocher --------- Signed-off-by: Glenn Jocher Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- models/common.py | 12 ++++++++++++ utils/general.py | 10 ++++++++++ 2 files changed, 22 insertions(+) diff --git a/models/common.py b/models/common.py index 4e03b9c7e47e..75cc4e97bbc7 100644 --- a/models/common.py +++ b/models/common.py @@ -23,6 +23,18 @@ import torch.nn as nn from PIL import Image from torch.cuda import amp + +# Import 'ultralytics' package or install if if missing +try: + import ultralytics + + assert hasattr(ultralytics, '__version__') # verify package is not directory +except (ImportError, AssertionError): + import os + + os.system('pip install -U ultralytics') + import ultralytics + from ultralytics.utils.plotting import Annotator, colors, save_one_box from utils import TryExcept diff --git a/utils/general.py b/utils/general.py index 8a265c741b91..017103752ef7 100644 --- a/utils/general.py +++ b/utils/general.py @@ -35,6 +35,16 @@ import torch import torchvision import yaml + +# Import 'ultralytics' package or install if if missing +try: + import ultralytics + + assert hasattr(ultralytics, '__version__') # verify package is not directory +except (ImportError, AssertionError): + os.system('pip install -U ultralytics') + import ultralytics + from ultralytics.utils.checks import check_requirements from utils import TryExcept, emojis