wandb 0.19.1rc1__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/data_types/helper_types/image_mask.py

@@ -0,0 +1,241 @@
+import numbers
+import os
+from typing import TYPE_CHECKING, Optional, Type, Union
+
+import wandb
+from wandb import util
+from wandb.sdk.lib import runid
+
+from .._private import MEDIA_TMP
+from ..base_types.media import Media
+
+if TYPE_CHECKING:  # pragma: no cover
+    from wandb.sdk.artifacts.artifact import Artifact
+
+    from ...wandb_run import Run as LocalRun
+
+
+class ImageMask(Media):
+    """Format image masks or overlays for logging to W&B.
+
+    Args:
+        val: (dictionary)
+            One of these two keys to represent the image:
+                mask_data : (2D numpy array) The mask containing an integer class label
+                    for each pixel in the image
+                path : (string) The path to a saved image file of the mask
+            class_labels : (dictionary of integers to strings, optional) A mapping of the
+                integer class labels in the mask to readable class names. These will default
+                to class_0, class_1, class_2, etc.
+
+        key: (string)
+            The readable name or id for this mask type (e.g. predictions, ground_truth)
+
+    Examples:
+        ### Logging a single masked image
+        <!--yeadoc-test:log-image-mask-->
+        ```python
+        import numpy as np
+        import wandb
+
+        wandb.init()
+        image = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8)
+        predicted_mask = np.empty((100, 100), dtype=np.uint8)
+        ground_truth_mask = np.empty((100, 100), dtype=np.uint8)
+
+        predicted_mask[:50, :50] = 0
+        predicted_mask[50:, :50] = 1
+        predicted_mask[:50, 50:] = 2
+        predicted_mask[50:, 50:] = 3
+
+        ground_truth_mask[:25, :25] = 0
+        ground_truth_mask[25:, :25] = 1
+        ground_truth_mask[:25, 25:] = 2
+        ground_truth_mask[25:, 25:] = 3
+
+        class_labels = {0: "person", 1: "tree", 2: "car", 3: "road"}
+
+        masked_image = wandb.Image(
+            image,
+            masks={
+                "predictions": {"mask_data": predicted_mask, "class_labels": class_labels},
+                "ground_truth": {
+                    "mask_data": ground_truth_mask,
+                    "class_labels": class_labels,
+                },
+            },
+        )
+        wandb.log({"img_with_masks": masked_image})
+        ```
+
+        ### Log a masked image inside a Table
+        <!--yeadoc-test:log-image-mask-table-->
+        ```python
+        import numpy as np
+        import wandb
+
+        wandb.init()
+        image = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8)
+        predicted_mask = np.empty((100, 100), dtype=np.uint8)
+        ground_truth_mask = np.empty((100, 100), dtype=np.uint8)
+
+        predicted_mask[:50, :50] = 0
+        predicted_mask[50:, :50] = 1
+        predicted_mask[:50, 50:] = 2
+        predicted_mask[50:, 50:] = 3
+
+        ground_truth_mask[:25, :25] = 0
+        ground_truth_mask[25:, :25] = 1
+        ground_truth_mask[:25, 25:] = 2
+        ground_truth_mask[25:, 25:] = 3
+
+        class_labels = {0: "person", 1: "tree", 2: "car", 3: "road"}
+
+        class_set = wandb.Classes(
+            [
+                {"name": "person", "id": 0},
+                {"name": "tree", "id": 1},
+                {"name": "car", "id": 2},
+                {"name": "road", "id": 3},
+            ]
+        )
+
+        masked_image = wandb.Image(
+            image,
+            masks={
+                "predictions": {"mask_data": predicted_mask, "class_labels": class_labels},
+                "ground_truth": {
+                    "mask_data": ground_truth_mask,
+                    "class_labels": class_labels,
+                },
+            },
+            classes=class_set,
+        )
+
+        table = wandb.Table(columns=["image"])
+        table.add_data(masked_image)
+        wandb.log({"random_field": table})
+        ```
+    """
+
+    _log_type = "mask"
+
+    def __init__(self, val: dict, key: str) -> None:
+        """Initialize an ImageMask object.
+
+        Args:
+            val: (dictionary) One of these two keys to represent the image:
+                mask_data : (2D numpy array) The mask containing an integer class label
+                    for each pixel in the image
+                path : (string) The path to a saved image file of the mask
+                class_labels : (dictionary of integers to strings, optional) A mapping
+                    of the integer class labels in the mask to readable class names.
+                    These will default to class_0, class_1, class_2, etc.
+
+        key: (string)
+            The readable name or id for this mask type (e.g. predictions, ground_truth)
+        """
+        super().__init__()
+
+        if "path" in val:
+            self._set_file(val["path"])
+        else:
+            np = util.get_module("numpy", required="Image mask support requires numpy")
+            # Add default class mapping
+            if "class_labels" not in val:
+                classes = np.unique(val["mask_data"]).astype(np.int32).tolist()
+                class_labels = {c: "class_" + str(c) for c in classes}
+                val["class_labels"] = class_labels
+
+            self.validate(val)
+            self._val = val
+            self._key = key
+
+            ext = "." + self.type_name() + ".png"
+            tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ext)
+
+            pil_image = util.get_module(
+                "PIL.Image",
+                required='wandb.Image needs the PIL package. To get it, run "pip install pillow".',
+            )
+            image = pil_image.fromarray(val["mask_data"].astype(np.int8), mode="L")
+
+            image.save(tmp_path, transparency=None)
+            self._set_file(tmp_path, is_tmp=True, extension=ext)
+
+    def bind_to_run(
+        self,
+        run: "LocalRun",
+        key: Union[int, str],
+        step: Union[int, str],
+        id_: Optional[Union[int, str]] = None,
+        ignore_copy_err: Optional[bool] = None,
+    ) -> None:
+        # bind_to_run key argument is the Image parent key
+        # the self._key value is the mask's sub key
+        super().bind_to_run(run, key, step, id_=id_, ignore_copy_err=ignore_copy_err)
+        if hasattr(self, "_val") and "class_labels" in self._val:
+            class_labels = self._val["class_labels"]
+
+            run._add_singleton(
+                "mask/class_labels",
+                str(key) + "_wandb_delimeter_" + self._key,
+                class_labels,
+            )
+
+    @classmethod
+    def get_media_subdir(cls: Type["ImageMask"]) -> str:
+        return os.path.join("media", "images", cls.type_name())
+
+    @classmethod
+    def from_json(
+        cls: Type["ImageMask"], json_obj: dict, source_artifact: "Artifact"
+    ) -> "ImageMask":
+        return cls(
+            {"path": source_artifact.get_entry(json_obj["path"]).download()},
+            key="",
+        )
+
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+        from wandb.sdk.wandb_run import Run
+
+        json_dict = super().to_json(run_or_artifact)
+
+        if isinstance(run_or_artifact, Run):
+            json_dict["_type"] = self.type_name()
+            return json_dict
+        elif isinstance(run_or_artifact, wandb.Artifact):
+            # Nothing special to add (used to add "digest", but no longer used.)
+            return json_dict
+        else:
+            raise ValueError("to_json accepts wandb_run.Run or wandb.Artifact")
+
+    @classmethod
+    def type_name(cls: Type["ImageMask"]) -> str:
+        return cls._log_type
+
+    def validate(self, val: dict) -> bool:
+        np = util.get_module("numpy", required="Image mask support requires numpy")
+        # 2D Make this work with all tensor(like) types
+        if "mask_data" not in val:
+            raise TypeError(
+                'Missing key "mask_data": An image mask requires mask data: a 2D array representing the predictions'
+            )
+        else:
+            error_str = "mask_data must be a 2D array"
+            shape = val["mask_data"].shape
+            if len(shape) != 2:
+                raise TypeError(error_str)
+            if not (
+                (val["mask_data"] >= 0).all() and (val["mask_data"] <= 255).all()
+            ) and issubclass(val["mask_data"].dtype.type, np.integer):
+                raise TypeError("Mask data must be integers between 0 and 255")
+
+        # Optional argument
+        if "class_labels" in val:
+            for k, v in list(val["class_labels"].items()):
+                if (not isinstance(k, numbers.Number)) or (not isinstance(v, str)):
+                    raise TypeError(
+                        "Class labels must be a dictionary of numbers to strings"
+                    )
+        return True

wandb/sdk/data_types/histogram.py

@@ -0,0 +1,94 @@
+import sys
+from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Union
+
+from wandb import util
+
+from .base_types.wb_value import WBValue
+
+if TYPE_CHECKING:  # pragma: no cover
+    import numpy as np
+
+    from wandb.sdk.artifacts.artifact import Artifact
+
+    from ..wandb_run import Run as LocalRun
+
+    NumpyHistogram = Tuple[np.ndarray, np.ndarray]
+
+
+class Histogram(WBValue):
+    """wandb class for histograms.
+
+    This object works just like numpy's histogram function
+    https://docs.scipy.org/doc/numpy/reference/generated/numpy.histogram.html
+
+    Examples:
+        Generate histogram from a sequence
+        ```python
+        wandb.Histogram([1, 2, 3])
+        ```
+
+        Efficiently initialize from np.histogram.
+        ```python
+        hist = np.histogram(data)
+        wandb.Histogram(np_histogram=hist)
+        ```
+
+    Args:
+        sequence: (array_like) input data for histogram
+        np_histogram: (numpy histogram) alternative input of a precomputed histogram
+        num_bins: (int) Number of bins for the histogram.  The default number of bins
+            is 64.  The maximum number of bins is 512
+
+    Attributes:
+        bins: ([float]) edges of bins
+        histogram: ([int]) number of elements falling in each bin
+    """
+
+    MAX_LENGTH: int = 512
+    _log_type = "histogram"
+
+    def __init__(
+        self,
+        sequence: Optional[Sequence] = None,
+        np_histogram: Optional["NumpyHistogram"] = None,
+        num_bins: int = 64,
+    ) -> None:
+        if np_histogram:
+            if len(np_histogram) == 2:
+                self.histogram = (
+                    np_histogram[0].tolist()
+                    if hasattr(np_histogram[0], "tolist")
+                    else np_histogram[0]
+                )
+                self.bins = (
+                    np_histogram[1].tolist()
+                    if hasattr(np_histogram[1], "tolist")
+                    else np_histogram[1]
+                )
+            else:
+                raise ValueError(
+                    "Expected np_histogram to be a tuple of (values, bin_edges) or sequence to be specified"
+                )
+        else:
+            np = util.get_module(
+                "numpy", required="Auto creation of histograms requires numpy"
+            )
+
+            histogram, bins = np.histogram(sequence, bins=num_bins)
+            self.histogram = histogram.tolist()
+            self.bins = bins.tolist()
+        if len(self.histogram) > self.MAX_LENGTH:
+            raise ValueError(f"The maximum length of a histogram is {self.MAX_LENGTH}")
+        if len(self.histogram) + 1 != len(self.bins):
+            raise ValueError("len(bins) must be len(histogram) + 1")
+
+    def to_json(self, run: Optional[Union["LocalRun", "Artifact"]] = None) -> dict:
+        return {"_type": self._log_type, "values": self.histogram, "bins": self.bins}
+
+    def __sizeof__(self) -> int:
+        """Estimated size in bytes.
+
+        Currently the factor of 1.7 is used to account for the JSON encoding. We use
+        this in tb_watcher.TBHistory.
+        """
+        return int((sys.getsizeof(self.histogram) + sys.getsizeof(self.bins)) * 1.7)

wandb/sdk/data_types/html.py

@@ -0,0 +1,115 @@
+import os
+from typing import TYPE_CHECKING, Sequence, Type, Union
+
+from wandb.sdk.lib import filesystem, runid
+
+from . import _dtypes
+from ._private import MEDIA_TMP
+from .base_types.media import BatchableMedia
+
+if TYPE_CHECKING:  # pragma: no cover
+    from typing import TextIO
+
+    from wandb.sdk.artifacts.artifact import Artifact
+
+    from ..wandb_run import Run as LocalRun
+
+
+class Html(BatchableMedia):
+    """Wandb class for arbitrary html.
+
+    Args:
+        data: (string or io object) HTML to display in wandb
+        inject: (boolean) Add a stylesheet to the HTML object.  If set
+            to False the HTML will pass through unchanged.
+    """
+
+    _log_type = "html-file"
+
+    def __init__(self, data: Union[str, "TextIO"], inject: bool = True) -> None:
+        super().__init__()
+        data_is_path = isinstance(data, str) and os.path.exists(data)
+        data_path = ""
+        if data_is_path:
+            assert isinstance(data, str)
+            data_path = data
+            with open(data_path, encoding="utf-8") as file:
+                self.html = file.read()
+        elif isinstance(data, str):
+            self.html = data
+        elif hasattr(data, "read"):
+            if hasattr(data, "seek"):
+                data.seek(0)
+            self.html = data.read()
+        else:
+            raise ValueError("data must be a string or an io object")
+
+        if inject:
+            self.inject_head()
+
+        if inject or not data_is_path:
+            tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".html")
+            with open(tmp_path, "w", encoding="utf-8") as out:
+                out.write(self.html)
+
+            self._set_file(tmp_path, is_tmp=True)
+        else:
+            self._set_file(data_path, is_tmp=False)
+
+    def inject_head(self) -> None:
+        join = ""
+        if "<head>" in self.html:
+            parts = self.html.split("<head>", 1)
+            parts[0] = parts[0] + "<head>"
+        elif "<html>" in self.html:
+            parts = self.html.split("<html>", 1)
+            parts[0] = parts[0] + "<html><head>"
+            parts[1] = "</head>" + parts[1]
+        else:
+            parts = ["", self.html]
+        parts.insert(
+            1,
+            '<base target="_blank"><link rel="stylesheet" type="text/css" href="https://app.wandb.ai/normalize.css" />',
+        )
+        self.html = join.join(parts).strip()
+
+    @classmethod
+    def get_media_subdir(cls: Type["Html"]) -> str:
+        return os.path.join("media", "html")
+
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+        json_dict = super().to_json(run_or_artifact)
+        json_dict["_type"] = self._log_type
+        return json_dict
+
+    @classmethod
+    def from_json(
+        cls: Type["Html"], json_obj: dict, source_artifact: "Artifact"
+    ) -> "Html":
+        return cls(source_artifact.get_entry(json_obj["path"]).download(), inject=False)
+
+    @classmethod
+    def seq_to_json(
+        cls: Type["Html"],
+        seq: Sequence["BatchableMedia"],
+        run: "LocalRun",
+        key: str,
+        step: Union[int, str],
+    ) -> dict:
+        base_path = os.path.join(run.dir, cls.get_media_subdir())
+        filesystem.mkdir_exists_ok(base_path)
+
+        meta = {
+            "_type": "html",
+            "count": len(seq),
+            "html": [h.to_json(run) for h in seq],
+        }
+        return meta
+
+
+class _HtmlFileType(_dtypes.Type):
+    name = "html-file"
+    types = [Html]
+
+
+_dtypes.TypeRegistry.add(_HtmlFileType)