wandb 0.19.1rc1__py3-none-musllinux_1_2_aarch64.whl

wandb/integration/ultralytics/mask_utils.py

@@ -0,0 +1,202 @@
+from typing import Dict, Optional, Tuple
+
+import cv2
+import numpy as np
+from tqdm.auto import tqdm
+from ultralytics.engine.results import Results
+from ultralytics.models.yolo.segment import SegmentationPredictor
+from ultralytics.utils.ops import scale_image
+
+import wandb
+from wandb.integration.ultralytics.bbox_utils import (
+    get_ground_truth_bbox_annotations,
+    get_mean_confidence_map,
+)
+
+
+def instance_mask_to_semantic_mask(instance_mask, class_indices):
+    height, width, num_instances = instance_mask.shape
+    semantic_mask = np.zeros((height, width), dtype=np.uint8)
+    for i in range(num_instances):
+        instance_map = instance_mask[:, :, i]
+        class_index = class_indices[i]
+        semantic_mask[instance_map == 1] = class_index
+    return semantic_mask
+
+
+def get_boxes_and_masks(result: Results) -> Tuple[Dict, Dict, Dict]:
+    boxes = result.boxes.xywh.long().numpy()
+    classes = result.boxes.cls.long().numpy()
+    confidence = result.boxes.conf.numpy()
+    class_id_to_label = {int(k): str(v) for k, v in result.names.items()}
+    class_id_to_label.update({len(result.names.items()): "background"})
+    mean_confidence_map = get_mean_confidence_map(
+        classes, confidence, class_id_to_label
+    )
+    masks = None
+    if result.masks is not None:
+        scaled_instance_mask = scale_image(
+            np.transpose(result.masks.data.numpy(), (1, 2, 0)),
+            result.orig_img[:, :, ::-1].shape,
+        )
+        scaled_semantic_mask = instance_mask_to_semantic_mask(
+            scaled_instance_mask, classes.tolist()
+        )
+        scaled_semantic_mask[scaled_semantic_mask == 0] = len(result.names.items())
+        masks = {
+            "predictions": {
+                "mask_data": scaled_semantic_mask,
+                "class_labels": class_id_to_label,
+            }
+        }
+    box_data, total_confidence = [], 0.0
+    for idx in range(len(boxes)):
+        box_data.append(
+            {
+                "position": {
+                    "middle": [int(boxes[idx][0]), int(boxes[idx][1])],
+                    "width": int(boxes[idx][2]),
+                    "height": int(boxes[idx][3]),
+                },
+                "domain": "pixel",
+                "class_id": int(classes[idx]),
+                "box_caption": class_id_to_label[int(classes[idx])],
+                "scores": {"confidence": float(confidence[idx])},
+            }
+        )
+        total_confidence += float(confidence[idx])
+
+    boxes = {
+        "predictions": {
+            "box_data": box_data,
+            "class_labels": class_id_to_label,
+        },
+    }
+    return boxes, masks, mean_confidence_map
+
+
+def plot_mask_predictions(
+    result: Results, model_name: str, table: Optional[wandb.Table] = None
+) -> Tuple[wandb.Image, Dict, Dict, Dict]:
+    result = result.to("cpu")
+    boxes, masks, mean_confidence_map = get_boxes_and_masks(result)
+    image = wandb.Image(result.orig_img[:, :, ::-1], boxes=boxes, masks=masks)
+    if table is not None:
+        table.add_data(
+            model_name,
+            image,
+            len(boxes["predictions"]["box_data"]),
+            mean_confidence_map,
+            result.speed,
+        )
+        return table
+    return image, masks, boxes["predictions"], mean_confidence_map
+
+
+def structure_prompts_and_image(image: np.array, prompt: Dict) -> Dict:
+    wb_box_data = []
+    if prompt["bboxes"] is not None:
+        wb_box_data.append(
+            {
+                "position": {
+                    "middle": [prompt["bboxes"][0], prompt["bboxes"][1]],
+                    "width": prompt["bboxes"][2],
+                    "height": prompt["bboxes"][3],
+                },
+                "domain": "pixel",
+                "class_id": 1,
+                "box_caption": "Prompt-Box",
+            }
+        )
+    if prompt["points"] is not None:
+        image = image.copy().astype(np.uint8)
+        image = cv2.circle(
+            image, tuple(prompt["points"]), 5, (0, 255, 0), -1, lineType=cv2.LINE_AA
+        )
+    wb_box_data = {
+        "prompts": {
+            "box_data": wb_box_data,
+            "class_labels": {1: "Prompt-Box"},
+        }
+    }
+    return image, wb_box_data
+
+
+def plot_sam_predictions(
+    result: Results, prompt: Dict, table: wandb.Table
+) -> wandb.Table:
+    result = result.to("cpu")
+    image = result.orig_img[:, :, ::-1]
+    image, wb_box_data = structure_prompts_and_image(image, prompt)
+    image = wandb.Image(
+        image,
+        boxes=wb_box_data,
+        masks={
+            "predictions": {
+                "mask_data": np.squeeze(result.masks.data.cpu().numpy().astype(int)),
+                "class_labels": {0: "Background", 1: "Prediction"},
+            }
+        },
+    )
+    table.add_data(image)
+    return table
+
+
+def plot_segmentation_validation_results(
+    dataloader,
+    class_label_map,
+    model_name: str,
+    predictor: SegmentationPredictor,
+    table: wandb.Table,
+    max_validation_batches: int,
+    epoch: Optional[int] = None,
+):
+    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):
+        prediction_results = predictor(batch["im_file"])
+        progress_bar_result_iterable = tqdm(
+            enumerate(prediction_results),
+            total=len(prediction_results),
+            desc=f"Generating Visualizations for batch-{batch_idx + 1}/{max_validation_batches}",
+        )
+        for img_idx, prediction_result in progress_bar_result_iterable:
+            prediction_result = prediction_result.to("cpu")
+            (
+                _,
+                prediction_mask_data,
+                prediction_box_data,
+                mean_confidence_map,
+            ) = plot_mask_predictions(prediction_result, model_name)
+            try:
+                ground_truth_data = get_ground_truth_bbox_annotations(
+                    img_idx, batch["im_file"][img_idx], batch, class_label_map
+                )
+                wandb_image = wandb.Image(
+                    batch["im_file"][img_idx],
+                    boxes={
+                        "ground-truth": {
+                            "box_data": ground_truth_data,
+                            "class_labels": class_label_map,
+                        },
+                        "predictions": prediction_box_data,
+                    },
+                    masks=prediction_mask_data,
+                )
+                table_rows = [
+                    data_idx,
+                    batch_idx,
+                    wandb_image,
+                    mean_confidence_map,
+                    prediction_result.speed,
+                ]
+                table_rows = [epoch] + table_rows if epoch is not None else table_rows
+                table_rows = [model_name] + table_rows
+                table.add_data(*table_rows)
+                data_idx += 1
+            except TypeError:
+                pass
+        if batch_idx + 1 == max_validation_batches:
+            break
+    return table

wandb/integration/ultralytics/pose_utils.py

@@ -0,0 +1,103 @@
+from typing import Any, Optional
+
+import numpy as np
+from PIL import Image
+from tqdm.auto import tqdm
+from ultralytics.engine.results import Results
+from ultralytics.models.yolo.pose import PosePredictor
+from ultralytics.utils.plotting import Annotator
+
+import wandb
+from wandb.integration.ultralytics.bbox_utils import (
+    get_boxes,
+    get_ground_truth_bbox_annotations,
+)
+
+
+def annotate_keypoint_results(result: Results, visualize_skeleton: bool):
+    annotator = Annotator(np.ascontiguousarray(result.orig_img[:, :, ::-1]))
+    key_points = result.keypoints.data.numpy()
+    for idx in range(key_points.shape[0]):
+        annotator.kpts(key_points[idx], kpt_line=visualize_skeleton)
+    return annotator.im
+
+
+def annotate_keypoint_batch(image_path: str, keypoints: Any, visualize_skeleton: bool):
+    with Image.open(image_path) as original_image:
+        original_image = np.ascontiguousarray(original_image)
+        annotator = Annotator(original_image)
+        annotator.kpts(keypoints.numpy(), kpt_line=visualize_skeleton)
+        return annotator.im
+
+
+def plot_pose_predictions(
+    result: Results,
+    model_name: str,
+    visualize_skeleton: bool,
+    table: Optional[wandb.Table] = None,
+):
+    result = result.to("cpu")
+    boxes, mean_confidence_map = get_boxes(result)
+    annotated_image = annotate_keypoint_results(result, visualize_skeleton)
+    prediction_image = wandb.Image(annotated_image, boxes=boxes)
+    table_row = [
+        model_name,
+        prediction_image,
+        len(boxes["predictions"]["box_data"]),
+        mean_confidence_map,
+        result.speed,
+    ]
+    if table is not None:
+        table.add_data(*table_row)
+        return table
+    return table_row
+
+
+def plot_pose_validation_results(
+    dataloader,
+    class_label_map,
+    model_name: str,
+    predictor: PosePredictor,
+    visualize_skeleton: bool,
+    table: wandb.Table,
+    max_validation_batches: int,
+    epoch: Optional[int] = None,
+) -> 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):
+        prediction_results = predictor(batch["im_file"])
+        progress_bar_result_iterable = tqdm(
+            enumerate(prediction_results),
+            total=len(prediction_results),
+            desc=f"Generating Visualizations for batch-{batch_idx + 1}/{max_validation_batches}",
+        )
+        for img_idx, prediction_result in progress_bar_result_iterable:
+            prediction_result = prediction_result.to("cpu")
+            table_row = plot_pose_predictions(
+                prediction_result, model_name, visualize_skeleton
+            )
+            ground_truth_image = wandb.Image(
+                annotate_keypoint_batch(
+                    batch["im_file"][img_idx],
+                    batch["keypoints"][img_idx],
+                    visualize_skeleton,
+                ),
+                boxes={
+                    "ground-truth": {
+                        "box_data": get_ground_truth_bbox_annotations(
+                            img_idx, batch["im_file"][img_idx], batch, class_label_map
+                        ),
+                        "class_labels": class_label_map,
+                    },
+                },
+            )
+            table_row = [data_idx, batch_idx, ground_truth_image] + table_row[1:]
+            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
+        if batch_idx + 1 == max_validation_batches:
+            break
+    return table

wandb/integration/xgboost/__init__.py

@@ -0,0 +1,11 @@
+"""W&B callback for xgboost.
+
+Simple callback to get logging for each tree
+
+Use the `wandb_callback` to add `wandb` logging to any `XGboost` model. However, it will
+be deprecated in favor of WandbCallback. Use it instead for more features.
+"""
+
+from .xgboost import WandbCallback, wandb_callback
+
+__all__ = ["wandb_callback", "WandbCallback"]

wandb/integration/xgboost/xgboost.py

@@ -0,0 +1,189 @@
+"""xgboost init!"""
+
+import json
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+import xgboost as xgb  # type: ignore
+from xgboost import Booster
+
+import wandb
+from wandb.sdk.lib import telemetry as wb_telemetry
+
+MINIMIZE_METRICS = [
+    "rmse",
+    "rmsle",
+    "mae",
+    "mape",
+    "mphe",
+    "logloss",
+    "error",
+    "error@t",
+    "merror",
+]
+
+MAXIMIZE_METRICS = ["auc", "aucpr", "ndcg", "map", "ndcg@n", "map@n"]
+
+
+if TYPE_CHECKING:
+    from typing import Callable, List, NamedTuple
+
+    class CallbackEnv(NamedTuple):
+        evaluation_result_list: List
+
+
+def wandb_callback() -> "Callable":
+    """Old style callback that will be deprecated in favor of WandbCallback. Please try the new logger for more features."""
+    warnings.warn(
+        "wandb_callback will be deprecated in favor of WandbCallback. Please use WandbCallback for more features.",
+        UserWarning,
+        stacklevel=2,
+    )
+
+    with wb_telemetry.context() as tel:
+        tel.feature.xgboost_old_wandb_callback = True
+
+    def callback(env: "CallbackEnv") -> None:
+        for k, v in env.evaluation_result_list:
+            wandb.log({k: v}, commit=False)
+        wandb.log({})
+
+    return callback
+
+
+class WandbCallback(xgb.callback.TrainingCallback):
+    """`WandbCallback` automatically integrates XGBoost with wandb.
+
+    Args:
+        log_model: (boolean) if True save and upload the model to Weights & Biases Artifacts
+        log_feature_importance: (boolean) if True log a feature importance bar plot
+        importance_type: (str) one of {weight, gain, cover, total_gain, total_cover} for tree model. weight for linear model.
+        define_metric: (boolean) if True (default) capture model performance at the best step, instead of the last step, of training in your `wandb.summary`.
+
+    Passing `WandbCallback` to XGBoost will:
+
+    - log the booster model configuration to Weights & Biases
+    - log evaluation metrics collected by XGBoost, such as rmse, accuracy etc. to Weights & Biases
+    - log training metric collected by XGBoost (if you provide training data to eval_set)
+    - log the best score and the best iteration
+    - save and upload your trained model to Weights & Biases Artifacts (when `log_model = True`)
+    - log feature importance plot when `log_feature_importance=True` (default).
+    - Capture the best eval metric in `wandb.summary` when `define_metric=True` (default).
+
+    Example:
+        ```python
+        bst_params = dict(
+            objective="reg:squarederror",
+            colsample_bytree=0.3,
+            learning_rate=0.1,
+            max_depth=5,
+            alpha=10,
+            n_estimators=10,
+            tree_method="hist",
+            callbacks=[WandbCallback()],
+        )
+
+        xg_reg = xgb.XGBRegressor(**bst_params)
+        xg_reg.fit(
+            X_train,
+            y_train,
+            eval_set=[(X_test, y_test)],
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        log_model: bool = False,
+        log_feature_importance: bool = True,
+        importance_type: str = "gain",
+        define_metric: bool = True,
+    ):
+        self.log_model: bool = log_model
+        self.log_feature_importance: bool = log_feature_importance
+        self.importance_type: str = importance_type
+        self.define_metric: bool = define_metric
+
+        if wandb.run is None:
+            raise wandb.Error("You must call wandb.init() before WandbCallback()")
+
+        with wb_telemetry.context() as tel:
+            tel.feature.xgboost_wandb_callback = True
+
+    def before_training(self, model: Booster) -> Booster:
+        """Run before training is finished."""
+        # Update W&B config
+        config = model.save_config()
+        wandb.config.update(json.loads(config))
+
+        return model
+
+    def after_training(self, model: Booster) -> Booster:
+        """Run after training is finished."""
+        # Log the booster model as artifacts
+        if self.log_model:
+            self._log_model_as_artifact(model)
+
+        # Plot feature importance
+        if self.log_feature_importance:
+            self._log_feature_importance(model)
+
+        # Log the best score and best iteration
+        if model.attr("best_score") is not None:
+            wandb.log(
+                {
+                    "best_score": float(cast(str, model.attr("best_score"))),
+                    "best_iteration": int(cast(str, model.attr("best_iteration"))),
+                }
+            )
+
+        return model
+
+    def after_iteration(self, model: Booster, epoch: int, evals_log: dict) -> bool:
+        """Run after each iteration. Return True when training should stop."""
+        # Log metrics
+        for data, metric in evals_log.items():
+            for metric_name, log in metric.items():
+                if self.define_metric:
+                    self._define_metric(data, metric_name)
+                    wandb.log({f"{data}-{metric_name}": log[-1]}, commit=False)
+                else:
+                    wandb.log({f"{data}-{metric_name}": log[-1]}, commit=False)
+
+        wandb.log({"epoch": epoch})
+
+        self.define_metric = False
+
+        return False
+
+    def _log_model_as_artifact(self, model: Booster) -> None:
+        model_name = f"{wandb.run.id}_model.json"  # type: ignore
+        model_path = Path(wandb.run.dir) / model_name  # type: ignore
+        model.save_model(str(model_path))
+
+        model_artifact = wandb.Artifact(name=model_name, type="model")
+        model_artifact.add_file(str(model_path))
+        wandb.log_artifact(model_artifact)
+
+    def _log_feature_importance(self, model: Booster) -> None:
+        fi = model.get_score(importance_type=self.importance_type)
+        fi_data = [[k, fi[k]] for k in fi]
+        table = wandb.Table(data=fi_data, columns=["Feature", "Importance"])
+        wandb.log(
+            {
+                "Feature Importance": wandb.plot.bar(
+                    table, "Feature", "Importance", title="Feature Importance"
+                )
+            }
+        )
+
+    def _define_metric(self, data: str, metric_name: str) -> None:
+        if "loss" in str.lower(metric_name):
+            wandb.define_metric(f"{data}-{metric_name}", summary="min")
+        elif str.lower(metric_name) in MINIMIZE_METRICS:
+            wandb.define_metric(f"{data}-{metric_name}", summary="min")
+        elif str.lower(metric_name) in MAXIMIZE_METRICS:
+            wandb.define_metric(f"{data}-{metric_name}", summary="max")
+        else:
+            pass