wandb 0.19.1__py3-none-musllinux_1_2_aarch64.whl

wandb/integration/ultralytics/callback.py

@@ -0,0 +1,524 @@
+import copy
+from datetime import datetime
+from typing import Callable, Dict, Optional, Union
+
+from packaging import version
+
+try:
+    import dill as pickle
+except ImportError:
+    import pickle
+
+import wandb
+from wandb.sdk.lib import telemetry
+
+try:
+    import torch
+    import ultralytics
+    from tqdm.auto import tqdm
+
+    if version.parse(ultralytics.__version__) > version.parse("8.0.238"):
+        wandb.termwarn(
+            """This integration is tested and supported for ultralytics v8.0.238 and below.
+            Please report any issues to https://github.com/wandb/wandb/issues with the tag `yolov8`.""",
+            repeat=False,
+        )
+
+    from ultralytics.models import YOLO
+    from ultralytics.models.sam.predict import Predictor as SAMPredictor
+    from ultralytics.models.yolo.classify import (
+        ClassificationPredictor,
+        ClassificationTrainer,
+        ClassificationValidator,
+    )
+    from ultralytics.models.yolo.detect import (
+        DetectionPredictor,
+        DetectionTrainer,
+        DetectionValidator,
+    )
+    from ultralytics.models.yolo.pose import PosePredictor, PoseTrainer, PoseValidator
+    from ultralytics.models.yolo.segment import (
+        SegmentationPredictor,
+        SegmentationTrainer,
+        SegmentationValidator,
+    )
+    from ultralytics.utils.torch_utils import de_parallel
+
+    try:
+        from ultralytics.yolo.utils import RANK, __version__
+    except ModuleNotFoundError:
+        from ultralytics.utils import RANK, __version__
+
+    from wandb.integration.ultralytics.bbox_utils import (
+        plot_bbox_predictions,
+        plot_detection_validation_results,
+    )
+    from wandb.integration.ultralytics.classification_utils import (
+        plot_classification_predictions,
+        plot_classification_validation_results,
+    )
+    from wandb.integration.ultralytics.mask_utils import (
+        plot_mask_predictions,
+        plot_sam_predictions,
+        plot_segmentation_validation_results,
+    )
+    from wandb.integration.ultralytics.pose_utils import (
+        plot_pose_predictions,
+        plot_pose_validation_results,
+    )
+except Exception as e:
+    wandb.Error(e)
+
+
+TRAINER_TYPE = Union[
+    ClassificationTrainer, DetectionTrainer, SegmentationTrainer, PoseTrainer
+]
+VALIDATOR_TYPE = Union[
+    ClassificationValidator, DetectionValidator, SegmentationValidator, PoseValidator
+]
+PREDICTOR_TYPE = Union[
+    ClassificationPredictor,
+    DetectionPredictor,
+    SegmentationPredictor,
+    PosePredictor,
+    SAMPredictor,
+]
+
+
+class WandBUltralyticsCallback:
+    """Stateful callback for logging to W&B.
+
+    In particular, it will log model checkpoints, predictions, and
+    ground-truth annotations with interactive overlays for bounding boxes
+    to Weights & Biases Tables during training, validation and prediction
+    for a `ultratytics` workflow.
+
+    Example:
+        ```python
+        from ultralytics.yolo.engine.model import YOLO
+        from wandb.yolov8 import add_wandb_callback
+
+        # initialize YOLO model
+        model = YOLO("yolov8n.pt")
+
+        # add wandb callback
+        add_wandb_callback(model, max_validation_batches=2, enable_model_checkpointing=True)
+
+        # train
+        model.train(data="coco128.yaml", epochs=5, imgsz=640)
+
+        # validate
+        model.val()
+
+        # perform inference
+        model(["img1.jpeg", "img2.jpeg"])
+        ```
+
+    Args:
+        model: (ultralytics.yolo.engine.model.YOLO) YOLO Model of type
+            `ultralytics.yolo.engine.model.YOLO`.
+        epoch_logging_interval: (int) interval to log the prediction visualizations
+            during training.
+        max_validation_batches: (int) maximum number of validation batches to log to
+            a table per epoch.
+        enable_model_checkpointing: (bool) enable logging model checkpoints as
+            artifacts at the end of eveny epoch if set to `True`.
+        visualize_skeleton: (bool) visualize pose skeleton by drawing lines connecting
+            keypoints for human pose.
+    """
+
+    def __init__(
+        self,
+        model: YOLO,
+        epoch_logging_interval: int = 1,
+        max_validation_batches: int = 1,
+        enable_model_checkpointing: bool = False,
+        visualize_skeleton: bool = False,
+    ) -> None:
+        self.epoch_logging_interval = epoch_logging_interval
+        self.max_validation_batches = max_validation_batches
+        self.enable_model_checkpointing = enable_model_checkpointing
+        self.visualize_skeleton = visualize_skeleton
+        self.task = model.task
+        self.task_map = model.task_map
+        self.model_name = (
+            model.overrides["model"].split(".")[0]
+            if "model" in model.overrides
+            else None
+        )
+        self._make_tables()
+        self._make_predictor(model)
+        self.supported_tasks = ["detect", "segment", "pose", "classify"]
+        self.prompts = None
+        self.run_id = None
+        self.train_epoch = None
+
+    def _make_tables(self):
+        if self.task in ["detect", "segment"]:
+            validation_columns = [
+                "Data-Index",
+                "Batch-Index",
+                "Image",
+                "Mean-Confidence",
+                "Speed",
+            ]
+            train_columns = ["Epoch"] + validation_columns
+            self.train_validation_table = wandb.Table(
+                columns=["Model-Name"] + train_columns
+            )
+            self.validation_table = wandb.Table(
+                columns=["Model-Name"] + validation_columns
+            )
+            self.prediction_table = wandb.Table(
+                columns=[
+                    "Model-Name",
+                    "Image",
+                    "Num-Objects",
+                    "Mean-Confidence",
+                    "Speed",
+                ]
+            )
+        elif self.task == "classify":
+            classification_columns = [
+                "Image",
+                "Predicted-Category",
+                "Prediction-Confidence",
+                "Top-5-Prediction-Categories",
+                "Top-5-Prediction-Confindence",
+                "Probabilities",
+                "Speed",
+            ]
+            validation_columns = ["Data-Index", "Batch-Index"] + classification_columns
+            validation_columns.insert(3, "Ground-Truth-Category")
+            self.train_validation_table = wandb.Table(
+                columns=["Model-Name", "Epoch"] + validation_columns
+            )
+            self.validation_table = wandb.Table(
+                columns=["Model-Name"] + validation_columns
+            )
+            self.prediction_table = wandb.Table(
+                columns=["Model-Name"] + classification_columns
+            )
+        elif self.task == "pose":
+            validation_columns = [
+                "Data-Index",
+                "Batch-Index",
+                "Image-Ground-Truth",
+                "Image-Prediction",
+                "Num-Instances",
+                "Mean-Confidence",
+                "Speed",
+            ]
+            train_columns = ["Epoch"] + validation_columns
+            self.train_validation_table = wandb.Table(
+                columns=["Model-Name"] + train_columns
+            )
+            self.validation_table = wandb.Table(
+                columns=["Model-Name"] + validation_columns
+            )
+            self.prediction_table = wandb.Table(
+                columns=[
+                    "Model-Name",
+                    "Image-Prediction",
+                    "Num-Instances",
+                    "Mean-Confidence",
+                    "Speed",
+                ]
+            )
+
+    def _make_predictor(self, model: YOLO):
+        overrides = copy.deepcopy(model.overrides)
+        overrides["conf"] = 0.1
+        self.predictor = self.task_map[self.task]["predictor"](overrides=overrides)
+        self.predictor.callbacks = {}
+        self.predictor.args.save = False
+        self.predictor.args.save_txt = False
+        self.predictor.args.save_crop = False
+        self.predictor.args.verbose = None
+
+    def _save_model(self, trainer: TRAINER_TYPE):
+        model_checkpoint_artifact = wandb.Artifact(f"run_{wandb.run.id}_model", "model")
+        checkpoint_dict = {
+            "epoch": trainer.epoch,
+            "best_fitness": trainer.best_fitness,
+            "model": copy.deepcopy(de_parallel(self.model)).half(),
+            "ema": copy.deepcopy(trainer.ema.ema).half(),
+            "updates": trainer.ema.updates,
+            "optimizer": trainer.optimizer.state_dict(),
+            "train_args": vars(trainer.args),
+            "date": datetime.now().isoformat(),
+            "version": __version__,
+        }
+        checkpoint_path = trainer.wdir / f"epoch{trainer.epoch}.pt"
+        torch.save(checkpoint_dict, checkpoint_path, pickle_module=pickle)
+        model_checkpoint_artifact.add_file(checkpoint_path)
+        wandb.log_artifact(
+            model_checkpoint_artifact, aliases=[f"epoch_{trainer.epoch}"]
+        )
+
+    def on_train_start(self, trainer: TRAINER_TYPE):
+        with telemetry.context(run=wandb.run) as tel:
+            tel.feature.ultralytics_yolov8 = True
+        wandb.config.train = vars(trainer.args)
+        self.run_id = wandb.run.id
+
+    @torch.no_grad()
+    def on_fit_epoch_end(self, trainer: DetectionTrainer):
+        if self.task in self.supported_tasks and self.train_epoch != trainer.epoch:
+            self.train_epoch = trainer.epoch
+            if (self.train_epoch + 1) % self.epoch_logging_interval == 0:
+                validator = trainer.validator
+                dataloader = validator.dataloader
+                class_label_map = validator.names
+                self.device = next(trainer.model.parameters()).device
+                if isinstance(trainer.model, torch.nn.parallel.DistributedDataParallel):
+                    model = trainer.model.module
+                else:
+                    model = trainer.model
+                self.model = copy.deepcopy(model).eval().to(self.device)
+                self.predictor.setup_model(model=self.model, verbose=False)
+                if self.task == "pose":
+                    self.train_validation_table = plot_pose_validation_results(
+                        dataloader=dataloader,
+                        class_label_map=class_label_map,
+                        model_name=self.model_name,
+                        predictor=self.predictor,
+                        visualize_skeleton=self.visualize_skeleton,
+                        table=self.train_validation_table,
+                        max_validation_batches=self.max_validation_batches,
+                        epoch=trainer.epoch,
+                    )
+                elif self.task == "segment":
+                    self.train_validation_table = plot_segmentation_validation_results(
+                        dataloader=dataloader,
+                        class_label_map=class_label_map,
+                        model_name=self.model_name,
+                        predictor=self.predictor,
+                        table=self.train_validation_table,
+                        max_validation_batches=self.max_validation_batches,
+                        epoch=trainer.epoch,
+                    )
+                elif self.task == "detect":
+                    self.train_validation_table = plot_detection_validation_results(
+                        dataloader=dataloader,
+                        class_label_map=class_label_map,
+                        model_name=self.model_name,
+                        predictor=self.predictor,
+                        table=self.train_validation_table,
+                        max_validation_batches=self.max_validation_batches,
+                        epoch=trainer.epoch,
+                    )
+                elif self.task == "classify":
+                    self.train_validation_table = (
+                        plot_classification_validation_results(
+                            dataloader=dataloader,
+                            model_name=self.model_name,
+                            predictor=self.predictor,
+                            table=self.train_validation_table,
+                            max_validation_batches=self.max_validation_batches,
+                            epoch=trainer.epoch,
+                        )
+                    )
+            if self.enable_model_checkpointing:
+                self._save_model(trainer)
+            trainer.model.to(self.device)
+
+    def on_train_end(self, trainer: TRAINER_TYPE):
+        if self.task in self.supported_tasks:
+            wandb.log({"Train-Table": self.train_validation_table}, commit=False)
+
+    def on_val_start(self, validator: VALIDATOR_TYPE):
+        wandb.run or wandb.init(
+            project=validator.args.project or "YOLOv8",
+            job_type="validation_" + validator.args.task,
+        )
+
+    @torch.no_grad()
+    def on_val_end(self, trainer: VALIDATOR_TYPE):
+        if self.task in self.supported_tasks:
+            validator = trainer
+            dataloader = validator.dataloader
+            class_label_map = validator.names
+            if self.task == "pose":
+                self.validation_table = plot_pose_validation_results(
+                    dataloader=dataloader,
+                    class_label_map=class_label_map,
+                    model_name=self.model_name,
+                    predictor=self.predictor,
+                    visualize_skeleton=self.visualize_skeleton,
+                    table=self.validation_table,
+                    max_validation_batches=self.max_validation_batches,
+                )
+            elif self.task == "segment":
+                self.validation_table = plot_segmentation_validation_results(
+                    dataloader=dataloader,
+                    class_label_map=class_label_map,
+                    model_name=self.model_name,
+                    predictor=self.predictor,
+                    table=self.validation_table,
+                    max_validation_batches=self.max_validation_batches,
+                )
+            elif self.task == "detect":
+                self.validation_table = plot_detection_validation_results(
+                    dataloader=dataloader,
+                    class_label_map=class_label_map,
+                    model_name=self.model_name,
+                    predictor=self.predictor,
+                    table=self.validation_table,
+                    max_validation_batches=self.max_validation_batches,
+                )
+            elif self.task == "classify":
+                self.validation_table = plot_classification_validation_results(
+                    dataloader=dataloader,
+                    model_name=self.model_name,
+                    predictor=self.predictor,
+                    table=self.validation_table,
+                    max_validation_batches=self.max_validation_batches,
+                )
+            wandb.log({"Validation-Table": self.validation_table}, commit=False)
+
+    def on_predict_start(self, predictor: PREDICTOR_TYPE):
+        wandb.run or wandb.init(
+            project=predictor.args.project or "YOLOv8",
+            config=vars(predictor.args),
+            job_type="prediction_" + predictor.args.task,
+        )
+        if isinstance(predictor, SAMPredictor):
+            self.prompts = copy.deepcopy(predictor.prompts)
+            self.prediction_table = wandb.Table(columns=["Image"])
+
+    def on_predict_end(self, predictor: PREDICTOR_TYPE):
+        wandb.config.prediction_configs = vars(predictor.args)
+        if self.task in self.supported_tasks:
+            for result in tqdm(predictor.results):
+                if self.task == "pose":
+                    self.prediction_table = plot_pose_predictions(
+                        result,
+                        self.model_name,
+                        self.visualize_skeleton,
+                        self.prediction_table,
+                    )
+                elif self.task == "segment":
+                    if isinstance(predictor, SegmentationPredictor):
+                        self.prediction_table = plot_mask_predictions(
+                            result, self.model_name, self.prediction_table
+                        )
+                    elif isinstance(predictor, SAMPredictor):
+                        self.prediction_table = plot_sam_predictions(
+                            result, self.prompts, self.prediction_table
+                        )
+                elif self.task == "detect":
+                    self.prediction_table = plot_bbox_predictions(
+                        result, self.model_name, self.prediction_table
+                    )
+                elif self.task == "classify":
+                    self.prediction_table = plot_classification_predictions(
+                        result, self.model_name, self.prediction_table
+                    )
+
+            wandb.log({"Prediction-Table": self.prediction_table}, commit=False)
+
+    @property
+    def callbacks(self) -> Dict[str, Callable]:
+        """Property contains all the relevant callbacks to add to the YOLO model for the Weights & Biases logging."""
+        return {
+            "on_train_start": self.on_train_start,
+            "on_fit_epoch_end": self.on_fit_epoch_end,
+            "on_train_end": self.on_train_end,
+            "on_val_start": self.on_val_start,
+            "on_val_end": self.on_val_end,
+            "on_predict_start": self.on_predict_start,
+            "on_predict_end": self.on_predict_end,
+        }
+
+
+# TODO: Add epoch interval
+def add_wandb_callback(
+    model: YOLO,
+    epoch_logging_interval: int = 1,
+    enable_model_checkpointing: bool = False,
+    enable_train_validation_logging: bool = True,
+    enable_validation_logging: bool = True,
+    enable_prediction_logging: bool = True,
+    max_validation_batches: Optional[int] = 1,
+    visualize_skeleton: Optional[bool] = True,
+):
+    """Function to add the `WandBUltralyticsCallback` callback to the `YOLO` model.
+
+    Example:
+        ```python
+        from ultralytics.yolo.engine.model import YOLO
+        from wandb.yolov8 import add_wandb_callback
+
+        # initialize YOLO model
+        model = YOLO("yolov8n.pt")
+
+        # add wandb callback
+        add_wandb_callback(model, max_validation_batches=2, enable_model_checkpointing=True)
+
+        # train
+        model.train(data="coco128.yaml", epochs=5, imgsz=640)
+
+        # validate
+        model.val()
+
+        # perform inference
+        model(["img1.jpeg", "img2.jpeg"])
+        ```
+
+    Args:
+        model: (ultralytics.yolo.engine.model.YOLO) YOLO Model of type
+            `ultralytics.yolo.engine.model.YOLO`.
+        epoch_logging_interval: (int) interval to log the prediction visualizations
+            during training.
+        enable_model_checkpointing: (bool) enable logging model checkpoints as
+            artifacts at the end of eveny epoch if set to `True`.
+        enable_train_validation_logging: (bool) enable logging the predictions and
+            ground-truths as interactive image overlays on the images from
+            the validation dataloader to a `wandb.Table` along with
+            mean-confidence of the predictions per-class at the end of each
+            training epoch.
+        enable_validation_logging: (bool) enable logging the predictions and
+            ground-truths as interactive image overlays on the images from the
+            validation dataloader to a `wandb.Table` along with
+            mean-confidence of the predictions per-class at the end of
+            validation.
+        enable_prediction_logging: (bool) enable logging the predictions and
+            ground-truths as interactive image overlays on the images from the
+            validation dataloader to a `wandb.Table` along with mean-confidence
+            of the predictions per-class at the end of each prediction.
+        max_validation_batches: (Optional[int]) maximum number of validation batches to log to
+            a table per epoch.
+        visualize_skeleton: (Optional[bool]) visualize pose skeleton by drawing lines connecting
+            keypoints for human pose.
+
+    Returns:
+        An instance of `ultralytics.yolo.engine.model.YOLO` with the `WandBUltralyticsCallback`.
+    """
+    if RANK in [-1, 0]:
+        wandb_callback = WandBUltralyticsCallback(
+            copy.deepcopy(model),
+            epoch_logging_interval,
+            max_validation_batches,
+            enable_model_checkpointing,
+            visualize_skeleton,
+        )
+        callbacks = wandb_callback.callbacks
+        if not enable_train_validation_logging:
+            _ = callbacks.pop("on_fit_epoch_end")
+            _ = callbacks.pop("on_train_end")
+        if not enable_validation_logging:
+            _ = callbacks.pop("on_val_start")
+            _ = callbacks.pop("on_val_end")
+        if not enable_prediction_logging:
+            _ = callbacks.pop("on_predict_start")
+            _ = callbacks.pop("on_predict_end")
+        for event, callback_fn in callbacks.items():
+            model.add_callback(event, callback_fn)
+    else:
+        wandb.termerror(
+            "The RANK of the process to add the callbacks was neither 0 or "
+            "-1. No Weights & Biases callbacks were added to this instance "
+            "of the YOLO model."
+        )
+    return model

wandb/integration/ultralytics/classification_utils.py

@@ -0,0 +1,83 @@
+from typing import Any, Optional
+
+import numpy as np
+from tqdm.auto import tqdm
+from ultralytics.engine.results import Results
+from ultralytics.models.yolo.classify import ClassificationPredictor
+
+import wandb
+
+
+def plot_classification_predictions(
+    result: Results,
+    model_name: str,
+    table: Optional[wandb.Table] = None,
+    original_image: Optional[np.array] = None,
+):
+    """Plot classification prediction results to a `wandb.Table` if the table is passed otherwise return the data."""
+    result = result.to("cpu")
+    probabilities = result.probs
+    probabilities_list = probabilities.data.numpy().tolist()
+    class_id_to_label = {int(k): str(v) for k, v in result.names.items()}
+    original_image = (
+        wandb.Image(original_image)
+        if original_image is not None
+        else wandb.Image(result.orig_img)
+    )
+    table_row = [
+        model_name,
+        original_image,
+        class_id_to_label[int(probabilities.top1)],
+        probabilities.top1conf,
+        [class_id_to_label[int(class_idx)] for class_idx in list(probabilities.top5)],
+        [probabilities_list[int(class_idx)] for class_idx in list(probabilities.top5)],
+        {
+            class_id_to_label[int(class_idx)]: probability
+            for class_idx, probability in enumerate(probabilities_list)
+        },
+        result.speed,
+    ]
+    if table is not None:
+        table.add_data(*table_row)
+        return table
+    return class_id_to_label, table_row
+
+
+def plot_classification_validation_results(
+    dataloader: Any,
+    model_name: str,
+    predictor: ClassificationPredictor,
+    table: wandb.Table,
+    max_validation_batches: int,
+    epoch: Optional[int] = None,
+) -> wandb.Table:
+    """Plot classification results to a `wandb.Table`."""
+    data_idx = 0
+    num_dataloader_batches = len(dataloader.dataset) // dataloader.batch_size
+    max_validation_batches = min(max_validation_batches, num_dataloader_batches)
+    for batch_idx, batch in enumerate(dataloader):
+        image_batch = [
+            image for image in np.transpose(batch["img"].numpy(), (0, 2, 3, 1))
+        ]
+        ground_truth = batch["cls"].numpy().tolist()
+        progress_bar_result_iterable = tqdm(
+            range(max_validation_batches),
+            desc=f"Generating Visualizations for batch-{batch_idx + 1}/{max_validation_batches}",
+        )
+        for img_idx in progress_bar_result_iterable:
+            try:
+                prediction_result = predictor(image_batch[img_idx])[0]
+                class_id_to_label, table_row = plot_classification_predictions(
+                    prediction_result, model_name, original_image=image_batch[img_idx]
+                )
+                table_row = [data_idx, batch_idx] + table_row[1:]
+                table_row.insert(3, class_id_to_label[ground_truth[img_idx]])
+                table_row = [epoch] + table_row if epoch is not None else table_row
+                table_row = [model_name] + table_row
+                table.add_data(*table_row)
+                data_idx += 1
+            except Exception:
+                pass
+        if batch_idx + 1 == max_validation_batches:
+            break
+    return table