ultralytics 8.0.197 save P, R, F1 curves to metrics (#5354)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: erminkev1 <83356055+erminkev1@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Andy <39454881+yermandy@users.noreply.github.com>

ultralytics/utils/callbacks/mlflow.py

@@ -38,7 +38,7 @@ def on_pretrain_routine_end(trainer):
             if not active_run:
                 active_run = mlflow.start_run(experiment_id=experiment.experiment_id, run_name=run_name)
             LOGGER.info(f'{prefix}Using run_id({active_run.info.run_id}) at {mlflow_location}')
-            run.log_params(vars(trainer.model.args))
+            run.log_params(trainer.args)
         except Exception as err:
             LOGGER.error(f'{prefix}Failing init - {repr(err)}')
             LOGGER.warning(f'{prefix}Continuing without Mlflow')

ultralytics/utils/callbacks/wb.py

@@ -8,7 +8,10 @@ try:
     assert SETTINGS['wandb'] is True  # verify integration is enabled
     import wandb as wb
 
-    assert hasattr(wb, '__version__')
+    assert hasattr(wb, '__version__')  # verify package is not directory
+
+    import numpy as np
+    import pandas as pd
 
     _processed_plots = {}
 
@@ -16,6 +19,83 @@ except (ImportError, AssertionError):
     wb = None
 
 
+def _custom_table(x, y, classes, title='Precision Recall Curve', x_axis_title='Recall', y_axis_title='Precision'):
+    """
+    Create and log a custom metric visualization to wandb.plot.pr_curve.
+
+    This function crafts a custom metric visualization that mimics the behavior of wandb's default precision-recall curve
+    while allowing for enhanced customization. The visual metric is useful for monitoring model performance across different classes.
+
+    Args:
+        x (List): Values for the x-axis; expected to have length N.
+        y (List): Corresponding values for the y-axis; also expected to have length N.
+        classes (List): Labels identifying the class of each point; length N.
+        title (str, optional): Title for the plot; defaults to 'Precision Recall Curve'.
+        x_title (str, optional): Label for the x-axis; defaults to 'Recall'.
+        y_title (str, optional): Label for the y-axis; defaults to 'Precision'.
+
+    Returns:
+        (wandb.Object): A wandb object suitable for logging, showcasing the crafted metric visualization.
+    """
+    df = pd.DataFrame({'class': classes, 'y': y, 'x': x}).round(3)
+    fields = {'x': 'x', 'y': 'y', 'class': 'class'}
+    string_fields = {'title': title, 'x-axis-title': x_axis_title, 'y-axis-title': y_axis_title}
+    return wb.plot_table('wandb/area-under-curve/v0',
+                         wb.Table(dataframe=df),
+                         fields=fields,
+                         string_fields=string_fields)
+
+
+def _plot_curve(x,
+                y,
+                names=None,
+                id='precision-recall',
+                title='Precision Recall Curve',
+                x_title='Recall',
+                y_title='Precision',
+                num_x=100,
+                only_mean=False):
+    """
+    Log a metric curve visualization.
+
+    This function generates a metric curve based on input data and logs the visualization to wandb.
+    The curve can represent aggregated data (mean) or individual class data, depending on the 'only_mean' flag.
+
+    Args:
+        x (np.ndarray): Data points for the x-axis with length N.
+        y (np.ndarray): Corresponding data points for the y-axis with shape CxN, where C represents the number of classes.
+        names (list, optional): Names of the classes corresponding to the y-axis data; length C. Defaults to an empty list.
+        id (str, optional): Unique identifier for the logged data in wandb. Defaults to 'precision-recall'.
+        title (str, optional): Title for the visualization plot. Defaults to 'Precision Recall Curve'.
+        x_title (str, optional): Label for the x-axis. Defaults to 'Recall'.
+        y_title (str, optional): Label for the y-axis. Defaults to 'Precision'.
+        num_x (int, optional): Number of interpolated data points for visualization. Defaults to 100.
+        only_mean (bool, optional): Flag to indicate if only the mean curve should be plotted. Defaults to True.
+
+    Note:
+        The function leverages the '_custom_table' function to generate the actual visualization.
+    """
+    # Create new x
+    if names is None:
+        names = []
+    x_new = np.linspace(x[0], x[-1], num_x).round(5)
+
+    # Create arrays for logging
+    x_log = x_new.tolist()
+    y_log = np.interp(x_new, x, np.mean(y, axis=0)).round(3).tolist()
+
+    if only_mean:
+        table = wb.Table(data=list(zip(x_log, y_log)), columns=[x_title, y_title])
+        wb.run.log({title: wb.plot.line(table, x_title, y_title, title=title)})
+    else:
+        classes = ['mean'] * len(x_log)
+        for i, yi in enumerate(y):
+            x_log.extend(x_new)  # add new x
+            y_log.extend(np.interp(x_new, x, yi))  # interpolate y to new x
+            classes.extend([names[i]] * len(x_new))  # add class names
+        wb.log({id: _custom_table(x_log, y_log, classes, title, x_title, y_title)}, commit=False)
+
+
 def _log_plots(plots, step):
     """Logs plots from the input dictionary if they haven't been logged already at the specified step."""
     for name, params in plots.items():
@@ -55,6 +135,17 @@ def on_train_end(trainer):
     if trainer.best.exists():
         art.add_file(trainer.best)
         wb.run.log_artifact(art, aliases=['best'])
+    for curve_name, curve_values in zip(trainer.validator.metrics.curves, trainer.validator.metrics.curves_results):
+        x, y, x_title, y_title = curve_values
+        _plot_curve(
+            x,
+            y,
+            names=list(trainer.validator.metrics.names.values()),
+            id=f'curves/{curve_name}',
+            title=curve_name,
+            x_title=x_title,
+            y_title=y_title,
+        )
     wb.run.finish()  # required or run continues on dashboard