wandb 0.15.9__py3-none-any.whl → 0.15.11__py3-none-any.whl

wandb/env.py

@@ -14,6 +14,7 @@ import json
 import os
 import sys
 from distutils.util import strtobool
+from pathlib import Path
 from typing import List, MutableMapping, Optional, Union
 
 import appdirs
@@ -50,6 +51,7 @@ RUN_GROUP = "WANDB_RUN_GROUP"
 RUN_DIR = "WANDB_RUN_DIR"
 SWEEP_ID = "WANDB_SWEEP_ID"
 HTTP_TIMEOUT = "WANDB_HTTP_TIMEOUT"
+FILE_PUSHER_TIMEOUT = "WANDB_FILE_PUSHER_TIMEOUT"
 API_KEY = "WANDB_API_KEY"
 JOB_TYPE = "WANDB_JOB_TYPE"
 DISABLE_CODE = "WANDB_DISABLE_CODE"
@@ -72,6 +74,7 @@ JUPYTER = "WANDB_JUPYTER"
 CONFIG_DIR = "WANDB_CONFIG_DIR"
 DATA_DIR = "WANDB_DATA_DIR"
 ARTIFACT_DIR = "WANDB_ARTIFACT_DIR"
+ARTIFACT_FETCH_FILE_URL_BATCH_SIZE = "WANDB_ARTIFACT_FETCH_FILE_URL_BATCH_SIZE"
 CACHE_DIR = "WANDB_CACHE_DIR"
 DISABLE_SSL = "WANDB_INSECURE_DISABLE_SSL"
 SERVICE = "WANDB_SERVICE"
@@ -120,6 +123,7 @@ def immutable_keys() -> List[str]:
         HOST,
         DATA_DIR,
         ARTIFACT_DIR,
+        ARTIFACT_FETCH_FILE_URL_BATCH_SIZE,
         CACHE_DIR,
         USE_V1_ARTIFACTS,
         DISABLE_SSL,
@@ -191,13 +195,24 @@ def get_docker(
     return env.get(DOCKER, default)
 
 
-def get_http_timeout(default: int = 10, env: Optional[Env] = None) -> int:
+def get_http_timeout(default: int = 20, env: Optional[Env] = None) -> int:
     if env is None:
         env = os.environ
 
     return int(env.get(HTTP_TIMEOUT, default))
 
 
+def get_file_pusher_timeout(
+    default: Optional[int] = None,
+    env: Optional[Env] = None,
+) -> Optional[int]:
+    if env is None:
+        env = os.environ
+
+    timeout = env.get(FILE_PUSHER_TIMEOUT, default)
+    return int(timeout) if timeout else None
+
+
 def get_ignore(
     default: Optional[List[str]] = None, env: Optional[Env] = None
 ) -> Optional[List[str]]:
@@ -367,14 +382,19 @@ def get_artifact_dir(env: Optional[Env] = None) -> str:
     return val
 
 
-def get_cache_dir(env: Optional[Env] = None) -> str:
-    default_dir = appdirs.user_cache_dir("wandb")
+def get_artifact_fetch_file_url_batch_size(env: Optional[Env] = None) -> int:
+    default_batch_size = 5000
     if env is None:
         env = os.environ
-    val = env.get(CACHE_DIR, default_dir)
+    val = int(env.get(ARTIFACT_FETCH_FILE_URL_BATCH_SIZE, default_batch_size))
     return val
 
 
+def get_cache_dir(env: Optional[Env] = None) -> Path:
+    env = env or os.environ
+    return Path(env.get(CACHE_DIR, appdirs.user_cache_dir("wandb")))
+
+
 def get_use_v1_artifacts(env: Optional[Env] = None) -> bool:
     if env is None:
         env = os.environ

wandb/filesync/step_checksum.py

@@ -16,8 +16,8 @@ if TYPE_CHECKING:
     import tempfile
 
     from wandb.filesync import stats
-    from wandb.sdk.artifacts import artifact_saver
     from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
+    from wandb.sdk.artifacts.artifact_saver import SaveFn, SaveFnAsync
     from wandb.sdk.internal import internal_api
 
 
@@ -30,8 +30,8 @@ class RequestUpload(NamedTuple):
 class RequestStoreManifestFiles(NamedTuple):
     manifest: "ArtifactManifest"
     artifact_id: str
-    save_fn: "artifact_saver.SaveFn"
-    save_fn_async: "artifact_saver.SaveFnAsync"
+    save_fn: "SaveFn"
+    save_fn_async: "SaveFnAsync"
 
 
 class RequestCommitArtifact(NamedTuple):

wandb/integration/openai/openai.py

@@ -13,6 +13,9 @@ autolog = AutologAPI(
         "Edit.create",
         "Completion.create",
         "ChatCompletion.create",
+        "Edit.acreate",
+        "Completion.acreate",
+        "ChatCompletion.acreate",
     ),
     resolver=OpenAIRequestResponseResolver(),
     telemetry_feature="openai_autolog",

wandb/integration/ultralytics/__init__.py

@@ -0,0 +1,9 @@
+"""Tools for integrating with [`ultralytics`](https://docs.ultralytics.com/).
+
+Ultralytics is a computer vision framework for training and deploying YOLOv8
+models.
+"""
+
+from wandb.integration.ultralytics.callback import add_wandb_callback
+
+__all__ = ("add_wandb_callback",)

wandb/integration/ultralytics/bbox_utils.py

@@ -0,0 +1,196 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from ultralytics.engine.results import Results
+from ultralytics.models.yolo.detect import DetectionPredictor
+from ultralytics.yolo.utils import ops
+
+import wandb
+
+
+def scale_bounding_box_to_original_image_shape(
+    box: torch.Tensor,
+    resized_image_shape: Tuple,
+    original_image_shape: Tuple,
+    ratio_pad: bool,
+) -> List[int]:
+    """YOLOv8 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.
+
+    Reference: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/yolo/utils/callbacks/comet.py#L105
+    """
+    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)
+    return box.tolist()
+
+
+def get_ground_truth_bbox_annotations(
+    img_idx: int, image_path: str, batch: Dict, class_name_map: Dict = None
+) -> List[Dict[str, Any]]:
+    """Get ground truth bounding box annotation data in the form required for `wandb.Image` overlay system."""
+    indices = batch["batch_idx"] == img_idx
+    bboxes = batch["bboxes"][indices]
+    cls_labels = batch["cls"][indices].squeeze(1).tolist()
+
+    class_name_map_reverse = {v: k for k, v in class_name_map.items()}
+
+    if len(bboxes) == 0:
+        wandb.termwarn(
+            f"Image: {image_path} has no bounding boxes labels", repeat=False
+        )
+        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(
+            {
+                "position": {
+                    "middle": [int(box[0]), int(box[1])],
+                    "width": int(box[2]),
+                    "height": int(box[3]),
+                },
+                "domain": "pixel",
+                "class_id": class_name_map_reverse[label],
+                "box_caption": label,
+            }
+        )
+
+    return data
+
+
+def get_mean_confidence_map(
+    classes: List, confidence: List, class_id_to_label: Dict
+) -> Dict[str, float]:
+    """Get Mean-confidence map from the predictions to be logged into a `wandb.Table`."""
+    confidence_map = {v: [] for _, v in class_id_to_label.items()}
+    for class_idx, confidence_value in zip(classes, confidence):
+        confidence_map[class_id_to_label[class_idx]].append(confidence_value)
+    updated_confidence_map = {}
+    for label, confidence_list in confidence_map.items():
+        if len(confidence_list) > 0:
+            updated_confidence_map[label] = sum(confidence_list) / len(confidence_list)
+        else:
+            updated_confidence_map[label] = 0
+    return updated_confidence_map
+
+
+def get_boxes(result: Results) -> Tuple[Dict, Dict]:
+    """Convert an ultralytics prediction result into metadata for the `wandb.Image` overlay system."""
+    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()}
+    mean_confidence_map = get_mean_confidence_map(
+        classes, confidence, class_id_to_label
+    )
+    box_data = []
+    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])},
+            }
+        )
+    boxes = {
+        "predictions": {
+            "box_data": box_data,
+            "class_labels": class_id_to_label,
+        },
+    }
+    return boxes, mean_confidence_map
+
+
+def plot_predictions(
+    result: Results, model_name: str, table: Optional[wandb.Table] = None
+) -> Union[wandb.Table, Tuple[wandb.Image, Dict, Dict]]:
+    """Plot the images with the W&B overlay system. The `wandb.Image` is either added to a `wandb.Table` or returned."""
+    result = result.to("cpu")
+    boxes, mean_confidence_map = get_boxes(result)
+    image = wandb.Image(result.orig_img[:, :, ::-1], boxes=boxes)
+    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, boxes["predictions"], mean_confidence_map
+
+
+def plot_validation_results(
+    dataloader: Any,
+    class_label_map: Dict,
+    model_name: str,
+    predictor: DetectionPredictor,
+    table: wandb.Table,
+    max_validation_batches: int,
+    epoch: Optional[int] = None,
+) -> wandb.Table:
+    """Plot validation results in a table."""
+    data_idx = 0
+    for batch_idx, batch in enumerate(dataloader):
+        for img_idx, image_path in enumerate(batch["im_file"]):
+            prediction_result = predictor(image_path)[0]
+            _, prediction_box_data, mean_confidence_map = plot_predictions(
+                prediction_result, model_name
+            )
+            try:
+                ground_truth_data = get_ground_truth_bbox_annotations(
+                    img_idx, image_path, batch, class_label_map
+                )
+                wandb_image = wandb.Image(
+                    image_path,
+                    boxes={
+                        "ground-truth": {
+                            "box_data": ground_truth_data,
+                            "class_labels": class_label_map,
+                        },
+                        "predictions": {
+                            "box_data": prediction_box_data["box_data"],
+                            "class_labels": class_label_map,
+                        },
+                    },
+                )
+                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