feat: initial HSAP platform
Huaxu Sentinel Active Safety Platform with embedded algorithm code, Docker Compose setup, and vendored dataset scaffolds for clone-and-run. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -0,0 +1,5 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from .base import add_integration_callbacks, default_callbacks, get_default_callbacks
|
||||
|
||||
__all__ = "add_integration_callbacks", "default_callbacks", "get_default_callbacks"
|
||||
233
algorithms/dms_yolo/code/ultralytics/utils/callbacks/base.py
Normal file
233
algorithms/dms_yolo/code/ultralytics/utils/callbacks/base.py
Normal file
@@ -0,0 +1,233 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
"""Base callbacks for Ultralytics training, validation, prediction, and export processes."""
|
||||
|
||||
from collections import defaultdict
|
||||
from copy import deepcopy
|
||||
|
||||
# Trainer callbacks ----------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer):
|
||||
"""Called before the pretraining routine starts."""
|
||||
pass
|
||||
|
||||
|
||||
def on_pretrain_routine_end(trainer):
|
||||
"""Called after the pretraining routine ends."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_start(trainer):
|
||||
"""Called when the training starts."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_epoch_start(trainer):
|
||||
"""Called at the start of each training epoch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_batch_start(trainer):
|
||||
"""Called at the start of each training batch."""
|
||||
pass
|
||||
|
||||
|
||||
def optimizer_step(trainer):
|
||||
"""Called when the optimizer takes a step."""
|
||||
pass
|
||||
|
||||
|
||||
def on_before_zero_grad(trainer):
|
||||
"""Called before the gradients are set to zero."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_batch_end(trainer):
|
||||
"""Called at the end of each training batch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_epoch_end(trainer):
|
||||
"""Called at the end of each training epoch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Called at the end of each fit epoch (train + val)."""
|
||||
pass
|
||||
|
||||
|
||||
def on_model_save(trainer):
|
||||
"""Called when the model is saved."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_end(trainer):
|
||||
"""Called when the training ends."""
|
||||
pass
|
||||
|
||||
|
||||
def on_params_update(trainer):
|
||||
"""Called when the model parameters are updated."""
|
||||
pass
|
||||
|
||||
|
||||
def teardown(trainer):
|
||||
"""Called during the teardown of the training process."""
|
||||
pass
|
||||
|
||||
|
||||
# Validator callbacks --------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def on_val_start(validator):
|
||||
"""Called when the validation starts."""
|
||||
pass
|
||||
|
||||
|
||||
def on_val_batch_start(validator):
|
||||
"""Called at the start of each validation batch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_val_batch_end(validator):
|
||||
"""Called at the end of each validation batch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_val_end(validator):
|
||||
"""Called when the validation ends."""
|
||||
pass
|
||||
|
||||
|
||||
# Predictor callbacks --------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def on_predict_start(predictor):
|
||||
"""Called when the prediction starts."""
|
||||
pass
|
||||
|
||||
|
||||
def on_predict_batch_start(predictor):
|
||||
"""Called at the start of each prediction batch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_predict_batch_end(predictor):
|
||||
"""Called at the end of each prediction batch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_predict_postprocess_end(predictor):
|
||||
"""Called after the post-processing of the prediction ends."""
|
||||
pass
|
||||
|
||||
|
||||
def on_predict_end(predictor):
|
||||
"""Called when the prediction ends."""
|
||||
pass
|
||||
|
||||
|
||||
# Exporter callbacks ---------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def on_export_start(exporter):
|
||||
"""Called when the model export starts."""
|
||||
pass
|
||||
|
||||
|
||||
def on_export_end(exporter):
|
||||
"""Called when the model export ends."""
|
||||
pass
|
||||
|
||||
|
||||
default_callbacks = {
|
||||
# Run in trainer
|
||||
"on_pretrain_routine_start": [on_pretrain_routine_start],
|
||||
"on_pretrain_routine_end": [on_pretrain_routine_end],
|
||||
"on_train_start": [on_train_start],
|
||||
"on_train_epoch_start": [on_train_epoch_start],
|
||||
"on_train_batch_start": [on_train_batch_start],
|
||||
"optimizer_step": [optimizer_step],
|
||||
"on_before_zero_grad": [on_before_zero_grad],
|
||||
"on_train_batch_end": [on_train_batch_end],
|
||||
"on_train_epoch_end": [on_train_epoch_end],
|
||||
"on_fit_epoch_end": [on_fit_epoch_end], # fit = train + val
|
||||
"on_model_save": [on_model_save],
|
||||
"on_train_end": [on_train_end],
|
||||
"on_params_update": [on_params_update],
|
||||
"teardown": [teardown],
|
||||
# Run in validator
|
||||
"on_val_start": [on_val_start],
|
||||
"on_val_batch_start": [on_val_batch_start],
|
||||
"on_val_batch_end": [on_val_batch_end],
|
||||
"on_val_end": [on_val_end],
|
||||
# Run in predictor
|
||||
"on_predict_start": [on_predict_start],
|
||||
"on_predict_batch_start": [on_predict_batch_start],
|
||||
"on_predict_postprocess_end": [on_predict_postprocess_end],
|
||||
"on_predict_batch_end": [on_predict_batch_end],
|
||||
"on_predict_end": [on_predict_end],
|
||||
# Run in exporter
|
||||
"on_export_start": [on_export_start],
|
||||
"on_export_end": [on_export_end],
|
||||
}
|
||||
|
||||
|
||||
def get_default_callbacks():
|
||||
"""Get the default callbacks for Ultralytics training, validation, prediction, and export processes.
|
||||
|
||||
Returns:
|
||||
(dict): Dictionary of default callbacks for various training events. Each key represents an event during the
|
||||
training process, and the corresponding value is a list of callback functions executed when that
|
||||
event occurs.
|
||||
|
||||
Examples:
|
||||
>>> callbacks = get_default_callbacks()
|
||||
>>> print(list(callbacks.keys())) # show all available callback events
|
||||
['on_pretrain_routine_start', 'on_pretrain_routine_end', ...]
|
||||
"""
|
||||
return defaultdict(list, deepcopy(default_callbacks))
|
||||
|
||||
|
||||
def add_integration_callbacks(instance):
|
||||
"""Add integration callbacks to the instance's callbacks dictionary.
|
||||
|
||||
This function loads and adds various integration callbacks to the provided instance. The specific callbacks added
|
||||
depend on the type of instance provided. All instances receive HUB callbacks, while Trainer instances also receive
|
||||
additional callbacks for various integrations like ClearML, Comet, DVC, MLflow, Neptune, Ray Tune, TensorBoard, and
|
||||
Weights & Biases.
|
||||
|
||||
Args:
|
||||
instance (Trainer | Predictor | Validator | Exporter): The object instance to which callbacks will be added. The
|
||||
type of instance determines which callbacks are loaded.
|
||||
|
||||
Examples:
|
||||
>>> from ultralytics.engine.trainer import BaseTrainer
|
||||
>>> trainer = BaseTrainer()
|
||||
>>> add_integration_callbacks(trainer)
|
||||
"""
|
||||
from .hub import callbacks as hub_cb
|
||||
from .platform import callbacks as platform_cb
|
||||
|
||||
# Load Ultralytics callbacks
|
||||
callbacks_list = [hub_cb, platform_cb]
|
||||
|
||||
# Load training callbacks
|
||||
if "Trainer" in instance.__class__.__name__:
|
||||
from .clearml import callbacks as clear_cb
|
||||
from .comet import callbacks as comet_cb
|
||||
from .dvc import callbacks as dvc_cb
|
||||
from .mlflow import callbacks as mlflow_cb
|
||||
from .neptune import callbacks as neptune_cb
|
||||
from .raytune import callbacks as tune_cb
|
||||
from .tensorboard import callbacks as tb_cb
|
||||
from .wb import callbacks as wb_cb
|
||||
|
||||
callbacks_list.extend([clear_cb, comet_cb, dvc_cb, mlflow_cb, neptune_cb, tune_cb, tb_cb, wb_cb])
|
||||
|
||||
# Add the callbacks to the callbacks dictionary
|
||||
for callbacks in callbacks_list:
|
||||
for k, v in callbacks.items():
|
||||
if v not in instance.callbacks[k]:
|
||||
instance.callbacks[k].append(v)
|
||||
146
algorithms/dms_yolo/code/ultralytics/utils/callbacks/clearml.py
Normal file
146
algorithms/dms_yolo/code/ultralytics/utils/callbacks/clearml.py
Normal file
@@ -0,0 +1,146 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from ultralytics.utils import LOGGER, SETTINGS, TESTS_RUNNING
|
||||
|
||||
try:
|
||||
assert not TESTS_RUNNING # do not log pytest
|
||||
assert SETTINGS["clearml"] is True # verify integration is enabled
|
||||
import clearml
|
||||
from clearml import Task
|
||||
|
||||
assert hasattr(clearml, "__version__") # verify package is not directory
|
||||
|
||||
except (ImportError, AssertionError):
|
||||
clearml = None
|
||||
|
||||
|
||||
def _log_debug_samples(files, title: str = "Debug Samples") -> None:
|
||||
"""Log files (images) as debug samples in the ClearML task.
|
||||
|
||||
Args:
|
||||
files (list[Path]): A list of file paths in PosixPath format.
|
||||
title (str): A title that groups together images with the same values.
|
||||
"""
|
||||
import re
|
||||
|
||||
if task := Task.current_task():
|
||||
for f in files:
|
||||
if f.exists():
|
||||
it = re.search(r"_batch(\d+)", f.name)
|
||||
iteration = int(it.groups()[0]) if it else 0
|
||||
task.get_logger().report_image(
|
||||
title=title, series=f.name.replace(it.group(), ""), local_path=str(f), iteration=iteration
|
||||
)
|
||||
|
||||
|
||||
def _log_plot(title: str, plot_path: str) -> None:
|
||||
"""Log an image as a plot in the plot section of ClearML.
|
||||
|
||||
Args:
|
||||
title (str): The title of the plot.
|
||||
plot_path (str): The path to the saved image file.
|
||||
"""
|
||||
import matplotlib.image as mpimg
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
img = mpimg.imread(plot_path)
|
||||
fig = plt.figure()
|
||||
ax = fig.add_axes([0, 0, 1, 1], frameon=False, aspect="auto", xticks=[], yticks=[]) # no ticks
|
||||
ax.imshow(img)
|
||||
|
||||
Task.current_task().get_logger().report_matplotlib_figure(
|
||||
title=title, series="", figure=fig, report_interactive=False
|
||||
)
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer) -> None:
|
||||
"""Initialize and connect ClearML task at the start of pretraining routine."""
|
||||
try:
|
||||
if task := Task.current_task():
|
||||
# WARNING: make sure the automatic pytorch and matplotlib bindings are disabled!
|
||||
# We are logging these plots and model files manually in the integration
|
||||
from clearml.binding.frameworks.pytorch_bind import PatchPyTorchModelIO
|
||||
from clearml.binding.matplotlib_bind import PatchedMatplotlib
|
||||
|
||||
PatchPyTorchModelIO.update_current_task(None)
|
||||
PatchedMatplotlib.update_current_task(None)
|
||||
else:
|
||||
task = Task.init(
|
||||
project_name=trainer.args.project or "Ultralytics",
|
||||
task_name=trainer.args.name,
|
||||
tags=["Ultralytics"],
|
||||
output_uri=True,
|
||||
reuse_last_task_id=False,
|
||||
auto_connect_frameworks={"pytorch": False, "matplotlib": False},
|
||||
)
|
||||
LOGGER.warning(
|
||||
"ClearML Initialized a new task. If you want to run remotely, "
|
||||
"please add clearml-init and connect your arguments before initializing YOLO."
|
||||
)
|
||||
task.connect(vars(trainer.args), name="General")
|
||||
except Exception as e:
|
||||
LOGGER.warning(f"ClearML installed but not initialized correctly, not logging this run. {e}")
|
||||
|
||||
|
||||
def on_train_epoch_end(trainer) -> None:
|
||||
"""Log debug samples for the first epoch and report current training progress."""
|
||||
if task := Task.current_task():
|
||||
# Log debug samples for first epoch only
|
||||
if trainer.epoch == 1:
|
||||
_log_debug_samples(sorted(trainer.save_dir.glob("train_batch*.jpg")), "Mosaic")
|
||||
# Report the current training progress
|
||||
for k, v in trainer.label_loss_items(trainer.tloss, prefix="train").items():
|
||||
task.get_logger().report_scalar("train", k, v, iteration=trainer.epoch)
|
||||
for k, v in trainer.lr.items():
|
||||
task.get_logger().report_scalar("lr", k, v, iteration=trainer.epoch)
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer) -> None:
|
||||
"""Report model information and metrics to logger at the end of an epoch."""
|
||||
if task := Task.current_task():
|
||||
# Report epoch time and validation metrics
|
||||
task.get_logger().report_scalar(
|
||||
title="Epoch Time", series="Epoch Time", value=trainer.epoch_time, iteration=trainer.epoch
|
||||
)
|
||||
for k, v in trainer.metrics.items():
|
||||
title = k.split("/")[0]
|
||||
task.get_logger().report_scalar(title, k, v, iteration=trainer.epoch)
|
||||
if trainer.epoch == 0:
|
||||
from ultralytics.utils.torch_utils import model_info_for_loggers
|
||||
|
||||
for k, v in model_info_for_loggers(trainer).items():
|
||||
task.get_logger().report_single_value(k, v)
|
||||
|
||||
|
||||
def on_val_end(validator) -> None:
|
||||
"""Log validation results including labels and predictions."""
|
||||
if Task.current_task():
|
||||
# Log validation labels and predictions
|
||||
_log_debug_samples(sorted(validator.save_dir.glob("val*.jpg")), "Validation")
|
||||
|
||||
|
||||
def on_train_end(trainer) -> None:
|
||||
"""Log final model and training results on training completion."""
|
||||
if task := Task.current_task():
|
||||
# Log final results, confusion matrix and PR plots
|
||||
for f in [*trainer.plots.keys(), *trainer.validator.plots.keys()]:
|
||||
if "batch" not in f.name:
|
||||
_log_plot(title=f.stem, plot_path=f)
|
||||
# Report final metrics
|
||||
for k, v in trainer.validator.metrics.results_dict.items():
|
||||
task.get_logger().report_single_value(k, v)
|
||||
# Log the final model
|
||||
task.update_output_model(model_path=str(trainer.best), model_name=trainer.args.name, auto_delete_file=False)
|
||||
|
||||
|
||||
callbacks = (
|
||||
{
|
||||
"on_pretrain_routine_start": on_pretrain_routine_start,
|
||||
"on_train_epoch_end": on_train_epoch_end,
|
||||
"on_fit_epoch_end": on_fit_epoch_end,
|
||||
"on_val_end": on_val_end,
|
||||
"on_train_end": on_train_end,
|
||||
}
|
||||
if clearml
|
||||
else {}
|
||||
)
|
||||
625
algorithms/dms_yolo/code/ultralytics/utils/callbacks/comet.py
Normal file
625
algorithms/dms_yolo/code/ultralytics/utils/callbacks/comet.py
Normal file
@@ -0,0 +1,625 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ultralytics.utils import LOGGER, RANK, SETTINGS, TESTS_RUNNING, ops
|
||||
from ultralytics.utils.metrics import ClassifyMetrics, DetMetrics, OBBMetrics, PoseMetrics, SegmentMetrics
|
||||
|
||||
try:
|
||||
assert not TESTS_RUNNING # do not log pytest
|
||||
assert SETTINGS["comet"] is True # verify integration is enabled
|
||||
import comet_ml
|
||||
|
||||
assert hasattr(comet_ml, "__version__") # verify package is not directory
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# Ensures certain logging functions only run for supported tasks
|
||||
COMET_SUPPORTED_TASKS = ["detect", "segment"]
|
||||
|
||||
# Names of plots created by Ultralytics that are logged to Comet
|
||||
CONFUSION_MATRIX_PLOT_NAMES = "confusion_matrix", "confusion_matrix_normalized"
|
||||
EVALUATION_PLOT_NAMES = "F1_curve", "P_curve", "R_curve", "PR_curve"
|
||||
LABEL_PLOT_NAMES = ["labels"]
|
||||
SEGMENT_METRICS_PLOT_PREFIX = "Box", "Mask"
|
||||
POSE_METRICS_PLOT_PREFIX = "Box", "Pose"
|
||||
DETECTION_METRICS_PLOT_PREFIX = ["Box"]
|
||||
RESULTS_TABLE_NAME = "results.csv"
|
||||
ARGS_YAML_NAME = "args.yaml"
|
||||
|
||||
_comet_image_prediction_count = 0
|
||||
|
||||
except (ImportError, AssertionError):
|
||||
comet_ml = None
|
||||
|
||||
|
||||
def _get_comet_mode() -> str:
|
||||
"""Return the Comet mode from environment variables, defaulting to 'online'."""
|
||||
comet_mode = os.getenv("COMET_MODE")
|
||||
if comet_mode is not None:
|
||||
LOGGER.warning(
|
||||
"The COMET_MODE environment variable is deprecated. "
|
||||
"Please use COMET_START_ONLINE to set the Comet experiment mode. "
|
||||
"To start an offline Comet experiment, use 'export COMET_START_ONLINE=0'. "
|
||||
"If COMET_START_ONLINE is not set or is set to '1', an online Comet experiment will be created."
|
||||
)
|
||||
return comet_mode
|
||||
|
||||
return "online"
|
||||
|
||||
|
||||
def _get_comet_model_name() -> str:
|
||||
"""Return the Comet model name from environment variable or default to 'Ultralytics'."""
|
||||
return os.getenv("COMET_MODEL_NAME", "Ultralytics")
|
||||
|
||||
|
||||
def _get_eval_batch_logging_interval() -> int:
|
||||
"""Get the evaluation batch logging interval from environment variable or use default value 1."""
|
||||
return int(os.getenv("COMET_EVAL_BATCH_LOGGING_INTERVAL", 1))
|
||||
|
||||
|
||||
def _get_max_image_predictions_to_log() -> int:
|
||||
"""Get the maximum number of image predictions to log from environment variables."""
|
||||
return int(os.getenv("COMET_MAX_IMAGE_PREDICTIONS", 100))
|
||||
|
||||
|
||||
def _scale_confidence_score(score: float) -> float:
|
||||
"""Scale the confidence score by a factor specified in environment variable."""
|
||||
scale = float(os.getenv("COMET_MAX_CONFIDENCE_SCORE", 100.0))
|
||||
return score * scale
|
||||
|
||||
|
||||
def _should_log_confusion_matrix() -> bool:
|
||||
"""Determine if the confusion matrix should be logged based on environment variable settings."""
|
||||
return os.getenv("COMET_EVAL_LOG_CONFUSION_MATRIX", "false").lower() == "true"
|
||||
|
||||
|
||||
def _should_log_image_predictions() -> bool:
|
||||
"""Determine whether to log image predictions based on environment variable."""
|
||||
return os.getenv("COMET_EVAL_LOG_IMAGE_PREDICTIONS", "true").lower() == "true"
|
||||
|
||||
|
||||
def _resume_or_create_experiment(args: SimpleNamespace) -> None:
|
||||
"""Resume CometML experiment or create a new experiment based on args.
|
||||
|
||||
Ensures that the experiment object is only created in a single process during distributed training.
|
||||
|
||||
Args:
|
||||
args (SimpleNamespace): Training arguments containing project configuration and other parameters.
|
||||
"""
|
||||
if RANK not in {-1, 0}:
|
||||
return
|
||||
|
||||
# Set environment variable (if not set by the user) to configure the Comet experiment's online mode under the hood.
|
||||
# IF COMET_START_ONLINE is set by the user it will override COMET_MODE value.
|
||||
if os.getenv("COMET_START_ONLINE") is None:
|
||||
comet_mode = _get_comet_mode()
|
||||
os.environ["COMET_START_ONLINE"] = "1" if comet_mode != "offline" else "0"
|
||||
|
||||
try:
|
||||
_project_name = os.getenv("COMET_PROJECT_NAME", args.project)
|
||||
experiment = comet_ml.start(project_name=_project_name)
|
||||
experiment.log_parameters(vars(args))
|
||||
experiment.log_others(
|
||||
{
|
||||
"eval_batch_logging_interval": _get_eval_batch_logging_interval(),
|
||||
"log_confusion_matrix_on_eval": _should_log_confusion_matrix(),
|
||||
"log_image_predictions": _should_log_image_predictions(),
|
||||
"max_image_predictions": _get_max_image_predictions_to_log(),
|
||||
}
|
||||
)
|
||||
experiment.log_other("Created from", "ultralytics")
|
||||
|
||||
except Exception as e:
|
||||
LOGGER.warning(f"Comet installed but not initialized correctly, not logging this run. {e}")
|
||||
|
||||
|
||||
def _fetch_trainer_metadata(trainer) -> dict:
|
||||
"""Return metadata for YOLO training including epoch and asset saving status.
|
||||
|
||||
Args:
|
||||
trainer (ultralytics.engine.trainer.BaseTrainer): The YOLO trainer object containing training state and config.
|
||||
|
||||
Returns:
|
||||
(dict): Dictionary containing current epoch, step, save assets flag, and final epoch flag.
|
||||
"""
|
||||
curr_epoch = trainer.epoch + 1
|
||||
|
||||
train_num_steps_per_epoch = len(trainer.train_loader.dataset) // trainer.batch_size
|
||||
curr_step = curr_epoch * train_num_steps_per_epoch
|
||||
final_epoch = curr_epoch == trainer.epochs
|
||||
|
||||
save = trainer.args.save
|
||||
save_period = trainer.args.save_period
|
||||
save_interval = curr_epoch % save_period == 0
|
||||
save_assets = save and save_period > 0 and save_interval and not final_epoch
|
||||
|
||||
return dict(curr_epoch=curr_epoch, curr_step=curr_step, save_assets=save_assets, final_epoch=final_epoch)
|
||||
|
||||
|
||||
def _scale_bounding_box_to_original_image_shape(
|
||||
box, resized_image_shape, original_image_shape, ratio_pad
|
||||
) -> list[float]:
|
||||
"""Scale bounding box from resized image coordinates to original image coordinates.
|
||||
|
||||
YOLO resizes images during training and the label values are normalized based on this resized shape. This function
|
||||
rescales the bounding box labels to the original image shape.
|
||||
|
||||
Args:
|
||||
box (torch.Tensor): Bounding box in normalized xywh format.
|
||||
resized_image_shape (tuple): Shape of the resized image (height, width).
|
||||
original_image_shape (tuple): Shape of the original image (height, width).
|
||||
ratio_pad (tuple): Ratio and padding information for scaling.
|
||||
|
||||
Returns:
|
||||
(list[float]): Scaled bounding box coordinates in xywh format with top-left corner adjustment.
|
||||
"""
|
||||
resized_image_height, resized_image_width = resized_image_shape
|
||||
|
||||
# Convert normalized xywh format predictions to xyxy in resized scale format
|
||||
box = ops.xywhn2xyxy(box, h=resized_image_height, w=resized_image_width)
|
||||
# Scale box predictions from resized image scale back to original image scale
|
||||
box = ops.scale_boxes(resized_image_shape, box, original_image_shape, ratio_pad)
|
||||
# Convert bounding box format from xyxy to xywh for Comet logging
|
||||
box = ops.xyxy2xywh(box)
|
||||
# Adjust xy center to correspond top-left corner
|
||||
box[:2] -= box[2:] / 2
|
||||
box = box.tolist()
|
||||
|
||||
return box
|
||||
|
||||
|
||||
def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, class_name_map=None) -> dict | None:
|
||||
"""Format ground truth annotations for object detection.
|
||||
|
||||
This function processes ground truth annotations from a batch of images for object detection tasks. It extracts
|
||||
bounding boxes, class labels, and other metadata for a specific image in the batch, and formats them for
|
||||
visualization or evaluation.
|
||||
|
||||
Args:
|
||||
img_idx (int): Index of the image in the batch to process.
|
||||
image_path (str | Path): Path to the image file.
|
||||
batch (dict): Batch dictionary containing detection data with keys:
|
||||
- 'batch_idx': Tensor of batch indices
|
||||
- 'bboxes': Tensor of bounding boxes in normalized xywh format
|
||||
- 'cls': Tensor of class labels
|
||||
- 'ori_shape': Original image shapes
|
||||
- 'resized_shape': Resized image shapes
|
||||
- 'ratio_pad': Ratio and padding information
|
||||
class_name_map (dict, optional): Mapping from class indices to class names.
|
||||
|
||||
Returns:
|
||||
(dict | None): Formatted ground truth annotations with the following structure:
|
||||
- 'boxes': List of box coordinates [x, y, width, height]
|
||||
- 'label': Label string with format "gt_{class_name}"
|
||||
- 'score': Confidence score (always 1.0, scaled by _scale_confidence_score)
|
||||
Returns None if no bounding boxes are found for the image.
|
||||
"""
|
||||
indices = batch["batch_idx"] == img_idx
|
||||
bboxes = batch["bboxes"][indices]
|
||||
if len(bboxes) == 0:
|
||||
LOGGER.debug(f"Comet Image: {image_path} has no bounding boxes labels")
|
||||
return None
|
||||
|
||||
cls_labels = batch["cls"][indices].squeeze(1).tolist()
|
||||
if class_name_map:
|
||||
cls_labels = [str(class_name_map[label]) for label in cls_labels]
|
||||
|
||||
original_image_shape = batch["ori_shape"][img_idx]
|
||||
resized_image_shape = batch["resized_shape"][img_idx]
|
||||
ratio_pad = batch["ratio_pad"][img_idx]
|
||||
|
||||
data = []
|
||||
for box, label in zip(bboxes, cls_labels):
|
||||
box = _scale_bounding_box_to_original_image_shape(box, resized_image_shape, original_image_shape, ratio_pad)
|
||||
data.append(
|
||||
{
|
||||
"boxes": [box],
|
||||
"label": f"gt_{label}",
|
||||
"score": _scale_confidence_score(1.0),
|
||||
}
|
||||
)
|
||||
|
||||
return {"name": "ground_truth", "data": data}
|
||||
|
||||
|
||||
def _format_prediction_annotations(image_path, metadata, class_label_map=None, class_map=None) -> dict | None:
|
||||
"""Format YOLO predictions for object detection visualization.
|
||||
|
||||
Args:
|
||||
image_path (Path): Path to the image file.
|
||||
metadata (dict): Prediction metadata containing bounding boxes and class information.
|
||||
class_label_map (dict, optional): Mapping from class indices to class names.
|
||||
class_map (dict, optional): Additional class mapping for label conversion.
|
||||
|
||||
Returns:
|
||||
(dict | None): Formatted prediction annotations or None if no predictions exist.
|
||||
"""
|
||||
stem = image_path.stem
|
||||
image_id = int(stem) if stem.isnumeric() else stem
|
||||
|
||||
predictions = metadata.get(image_id)
|
||||
if not predictions:
|
||||
LOGGER.debug(f"Comet Image: {image_path} has no bounding boxes predictions")
|
||||
return None
|
||||
|
||||
# apply the mapping that was used to map the predicted classes when the JSON was created
|
||||
if class_label_map and class_map:
|
||||
class_label_map = {class_map[k]: v for k, v in class_label_map.items()}
|
||||
try:
|
||||
# import pycotools utilities to decompress annotations for various tasks, e.g. segmentation
|
||||
from faster_coco_eval.core.mask import decode
|
||||
except ImportError:
|
||||
decode = None
|
||||
|
||||
data = []
|
||||
for prediction in predictions:
|
||||
boxes = prediction["bbox"]
|
||||
score = _scale_confidence_score(prediction["score"])
|
||||
cls_label = prediction["category_id"]
|
||||
if class_label_map:
|
||||
cls_label = str(class_label_map[cls_label])
|
||||
|
||||
annotation_data = {"boxes": [boxes], "label": cls_label, "score": score}
|
||||
|
||||
if decode is not None:
|
||||
# do segmentation processing only if we are able to decode it
|
||||
segments = prediction.get("segmentation", None)
|
||||
if segments is not None:
|
||||
segments = _extract_segmentation_annotation(segments, decode)
|
||||
if segments is not None:
|
||||
annotation_data["points"] = segments
|
||||
|
||||
data.append(annotation_data)
|
||||
|
||||
return {"name": "prediction", "data": data}
|
||||
|
||||
|
||||
def _extract_segmentation_annotation(segmentation_raw: str, decode: Callable) -> list[list[Any]] | None:
|
||||
"""Extract segmentation annotation from compressed segmentations as list of polygons.
|
||||
|
||||
Args:
|
||||
segmentation_raw (str): Raw segmentation data in compressed format.
|
||||
decode (Callable): Function to decode the compressed segmentation data.
|
||||
|
||||
Returns:
|
||||
(list[list[Any]] | None): List of polygon points or None if extraction fails.
|
||||
"""
|
||||
try:
|
||||
mask = decode(segmentation_raw)
|
||||
contours, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
|
||||
annotations = [np.array(polygon).squeeze() for polygon in contours if len(polygon) >= 3]
|
||||
return [annotation.ravel().tolist() for annotation in annotations]
|
||||
except Exception as e:
|
||||
LOGGER.warning(f"Comet Failed to extract segmentation annotation: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _fetch_annotations(img_idx, image_path, batch, prediction_metadata_map, class_label_map, class_map) -> list | None:
|
||||
"""Join the ground truth and prediction annotations if they exist.
|
||||
|
||||
Args:
|
||||
img_idx (int): Index of the image in the batch.
|
||||
image_path (Path): Path to the image file.
|
||||
batch (dict): Batch data containing ground truth annotations.
|
||||
prediction_metadata_map (dict): Map of prediction metadata by image ID.
|
||||
class_label_map (dict): Mapping from class indices to class names.
|
||||
class_map (dict): Additional class mapping for label conversion.
|
||||
|
||||
Returns:
|
||||
(list | None): List of annotation dictionaries or None if no annotations exist.
|
||||
"""
|
||||
ground_truth_annotations = _format_ground_truth_annotations_for_detection(
|
||||
img_idx, image_path, batch, class_label_map
|
||||
)
|
||||
prediction_annotations = _format_prediction_annotations(
|
||||
image_path, prediction_metadata_map, class_label_map, class_map
|
||||
)
|
||||
|
||||
annotations = [
|
||||
annotation for annotation in [ground_truth_annotations, prediction_annotations] if annotation is not None
|
||||
]
|
||||
return [annotations] if annotations else None
|
||||
|
||||
|
||||
def _create_prediction_metadata_map(model_predictions) -> dict:
|
||||
"""Create metadata map for model predictions by grouping them based on image ID."""
|
||||
pred_metadata_map = {}
|
||||
for prediction in model_predictions:
|
||||
pred_metadata_map.setdefault(prediction["image_id"], [])
|
||||
pred_metadata_map[prediction["image_id"]].append(prediction)
|
||||
|
||||
return pred_metadata_map
|
||||
|
||||
|
||||
def _log_confusion_matrix(experiment, trainer, curr_step, curr_epoch) -> None:
|
||||
"""Log the confusion matrix to Comet experiment."""
|
||||
conf_mat = trainer.validator.confusion_matrix.matrix
|
||||
names = [*list(trainer.data["names"].values()), "background"]
|
||||
experiment.log_confusion_matrix(
|
||||
matrix=conf_mat, labels=names, max_categories=len(names), epoch=curr_epoch, step=curr_step
|
||||
)
|
||||
|
||||
|
||||
def _log_images(experiment, image_paths, curr_step: int | None, annotations=None) -> None:
|
||||
"""Log images to the experiment with optional annotations.
|
||||
|
||||
This function logs images to a Comet ML experiment, optionally including annotation data for visualization such as
|
||||
bounding boxes or segmentation masks.
|
||||
|
||||
Args:
|
||||
experiment (comet_ml.CometExperiment): The Comet ML experiment to log images to.
|
||||
image_paths (list[Path]): List of paths to images that will be logged.
|
||||
curr_step (int): Current training step/iteration for tracking in the experiment timeline.
|
||||
annotations (list[list[dict]], optional): Nested list of annotation dictionaries for each image. Each annotation
|
||||
contains visualization data like bounding boxes, labels, and confidence scores.
|
||||
"""
|
||||
if annotations:
|
||||
for image_path, annotation in zip(image_paths, annotations):
|
||||
experiment.log_image(image_path, name=image_path.stem, step=curr_step, annotations=annotation)
|
||||
|
||||
else:
|
||||
for image_path in image_paths:
|
||||
experiment.log_image(image_path, name=image_path.stem, step=curr_step)
|
||||
|
||||
|
||||
def _log_image_predictions(experiment, validator, curr_step) -> None:
|
||||
"""Log predicted boxes for a single image during training.
|
||||
|
||||
This function logs image predictions to a Comet ML experiment during model validation. It processes validation data
|
||||
and formats both ground truth and prediction annotations for visualization in the Comet
|
||||
dashboard. The function respects configured limits on the number of images to log.
|
||||
|
||||
Args:
|
||||
experiment (comet_ml.CometExperiment): The Comet ML experiment to log to.
|
||||
validator (BaseValidator): The validator instance containing validation data and predictions.
|
||||
curr_step (int): The current training step for logging timeline.
|
||||
|
||||
Notes:
|
||||
This function uses global state to track the number of logged predictions across calls.
|
||||
It only logs predictions for supported tasks defined in COMET_SUPPORTED_TASKS.
|
||||
The number of logged images is limited by the COMET_MAX_IMAGE_PREDICTIONS environment variable.
|
||||
"""
|
||||
global _comet_image_prediction_count
|
||||
|
||||
task = validator.args.task
|
||||
if task not in COMET_SUPPORTED_TASKS:
|
||||
return
|
||||
|
||||
jdict = validator.jdict
|
||||
if not jdict:
|
||||
return
|
||||
|
||||
predictions_metadata_map = _create_prediction_metadata_map(jdict)
|
||||
dataloader = validator.dataloader
|
||||
class_label_map = validator.names
|
||||
class_map = getattr(validator, "class_map", None)
|
||||
|
||||
batch_logging_interval = _get_eval_batch_logging_interval()
|
||||
max_image_predictions = _get_max_image_predictions_to_log()
|
||||
|
||||
for batch_idx, batch in enumerate(dataloader):
|
||||
if (batch_idx + 1) % batch_logging_interval != 0:
|
||||
continue
|
||||
|
||||
image_paths = batch["im_file"]
|
||||
for img_idx, image_path in enumerate(image_paths):
|
||||
if _comet_image_prediction_count >= max_image_predictions:
|
||||
return
|
||||
|
||||
image_path = Path(image_path)
|
||||
annotations = _fetch_annotations(
|
||||
img_idx,
|
||||
image_path,
|
||||
batch,
|
||||
predictions_metadata_map,
|
||||
class_label_map,
|
||||
class_map=class_map,
|
||||
)
|
||||
_log_images(
|
||||
experiment,
|
||||
[image_path],
|
||||
curr_step,
|
||||
annotations=annotations,
|
||||
)
|
||||
_comet_image_prediction_count += 1
|
||||
|
||||
|
||||
def _log_plots(experiment, trainer) -> None:
|
||||
"""Log evaluation plots and label plots for the experiment.
|
||||
|
||||
This function logs various evaluation plots and confusion matrices to the experiment tracking system. It handles
|
||||
different types of metrics (SegmentMetrics, PoseMetrics, DetMetrics, OBBMetrics) and logs the appropriate plots for
|
||||
each type.
|
||||
|
||||
Args:
|
||||
experiment (comet_ml.CometExperiment): The Comet ML experiment to log plots to.
|
||||
trainer (ultralytics.engine.trainer.BaseTrainer): The trainer object containing validation metrics and save
|
||||
directory information.
|
||||
|
||||
Examples:
|
||||
>>> from ultralytics.utils.callbacks.comet import _log_plots
|
||||
>>> _log_plots(experiment, trainer)
|
||||
"""
|
||||
plot_filenames = None
|
||||
if isinstance(trainer.validator.metrics, SegmentMetrics):
|
||||
plot_filenames = [
|
||||
trainer.save_dir / f"{prefix}{plots}.png"
|
||||
for plots in EVALUATION_PLOT_NAMES
|
||||
for prefix in SEGMENT_METRICS_PLOT_PREFIX
|
||||
]
|
||||
elif isinstance(trainer.validator.metrics, PoseMetrics):
|
||||
plot_filenames = [
|
||||
trainer.save_dir / f"{prefix}{plots}.png"
|
||||
for plots in EVALUATION_PLOT_NAMES
|
||||
for prefix in POSE_METRICS_PLOT_PREFIX
|
||||
]
|
||||
elif isinstance(trainer.validator.metrics, (DetMetrics, OBBMetrics)):
|
||||
plot_filenames = [
|
||||
trainer.save_dir / f"{prefix}{plots}.png"
|
||||
for plots in EVALUATION_PLOT_NAMES
|
||||
for prefix in DETECTION_METRICS_PLOT_PREFIX
|
||||
]
|
||||
|
||||
if plot_filenames is not None:
|
||||
_log_images(experiment, plot_filenames, None)
|
||||
|
||||
confusion_matrix_filenames = [trainer.save_dir / f"{plots}.png" for plots in CONFUSION_MATRIX_PLOT_NAMES]
|
||||
_log_images(experiment, confusion_matrix_filenames, None)
|
||||
|
||||
if not isinstance(trainer.validator.metrics, ClassifyMetrics):
|
||||
label_plot_filenames = [trainer.save_dir / f"{labels}.jpg" for labels in LABEL_PLOT_NAMES]
|
||||
_log_images(experiment, label_plot_filenames, None)
|
||||
|
||||
|
||||
def _log_model(experiment, trainer) -> None:
|
||||
"""Log the best-trained model to Comet.ml."""
|
||||
model_name = _get_comet_model_name()
|
||||
experiment.log_model(model_name, file_or_folder=str(trainer.best), file_name="best.pt", overwrite=True)
|
||||
|
||||
|
||||
def _log_image_batches(experiment, trainer, curr_step: int) -> None:
|
||||
"""Log samples of image batches for train, validation, and test."""
|
||||
_log_images(experiment, trainer.save_dir.glob("train_batch*.jpg"), curr_step)
|
||||
_log_images(experiment, trainer.save_dir.glob("val_batch*.jpg"), curr_step)
|
||||
|
||||
|
||||
def _log_asset(experiment, asset_path) -> None:
|
||||
"""Logs a specific asset file to the given experiment.
|
||||
|
||||
This function facilitates logging an asset, such as a file, to the provided
|
||||
experiment. It enables integration with experiment tracking platforms.
|
||||
|
||||
Args:
|
||||
experiment (comet_ml.CometExperiment): The experiment instance to which the asset will be logged.
|
||||
asset_path (Path): The file path of the asset to log.
|
||||
"""
|
||||
experiment.log_asset(asset_path)
|
||||
|
||||
|
||||
def _log_table(experiment, table_path) -> None:
|
||||
"""Logs a table to the provided experiment.
|
||||
|
||||
This function is used to log a table file to the given experiment. The table is identified by its file path.
|
||||
|
||||
Args:
|
||||
experiment (comet_ml.CometExperiment): The experiment object where the table file will be logged.
|
||||
table_path (Path): The file path of the table to be logged.
|
||||
"""
|
||||
experiment.log_table(str(table_path))
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer) -> None:
|
||||
"""Create or resume a CometML experiment at the start of a YOLO pre-training routine."""
|
||||
_resume_or_create_experiment(trainer.args)
|
||||
|
||||
|
||||
def on_train_epoch_end(trainer) -> None:
|
||||
"""Log metrics and save batch images at the end of training epochs."""
|
||||
experiment = comet_ml.get_running_experiment()
|
||||
if not experiment:
|
||||
return
|
||||
|
||||
metadata = _fetch_trainer_metadata(trainer)
|
||||
curr_epoch = metadata["curr_epoch"]
|
||||
curr_step = metadata["curr_step"]
|
||||
|
||||
experiment.log_metrics(trainer.label_loss_items(trainer.tloss, prefix="train"), step=curr_step, epoch=curr_epoch)
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer) -> None:
|
||||
"""Log model assets at the end of each epoch during training.
|
||||
|
||||
This function is called at the end of each training epoch to log metrics, learning rates, and model information to a
|
||||
Comet ML experiment. It also logs model assets, confusion matrices, and image predictions based on configuration
|
||||
settings.
|
||||
|
||||
The function retrieves the current Comet ML experiment and logs various training metrics. If it's the first epoch,
|
||||
it also logs model information. On specified save intervals, it logs the model, confusion matrix (if enabled), and
|
||||
image predictions (if enabled).
|
||||
|
||||
Args:
|
||||
trainer (BaseTrainer): The YOLO trainer object containing training state, metrics, and configuration.
|
||||
|
||||
Examples:
|
||||
>>> # Inside a training loop
|
||||
>>> on_fit_epoch_end(trainer) # Log metrics and assets to Comet ML
|
||||
"""
|
||||
experiment = comet_ml.get_running_experiment()
|
||||
if not experiment:
|
||||
return
|
||||
|
||||
metadata = _fetch_trainer_metadata(trainer)
|
||||
curr_epoch = metadata["curr_epoch"]
|
||||
curr_step = metadata["curr_step"]
|
||||
save_assets = metadata["save_assets"]
|
||||
|
||||
experiment.log_metrics(trainer.metrics, step=curr_step, epoch=curr_epoch)
|
||||
experiment.log_metrics(trainer.lr, step=curr_step, epoch=curr_epoch)
|
||||
if curr_epoch == 1:
|
||||
from ultralytics.utils.torch_utils import model_info_for_loggers
|
||||
|
||||
experiment.log_metrics(model_info_for_loggers(trainer), step=curr_step, epoch=curr_epoch)
|
||||
|
||||
if not save_assets:
|
||||
return
|
||||
|
||||
_log_model(experiment, trainer)
|
||||
if _should_log_confusion_matrix():
|
||||
_log_confusion_matrix(experiment, trainer, curr_step, curr_epoch)
|
||||
if _should_log_image_predictions():
|
||||
_log_image_predictions(experiment, trainer.validator, curr_step)
|
||||
|
||||
|
||||
def on_train_end(trainer) -> None:
|
||||
"""Perform operations at the end of training."""
|
||||
experiment = comet_ml.get_running_experiment()
|
||||
if not experiment:
|
||||
return
|
||||
|
||||
metadata = _fetch_trainer_metadata(trainer)
|
||||
curr_epoch = metadata["curr_epoch"]
|
||||
curr_step = metadata["curr_step"]
|
||||
plots = trainer.args.plots
|
||||
|
||||
_log_model(experiment, trainer)
|
||||
if plots:
|
||||
_log_plots(experiment, trainer)
|
||||
|
||||
_log_confusion_matrix(experiment, trainer, curr_step, curr_epoch)
|
||||
_log_image_predictions(experiment, trainer.validator, curr_step)
|
||||
_log_image_batches(experiment, trainer, curr_step)
|
||||
# log results table
|
||||
table_path = trainer.save_dir / RESULTS_TABLE_NAME
|
||||
if table_path.exists():
|
||||
_log_table(experiment, table_path)
|
||||
|
||||
# log arguments YAML
|
||||
args_path = trainer.save_dir / ARGS_YAML_NAME
|
||||
if args_path.exists():
|
||||
_log_asset(experiment, args_path)
|
||||
|
||||
experiment.end()
|
||||
|
||||
global _comet_image_prediction_count
|
||||
_comet_image_prediction_count = 0
|
||||
|
||||
|
||||
callbacks = (
|
||||
{
|
||||
"on_pretrain_routine_start": on_pretrain_routine_start,
|
||||
"on_train_epoch_end": on_train_epoch_end,
|
||||
"on_fit_epoch_end": on_fit_epoch_end,
|
||||
"on_train_end": on_train_end,
|
||||
}
|
||||
if comet_ml
|
||||
else {}
|
||||
)
|
||||
197
algorithms/dms_yolo/code/ultralytics/utils/callbacks/dvc.py
Normal file
197
algorithms/dms_yolo/code/ultralytics/utils/callbacks/dvc.py
Normal file
@@ -0,0 +1,197 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from ultralytics.utils import LOGGER, SETTINGS, TESTS_RUNNING, checks
|
||||
|
||||
try:
|
||||
assert not TESTS_RUNNING # do not log pytest
|
||||
assert SETTINGS["dvc"] is True # verify integration is enabled
|
||||
import dvclive
|
||||
|
||||
assert checks.check_version("dvclive", "2.11.0", verbose=True)
|
||||
|
||||
import os
|
||||
import re
|
||||
|
||||
# DVCLive logger instance
|
||||
live = None
|
||||
_processed_plots = {}
|
||||
|
||||
# `on_fit_epoch_end` is called on final validation (probably need to be fixed) for now this is the way we
|
||||
# distinguish final evaluation of the best model vs last epoch validation
|
||||
_training_epoch = False
|
||||
|
||||
except (ImportError, AssertionError, TypeError):
|
||||
dvclive = None
|
||||
|
||||
|
||||
def _log_images(path: Path, prefix: str = "") -> None:
|
||||
"""Log images at specified path with an optional prefix using DVCLive.
|
||||
|
||||
This function logs images found at the given path to DVCLive, organizing them by batch to enable slider
|
||||
functionality in the UI. It processes image filenames to extract batch information and restructures the path
|
||||
accordingly.
|
||||
|
||||
Args:
|
||||
path (Path): Path to the image file to be logged.
|
||||
prefix (str, optional): Optional prefix to add to the image name when logging.
|
||||
|
||||
Examples:
|
||||
>>> from pathlib import Path
|
||||
>>> _log_images(Path("runs/train/exp/val_batch0_pred.jpg"), prefix="validation")
|
||||
"""
|
||||
if live:
|
||||
name = path.name
|
||||
|
||||
# Group images by batch to enable sliders in UI
|
||||
if m := re.search(r"_batch(\d+)", name):
|
||||
ni = m[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: dict, prefix: str = "") -> None:
|
||||
"""Log plot images for training progress if they have not been previously processed.
|
||||
|
||||
Args:
|
||||
plots (dict): Dictionary containing plot information with timestamps.
|
||||
prefix (str, optional): Optional prefix to add to the logged image paths.
|
||||
"""
|
||||
for name, params in plots.items():
|
||||
timestamp = params["timestamp"]
|
||||
if _processed_plots.get(name) != timestamp:
|
||||
_log_images(name, prefix)
|
||||
_processed_plots[name] = timestamp
|
||||
|
||||
|
||||
def _log_confusion_matrix(validator) -> None:
|
||||
"""Log confusion matrix for a validator using DVCLive.
|
||||
|
||||
This function processes the confusion matrix from a validator object and logs it to DVCLive by converting the matrix
|
||||
into lists of target and prediction labels.
|
||||
|
||||
Args:
|
||||
validator (BaseValidator): The validator object containing the confusion matrix and class names. Must have
|
||||
attributes confusion_matrix.matrix, confusion_matrix.task, and names.
|
||||
"""
|
||||
targets = []
|
||||
preds = []
|
||||
matrix = validator.confusion_matrix.matrix
|
||||
names = list(validator.names.values())
|
||||
if validator.confusion_matrix.task == "detect":
|
||||
names += ["background"]
|
||||
|
||||
for ti, pred in enumerate(matrix.T.astype(int)):
|
||||
for pi, num in enumerate(pred):
|
||||
targets.extend([names[ti]] * num)
|
||||
preds.extend([names[pi]] * num)
|
||||
|
||||
live.log_sklearn_plot("confusion_matrix", targets, preds, name="cf.json", normalized=True)
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer) -> None:
|
||||
"""Initialize DVCLive logger for training metadata during pre-training routine."""
|
||||
try:
|
||||
global live
|
||||
live = dvclive.Live(save_dvc_exp=True, cache_images=True)
|
||||
LOGGER.info("DVCLive is detected and auto logging is enabled (run 'yolo settings dvc=False' to disable).")
|
||||
except Exception as e:
|
||||
LOGGER.warning(f"DVCLive installed but not initialized correctly, not logging this run. {e}")
|
||||
|
||||
|
||||
def on_pretrain_routine_end(trainer) -> None:
|
||||
"""Log plots related to the training process at the end of the pretraining routine."""
|
||||
_log_plots(trainer.plots, "train")
|
||||
|
||||
|
||||
def on_train_start(trainer) -> None:
|
||||
"""Log the training parameters if DVCLive logging is active."""
|
||||
if live:
|
||||
live.log_params(trainer.args)
|
||||
|
||||
|
||||
def on_train_epoch_start(trainer) -> None:
|
||||
"""Set the global variable _training_epoch value to True at the start of training each epoch."""
|
||||
global _training_epoch
|
||||
_training_epoch = True
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer) -> None:
|
||||
"""Log training metrics, model info, and advance to next step at the end of each fit epoch.
|
||||
|
||||
This function is called at the end of each fit epoch during training. It logs various metrics including training
|
||||
loss items, validation metrics, and learning rates. On the first epoch, it also logs model
|
||||
information. Additionally, it logs training and validation plots and advances the DVCLive step counter.
|
||||
|
||||
Args:
|
||||
trainer (BaseTrainer): The trainer object containing training state, metrics, and plots.
|
||||
|
||||
Notes:
|
||||
This function only performs logging operations when DVCLive logging is active and during a training epoch.
|
||||
The global variable _training_epoch is used to track whether the current epoch is a training epoch.
|
||||
"""
|
||||
global _training_epoch
|
||||
if live and _training_epoch:
|
||||
all_metrics = {**trainer.label_loss_items(trainer.tloss, prefix="train"), **trainer.metrics, **trainer.lr}
|
||||
for metric, value in all_metrics.items():
|
||||
live.log_metric(metric, value)
|
||||
|
||||
if trainer.epoch == 0:
|
||||
from ultralytics.utils.torch_utils import model_info_for_loggers
|
||||
|
||||
for metric, value in model_info_for_loggers(trainer).items():
|
||||
live.log_metric(metric, value, plot=False)
|
||||
|
||||
_log_plots(trainer.plots, "train")
|
||||
_log_plots(trainer.validator.plots, "val")
|
||||
|
||||
live.next_step()
|
||||
_training_epoch = False
|
||||
|
||||
|
||||
def on_train_end(trainer) -> None:
|
||||
"""Log best metrics, plots, and confusion matrix at the end of training.
|
||||
|
||||
This function is called at the conclusion of the training process to log final metrics, visualizations, and model
|
||||
artifacts if DVCLive logging is active. It captures the best model performance metrics, training plots, validation
|
||||
plots, and confusion matrix for later analysis.
|
||||
|
||||
Args:
|
||||
trainer (BaseTrainer): The trainer object containing training state, metrics, and validation results.
|
||||
|
||||
Examples:
|
||||
>>> # Inside a custom training loop
|
||||
>>> from ultralytics.utils.callbacks.dvc import on_train_end
|
||||
>>> on_train_end(trainer) # Log final metrics and artifacts
|
||||
"""
|
||||
if live:
|
||||
# At the end log the best metrics. It runs validator on the best model internally.
|
||||
all_metrics = {**trainer.label_loss_items(trainer.tloss, prefix="train"), **trainer.metrics, **trainer.lr}
|
||||
for metric, value in all_metrics.items():
|
||||
live.log_metric(metric, value, plot=False)
|
||||
|
||||
_log_plots(trainer.plots, "val")
|
||||
_log_plots(trainer.validator.plots, "val")
|
||||
_log_confusion_matrix(trainer.validator)
|
||||
|
||||
if trainer.best.exists():
|
||||
live.log_artifact(trainer.best, copy=True, type="model")
|
||||
|
||||
live.end()
|
||||
|
||||
|
||||
callbacks = (
|
||||
{
|
||||
"on_pretrain_routine_start": on_pretrain_routine_start,
|
||||
"on_pretrain_routine_end": on_pretrain_routine_end,
|
||||
"on_train_start": on_train_start,
|
||||
"on_train_epoch_start": on_train_epoch_start,
|
||||
"on_fit_epoch_end": on_fit_epoch_end,
|
||||
"on_train_end": on_train_end,
|
||||
}
|
||||
if dvclive
|
||||
else {}
|
||||
)
|
||||
110
algorithms/dms_yolo/code/ultralytics/utils/callbacks/hub.py
Normal file
110
algorithms/dms_yolo/code/ultralytics/utils/callbacks/hub.py
Normal file
@@ -0,0 +1,110 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
import json
|
||||
from time import time
|
||||
|
||||
from ultralytics.hub import HUB_WEB_ROOT, PREFIX, HUBTrainingSession
|
||||
from ultralytics.utils import LOGGER, RANK, SETTINGS
|
||||
from ultralytics.utils.events import events
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer):
|
||||
"""Create a remote Ultralytics HUB session to log local model training."""
|
||||
if RANK in {-1, 0} and SETTINGS["hub"] is True and SETTINGS["api_key"] and trainer.hub_session is None:
|
||||
trainer.hub_session = HUBTrainingSession.create_session(trainer.args.model, trainer.args)
|
||||
|
||||
|
||||
def on_pretrain_routine_end(trainer):
|
||||
"""Initialize timers for upload rate limiting before training begins."""
|
||||
if session := getattr(trainer, "hub_session", None):
|
||||
# Start timer for upload rate limit
|
||||
session.timers = {"metrics": time(), "ckpt": time()} # start timer for session rate limiting
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Upload training progress metrics to Ultralytics HUB at the end of each epoch."""
|
||||
if session := getattr(trainer, "hub_session", None):
|
||||
# Upload metrics after validation ends
|
||||
all_plots = {
|
||||
**trainer.label_loss_items(trainer.tloss, prefix="train"),
|
||||
**trainer.metrics,
|
||||
}
|
||||
if trainer.epoch == 0:
|
||||
from ultralytics.utils.torch_utils import model_info_for_loggers
|
||||
|
||||
all_plots = {**all_plots, **model_info_for_loggers(trainer)}
|
||||
|
||||
session.metrics_queue[trainer.epoch] = json.dumps(all_plots)
|
||||
|
||||
# If any metrics failed to upload previously, add them to the queue to attempt uploading again
|
||||
if session.metrics_upload_failed_queue:
|
||||
session.metrics_queue.update(session.metrics_upload_failed_queue)
|
||||
|
||||
if time() - session.timers["metrics"] > session.rate_limits["metrics"]:
|
||||
session.upload_metrics()
|
||||
session.timers["metrics"] = time() # reset timer
|
||||
session.metrics_queue = {} # reset queue
|
||||
|
||||
|
||||
def on_model_save(trainer):
|
||||
"""Upload model checkpoints to Ultralytics HUB with rate limiting."""
|
||||
if session := getattr(trainer, "hub_session", None):
|
||||
# Upload checkpoints with rate limiting
|
||||
is_best = trainer.best_fitness == trainer.fitness
|
||||
if time() - session.timers["ckpt"] > session.rate_limits["ckpt"]:
|
||||
LOGGER.info(f"{PREFIX}Uploading checkpoint {HUB_WEB_ROOT}/models/{session.model.id}")
|
||||
session.upload_model(trainer.epoch, trainer.last, is_best)
|
||||
session.timers["ckpt"] = time() # reset timer
|
||||
|
||||
|
||||
def on_train_end(trainer):
|
||||
"""Upload final model and metrics to Ultralytics HUB at the end of training."""
|
||||
if session := getattr(trainer, "hub_session", None):
|
||||
# Upload final model and metrics with exponential standoff
|
||||
LOGGER.info(f"{PREFIX}Syncing final model...")
|
||||
session.upload_model(
|
||||
trainer.epoch,
|
||||
trainer.best,
|
||||
map=trainer.metrics.get("metrics/mAP50-95(B)", 0),
|
||||
final=True,
|
||||
)
|
||||
session.alive = False # stop heartbeats
|
||||
LOGGER.info(f"{PREFIX}Done ✅\n{PREFIX}View model at {session.model_url} 🚀")
|
||||
|
||||
|
||||
def on_train_start(trainer):
|
||||
"""Run events on train start."""
|
||||
events(trainer.args, trainer.device)
|
||||
|
||||
|
||||
def on_val_start(validator):
|
||||
"""Run events on validation start."""
|
||||
if not validator.training:
|
||||
events(validator.args, validator.device)
|
||||
|
||||
|
||||
def on_predict_start(predictor):
|
||||
"""Run events on predict start."""
|
||||
events(predictor.args, predictor.device)
|
||||
|
||||
|
||||
def on_export_start(exporter):
|
||||
"""Run events on export start."""
|
||||
events(exporter.args, exporter.device)
|
||||
|
||||
|
||||
callbacks = (
|
||||
{
|
||||
"on_pretrain_routine_start": on_pretrain_routine_start,
|
||||
"on_pretrain_routine_end": on_pretrain_routine_end,
|
||||
"on_fit_epoch_end": on_fit_epoch_end,
|
||||
"on_model_save": on_model_save,
|
||||
"on_train_end": on_train_end,
|
||||
"on_train_start": on_train_start,
|
||||
"on_val_start": on_val_start,
|
||||
"on_predict_start": on_predict_start,
|
||||
"on_export_start": on_export_start,
|
||||
}
|
||||
if SETTINGS["hub"] is True
|
||||
else {}
|
||||
)
|
||||
134
algorithms/dms_yolo/code/ultralytics/utils/callbacks/mlflow.py
Normal file
134
algorithms/dms_yolo/code/ultralytics/utils/callbacks/mlflow.py
Normal file
@@ -0,0 +1,134 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
"""
|
||||
MLflow Logging for Ultralytics YOLO.
|
||||
|
||||
This module enables MLflow logging for Ultralytics YOLO. It logs metrics, parameters, and model artifacts.
|
||||
For setting up, a tracking URI should be specified. The logging can be customized using environment variables.
|
||||
|
||||
Commands:
|
||||
1. To set a project name:
|
||||
`export MLFLOW_EXPERIMENT_NAME=<your_experiment_name>` or use the project=<project> argument
|
||||
|
||||
2. To set a run name:
|
||||
`export MLFLOW_RUN=<your_run_name>` or use the name=<name> argument
|
||||
|
||||
3. To start a local MLflow server:
|
||||
mlflow server --backend-store-uri runs/mlflow
|
||||
It will by default start a local server at http://127.0.0.1:5000.
|
||||
To specify a different URI, set the MLFLOW_TRACKING_URI environment variable.
|
||||
|
||||
4. To kill all running MLflow server instances:
|
||||
ps aux | grep 'mlflow' | grep -v 'grep' | awk '{print $2}' | xargs kill -9
|
||||
"""
|
||||
|
||||
from ultralytics.utils import LOGGER, RUNS_DIR, SETTINGS, TESTS_RUNNING, colorstr
|
||||
|
||||
try:
|
||||
import os
|
||||
|
||||
assert not TESTS_RUNNING or "test_mlflow" in os.environ.get("PYTEST_CURRENT_TEST", "") # do not log pytest
|
||||
assert SETTINGS["mlflow"] is True # verify integration is enabled
|
||||
import mlflow
|
||||
|
||||
assert hasattr(mlflow, "__version__") # verify package is not directory
|
||||
from pathlib import Path
|
||||
|
||||
PREFIX = colorstr("MLflow: ")
|
||||
|
||||
except (ImportError, AssertionError):
|
||||
mlflow = None
|
||||
|
||||
|
||||
def sanitize_dict(x: dict) -> dict:
|
||||
"""Sanitize dictionary keys by removing parentheses and converting values to floats."""
|
||||
return {k.replace("(", "").replace(")", ""): float(v) for k, v in x.items()}
|
||||
|
||||
|
||||
def on_pretrain_routine_end(trainer):
|
||||
"""Log training parameters to MLflow at the end of the pretraining routine.
|
||||
|
||||
This function sets up MLflow logging based on environment variables and trainer arguments. It sets the tracking URI,
|
||||
experiment name, and run name, then starts the MLflow run if not already active. It finally logs the parameters from
|
||||
the trainer.
|
||||
|
||||
Args:
|
||||
trainer (ultralytics.engine.trainer.BaseTrainer): The training object with arguments and parameters to log.
|
||||
|
||||
Notes:
|
||||
MLFLOW_TRACKING_URI: The URI for MLflow tracking. If not set, defaults to 'runs/mlflow'.
|
||||
MLFLOW_EXPERIMENT_NAME: The name of the MLflow experiment. If not set, defaults to trainer.args.project.
|
||||
MLFLOW_RUN: The name of the MLflow run. If not set, defaults to trainer.args.name.
|
||||
MLFLOW_KEEP_RUN_ACTIVE: Boolean indicating whether to keep the MLflow run active after training ends.
|
||||
"""
|
||||
global mlflow
|
||||
|
||||
uri = os.environ.get("MLFLOW_TRACKING_URI") or str(RUNS_DIR / "mlflow")
|
||||
LOGGER.debug(f"{PREFIX} tracking uri: {uri}")
|
||||
mlflow.set_tracking_uri(uri)
|
||||
|
||||
# Set experiment and run names
|
||||
experiment_name = os.environ.get("MLFLOW_EXPERIMENT_NAME") or trainer.args.project or "/Shared/Ultralytics"
|
||||
run_name = os.environ.get("MLFLOW_RUN") or trainer.args.name
|
||||
mlflow.set_experiment(experiment_name)
|
||||
|
||||
mlflow.autolog()
|
||||
try:
|
||||
active_run = mlflow.active_run() or mlflow.start_run(run_name=run_name)
|
||||
LOGGER.info(f"{PREFIX}logging run_id({active_run.info.run_id}) to {uri}")
|
||||
if Path(uri).is_dir():
|
||||
LOGGER.info(f"{PREFIX}view at http://127.0.0.1:5000 with 'mlflow server --backend-store-uri {uri}'")
|
||||
LOGGER.info(f"{PREFIX}disable with 'yolo settings mlflow=False'")
|
||||
mlflow.log_params(dict(trainer.args))
|
||||
except Exception as e:
|
||||
LOGGER.warning(f"{PREFIX}Failed to initialize: {e}")
|
||||
LOGGER.warning(f"{PREFIX}Not tracking this run")
|
||||
|
||||
|
||||
def on_train_epoch_end(trainer):
|
||||
"""Log training metrics at the end of each train epoch to MLflow."""
|
||||
if mlflow:
|
||||
mlflow.log_metrics(
|
||||
metrics={
|
||||
**sanitize_dict(trainer.lr),
|
||||
**sanitize_dict(trainer.label_loss_items(trainer.tloss, prefix="train")),
|
||||
},
|
||||
step=trainer.epoch,
|
||||
)
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Log training metrics at the end of each fit epoch to MLflow."""
|
||||
if mlflow:
|
||||
mlflow.log_metrics(metrics=sanitize_dict(trainer.metrics), step=trainer.epoch)
|
||||
|
||||
|
||||
def on_train_end(trainer):
|
||||
"""Log model artifacts at the end of training."""
|
||||
if not mlflow:
|
||||
return
|
||||
mlflow.log_artifact(str(trainer.best.parent)) # log save_dir/weights directory with best.pt and last.pt
|
||||
for f in trainer.save_dir.glob("*"): # log all other files in save_dir
|
||||
if f.suffix in {".png", ".jpg", ".csv", ".pt", ".yaml"}:
|
||||
mlflow.log_artifact(str(f))
|
||||
keep_run_active = os.environ.get("MLFLOW_KEEP_RUN_ACTIVE", "False").lower() == "true"
|
||||
if keep_run_active:
|
||||
LOGGER.info(f"{PREFIX}mlflow run still alive, remember to close it using mlflow.end_run()")
|
||||
else:
|
||||
mlflow.end_run()
|
||||
LOGGER.debug(f"{PREFIX}mlflow run ended")
|
||||
|
||||
LOGGER.info(
|
||||
f"{PREFIX}results logged to {mlflow.get_tracking_uri()}\n{PREFIX}disable with 'yolo settings mlflow=False'"
|
||||
)
|
||||
|
||||
|
||||
callbacks = (
|
||||
{
|
||||
"on_pretrain_routine_end": on_pretrain_routine_end,
|
||||
"on_train_epoch_end": on_train_epoch_end,
|
||||
"on_fit_epoch_end": on_fit_epoch_end,
|
||||
"on_train_end": on_train_end,
|
||||
}
|
||||
if mlflow
|
||||
else {}
|
||||
)
|
||||
126
algorithms/dms_yolo/code/ultralytics/utils/callbacks/neptune.py
Normal file
126
algorithms/dms_yolo/code/ultralytics/utils/callbacks/neptune.py
Normal file
@@ -0,0 +1,126 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from ultralytics.utils import LOGGER, SETTINGS, TESTS_RUNNING
|
||||
|
||||
try:
|
||||
assert not TESTS_RUNNING # do not log pytest
|
||||
assert SETTINGS["neptune"] is True # verify integration is enabled
|
||||
|
||||
import neptune
|
||||
from neptune.types import File
|
||||
|
||||
assert hasattr(neptune, "__version__")
|
||||
|
||||
run = None # NeptuneAI experiment logger instance
|
||||
|
||||
except (ImportError, AssertionError):
|
||||
neptune = None
|
||||
|
||||
|
||||
def _log_scalars(scalars: dict, step: int = 0) -> None:
|
||||
"""Log scalars to the NeptuneAI experiment logger.
|
||||
|
||||
Args:
|
||||
scalars (dict): Dictionary of scalar values to log to NeptuneAI.
|
||||
step (int, optional): The current step or iteration number for logging.
|
||||
|
||||
Examples:
|
||||
>>> metrics = {"mAP": 0.85, "loss": 0.32}
|
||||
>>> _log_scalars(metrics, step=100)
|
||||
"""
|
||||
if run:
|
||||
for k, v in scalars.items():
|
||||
run[k].append(value=v, step=step)
|
||||
|
||||
|
||||
def _log_images(imgs_dict: dict, group: str = "") -> None:
|
||||
"""Log images to the NeptuneAI experiment logger.
|
||||
|
||||
This function logs image data to Neptune.ai when a valid Neptune run is active. Images are organized under the
|
||||
specified group name.
|
||||
|
||||
Args:
|
||||
imgs_dict (dict): Dictionary of images to log, with keys as image names and values as image data.
|
||||
group (str, optional): Group name to organize images under in the Neptune UI.
|
||||
|
||||
Examples:
|
||||
>>> # Log validation images
|
||||
>>> _log_images({"val_batch": img_tensor}, group="validation")
|
||||
"""
|
||||
if run:
|
||||
for k, v in imgs_dict.items():
|
||||
run[f"{group}/{k}"].upload(File(v))
|
||||
|
||||
|
||||
def _log_plot(title: str, plot_path: str) -> None:
|
||||
"""Log plots to the NeptuneAI experiment logger."""
|
||||
import matplotlib.image as mpimg
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
img = mpimg.imread(plot_path)
|
||||
fig = plt.figure()
|
||||
ax = fig.add_axes([0, 0, 1, 1], frameon=False, aspect="auto", xticks=[], yticks=[]) # no ticks
|
||||
ax.imshow(img)
|
||||
run[f"Plots/{title}"].upload(fig)
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer) -> None:
|
||||
"""Initialize NeptuneAI run and log hyperparameters before training starts."""
|
||||
try:
|
||||
global run
|
||||
run = neptune.init_run(
|
||||
project=trainer.args.project or "Ultralytics",
|
||||
name=trainer.args.name,
|
||||
tags=["Ultralytics"],
|
||||
)
|
||||
run["Configuration/Hyperparameters"] = {k: "" if v is None else v for k, v in vars(trainer.args).items()}
|
||||
except Exception as e:
|
||||
LOGGER.warning(f"NeptuneAI installed but not initialized correctly, not logging this run. {e}")
|
||||
|
||||
|
||||
def on_train_epoch_end(trainer) -> None:
|
||||
"""Log training metrics and learning rate at the end of each training epoch."""
|
||||
_log_scalars(trainer.label_loss_items(trainer.tloss, prefix="train"), trainer.epoch + 1)
|
||||
_log_scalars(trainer.lr, trainer.epoch + 1)
|
||||
if trainer.epoch == 1:
|
||||
_log_images({f.stem: str(f) for f in trainer.save_dir.glob("train_batch*.jpg")}, "Mosaic")
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer) -> None:
|
||||
"""Log model info and validation metrics at the end of each fit epoch."""
|
||||
if run and trainer.epoch == 0:
|
||||
from ultralytics.utils.torch_utils import model_info_for_loggers
|
||||
|
||||
run["Configuration/Model"] = model_info_for_loggers(trainer)
|
||||
_log_scalars(trainer.metrics, trainer.epoch + 1)
|
||||
|
||||
|
||||
def on_val_end(validator) -> None:
|
||||
"""Log validation images at the end of validation."""
|
||||
if run:
|
||||
# Log val_labels and val_pred
|
||||
_log_images({f.stem: str(f) for f in validator.save_dir.glob("val*.jpg")}, "Validation")
|
||||
|
||||
|
||||
def on_train_end(trainer) -> None:
|
||||
"""Log final results, plots, and model weights at the end of training."""
|
||||
if run:
|
||||
# Log final results, CM matrix + PR plots
|
||||
for f in [*trainer.plots.keys(), *trainer.validator.plots.keys()]:
|
||||
if "batch" not in f.name:
|
||||
_log_plot(title=f.stem, plot_path=f)
|
||||
# Log the final model
|
||||
run[f"weights/{trainer.args.name or trainer.args.task}/{trainer.best.name}"].upload(File(str(trainer.best)))
|
||||
|
||||
|
||||
callbacks = (
|
||||
{
|
||||
"on_pretrain_routine_start": on_pretrain_routine_start,
|
||||
"on_train_epoch_end": on_train_epoch_end,
|
||||
"on_fit_epoch_end": on_fit_epoch_end,
|
||||
"on_val_end": on_val_end,
|
||||
"on_train_end": on_train_end,
|
||||
}
|
||||
if neptune
|
||||
else {}
|
||||
)
|
||||
458
algorithms/dms_yolo/code/ultralytics/utils/callbacks/platform.py
Normal file
458
algorithms/dms_yolo/code/ultralytics/utils/callbacks/platform.py
Normal file
@@ -0,0 +1,458 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
import os
|
||||
import platform
|
||||
import re
|
||||
import socket
|
||||
import sys
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from time import time
|
||||
|
||||
from ultralytics.utils import ENVIRONMENT, GIT, LOGGER, PYTHON_VERSION, RANK, SETTINGS, TESTS_RUNNING, colorstr
|
||||
|
||||
PREFIX = colorstr("Platform: ")
|
||||
|
||||
# Configurable platform URL for debugging (e.g. ULTRALYTICS_PLATFORM_URL=http://localhost:3000)
|
||||
PLATFORM_URL = os.getenv("ULTRALYTICS_PLATFORM_URL", "https://platform.ultralytics.com").rstrip("/")
|
||||
PLATFORM_API_URL = f"{PLATFORM_URL}/api/webhooks"
|
||||
|
||||
|
||||
def slugify(text):
|
||||
"""Convert text to URL-safe slug (e.g., 'My Project 1' -> 'my-project-1')."""
|
||||
if not text:
|
||||
return text
|
||||
return re.sub(r"-+", "-", re.sub(r"[^a-z0-9\s-]", "", str(text).lower()).replace(" ", "-")).strip("-")[:128]
|
||||
|
||||
|
||||
try:
|
||||
assert not TESTS_RUNNING # do not log pytest
|
||||
assert SETTINGS.get("platform", False) is True or os.getenv("ULTRALYTICS_API_KEY") or SETTINGS.get("api_key")
|
||||
_api_key = os.getenv("ULTRALYTICS_API_KEY") or SETTINGS.get("api_key")
|
||||
assert _api_key # verify API key is present
|
||||
|
||||
import requests
|
||||
|
||||
from ultralytics.utils.logger import ConsoleLogger, SystemLogger
|
||||
from ultralytics.utils.torch_utils import model_info_for_loggers
|
||||
|
||||
_executor = ThreadPoolExecutor(max_workers=10) # Bounded thread pool for async operations
|
||||
|
||||
except (AssertionError, ImportError):
|
||||
_api_key = None
|
||||
|
||||
|
||||
def resolve_platform_uri(uri, hard=True):
|
||||
"""Resolve ul:// URIs to signed URLs by authenticating with Ultralytics Platform.
|
||||
|
||||
Formats:
|
||||
ul://username/datasets/slug -> Returns signed URL to NDJSON file
|
||||
ul://username/project/model -> Returns signed URL to .pt file
|
||||
|
||||
Args:
|
||||
uri (str): Platform URI starting with "ul://".
|
||||
hard (bool): Whether to raise an error if resolution fails (FileNotFoundError only).
|
||||
|
||||
Returns:
|
||||
(str | None): Signed URL on success, None if not found and hard=False.
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is missing/invalid or URI format is wrong.
|
||||
PermissionError: If access is denied.
|
||||
RuntimeError: If resource is not ready (e.g., dataset still processing).
|
||||
FileNotFoundError: If resource not found and hard=True.
|
||||
ConnectionError: If network request fails and hard=True.
|
||||
"""
|
||||
import requests
|
||||
|
||||
path = uri[5:] # Remove "ul://"
|
||||
parts = path.split("/")
|
||||
|
||||
api_key = os.getenv("ULTRALYTICS_API_KEY") or SETTINGS.get("api_key")
|
||||
if not api_key:
|
||||
raise ValueError(f"ULTRALYTICS_API_KEY required for '{uri}'. Get key at {PLATFORM_URL}/settings")
|
||||
|
||||
base = PLATFORM_API_URL
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
|
||||
# ul://username/datasets/slug
|
||||
if len(parts) == 3 and parts[1] == "datasets":
|
||||
username, _, slug = parts
|
||||
url = f"{base}/datasets/{username}/{slug}/export"
|
||||
|
||||
# ul://username/project/model
|
||||
elif len(parts) == 3:
|
||||
username, project, model = parts
|
||||
url = f"{base}/models/{username}/{project}/{model}/download"
|
||||
|
||||
else:
|
||||
raise ValueError(f"Invalid platform URI: {uri}. Use ul://user/datasets/name or ul://user/project/model")
|
||||
|
||||
try:
|
||||
timeout = 3600 if "/datasets/" in url else 90 # NDJSON generation can be slow for large datasets
|
||||
r = requests.head(url, headers=headers, allow_redirects=False, timeout=timeout)
|
||||
|
||||
# Handle redirect responses (301, 302, 303, 307, 308)
|
||||
if 300 <= r.status_code < 400 and "location" in r.headers:
|
||||
return r.headers["location"] # Return signed URL
|
||||
|
||||
# Handle error responses
|
||||
if r.status_code == 401:
|
||||
raise ValueError(f"Invalid ULTRALYTICS_API_KEY for '{uri}'")
|
||||
if r.status_code == 403:
|
||||
raise PermissionError(f"Access denied for '{uri}'. Check dataset/model visibility settings.")
|
||||
if r.status_code == 404:
|
||||
if hard:
|
||||
raise FileNotFoundError(f"Not found on platform: {uri}")
|
||||
LOGGER.warning(f"Not found on platform: {uri}")
|
||||
return None
|
||||
if r.status_code == 409:
|
||||
raise RuntimeError(f"Resource not ready: {uri}. Dataset may still be processing.")
|
||||
|
||||
# Unexpected response
|
||||
r.raise_for_status()
|
||||
raise RuntimeError(f"Unexpected response from platform for '{uri}': {r.status_code}")
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
if hard:
|
||||
raise ConnectionError(f"Failed to resolve {uri}: {e}") from e
|
||||
LOGGER.warning(f"Failed to resolve {uri}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _interp_plot(plot, n=101):
|
||||
"""Interpolate plot curve data from 1000 to n points to reduce storage size."""
|
||||
import numpy as np
|
||||
|
||||
if not plot.get("x") or not plot.get("y"):
|
||||
return plot # No interpolation needed (e.g., confusion_matrix)
|
||||
|
||||
x, y = np.array(plot["x"]), np.array(plot["y"])
|
||||
if len(x) <= n:
|
||||
return plot # Already small enough
|
||||
|
||||
# New x values (101 points gives clean 0.01 increments: 0, 0.01, 0.02, ..., 1.0)
|
||||
x_new = np.linspace(x[0], x[-1], n)
|
||||
|
||||
# Interpolate y values (handle both 1D and 2D arrays)
|
||||
if y.ndim == 1:
|
||||
y_new = np.interp(x_new, x, y)
|
||||
else:
|
||||
y_new = np.array([np.interp(x_new, x, yi) for yi in y])
|
||||
|
||||
# Also interpolate ap if present (for PR curves)
|
||||
result = {**plot, "x": x_new.tolist(), "y": y_new.tolist()}
|
||||
if "ap" in plot:
|
||||
result["ap"] = plot["ap"] # Keep AP values as-is (per-class scalars)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _send(event, data, project, name, model_id=None):
|
||||
"""Send event to Platform endpoint. Returns response JSON on success."""
|
||||
try:
|
||||
payload = {"event": event, "project": project, "name": name, "data": data}
|
||||
if model_id:
|
||||
payload["modelId"] = model_id
|
||||
r = requests.post(
|
||||
f"{PLATFORM_API_URL}/training/metrics",
|
||||
json=payload,
|
||||
headers={"Authorization": f"Bearer {_api_key}"},
|
||||
timeout=10,
|
||||
)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
except Exception as e:
|
||||
LOGGER.debug(f"Platform: Failed to send {event}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _send_async(event, data, project, name, model_id=None):
|
||||
"""Send event asynchronously using bounded thread pool."""
|
||||
_executor.submit(_send, event, data, project, name, model_id)
|
||||
|
||||
|
||||
def _upload_model(model_path, project, name):
|
||||
"""Upload model checkpoint to Platform via signed URL."""
|
||||
try:
|
||||
model_path = Path(model_path)
|
||||
if not model_path.exists():
|
||||
return None
|
||||
|
||||
# Get signed upload URL
|
||||
response = requests.post(
|
||||
f"{PLATFORM_API_URL}/models/upload",
|
||||
json={"project": project, "name": name, "filename": model_path.name},
|
||||
headers={"Authorization": f"Bearer {_api_key}"},
|
||||
timeout=10,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
# Upload to GCS
|
||||
with open(model_path, "rb") as f:
|
||||
requests.put(
|
||||
data["uploadUrl"],
|
||||
data=f,
|
||||
headers={"Content-Type": "application/octet-stream"},
|
||||
timeout=600, # 10 min timeout for large models
|
||||
).raise_for_status()
|
||||
|
||||
# url = f"{PLATFORM_URL}/{project}/{name}"
|
||||
# LOGGER.info(f"{PREFIX}Model uploaded to {url}")
|
||||
return data.get("gcsPath")
|
||||
|
||||
except Exception as e:
|
||||
LOGGER.debug(f"Platform: Failed to upload model: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _upload_model_async(model_path, project, name):
|
||||
"""Upload model asynchronously using bounded thread pool."""
|
||||
_executor.submit(_upload_model, model_path, project, name)
|
||||
|
||||
|
||||
def _get_environment_info():
|
||||
"""Collect comprehensive environment info using existing ultralytics utilities."""
|
||||
import shutil
|
||||
|
||||
import psutil
|
||||
import torch
|
||||
|
||||
from ultralytics import __version__
|
||||
from ultralytics.utils.torch_utils import get_cpu_info, get_gpu_info
|
||||
|
||||
# Get RAM and disk totals
|
||||
memory = psutil.virtual_memory()
|
||||
disk_usage = shutil.disk_usage("/")
|
||||
|
||||
env = {
|
||||
"ultralyticsVersion": __version__,
|
||||
"hostname": socket.gethostname(),
|
||||
"os": platform.platform(),
|
||||
"environment": ENVIRONMENT,
|
||||
"pythonVersion": PYTHON_VERSION,
|
||||
"pythonExecutable": sys.executable,
|
||||
"cpuCount": os.cpu_count() or 0,
|
||||
"cpu": get_cpu_info(),
|
||||
"command": " ".join(sys.argv),
|
||||
"totalRamGb": round(memory.total / (1 << 30), 1), # Total RAM in GB
|
||||
"totalDiskGb": round(disk_usage.total / (1 << 30), 1), # Total disk in GB
|
||||
}
|
||||
|
||||
# Git info using cached GIT singleton (no subprocess calls)
|
||||
try:
|
||||
if GIT.is_repo:
|
||||
if GIT.origin:
|
||||
env["gitRepository"] = GIT.origin
|
||||
if GIT.branch:
|
||||
env["gitBranch"] = GIT.branch
|
||||
if GIT.commit:
|
||||
env["gitCommit"] = GIT.commit[:12] # Short hash
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# GPU info
|
||||
try:
|
||||
if torch.cuda.is_available():
|
||||
env["gpuCount"] = torch.cuda.device_count()
|
||||
env["gpuType"] = get_gpu_info(0) if torch.cuda.device_count() > 0 else None
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return env
|
||||
|
||||
|
||||
def _get_project_name(trainer):
|
||||
"""Get slugified project and name from trainer args."""
|
||||
raw = str(trainer.args.project)
|
||||
parts = raw.split("/", 1)
|
||||
project = f"{parts[0]}/{slugify(parts[1])}" if len(parts) == 2 else slugify(raw)
|
||||
return project, slugify(str(trainer.args.name or "train"))
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer):
|
||||
"""Initialize Platform logging at training start."""
|
||||
if RANK not in {-1, 0} or not trainer.args.project:
|
||||
return
|
||||
|
||||
# Per-trainer state to isolate concurrent training runs
|
||||
trainer._platform_model_id = None
|
||||
trainer._platform_last_upload = time()
|
||||
|
||||
project, name = _get_project_name(trainer)
|
||||
url = f"{PLATFORM_URL}/{project}/{name}"
|
||||
LOGGER.info(f"{PREFIX}Streaming to {url}")
|
||||
|
||||
# Create callback to send console output to Platform
|
||||
def send_console_output(content, line_count, chunk_id):
|
||||
"""Send batched console output to Platform webhook."""
|
||||
_send_async(
|
||||
"console_output",
|
||||
{"chunkId": chunk_id, "content": content, "lineCount": line_count},
|
||||
project,
|
||||
name,
|
||||
getattr(trainer, "_platform_model_id", None),
|
||||
)
|
||||
|
||||
# Start console capture with batching (5 lines or 5 seconds)
|
||||
trainer._platform_console_logger = ConsoleLogger(batch_size=5, flush_interval=5.0, on_flush=send_console_output)
|
||||
trainer._platform_console_logger.start_capture()
|
||||
|
||||
# Collect environment info (W&B-style metadata)
|
||||
environment = _get_environment_info()
|
||||
|
||||
# Build trainArgs - callback runs before get_dataset() so args.data is still original (e.g., ul:// URIs)
|
||||
# Note: model_info is sent later in on_fit_epoch_end (epoch 0) when the model is actually loaded
|
||||
train_args = {k: str(v) for k, v in vars(trainer.args).items()}
|
||||
|
||||
# Send synchronously to get modelId for subsequent webhooks
|
||||
response = _send(
|
||||
"training_started",
|
||||
{
|
||||
"trainArgs": train_args,
|
||||
"epochs": trainer.epochs,
|
||||
"device": str(trainer.device),
|
||||
"environment": environment,
|
||||
},
|
||||
project,
|
||||
name,
|
||||
)
|
||||
if response and response.get("modelId"):
|
||||
trainer._platform_model_id = response["modelId"]
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Log training and system metrics at epoch end."""
|
||||
if RANK not in {-1, 0} or not trainer.args.project:
|
||||
return
|
||||
|
||||
project, name = _get_project_name(trainer)
|
||||
metrics = {**trainer.label_loss_items(trainer.tloss, prefix="train"), **trainer.metrics}
|
||||
|
||||
if trainer.optimizer and trainer.optimizer.param_groups:
|
||||
metrics["lr"] = trainer.optimizer.param_groups[0]["lr"]
|
||||
|
||||
# Extract model info at epoch 0 (sent as separate field, not in metrics)
|
||||
model_info = None
|
||||
if trainer.epoch == 0:
|
||||
try:
|
||||
info = model_info_for_loggers(trainer)
|
||||
model_info = {
|
||||
"parameters": info.get("model/parameters", 0),
|
||||
"gflops": info.get("model/GFLOPs", 0),
|
||||
"speedMs": info.get("model/speed_PyTorch(ms)", 0),
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Get system metrics (cache SystemLogger on trainer for efficiency)
|
||||
system = {}
|
||||
try:
|
||||
if not hasattr(trainer, "_platform_system_logger"):
|
||||
trainer._platform_system_logger = SystemLogger()
|
||||
system = trainer._platform_system_logger.get_metrics(rates=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
payload = {
|
||||
"epoch": trainer.epoch,
|
||||
"metrics": metrics,
|
||||
"system": system,
|
||||
"fitness": trainer.fitness,
|
||||
"best_fitness": trainer.best_fitness,
|
||||
}
|
||||
if model_info:
|
||||
payload["modelInfo"] = model_info
|
||||
|
||||
_send_async(
|
||||
"epoch_end",
|
||||
payload,
|
||||
project,
|
||||
name,
|
||||
getattr(trainer, "_platform_model_id", None),
|
||||
)
|
||||
|
||||
|
||||
def on_model_save(trainer):
|
||||
"""Upload model checkpoint (rate limited to every 15 min)."""
|
||||
if RANK not in {-1, 0} or not trainer.args.project:
|
||||
return
|
||||
|
||||
# Rate limit to every 15 minutes (900 seconds)
|
||||
if time() - getattr(trainer, "_platform_last_upload", 0) < 900:
|
||||
return
|
||||
|
||||
model_path = trainer.best if trainer.best and Path(trainer.best).exists() else trainer.last
|
||||
if not model_path:
|
||||
return
|
||||
|
||||
project, name = _get_project_name(trainer)
|
||||
_upload_model_async(model_path, project, name)
|
||||
trainer._platform_last_upload = time()
|
||||
|
||||
|
||||
def on_train_end(trainer):
|
||||
"""Log final results, upload best model, and send validation plot data."""
|
||||
if RANK not in {-1, 0} or not trainer.args.project:
|
||||
return
|
||||
|
||||
project, name = _get_project_name(trainer)
|
||||
|
||||
# Stop console capture
|
||||
if hasattr(trainer, "_platform_console_logger") and trainer._platform_console_logger:
|
||||
trainer._platform_console_logger.stop_capture()
|
||||
trainer._platform_console_logger = None
|
||||
|
||||
# Upload best model (blocking to ensure it completes)
|
||||
model_path = None
|
||||
model_size = None
|
||||
if trainer.best and Path(trainer.best).exists():
|
||||
model_size = Path(trainer.best).stat().st_size
|
||||
model_path = _upload_model(trainer.best, project, name)
|
||||
|
||||
# Collect plots from trainer and validator, deduplicating by type
|
||||
plots_by_type = {}
|
||||
for info in getattr(trainer, "plots", {}).values():
|
||||
if info.get("data") and info["data"].get("type"):
|
||||
plots_by_type[info["data"]["type"]] = info["data"]
|
||||
for info in getattr(getattr(trainer, "validator", None), "plots", {}).values():
|
||||
if info.get("data") and info["data"].get("type"):
|
||||
plots_by_type.setdefault(info["data"]["type"], info["data"]) # Don't overwrite trainer plots
|
||||
plots = [_interp_plot(p) for p in plots_by_type.values()] # Interpolate curves to reduce size
|
||||
|
||||
# Get class names
|
||||
names = getattr(getattr(trainer, "validator", None), "names", None) or (trainer.data or {}).get("names")
|
||||
class_names = list(names.values()) if isinstance(names, dict) else list(names) if names else None
|
||||
|
||||
_send(
|
||||
"training_complete",
|
||||
{
|
||||
"results": {
|
||||
"metrics": {**trainer.metrics, "fitness": trainer.fitness},
|
||||
"bestEpoch": getattr(trainer, "best_epoch", trainer.epoch),
|
||||
"bestFitness": trainer.best_fitness,
|
||||
"modelPath": model_path or (str(trainer.best) if trainer.best else None),
|
||||
"modelSize": model_size,
|
||||
},
|
||||
"classNames": class_names,
|
||||
"plots": plots,
|
||||
},
|
||||
project,
|
||||
name,
|
||||
getattr(trainer, "_platform_model_id", None),
|
||||
)
|
||||
url = f"{PLATFORM_URL}/{project}/{name}"
|
||||
LOGGER.info(f"{PREFIX}View results at {url}")
|
||||
|
||||
|
||||
callbacks = (
|
||||
{
|
||||
"on_pretrain_routine_start": on_pretrain_routine_start,
|
||||
"on_fit_epoch_end": on_fit_epoch_end,
|
||||
"on_model_save": on_model_save,
|
||||
"on_train_end": on_train_end,
|
||||
}
|
||||
if _api_key
|
||||
else {}
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from ultralytics.utils import SETTINGS
|
||||
|
||||
try:
|
||||
assert SETTINGS["raytune"] is True # verify integration is enabled
|
||||
import ray
|
||||
from ray import tune
|
||||
from ray.air import session
|
||||
|
||||
except (ImportError, AssertionError):
|
||||
tune = None
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Report training metrics to Ray Tune at epoch end when a Ray session is active.
|
||||
|
||||
Captures metrics from the trainer object and sends them to Ray Tune with the current epoch number, enabling
|
||||
hyperparameter tuning optimization. Only executes when within an active Ray Tune session.
|
||||
|
||||
Args:
|
||||
trainer (ultralytics.engine.trainer.BaseTrainer): The Ultralytics trainer object containing metrics and epochs.
|
||||
|
||||
Examples:
|
||||
>>> # Called automatically by the Ultralytics training loop
|
||||
>>> on_fit_epoch_end(trainer)
|
||||
|
||||
References:
|
||||
Ray Tune docs: https://docs.ray.io/en/latest/tune/index.html
|
||||
"""
|
||||
if ray.train._internal.session.get_session(): # check if Ray Tune session is active
|
||||
metrics = trainer.metrics
|
||||
session.report({**metrics, **{"epoch": trainer.epoch + 1}})
|
||||
|
||||
|
||||
callbacks = (
|
||||
{
|
||||
"on_fit_epoch_end": on_fit_epoch_end,
|
||||
}
|
||||
if tune
|
||||
else {}
|
||||
)
|
||||
@@ -0,0 +1,125 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from ultralytics.utils import LOGGER, SETTINGS, TESTS_RUNNING, colorstr, torch_utils
|
||||
from ultralytics.utils.torch_utils import smart_inference_mode
|
||||
|
||||
try:
|
||||
assert not TESTS_RUNNING # do not log pytest
|
||||
assert SETTINGS["tensorboard"] is True # verify integration is enabled
|
||||
WRITER = None # TensorBoard SummaryWriter instance
|
||||
PREFIX = colorstr("TensorBoard: ")
|
||||
|
||||
# Imports below only required if TensorBoard enabled
|
||||
from copy import deepcopy
|
||||
|
||||
import torch
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
except (ImportError, AssertionError, TypeError, AttributeError):
|
||||
# TypeError for handling 'Descriptors cannot not be created directly.' protobuf errors in Windows
|
||||
# AttributeError: module 'tensorflow' has no attribute 'io' if 'tensorflow' not installed
|
||||
SummaryWriter = None
|
||||
|
||||
|
||||
def _log_scalars(scalars: dict, step: int = 0) -> None:
|
||||
"""Log scalar values to TensorBoard.
|
||||
|
||||
Args:
|
||||
scalars (dict): Dictionary of scalar values to log to TensorBoard. Keys are scalar names and values are the
|
||||
corresponding scalar values.
|
||||
step (int): Global step value to record with the scalar values. Used for x-axis in TensorBoard graphs.
|
||||
|
||||
Examples:
|
||||
Log training metrics
|
||||
>>> metrics = {"loss": 0.5, "accuracy": 0.95}
|
||||
>>> _log_scalars(metrics, step=100)
|
||||
"""
|
||||
if WRITER:
|
||||
for k, v in scalars.items():
|
||||
WRITER.add_scalar(k, v, step)
|
||||
|
||||
|
||||
@smart_inference_mode()
|
||||
def _log_tensorboard_graph(trainer) -> None:
|
||||
"""Log model graph to TensorBoard.
|
||||
|
||||
This function attempts to visualize the model architecture in TensorBoard by tracing the model with a dummy input
|
||||
tensor. It first tries a simple method suitable for YOLO models, and if that fails, falls back to a more complex
|
||||
approach for models like RTDETR that may require special handling.
|
||||
|
||||
Args:
|
||||
trainer (ultralytics.engine.trainer.BaseTrainer): The trainer object containing the model to visualize. Must
|
||||
have attributes model and args with imgsz.
|
||||
|
||||
Notes:
|
||||
This function requires TensorBoard integration to be enabled and the global WRITER to be initialized.
|
||||
It handles potential warnings from the PyTorch JIT tracer and attempts to gracefully handle different
|
||||
model architectures.
|
||||
"""
|
||||
# Input image
|
||||
imgsz = trainer.args.imgsz
|
||||
imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else imgsz
|
||||
p = next(trainer.model.parameters()) # for device, type
|
||||
im = torch.zeros((1, 3, *imgsz), device=p.device, dtype=p.dtype) # input image (must be zeros, not empty)
|
||||
|
||||
# Try simple method first (YOLO)
|
||||
try:
|
||||
trainer.model.eval() # place in .eval() mode to avoid BatchNorm statistics changes
|
||||
WRITER.add_graph(torch.jit.trace(torch_utils.unwrap_model(trainer.model), im, strict=False), [])
|
||||
LOGGER.info(f"{PREFIX}model graph visualization added ✅")
|
||||
return
|
||||
except Exception as e1:
|
||||
# Fallback to TorchScript export steps (RTDETR)
|
||||
try:
|
||||
model = deepcopy(torch_utils.unwrap_model(trainer.model))
|
||||
model.eval()
|
||||
model = model.fuse(verbose=False)
|
||||
for m in model.modules():
|
||||
if hasattr(m, "export"): # Detect, RTDETRDecoder (Segment and Pose use Detect base class)
|
||||
m.export = True
|
||||
m.format = "torchscript"
|
||||
model(im) # dry run
|
||||
WRITER.add_graph(torch.jit.trace(model, im, strict=False), [])
|
||||
LOGGER.info(f"{PREFIX}model graph visualization added ✅")
|
||||
except Exception as e2:
|
||||
LOGGER.warning(f"{PREFIX}TensorBoard graph visualization failure: {e1} -> {e2}")
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer) -> None:
|
||||
"""Initialize TensorBoard logging with SummaryWriter."""
|
||||
if SummaryWriter:
|
||||
try:
|
||||
global WRITER
|
||||
WRITER = SummaryWriter(str(trainer.save_dir))
|
||||
LOGGER.info(f"{PREFIX}Start with 'tensorboard --logdir {trainer.save_dir}', view at http://localhost:6006/")
|
||||
except Exception as e:
|
||||
LOGGER.warning(f"{PREFIX}TensorBoard not initialized correctly, not logging this run. {e}")
|
||||
|
||||
|
||||
def on_train_start(trainer) -> None:
|
||||
"""Log TensorBoard graph."""
|
||||
if WRITER:
|
||||
_log_tensorboard_graph(trainer)
|
||||
|
||||
|
||||
def on_train_epoch_end(trainer) -> None:
|
||||
"""Log scalar statistics at the end of a training epoch."""
|
||||
_log_scalars(trainer.label_loss_items(trainer.tloss, prefix="train"), trainer.epoch + 1)
|
||||
_log_scalars(trainer.lr, trainer.epoch + 1)
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer) -> None:
|
||||
"""Log epoch metrics at end of training epoch."""
|
||||
_log_scalars(trainer.metrics, trainer.epoch + 1)
|
||||
|
||||
|
||||
callbacks = (
|
||||
{
|
||||
"on_pretrain_routine_start": on_pretrain_routine_start,
|
||||
"on_train_start": on_train_start,
|
||||
"on_fit_epoch_end": on_fit_epoch_end,
|
||||
"on_train_epoch_end": on_train_epoch_end,
|
||||
}
|
||||
if SummaryWriter
|
||||
else {}
|
||||
)
|
||||
193
algorithms/dms_yolo/code/ultralytics/utils/callbacks/wb.py
Normal file
193
algorithms/dms_yolo/code/ultralytics/utils/callbacks/wb.py
Normal file
@@ -0,0 +1,193 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from ultralytics.utils import SETTINGS, TESTS_RUNNING
|
||||
from ultralytics.utils.torch_utils import model_info_for_loggers
|
||||
|
||||
try:
|
||||
assert not TESTS_RUNNING # do not log pytest
|
||||
assert SETTINGS["wandb"] is True # verify integration is enabled
|
||||
import wandb as wb
|
||||
|
||||
assert hasattr(wb, "__version__") # verify package is not directory
|
||||
_processed_plots = {}
|
||||
|
||||
except (ImportError, AssertionError):
|
||||
wb = None
|
||||
|
||||
|
||||
def _custom_table(x, y, classes, title="Precision Recall Curve", x_title="Recall", y_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 the default wandb 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.
|
||||
x_title (str, optional): Label for the x-axis.
|
||||
y_title (str, optional): Label for the y-axis.
|
||||
|
||||
Returns:
|
||||
(wandb.Object): A wandb object suitable for logging, showcasing the crafted metric visualization.
|
||||
"""
|
||||
import polars as pl # scope for faster 'import ultralytics'
|
||||
import polars.selectors as cs
|
||||
|
||||
df = pl.DataFrame({"class": classes, "y": y, "x": x}).with_columns(cs.numeric().round(3))
|
||||
data = df.select(["class", "y", "x"]).rows()
|
||||
|
||||
fields = {"x": "x", "y": "y", "class": "class"}
|
||||
string_fields = {"title": title, "x-axis-title": x_title, "y-axis-title": y_title}
|
||||
return wb.plot_table(
|
||||
"wandb/area-under-curve/v0",
|
||||
wb.Table(data=data, columns=["class", "y", "x"]),
|
||||
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 (C, N), where C is the number of classes.
|
||||
names (list, optional): Names of the classes corresponding to the y-axis data; length C.
|
||||
id (str, optional): Unique identifier for the logged data in wandb.
|
||||
title (str, optional): Title for the visualization plot.
|
||||
x_title (str, optional): Label for the x-axis.
|
||||
y_title (str, optional): Label for the y-axis.
|
||||
num_x (int, optional): Number of interpolated data points for visualization.
|
||||
only_mean (bool, optional): Flag to indicate if only the mean curve should be plotted.
|
||||
|
||||
Notes:
|
||||
The function leverages the '_custom_table' function to generate the actual visualization.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
# 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):
|
||||
"""Log plots to WandB at a specific step if they haven't been logged already.
|
||||
|
||||
This function checks each plot in the input dictionary against previously processed plots and logs new or updated
|
||||
plots to WandB at the specified step.
|
||||
|
||||
Args:
|
||||
plots (dict): Dictionary of plots to log, where keys are plot names and values are dictionaries containing plot
|
||||
metadata including timestamps.
|
||||
step (int): The step/epoch at which to log the plots in the WandB run.
|
||||
|
||||
Notes:
|
||||
The function uses a shallow copy of the plots dictionary to prevent modification during iteration.
|
||||
Plots are identified by their stem name (filename without extension).
|
||||
Each plot is logged as a WandB Image object.
|
||||
"""
|
||||
for name, params in plots.copy().items(): # shallow copy to prevent plots dict changing during iteration
|
||||
timestamp = params["timestamp"]
|
||||
if _processed_plots.get(name) != timestamp:
|
||||
wb.run.log({name.stem: wb.Image(str(name))}, step=step)
|
||||
_processed_plots[name] = timestamp
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer):
|
||||
"""Initialize and start wandb project if module is present."""
|
||||
if not wb.run:
|
||||
from datetime import datetime
|
||||
|
||||
name = str(trainer.args.name).replace("/", "-").replace(" ", "_")
|
||||
wb.init(
|
||||
project=str(trainer.args.project).replace("/", "-") if trainer.args.project else "Ultralytics",
|
||||
name=name,
|
||||
config=vars(trainer.args),
|
||||
id=f"{name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}", # add unique id
|
||||
dir=str(trainer.save_dir),
|
||||
)
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Log training metrics and model information at the end of an epoch."""
|
||||
_log_plots(trainer.plots, step=trainer.epoch + 1)
|
||||
_log_plots(trainer.validator.plots, step=trainer.epoch + 1)
|
||||
if trainer.epoch == 0:
|
||||
wb.run.log(model_info_for_loggers(trainer), step=trainer.epoch + 1)
|
||||
wb.run.log(trainer.metrics, step=trainer.epoch + 1, commit=True) # commit forces sync
|
||||
|
||||
|
||||
def on_train_epoch_end(trainer):
|
||||
"""Log metrics and save images at the end of each training epoch."""
|
||||
wb.run.log(trainer.label_loss_items(trainer.tloss, prefix="train"), step=trainer.epoch + 1)
|
||||
wb.run.log(trainer.lr, step=trainer.epoch + 1)
|
||||
if trainer.epoch == 1:
|
||||
_log_plots(trainer.plots, step=trainer.epoch + 1)
|
||||
|
||||
|
||||
def on_train_end(trainer):
|
||||
"""Save the best model as an artifact and log final plots at the end of training."""
|
||||
_log_plots(trainer.validator.plots, step=trainer.epoch + 1)
|
||||
_log_plots(trainer.plots, step=trainer.epoch + 1)
|
||||
art = wb.Artifact(type="model", name=f"run_{wb.run.id}_model")
|
||||
if trainer.best.exists():
|
||||
art.add_file(trainer.best)
|
||||
wb.run.log_artifact(art, aliases=["best"])
|
||||
# Check if we actually have plots to save
|
||||
if trainer.args.plots and hasattr(trainer.validator.metrics, "curves_results"):
|
||||
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
|
||||
|
||||
|
||||
callbacks = (
|
||||
{
|
||||
"on_pretrain_routine_start": on_pretrain_routine_start,
|
||||
"on_train_epoch_end": on_train_epoch_end,
|
||||
"on_fit_epoch_end": on_fit_epoch_end,
|
||||
"on_train_end": on_train_end,
|
||||
}
|
||||
if wb
|
||||
else {}
|
||||
)
|
||||
Reference in New Issue
Block a user