ultralytics 8.0.71 updates and fixes (#1907)

Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Pavel Bugneac <50273042+pavelbugneac@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, GPL-3.0 license
 
-__version__ = '8.0.70'
+__version__ = '8.0.71'
 
 from ultralytics.hub import start
 from ultralytics.yolo.engine.model import YOLO

ultralytics/tracker/track.py

@@ -1,5 +1,7 @@
 # Ultralytics YOLO 🚀, GPL-3.0 license
 
+from functools import partial
+
 import torch
 
 from ultralytics.yolo.utils import IterableSimpleNamespace, yaml_load
@@ -10,7 +12,19 @@ from .trackers import BOTSORT, BYTETracker
 TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT}
 
 
-def on_predict_start(predictor):
+def on_predict_start(predictor, persist=False):
+    """
+    Initialize trackers for object tracking during prediction.
+
+    Args:
+        predictor (object): The predictor object to initialize trackers for.
+        persist (bool, optional): Whether to persist the trackers if they already exist. Defaults to False.
+
+    Raises:
+        AssertionError: If the tracker_type is not 'bytetrack' or 'botsort'.
+    """
+    if hasattr(predictor, 'trackers') and persist:
+        return
     tracker = check_yaml(predictor.args.tracker)
     cfg = IterableSimpleNamespace(**yaml_load(tracker))
     assert cfg.tracker_type in ['bytetrack', 'botsort'], \
@@ -38,6 +52,14 @@ def on_predict_postprocess_end(predictor):
         predictor.results[i].update(boxes=torch.as_tensor(tracks[:, :-1]))
 
 
-def register_tracker(model):
-    model.add_callback('on_predict_start', on_predict_start)
+def register_tracker(model, persist):
+    """
+    Register tracking callbacks to the model for object tracking during prediction.
+
+    Args:
+        model (object): The model object to register tracking callbacks for.
+        persist (bool): Whether to persist the trackers if they already exist.
+
+    """
+    model.add_callback('on_predict_start', partial(on_predict_start, persist=persist))
     model.add_callback('on_predict_postprocess_end', on_predict_postprocess_end)

ultralytics/tracker/trackers/byte_tracker.py

@@ -277,12 +277,13 @@ class BYTETracker:
         self.lost_stracks = self.sub_stracks(self.lost_stracks, self.tracked_stracks)
         self.lost_stracks.extend(lost_stracks)
         self.lost_stracks = self.sub_stracks(self.lost_stracks, self.removed_stracks)
-        self.removed_stracks.extend(removed_stracks)
         self.tracked_stracks, self.lost_stracks = self.remove_duplicate_stracks(self.tracked_stracks, self.lost_stracks)
-        output = [
-            track.tlbr.tolist() + [track.track_id, track.score, track.cls, track.idx] for track in self.tracked_stracks
-            if track.is_activated]
-        return np.asarray(output, dtype=np.float32)
+        self.removed_stracks.extend(removed_stracks)
+        if len(self.removed_stracks) > 1000:
+            self.removed_stracks = self.removed_stracks[-999:]  # clip remove stracks to 1000 maximum
+        return np.asarray(
+            [x.tlbr.tolist() + [x.track_id, x.score, x.cls, x.idx] for x in self.tracked_stracks if x.is_activated],
+            dtype=np.float32)
 
     def get_kalmanfilter(self):
         return KalmanFilterXYAH()
@@ -319,12 +320,16 @@ class BYTETracker:
 
     @staticmethod
     def sub_stracks(tlista, tlistb):
+        """ DEPRECATED CODE in https://github.com/ultralytics/ultralytics/pull/1890/
         stracks = {t.track_id: t for t in tlista}
         for t in tlistb:
             tid = t.track_id
             if stracks.get(tid, 0):
                 del stracks[tid]
         return list(stracks.values())
+        """
+        track_ids_b = {t.track_id for t in tlistb}
+        return [t for t in tlista if t.track_id not in track_ids_b]
 
     @staticmethod
     def remove_duplicate_stracks(stracksa, stracksb):

ultralytics/yolo/cfg/__init__.py

@@ -63,7 +63,7 @@ CLI_HELP_MSG = \
     """
 
 # Define keys for arg type checks
-CFG_FLOAT_KEYS = 'warmup_epochs', 'box', 'cls', 'dfl', 'degrees', 'shear', 'fl_gamma'
+CFG_FLOAT_KEYS = 'warmup_epochs', 'box', 'cls', 'dfl', 'degrees', 'shear'
 CFG_FRACTION_KEYS = ('dropout', 'iou', 'lr0', 'lrf', 'momentum', 'weight_decay', 'warmup_momentum', 'warmup_bias_lr',
                      'label_smoothing', 'hsv_h', 'hsv_s', 'hsv_v', 'translate', 'scale', 'perspective', 'flipud',
                      'fliplr', 'mosaic', 'mixup', 'copy_paste', 'conf', 'iou')  # fractional floats limited to 0.0 - 1.0

ultralytics/yolo/cfg/default.yaml

@@ -90,7 +90,6 @@ cls: 0.5  # cls loss gain (scale with pixels)
 dfl: 1.5  # dfl loss gain
 pose: 12.0  # pose loss gain
 kobj: 1.0  # keypoint obj loss gain
-fl_gamma: 0.0  # focal loss gamma (efficientDet default gamma=1.5)
 label_smoothing: 0.0  # label smoothing (fraction)
 nbs: 64  # nominal batch size
 hsv_h: 0.015  # image HSV-Hue augmentation (fraction)

ultralytics/yolo/data/build.py

@@ -93,15 +93,17 @@ def build_dataloader(cfg, batch, img_path, data_info, stride=32, rect=False, ran
     loader = DataLoader if cfg.image_weights or cfg.close_mosaic else InfiniteDataLoader  # allow attribute updates
     generator = torch.Generator()
     generator.manual_seed(6148914691236517205 + RANK)
-    return loader(dataset=dataset,
-                  batch_size=batch,
-                  shuffle=shuffle and sampler is None,
-                  num_workers=nw,
-                  sampler=sampler,
-                  pin_memory=PIN_MEMORY,
-                  collate_fn=getattr(dataset, 'collate_fn', None),
-                  worker_init_fn=seed_worker,
-                  generator=generator), dataset
+    return loader(
+        dataset=dataset,
+        batch_size=batch,
+        shuffle=shuffle and sampler is None,
+        num_workers=nw,
+        sampler=sampler,
+        pin_memory=PIN_MEMORY,
+        collate_fn=getattr(dataset, 'collate_fn', None),
+        worker_init_fn=seed_worker,
+        persistent_workers=(nw > 0) and (loader == DataLoader),  # persist workers if using default PyTorch DataLoader
+        generator=generator), dataset
 
 
 # build classification

ultralytics/yolo/data/dataset.py

@@ -37,7 +37,7 @@ class YOLODataset(BaseDataset):
         single_cls (bool): if True, single class training is used (default: False).
         use_segments (bool): if True, segmentation masks are used as labels (default: False).
         use_keypoints (bool): if True, keypoints are used as labels (default: False).
-        names (list): class names (default: None).
+        names (dict): A dictionary of class names. (default: None).
 
     Returns:
         A PyTorch dataset object that can be used for training an object detection or segmentation model.

ultralytics/yolo/engine/exporter.py

@@ -138,7 +138,7 @@ class Exporter:
             overrides (dict, optional): Configuration overrides. Defaults to None.
         """
         self.args = get_cfg(cfg, overrides)
-        self.callbacks = _callbacks if _callbacks else callbacks.get_default_callbacks()
+        self.callbacks = _callbacks or callbacks.get_default_callbacks()
         callbacks.add_integration_callbacks(self)
 
     @smart_inference_mode()
@@ -379,6 +379,7 @@ class Exporter:
         yaml_save(Path(f) / 'metadata.yaml', self.metadata)  # add metadata.yaml
         return f, None
 
+    @try_export
     def _export_coreml(self, prefix=colorstr('CoreML:')):
         # YOLOv8 CoreML export
         check_requirements('coremltools>=6.0')

ultralytics/yolo/engine/model.py

@@ -235,7 +235,8 @@ class YOLO:
         overrides.update(kwargs)  # prefer kwargs
         overrides['mode'] = kwargs.get('mode', 'predict')
         assert overrides['mode'] in ['track', 'predict']
-        overrides['save'] = kwargs.get('save', False)  # not save files by default
+        if not is_cli:
+            overrides['save'] = kwargs.get('save', False)  # do not save by default if called in Python
         if not self.predictor:
             self.task = overrides.get('task') or self.task
             self.predictor = TASK_MAP[self.task][3](overrides=overrides, _callbacks=self.callbacks)
@@ -244,10 +245,23 @@ class YOLO:
             self.predictor.args = get_cfg(self.predictor.args, overrides)
         return self.predictor.predict_cli(source=source) if is_cli else self.predictor(source=source, stream=stream)
 
-    def track(self, source=None, stream=False, **kwargs):
+    def track(self, source=None, stream=False, persist=False, **kwargs):
+        """
+        Perform object tracking on the input source using the registered trackers.
+
+        Args:
+            source (str, optional): The input source for object tracking. Can be a file path or a video stream.
+            stream (bool, optional): Whether the input source is a video stream. Defaults to False.
+            persist (bool, optional): Whether to persist the trackers if they already exist. Defaults to False.
+            **kwargs: Additional keyword arguments for the tracking process.
+
+        Returns:
+            object: The tracking results.
+
+        """
         if not hasattr(self.predictor, 'trackers'):
             from ultralytics.tracker import register_tracker
-            register_tracker(self)
+            register_tracker(self, persist)
         # ByteTrack-based method needs low confidence predictions as input
         conf = kwargs.get('conf') or 0.1
         kwargs['conf'] = conf

ultralytics/yolo/engine/predictor.py

@@ -103,7 +103,7 @@ class BasePredictor:
         self.data_path = None
         self.source_type = None
         self.batch = None
-        self.callbacks = _callbacks if _callbacks else callbacks.get_default_callbacks()
+        self.callbacks = _callbacks or callbacks.get_default_callbacks()
         callbacks.add_integration_callbacks(self)
 
     def preprocess(self, img):

ultralytics/yolo/engine/results.py

@@ -70,10 +70,12 @@ class Results(SimpleClass):
     Args:
         orig_img (numpy.ndarray): The original image as a numpy array.
         path (str): The path to the image file.
-        names (List[str]): A list of class names.
+        names (dict): A dictionary of class names.
         boxes (List[List[float]], optional): A list of bounding box coordinates for each detection.
         masks (numpy.ndarray, optional): A 3D numpy array of detection masks, where each mask is a binary image.
         probs (numpy.ndarray, optional): A 2D numpy array of detection probabilities for each class.
+        keypoints (List[List[float]], optional): A list of detected keypoints for each object.
+
 
     Attributes:
         orig_img (numpy.ndarray): The original image as a numpy array.
@@ -81,9 +83,12 @@ class Results(SimpleClass):
         boxes (Boxes, optional): A Boxes object containing the detection bounding boxes.
         masks (Masks, optional): A Masks object containing the detection masks.
         probs (numpy.ndarray, optional): A 2D numpy array of detection probabilities for each class.
-        names (List[str]): A list of class names.
+        names (dict): A dictionary of class names.
         path (str): The path to the image file.
+        keypoints (List[List[float]], optional): A list of detected keypoints for each object.
+        speed (dict): A dictionary of preprocess, inference and postprocess speeds in milliseconds per image.
         _keys (tuple): A tuple of attribute names for non-empty attributes.
+
     """
 
     def __init__(self, orig_img, path, names, boxes=None, masks=None, probs=None, keypoints=None) -> None:
@@ -93,6 +98,7 @@ class Results(SimpleClass):
         self.masks = Masks(masks, self.orig_shape) if masks is not None else None  # native size or imgsz masks
         self.probs = probs if probs is not None else None
         self.keypoints = keypoints if keypoints is not None else None
+        self.speed = {'preprocess': None, 'inference': None, 'postprocess': None}  # milliseconds per image
         self.names = names
         self.path = path
         self._keys = ('boxes', 'masks', 'probs', 'keypoints')
@@ -203,7 +209,7 @@ class Results(SimpleClass):
         keypoints = self.keypoints
         if pred_masks and show_masks:
             if img_gpu is None:
-                img = LetterBox(pred_masks.shape[1:])(image=annotator.im)
+                img = LetterBox(pred_masks.shape[1:])(image=annotator.result())
                 img_gpu = torch.as_tensor(img, dtype=torch.float16, device=pred_masks.masks.device).permute(
                     2, 0, 1).flip(0).contiguous() / 255
             annotator.masks(pred_masks.data, colors=[colors(x, True) for x in pred_boxes.cls], im_gpu=img_gpu)

ultralytics/yolo/engine/trainer.py

@@ -142,7 +142,7 @@ class BaseTrainer:
         self.plot_idx = [0, 1, 2]
 
         # Callbacks
-        self.callbacks = _callbacks if _callbacks else callbacks.get_default_callbacks()
+        self.callbacks = _callbacks or callbacks.get_default_callbacks()
         if RANK in (-1, 0):
             callbacks.add_integration_callbacks(self)
 

ultralytics/yolo/engine/validator.py

@@ -84,7 +84,7 @@ class BaseValidator:
         if self.args.conf is None:
             self.args.conf = 0.001  # default conf=0.001
 
-        self.callbacks = _callbacks if _callbacks else callbacks.get_default_callbacks()
+        self.callbacks = _callbacks or callbacks.get_default_callbacks()
 
     @smart_inference_mode()
     def __call__(self, trainer=None, model=None):