Clean and bump dvc callback: settings, stacked images (#4343)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

docs/quickstart.md

@@ -259,7 +259,7 @@ The table below provides an overview of the settings available for adjustment wi
 | `api_key`          | `''`                  | `str`     | Ultralytics HUB [API Key](https://hub.ultralytics.com/settings?tab=api+keys)                                     |
 | `clearml`          | `True`                | `bool`    | Whether to use ClearML logging                                                                                   |
 | `comet`            | `True`                | `bool`    | Whether to use [Comet ML](https://bit.ly/yolov8-readme-comet) for experiment tracking and visualization          |
-| `dvc`              | `True`                | `bool`    | Whether to use DVC for version control                                                                           |
+| `dvc`              | `True`                | `bool`    | Whether to use [DVC for experiment tracking](https://dvc.org/doc/dvclive/ml-frameworks/yolo) and version control                                                                           |
 | `hub`              | `True`                | `bool`    | Whether to use [Ultralytics HUB](https://hub.ultralytics.com) integration                                        |
 | `mlflow`           | `True`                | `bool`    | Whether to use MLFlow for experiment tracking                                                                    |
 | `neptune`          | `True`                | `bool`    | Whether to use Neptune for experiment tracking                                                                   |

ultralytics/engine/trainer.py

@@ -543,7 +543,8 @@ class BaseTrainer:
 
     def on_plot(self, name, data=None):
         """Registers plots (e.g. to be consumed in callbacks)"""
-        self.plots[name] = {'data': data, 'timestamp': time.time()}
+        path = Path(name)
+        self.plots[path] = {'data': data, 'timestamp': time.time()}
 
     def final_eval(self):
         """Performs final evaluation and validation for object detection YOLO model."""

ultralytics/engine/validator.py

@@ -303,7 +303,8 @@ class BaseValidator:
 
     def on_plot(self, name, data=None):
         """Registers plots (e.g. to be consumed in callbacks)"""
-        self.plots[name] = {'data': data, 'timestamp': time.time()}
+        path = Path(name)
+        self.plots[path] = {'data': data, 'timestamp': time.time()}
 
     # TODO: may need to put these following functions into callback
     def plot_val_samples(self, batch, ni):

ultralytics/utils/callbacks/dvc.py

@@ -1,6 +1,8 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
 import os
+import re
+from pathlib import Path
 
 import pkg_resources as pkg
 
@@ -32,13 +34,17 @@ _processed_plots = {}
 _training_epoch = False
 
 
-def _logger_disabled():
+def _log_images(path, prefix=''):
-    return os.getenv('ULTRALYTICS_DVC_DISABLED', 'false').lower() == 'true'
+    if live:
+        name = path.name
 
+        # Group images by batch to enable sliders in UI
+        if m := re.search(r'_batch(\d+)', name):
+            ni = m.group(1)
+            new_stem = re.sub(r'_batch(\d+)', '_batch', path.stem)
+            name = (Path(new_stem) / ni).with_suffix(path.suffix)
 
-def _log_images(image_path, prefix=''):
+        live.log_image(os.path.join(prefix, name), path)
-    if live:
-        live.log_image(os.path.join(prefix, image_path.name), image_path)
 
 
 def _log_plots(plots, prefix=''):
@@ -68,14 +74,10 @@ def _log_confusion_matrix(validator):
 def on_pretrain_routine_start(trainer):
     try:
         global live
-        if not _logger_disabled():
+        live = dvclive.Live(save_dvc_exp=True, cache_images=True)
-            live = dvclive.Live(save_dvc_exp=True, cache_images=True)
+        LOGGER.info(
-            LOGGER.info(
+            f'DVCLive is detected and auto logging is enabled (can be disabled in the {SETTINGS.file} with `dvc: false`).'
-                'DVCLive is detected and auto logging is enabled (can be disabled with `ULTRALYTICS_DVC_DISABLED=true`).'
+        )
-            )
-        else:
-            LOGGER.debug('DVCLive is detected and auto logging is disabled via `ULTRALYTICS_DVC_DISABLED`.')
-            live = None
     except Exception as e:
         LOGGER.warning(f'WARNING ⚠️ DVCLive installed but not initialized correctly, not logging this run. {e}')