ultralytics 8.0.106 (#2736)

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: vyskocj <whiskey1939@seznam.cz>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: triple Mu <gpu@163.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>

ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = '8.0.105'
+__version__ = '8.0.106'
 
 from ultralytics.hub import start
 from ultralytics.vit.rtdetr import RTDETR

ultralytics/nn/modules/head.py

@@ -127,7 +127,7 @@ class Pose(Detect):
             y = kpts.view(bs, *self.kpt_shape, -1)
             a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * self.strides
             if ndim == 3:
-                a = torch.cat((a, y[:, :, 1:2].sigmoid()), 2)
+                a = torch.cat((a, y[:, :, 2:3].sigmoid()), 2)
             return a.view(bs, self.nk, -1)
         else:
             y = kpts.clone()

ultralytics/vit/rtdetr/predict.py

@@ -17,9 +17,11 @@ class RTDETRPredictor(BasePredictor):
         results = []
         for i, bbox in enumerate(bboxes):  # (300, 4)
             bbox = ops.xywh2xyxy(bbox)
-            score, cls = scores[i].max(-1)  # (300, )
-            idx = score > self.args.conf
-            pred = torch.cat([bbox, score[..., None], cls[..., None]], dim=-1)[idx]  # filter
+            score, cls = scores[i].max(-1, keepdim=True)  # (300, 1)
+            idx = score.squeeze(-1) > self.args.conf  # (300, )
+            if self.args.classes is not None:
+                idx = (cls == torch.tensor(self.args.classes, device=cls.device)).any(1) & idx
+            pred = torch.cat([bbox, score, cls], dim=-1)[idx]  # filter
             orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
             oh, ow = orig_img.shape[:2]
             if not isinstance(orig_imgs, torch.Tensor):

ultralytics/yolo/data/dataset.py

@@ -206,8 +206,6 @@ class ClassificationDataset(torchvision.datasets.ImageFolder):
 
     Args:
         root (str): Dataset path.
-        transform (callable, optional): torchvision transforms, used by default.
-        album_transform (callable, optional): Albumentations transforms, used if installed.
 
     Attributes:
         cache_ram (bool): True if images should be cached in RAM, False otherwise.

ultralytics/yolo/engine/trainer.py

@@ -414,12 +414,18 @@ class BaseTrainer:
             'date': datetime.now().isoformat(),
             'version': __version__}
 
+        # Use dill (if exists) to serialize the lambda functions where pickle does not do this
+        try:
+            import dill as pickle
+        except ImportError:
+            import pickle
+
         # Save last, best and delete
-        torch.save(ckpt, self.last)
+        torch.save(ckpt, self.last, pickle_module=pickle)
         if self.best_fitness == self.fitness:
-            torch.save(ckpt, self.best)
+            torch.save(ckpt, self.best, pickle_module=pickle)
         if (self.epoch > 0) and (self.save_period > 0) and (self.epoch % self.save_period == 0):
-            torch.save(ckpt, self.wdir / f'epoch{self.epoch}.pt')
+            torch.save(ckpt, self.wdir / f'epoch{self.epoch}.pt', pickle_module=pickle)
         del ckpt
 
     @staticmethod

ultralytics/yolo/utils/__init__.py

@@ -754,25 +754,9 @@ ENVIRONMENT = 'Colab' if is_colab() else 'Kaggle' if is_kaggle() else 'Jupyter'
 TESTS_RUNNING = is_pytest_running() or is_github_actions_ci()
 set_sentry()
 
-# OpenCV Multilanguage-friendly functions ------------------------------------------------------------------------------
-imshow_ = cv2.imshow  # copy to avoid recursion errors
-
-
-def imread(filename, flags=cv2.IMREAD_COLOR):
-    return cv2.imdecode(np.fromfile(filename, np.uint8), flags)
-
-
-def imwrite(filename, img):
-    try:
-        cv2.imencode(Path(filename).suffix, img)[1].tofile(filename)
-        return True
-    except Exception:
-        return False
-
-
-def imshow(path, im):
-    imshow_(path.encode('unicode_escape').decode(), im)
-
+# Apply monkey patches if the script is being run from within the parent directory of the script's location
+from .patches import imread, imshow, imwrite
 
+# torch.save = torch_save
 if Path(inspect.stack()[0].filename).parent.parent.as_posix() in inspect.stack()[-1].filename:
-    cv2.imread, cv2.imwrite, cv2.imshow = imread, imwrite, imshow  # redefine
+    cv2.imread, cv2.imwrite, cv2.imshow = imread, imwrite, imshow

ultralytics/yolo/utils/ops.py

@@ -1,3 +1,5 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
 import contextlib
 import math
 import re

ultralytics/yolo/utils/patches.py

@@ -0,0 +1,45 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+"""
+Monkey patches to update/extend functionality of existing functions
+"""
+
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+
+# OpenCV Multilanguage-friendly functions ------------------------------------------------------------------------------
+_imshow = cv2.imshow  # copy to avoid recursion errors
+
+
+def imread(filename, flags=cv2.IMREAD_COLOR):
+    return cv2.imdecode(np.fromfile(filename, np.uint8), flags)
+
+
+def imwrite(filename, img):
+    try:
+        cv2.imencode(Path(filename).suffix, img)[1].tofile(filename)
+        return True
+    except Exception:
+        return False
+
+
+def imshow(path, im):
+    _imshow(path.encode('unicode_escape').decode(), im)
+
+
+# PyTorch functions ----------------------------------------------------------------------------------------------------
+_torch_save = torch.save  # copy to avoid recursion errors
+
+
+def torch_save(*args, **kwargs):
+    # Use dill (if exists) to serialize the lambda functions where pickle does not do this
+    try:
+        import dill as pickle
+    except ImportError:
+        import pickle
+
+    if 'pickle_module' not in kwargs:
+        kwargs['pickle_module'] = pickle
+    return _torch_save(*args, **kwargs)

ultralytics/yolo/utils/torch_utils.py

@@ -341,6 +341,12 @@ def strip_optimizer(f: Union[str, Path] = 'best.pt', s: str = '') -> None:
         for f in Path('/Users/glennjocher/Downloads/weights').rglob('*.pt'):
             strip_optimizer(f)
     """
+    # Use dill (if exists) to serialize the lambda functions where pickle does not do this
+    try:
+        import dill as pickle
+    except ImportError:
+        import pickle
+
     x = torch.load(f, map_location=torch.device('cpu'))
     args = {**DEFAULT_CFG_DICT, **x['train_args']}  # combine model args with default args, preferring model args
     if x.get('ema'):
@@ -353,7 +359,7 @@ def strip_optimizer(f: Union[str, Path] = 'best.pt', s: str = '') -> None:
         p.requires_grad = False
     x['train_args'] = {k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS}  # strip non-default keys
     # x['model'].args = x['train_args']
-    torch.save(x, s or f)
+    torch.save(x, s or f, pickle_module=pickle)
     mb = os.path.getsize(s or f) / 1E6  # filesize
     LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB")
 

ultralytics/yolo/utils/tuner.py

@@ -1,3 +1,5 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
 from ultralytics.yolo.utils import LOGGER
 
 try:

ultralytics/yolo/v8/detect/val.py

@@ -62,14 +62,13 @@ class DetectionValidator(BaseValidator):
 
     def postprocess(self, preds):
         """Apply Non-maximum suppression to prediction outputs."""
-        preds = ops.non_max_suppression(preds,
-                                        self.args.conf,
-                                        self.args.iou,
-                                        labels=self.lb,
-                                        multi_label=True,
-                                        agnostic=self.args.single_cls,
-                                        max_det=self.args.max_det)
-        return preds
+        return ops.non_max_suppression(preds,
+                                       self.args.conf,
+                                       self.args.iou,
+                                       labels=self.lb,
+                                       multi_label=True,
+                                       agnostic=self.args.single_cls,
+                                       max_det=self.args.max_det)
 
     def update_metrics(self, preds, batch):
         """Metrics."""

ultralytics/yolo/v8/pose/val.py

@@ -33,15 +33,14 @@ class PoseValidator(DetectionValidator):
 
     def postprocess(self, preds):
         """Apply non-maximum suppression and return detections with high confidence scores."""
-        preds = ops.non_max_suppression(preds,
-                                        self.args.conf,
-                                        self.args.iou,
-                                        labels=self.lb,
-                                        multi_label=True,
-                                        agnostic=self.args.single_cls,
-                                        max_det=self.args.max_det,
-                                        nc=self.nc)
-        return preds
+        return ops.non_max_suppression(preds,
+                                       self.args.conf,
+                                       self.args.iou,
+                                       labels=self.lb,
+                                       multi_label=True,
+                                       agnostic=self.args.single_cls,
+                                       max_det=self.args.max_det,
+                                       nc=self.nc)
 
     def init_metrics(self, model):
         """Initiate pose estimation metrics for YOLO model."""