wandb 0.19.12rc1__py3-none-macosx_11_0_arm64.whl → 0.20.0__py3-none-macosx_11_0_arm64.whl

wandb/sdk/data_types/image.py

@@ -1,10 +1,13 @@
 import hashlib
 import logging
 import os
+import pathlib
 from io import BytesIO
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Type, Union, cast
 from urllib import parse
 
+from packaging.version import parse as parse_version
+
 import wandb
 from wandb import util
 from wandb.sdk.lib import hashutil, runid
@@ -30,10 +33,60 @@ if TYPE_CHECKING:  # pragma: no cover
     ImageDataType = Union[
         "matplotlib.artist.Artist", "PILImage", "TorchTensorType", "np.ndarray"
     ]
-    ImageDataOrPathType = Union[str, "Image", ImageDataType]
+    ImageDataOrPathType = Union[str, pathlib.Path, "Image", ImageDataType]
     TorchTensorType = Union["torch.Tensor", "torch.Variable"]
 
 
+def _warn_on_invalid_data_range(
+    data: "np.ndarray",
+    normalize: bool = True,
+) -> None:
+    if not normalize:
+        return
+
+    np = util.get_module(
+        "numpy",
+        required="wandb.Image requires numpy if not supplying PIL Images: pip install numpy",
+    )
+
+    if np.min(data) < 0 or np.max(data) > 255:
+        wandb.termwarn(
+            "Data passed to `wandb.Image` should consist of values in the range [0, 255], "
+            "image data will be normalized to this range, "
+            "but behavior will be removed in a future version of wandb.",
+            repeat=False,
+        )
+
+
+def _normalize(data: "np.ndarray") -> "np.ndarray":
+    """Normalizes and converts image pixel values to uint8 in the range [0, 255]."""
+    np = util.get_module(
+        "numpy",
+        required="wandb.Image requires numpy if not supplying PIL Images: pip install numpy",
+    )
+
+    # if an image has negative values, set all values to be in the range [0, 1]
+    # This can lead to inconsistent behavior when an image has only a single negative value
+    if np.min(data) < 0:
+        data = data - np.min(data)
+
+    if np.ptp(data) != 0:
+        data = data / np.ptp(data)
+
+    if np.max(data) <= 1.0:
+        data = (255 * data).astype(np.int32)
+
+    return data.clip(0, 255)
+
+
+def _convert_to_uint8(data: "np.ndarray") -> "np.ndarray":
+    np = util.get_module(
+        "numpy",
+        required="wandb.Image requires numpy if not supplying PIL Images: pip install numpy",
+    )
+    return data.astype(np.uint8)
+
+
 def _server_accepts_image_filenames(run: "LocalRun") -> bool:
     if run.offline:
         return True
@@ -43,10 +96,9 @@ def _server_accepts_image_filenames(run: "LocalRun") -> bool:
     max_cli_version = util._get_max_cli_version()
     if max_cli_version is None:
         return False
-    from wandb.util import parse_version
 
-    accepts_image_filenames: bool = parse_version("0.12.10") <= parse_version(
-        max_cli_version
+    accepts_image_filenames: bool = parse_version(max_cli_version) >= parse_version(
+        "0.12.10"
     )
     return accepts_image_filenames
 
@@ -59,69 +111,11 @@ def _server_accepts_artifact_path(run: "LocalRun") -> bool:
     if max_cli_version is None:
         return False
 
-    return util.parse_version("0.12.14") <= util.parse_version(max_cli_version)
+    return parse_version(max_cli_version) >= parse_version("0.12.14")
 
 
 class Image(BatchableMedia):
-    """Format images for logging to W&B.
-
-    Args:
-        data_or_path: (numpy array, string, io) Accepts numpy array of
-            image data, or a PIL image. The class attempts to infer
-            the data format and converts it.
-        mode: (string) The PIL mode for an image. Most common are "L", "RGB",
-            "RGBA". Full explanation at https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes
-        caption: (string) Label for display of image.
-
-    Note : When logging a `torch.Tensor` as a `wandb.Image`, images are normalized. If you do not want to normalize your images, please convert your tensors to a PIL Image.
-
-    Examples:
-        ### Create a wandb.Image from a numpy array
-        ```python
-        import numpy as np
-        import wandb
-
-        with wandb.init() as run:
-            examples = []
-            for i in range(3):
-                pixels = np.random.randint(low=0, high=256, size=(100, 100, 3))
-                image = wandb.Image(pixels, caption=f"random field {i}")
-                examples.append(image)
-            run.log({"examples": examples})
-        ```
-
-        ### Create a wandb.Image from a PILImage
-        ```python
-        import numpy as np
-        from PIL import Image as PILImage
-        import wandb
-
-        with wandb.init() as run:
-            examples = []
-            for i in range(3):
-                pixels = np.random.randint(
-                    low=0, high=256, size=(100, 100, 3), dtype=np.uint8
-                )
-                pil_image = PILImage.fromarray(pixels, mode="RGB")
-                image = wandb.Image(pil_image, caption=f"random field {i}")
-                examples.append(image)
-            run.log({"examples": examples})
-        ```
-
-        ### log .jpg rather than .png (default)
-        ```python
-        import numpy as np
-        import wandb
-
-        with wandb.init() as run:
-            examples = []
-            for i in range(3):
-                pixels = np.random.randint(low=0, high=256, size=(100, 100, 3))
-                image = wandb.Image(pixels, caption=f"random field {i}", file_type="jpg")
-                examples.append(image)
-            run.log({"examples": examples})
-        ```
-    """
+    """A class for logging images to W&B."""
 
     MAX_ITEMS = 108
 
@@ -151,7 +145,85 @@ class Image(BatchableMedia):
         boxes: Optional[Union[Dict[str, "BoundingBoxes2D"], Dict[str, dict]]] = None,
         masks: Optional[Union[Dict[str, "ImageMask"], Dict[str, dict]]] = None,
         file_type: Optional[str] = None,
+        normalize: bool = True,
     ) -> None:
+        """Initialize a wandb.Image object.
+
+        Args:
+            data_or_path: Accepts numpy array/pytorch tensor of image data,
+                a PIL image object, or a path to an image file.
+
+                If a numpy array or pytorch tensor is provided,
+                the image data will be saved to the given file type.
+                If the values are not in the range [0, 255] or all values are in the range [0, 1],
+                the image pixel values will be normalized to the range [0, 255]
+                unless `normalize` is set to False.
+                - pytorch tensor should be in the format (channel, height, width)
+                - numpy array should be in the format (height, width, channel)
+            mode: The PIL mode for an image. Most common are "L", "RGB",
+                "RGBA". Full explanation at https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes
+            caption: Label for display of image.
+            grouping: The grouping number for the image.
+            classes: A list of class information for the image,
+                used for labeling bounding boxes, and image masks.
+            boxes: A dictionary containing bounding box information for the image.
+                see: https://docs.wandb.ai/ref/python/data-types/boundingboxes2d/
+            masks: A dictionary containing mask information for the image.
+                see: https://docs.wandb.ai/ref/python/data-types/imagemask/
+            file_type: The file type to save the image as.
+                This parameter has no effect if data_or_path is a path to an image file.
+            normalize: If True, normalize the image pixel values to fall within the range of [0, 255].
+                Normalize is only applied if data_or_path is a numpy array or pytorch tensor.
+
+        Examples:
+            ### Create a wandb.Image from a numpy array
+            ```python
+            import numpy as np
+            import wandb
+
+            with wandb.init() as run:
+                examples = []
+                for i in range(3):
+                    pixels = np.random.randint(low=0, high=256, size=(100, 100, 3))
+                    image = wandb.Image(pixels, caption=f"random field {i}")
+                    examples.append(image)
+                run.log({"examples": examples})
+            ```
+
+            ### Create a wandb.Image from a PILImage
+            ```python
+            import numpy as np
+            from PIL import Image as PILImage
+            import wandb
+
+            with wandb.init() as run:
+                examples = []
+                for i in range(3):
+                    pixels = np.random.randint(
+                        low=0, high=256, size=(100, 100, 3), dtype=np.uint8
+                    )
+                    pil_image = PILImage.fromarray(pixels, mode="RGB")
+                    image = wandb.Image(pil_image, caption=f"random field {i}")
+                    examples.append(image)
+                run.log({"examples": examples})
+            ```
+
+            ### log .jpg rather than .png (default)
+            ```python
+            import numpy as np
+            import wandb
+
+            with wandb.init() as run:
+                examples = []
+                for i in range(3):
+                    pixels = np.random.randint(low=0, high=256, size=(100, 100, 3))
+                    image = wandb.Image(
+                        pixels, caption=f"random field {i}", file_type="jpg"
+                    )
+                    examples.append(image)
+                run.log({"examples": examples})
+            ```
+        """
         super().__init__(caption=caption)
         # TODO: We should remove grouping, it's a terrible name and I don't
         # think anyone uses it.
@@ -169,13 +241,15 @@ class Image(BatchableMedia):
         # only overriding additional metadata passed in. If this pattern is compelling, we can generalize.
         if isinstance(data_or_path, Image):
             self._initialize_from_wbimage(data_or_path)
-        elif isinstance(data_or_path, str):
+        elif isinstance(data_or_path, (str, pathlib.Path)):
+            data_or_path = str(data_or_path)
+
             if self.path_is_reference(data_or_path):
                 self._initialize_from_reference(data_or_path)
             else:
                 self._initialize_from_path(data_or_path)
         else:
-            self._initialize_from_data(data_or_path, mode, file_type)
+            self._initialize_from_data(data_or_path, mode, file_type, normalize)
         self._set_initialization_meta(
             grouping, caption, classes, boxes, masks, file_type
         )
@@ -288,6 +362,7 @@ class Image(BatchableMedia):
         data: "ImageDataType",
         mode: Optional[str] = None,
         file_type: Optional[str] = None,
+        normalize: bool = True,
     ) -> None:
         pil_image = util.get_module(
             "PIL.Image",
@@ -309,28 +384,39 @@ class Image(BatchableMedia):
         elif isinstance(data, pil_image.Image):
             self._image = data
         elif util.is_pytorch_tensor_typename(util.get_full_typename(data)):
-            vis_util = util.get_module(
-                "torchvision.utils", "torchvision is required to render images"
-            )
             if hasattr(data, "requires_grad") and data.requires_grad:
                 data = data.detach()  # type: ignore
             if hasattr(data, "dtype") and str(data.dtype) == "torch.uint8":
-                data = data.to(float)
-            data = vis_util.make_grid(data, normalize=True)
+                data = data.to(float)  # type: ignore [union-attr]
             mode = mode or self.guess_mode(data, file_type)
+            data = data.permute(1, 2, 0).cpu().numpy()  # type: ignore [union-attr]
+
+            _warn_on_invalid_data_range(data, normalize)
+
+            data = _normalize(data) if normalize else data  # type: ignore [arg-type]
+            data = _convert_to_uint8(data)
+
+            if data.ndim > 2:
+                data = data.squeeze()
+
             self._image = pil_image.fromarray(
-                data.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy(),
+                data,
                 mode=mode,
             )
         else:
             if hasattr(data, "numpy"):  # TF data eager tensors
                 data = data.numpy()
-            if data.ndim > 2:
-                data = data.squeeze()  # get rid of trivial dimensions as a convenience
+            if data.ndim > 2:  # type: ignore [union-attr]
+                # get rid of trivial dimensions as a convenience
+                data = data.squeeze()  # type: ignore [union-attr]
+
+            _warn_on_invalid_data_range(data, normalize)  # type: ignore [arg-type]
 
             mode = mode or self.guess_mode(data, file_type)
+            data = _normalize(data) if normalize else data  # type: ignore [arg-type]
+            data = _convert_to_uint8(data)  # type: ignore [arg-type]
             self._image = pil_image.fromarray(
-                self.to_uint8(data),
+                data,
                 mode=mode,
             )
 
@@ -459,7 +545,7 @@ class Image(BatchableMedia):
                 }
 
         elif not isinstance(run_or_artifact, Run):
-            raise ValueError("to_json accepts wandb_run.Run or wandb_artifact.Artifact")
+            raise TypeError("to_json accepts wandb_run.Run or wandb_artifact.Artifact")
 
         if self._boxes:
             json_dict["boxes"] = {
@@ -485,7 +571,7 @@ class Image(BatchableMedia):
         else:
             num_channels = data.shape[-1]
 
-        if ndims == 2:
+        if ndims == 2 or num_channels == 1:
             return "L"
         elif num_channels == 3:
             return "RGB"
@@ -501,34 +587,9 @@ class Image(BatchableMedia):
                 return "RGBA"
         else:
             raise ValueError(
-                "Un-supported shape for image conversion {}".format(list(data.shape))
+                f"Un-supported shape for image conversion {list(data.shape)}"
             )
 
-    @classmethod
-    def to_uint8(cls, data: "np.ndarray") -> "np.ndarray":
-        """Convert image data to uint8.
-
-        Convert floating point image on the range [0,1] and integer images on the range
-        [0,255] to uint8, clipping if necessary.
-        """
-        np = util.get_module(
-            "numpy",
-            required="wandb.Image requires numpy if not supplying PIL Images: pip install numpy",
-        )
-
-        # I think it's better to check the image range vs the data type, since many
-        # image libraries will return floats between 0 and 255
-
-        # some images have range -1...1 or 0-1
-        dmin = np.min(data)
-        if dmin < 0:
-            data = (data - np.min(data)) / np.ptp(data)
-        if np.max(data) <= 1.0:
-            data = (data * 255).astype(np.int32)
-
-        # assert issubclass(data.dtype.type, np.integer), 'Illegal image format.'
-        return data.clip(0, 255).astype(np.uint8)
-
     @classmethod
     def seq_to_json(
         cls: Type["Image"],

wandb/sdk/data_types/molecule.py

@@ -26,7 +26,7 @@ class Molecule(BatchableMedia):
     """Wandb class for 3D Molecular data.
 
     Args:
-        data_or_path: (string, io)
+        data_or_path: (pathlib.Path, string, io)
             Molecule can be initialized from a file name or an io object.
         caption: (string)
             Caption associated with the molecule for display.
@@ -49,7 +49,7 @@ class Molecule(BatchableMedia):
 
     def __init__(
         self,
-        data_or_path: Union[str, "TextIO"],
+        data_or_path: Union[str, pathlib.Path, "TextIO"],
         caption: Optional[str] = None,
         **kwargs: str,
     ) -> None:
@@ -82,7 +82,9 @@ class Molecule(BatchableMedia):
                 f.write(molecule)
 
             self._set_file(tmp_path, is_tmp=True)
-        elif isinstance(data_or_path, str):
+        elif isinstance(data_or_path, (str, pathlib.Path)):
+            data_or_path = str(data_or_path)
+
             extension = os.path.splitext(data_or_path)[1][1:]
             if extension not in Molecule.SUPPORTED_TYPES:
                 raise ValueError(
@@ -144,9 +146,7 @@ class Molecule(BatchableMedia):
         elif isinstance(data_or_path, rdkit_chem.rdchem.Mol):
             molecule = data_or_path
         else:
-            raise ValueError(
-                "Data must be file name or an rdkit.Chem.rdchem.Mol object"
-            )
+            raise TypeError("Data must be file name or an rdkit.Chem.rdchem.Mol object")
 
         if convert_to_3d_and_optimize:
             molecule = rdkit_chem.AddHs(molecule)

wandb/sdk/data_types/object_3d.py

@@ -2,6 +2,7 @@ import codecs
 import itertools
 import json
 import os
+import pathlib
 from typing import (
     TYPE_CHECKING,
     ClassVar,
@@ -187,7 +188,7 @@ class Object3D(BatchableMedia):
     """Wandb class for 3D point clouds.
 
     Args:
-        data_or_path: (numpy array, string, io)
+        data_or_path: (numpy array, pathlib.Path, string, io)
             Object3D can be initialized from a file or a numpy array.
 
             You can pass a path to a file or an io object and a file_type
@@ -214,14 +215,16 @@ class Object3D(BatchableMedia):
 
     def __init__(
         self,
-        data_or_path: Union["np.ndarray", str, "TextIO", dict],
+        data_or_path: Union["np.ndarray", str, pathlib.Path, "TextIO", dict],
         caption: Optional[str] = None,
         **kwargs: Optional[Union[str, "FileFormat3D"]],
     ) -> None:
         super().__init__(caption=caption)
 
-        if hasattr(data_or_path, "name"):
-            # if the file has a path, we just detect the type and copy it from there
+        if hasattr(data_or_path, "name") and not isinstance(data_or_path, pathlib.Path):
+            # if the file has a path, we just detect the type and copy it from there.
+            # this does not work for pathlib.Path objects,
+            # where `.name` returns the last directory in the path.
             data_or_path = data_or_path.name
 
         if hasattr(data_or_path, "read"):
@@ -247,7 +250,9 @@ class Object3D(BatchableMedia):
                 f.write(object_3d)
 
             self._set_file(tmp_path, is_tmp=True, extension=extension)
-        elif isinstance(data_or_path, str):
+        elif isinstance(data_or_path, (str, pathlib.Path)):
+            data_or_path = str(data_or_path)
+
             path = data_or_path
             extension = None
             for supported_type in Object3D.SUPPORTED_TYPES:

wandb/sdk/data_types/saved_model.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import pathlib
 import shutil
 import sys
 from types import ModuleType
@@ -72,10 +73,12 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
 
     _model_obj: SavedModelObjType | None
     _path: str | None
-    _input_obj_or_path: SavedModelObjType | str
+    _input_obj_or_path: SavedModelObjType | str | pathlib.Path
 
     # Public Methods
-    def __init__(self, obj_or_path: SavedModelObjType | str, **kwargs: Any) -> None:
+    def __init__(
+        self, obj_or_path: SavedModelObjType | str | pathlib.Path, **kwargs: Any
+    ) -> None:
         super().__init__()
         if self.__class__ == _SavedModel:
             raise TypeError(
@@ -84,9 +87,11 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
         self._model_obj = None
         self._path = None
         self._input_obj_or_path = obj_or_path
-        input_is_path = isinstance(obj_or_path, str) and os.path.exists(obj_or_path)
+        input_is_path = isinstance(obj_or_path, (str, pathlib.Path)) and os.path.exists(
+            obj_or_path
+        )
         if input_is_path:
-            assert isinstance(obj_or_path, str)  # mypy
+            obj_or_path = str(obj_or_path)
             self._set_obj(self._deserialize(obj_or_path))
         else:
             self._set_obj(obj_or_path)
@@ -140,7 +145,7 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
         from wandb.sdk.wandb_run import Run
 
         if isinstance(run_or_artifact, Run):
-            raise ValueError("SavedModel cannot be added to run - must use artifact")
+            raise TypeError("SavedModel cannot be added to run - must use artifact")
         artifact = run_or_artifact
         json_obj = {
             "type": self._log_type,
@@ -280,7 +285,7 @@ class _PicklingSavedModel(_SavedModel[SavedModelObjType]):
 
     def __init__(
         self,
-        obj_or_path: SavedModelObjType | str,
+        obj_or_path: SavedModelObjType | str | pathlib.Path,
         dep_py_files: list[str] | None = None,
     ):
         super().__init__(obj_or_path)