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

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

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():
-    return os.getenv('ULTRALYTICS_DVC_DISABLED', 'false').lower() == 'true'
-
-
-def _log_images(image_path, prefix=''):
+def _log_images(path, prefix=''):
     if live:
-        live.log_image(os.path.join(prefix, image_path.name), image_path)
+        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)
+
+        live.log_image(os.path.join(prefix, name), 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)
-            LOGGER.info(
-                '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
+        live = dvclive.Live(save_dvc_exp=True, cache_images=True)
+        LOGGER.info(
+            f'DVCLive is detected and auto logging is enabled (can be disabled in the {SETTINGS.file} with `dvc: false`).'
+        )
     except Exception as e:
         LOGGER.warning(f'WARNING ⚠️ DVCLive installed but not initialized correctly, not logging this run. {e}')