wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/data_types/graph.py

@@ -0,0 +1,405 @@
+import codecs
+import os
+import pprint
+
+from wandb import util
+from wandb.sdk.data_types._private import MEDIA_TMP
+from wandb.sdk.data_types.base_types.media import Media, _numpy_arrays_to_lists
+from wandb.sdk.data_types.base_types.wb_value import WBValue
+from wandb.sdk.lib import runid
+
+
+def _nest(thing):
+    # Use tensorflows nest function if available, otherwise just wrap object in an array"""
+
+    tfutil = util.get_module("tensorflow.python.util")
+    if tfutil:
+        return tfutil.nest.flatten(thing)
+    else:
+        return [thing]
+
+
+class Edge(WBValue):
+    """Edge used in `Graph`."""
+
+    def __init__(self, from_node, to_node):
+        self._attributes = {}
+        self.from_node = from_node
+        self.to_node = to_node
+
+    def __repr__(self):
+        temp_attr = dict(self._attributes)
+        del temp_attr["from_node"]
+        del temp_attr["to_node"]
+        temp_attr["from_id"] = self.from_node.id
+        temp_attr["to_id"] = self.to_node.id
+        return str(temp_attr)
+
+    def to_json(self, run=None):
+        return [self.from_node.id, self.to_node.id]
+
+    @property
+    def name(self):
+        """Optional, not necessarily unique."""
+        return self._attributes.get("name")
+
+    @name.setter
+    def name(self, val):
+        self._attributes["name"] = val
+        return val
+
+    @property
+    def from_node(self):
+        return self._attributes.get("from_node")
+
+    @from_node.setter
+    def from_node(self, val):
+        self._attributes["from_node"] = val
+        return val
+
+    @property
+    def to_node(self):
+        return self._attributes.get("to_node")
+
+    @to_node.setter
+    def to_node(self, val):
+        self._attributes["to_node"] = val
+        return val
+
+
+class Node(WBValue):
+    """Node used in `Graph`."""
+
+    def __init__(
+        self,
+        id=None,
+        name=None,
+        class_name=None,
+        size=None,
+        parameters=None,
+        output_shape=None,
+        is_output=None,
+        num_parameters=None,
+        node=None,
+    ):
+        self._attributes = {"name": None}
+        self.in_edges = {}  # indexed by source node id
+        self.out_edges = {}  # indexed by dest node id
+        # optional object (e.g. PyTorch Parameter or Module) that this Node represents
+        self.obj = None
+
+        if node is not None:
+            self._attributes.update(node._attributes)
+            del self._attributes["id"]
+            self.obj = node.obj
+
+        if id is not None:
+            self.id = id
+        if name is not None:
+            self.name = name
+        if class_name is not None:
+            self.class_name = class_name
+        if size is not None:
+            self.size = size
+        if parameters is not None:
+            self.parameters = parameters
+        if output_shape is not None:
+            self.output_shape = output_shape
+        if is_output is not None:
+            self.is_output = is_output
+        if num_parameters is not None:
+            self.num_parameters = num_parameters
+
+    def to_json(self, run=None):
+        return self._attributes
+
+    def __repr__(self):
+        return repr(self._attributes)
+
+    @property
+    def id(self):
+        """Must be unique in the graph."""
+        return self._attributes.get("id")
+
+    @id.setter
+    def id(self, val):
+        self._attributes["id"] = val
+        return val
+
+    @property
+    def name(self):
+        """Usually the type of layer or sublayer."""
+        return self._attributes.get("name")
+
+    @name.setter
+    def name(self, val):
+        self._attributes["name"] = val
+        return val
+
+    @property
+    def class_name(self):
+        """Usually the type of layer or sublayer."""
+        return self._attributes.get("class_name")
+
+    @class_name.setter
+    def class_name(self, val):
+        self._attributes["class_name"] = val
+        return val
+
+    @property
+    def functions(self):
+        return self._attributes.get("functions", [])
+
+    @functions.setter
+    def functions(self, val):
+        self._attributes["functions"] = val
+        return val
+
+    @property
+    def parameters(self):
+        return self._attributes.get("parameters", [])
+
+    @parameters.setter
+    def parameters(self, val):
+        self._attributes["parameters"] = val
+        return val
+
+    @property
+    def size(self):
+        return self._attributes.get("size")
+
+    @size.setter
+    def size(self, val):
+        """Tensor size."""
+        self._attributes["size"] = tuple(val)
+        return val
+
+    @property
+    def output_shape(self):
+        return self._attributes.get("output_shape")
+
+    @output_shape.setter
+    def output_shape(self, val):
+        """Tensor output_shape."""
+        self._attributes["output_shape"] = val
+        return val
+
+    @property
+    def is_output(self):
+        return self._attributes.get("is_output")
+
+    @is_output.setter
+    def is_output(self, val):
+        """Tensor is_output."""
+        self._attributes["is_output"] = val
+        return val
+
+    @property
+    def num_parameters(self):
+        return self._attributes.get("num_parameters")
+
+    @num_parameters.setter
+    def num_parameters(self, val):
+        """Tensor num_parameters."""
+        self._attributes["num_parameters"] = val
+        return val
+
+    @property
+    def child_parameters(self):
+        return self._attributes.get("child_parameters")
+
+    @child_parameters.setter
+    def child_parameters(self, val):
+        """Tensor child_parameters."""
+        self._attributes["child_parameters"] = val
+        return val
+
+    @property
+    def is_constant(self):
+        return self._attributes.get("is_constant")
+
+    @is_constant.setter
+    def is_constant(self, val):
+        """Tensor is_constant."""
+        self._attributes["is_constant"] = val
+        return val
+
+    @classmethod
+    def from_keras(cls, layer):
+        node = cls()
+
+        try:
+            output_shape = layer.output_shape
+        except AttributeError:
+            output_shape = ["multiple"]
+
+        node.id = layer.name
+        node.name = layer.name
+        node.class_name = layer.__class__.__name__
+        node.output_shape = output_shape
+        node.num_parameters = layer.count_params()
+
+        return node
+
+
+class Graph(Media):
+    """Wandb class for graphs.
+
+    This class is typically used for saving and displaying neural net models.  It
+    represents the graph as an array of nodes and edges.  The nodes can have
+    labels that can be visualized by wandb.
+
+    Examples:
+        Import a keras model:
+        ```
+        Graph.from_keras(keras_model)
+        ```
+
+    Attributes:
+        format (string): Format to help wandb display the graph nicely.
+        nodes ([wandb.Node]): List of wandb.Nodes
+        nodes_by_id (dict): dict of ids -> nodes
+        edges ([(wandb.Node, wandb.Node)]): List of pairs of nodes interpreted as edges
+        loaded (boolean): Flag to tell whether the graph is completely loaded
+        root (wandb.Node): root node of the graph
+    """
+
+    _log_type = "graph-file"
+
+    def __init__(self, format="keras"):
+        super().__init__()
+        # LB: TODO: I think we should factor criterion and criterion_passed out
+        self.format = format
+        self.nodes = []
+        self.nodes_by_id = {}
+        self.edges = []
+        self.loaded = False
+        self.criterion = None
+        self.criterion_passed = False
+        self.root = None  # optional root Node if applicable
+
+    def _to_graph_json(self, run=None):
+        # Needs to be its own function for tests
+        return {
+            "format": self.format,
+            "nodes": [node.to_json() for node in self.nodes],
+            "edges": [edge.to_json() for edge in self.edges],
+        }
+
+    def bind_to_run(self, *args, **kwargs):
+        data = self._to_graph_json()
+        tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".graph.json")
+        data = _numpy_arrays_to_lists(data)
+        with codecs.open(tmp_path, "w", encoding="utf-8") as fp:
+            util.json_dump_safer(data, fp)
+        self._set_file(tmp_path, is_tmp=True, extension=".graph.json")
+        if self.is_bound():
+            return
+        super().bind_to_run(*args, **kwargs)
+
+    @classmethod
+    def get_media_subdir(cls):
+        return os.path.join("media", "graph")
+
+    def to_json(self, run):
+        json_dict = super().to_json(run)
+        json_dict["_type"] = self._log_type
+        return json_dict
+
+    def __getitem__(self, nid):
+        return self.nodes_by_id[nid]
+
+    def pprint(self):
+        for edge in self.edges:
+            pprint.pprint(edge.attributes)
+        for node in self.nodes:
+            pprint.pprint(node.attributes)
+
+    def add_node(self, node=None, **node_kwargs):
+        if node is None:
+            node = Node(**node_kwargs)
+        elif node_kwargs:
+            raise ValueError(
+                f"Only pass one of either node ({node}) or other keyword arguments ({node_kwargs})"
+            )
+        self.nodes.append(node)
+        self.nodes_by_id[node.id] = node
+
+        return node
+
+    def add_edge(self, from_node, to_node):
+        edge = Edge(from_node, to_node)
+        self.edges.append(edge)
+
+        return edge
+
+    @classmethod
+    def from_keras(cls, model):
+        # TODO: his method requires a refactor to work with the keras 3.
+        graph = cls()
+        # Shamelessly copied (then modified) from keras/keras/utils/layer_utils.py
+        sequential_like = cls._is_sequential(model)
+
+        relevant_nodes = None
+        if not sequential_like:
+            relevant_nodes = []
+            for v in model._nodes_by_depth.values():
+                relevant_nodes += v
+
+        layers = model.layers
+        for i in range(len(layers)):
+            node = Node.from_keras(layers[i])
+            if hasattr(layers[i], "_inbound_nodes"):
+                for in_node in layers[i]._inbound_nodes:
+                    if relevant_nodes and in_node not in relevant_nodes:
+                        # node is not part of the current network
+                        continue
+                    for in_layer in _nest(in_node.inbound_layers):
+                        inbound_keras_node = Node.from_keras(in_layer)
+
+                        if inbound_keras_node.id not in graph.nodes_by_id:
+                            graph.add_node(inbound_keras_node)
+                        inbound_node = graph.nodes_by_id[inbound_keras_node.id]
+
+                        graph.add_edge(inbound_node, node)
+            graph.add_node(node)
+        return graph
+
+    @classmethod
+    def _is_sequential(cls, model):
+        sequential_like = True
+
+        if (
+            model.__class__.__name__ != "Sequential"
+            and hasattr(model, "_is_graph_network")
+            and model._is_graph_network
+        ):
+            nodes_by_depth = model._nodes_by_depth.values()
+            nodes = []
+            for v in nodes_by_depth:
+                # TensorFlow2 doesn't insure inbound is always a list
+                inbound = v[0].inbound_layers
+                if not hasattr(inbound, "__len__"):
+                    inbound = [inbound]
+                if (len(v) > 1) or (len(v) == 1 and len(inbound) > 1):
+                    # if the model has multiple nodes
+                    # or if the nodes have multiple inbound_layers
+                    # the model is no longer sequential
+                    sequential_like = False
+                    break
+                nodes += v
+            if sequential_like:
+                # search for shared layers
+                for layer in model.layers:
+                    flag = False
+                    if hasattr(layer, "_inbound_nodes"):
+                        for node in layer._inbound_nodes:
+                            if node in nodes:
+                                if flag:
+                                    sequential_like = False
+                                    break
+                                else:
+                                    flag = True
+                    if not sequential_like:
+                        break
+        return sequential_like

wandb/sdk/data_types/helper_types/bounding_boxes_2d.py

@@ -0,0 +1,295 @@
+import numbers
+from typing import TYPE_CHECKING, Optional, Type, Union
+
+import wandb
+from wandb import util
+from wandb.util import has_num
+
+from ..base_types.json_metadata import JSONMetadata
+
+if TYPE_CHECKING:  # pragma: no cover
+    from wandb.sdk.artifacts.artifact import Artifact
+
+    from ...wandb_run import Run as LocalRun
+
+
+class BoundingBoxes2D(JSONMetadata):
+    """Format images with 2D bounding box overlays for logging to W&B.
+
+    Arguments:
+        val: (dictionary) A dictionary of the following form:
+            box_data: (list of dictionaries) One dictionary for each bounding box, containing:
+                position: (dictionary) the position and size of the bounding box, in one of two formats
+                    Note that boxes need not all use the same format.
+                    {"minX", "minY", "maxX", "maxY"}: (dictionary) A set of coordinates defining
+                        the upper and lower bounds of the box (the bottom left and top right corners)
+                    {"middle", "width", "height"}: (dictionary) A set of coordinates defining the
+                        center and dimensions of the box, with "middle" as a list [x, y] for the
+                        center point and "width" and "height" as numbers
+                domain: (string) One of two options for the bounding box coordinate domain
+                    null: By default, or if no argument is passed, the coordinate domain
+                        is assumed to be relative to the original image, expressing this box as a fraction
+                        or percentage of the original image. This means all coordinates and dimensions
+                        passed into the "position" argument are floating point numbers between 0 and 1.
+                    "pixel": (string literal) The coordinate domain is set to the pixel space. This means all
+                        coordinates and dimensions passed into "position" are integers within the bounds
+                        of the image dimensions.
+                class_id: (integer) The class label id for this box
+                scores: (dictionary of string to number, optional) A mapping of named fields
+                        to numerical values (float or int), can be used for filtering boxes in the UI
+                        based on a range of values for the corresponding field
+                box_caption: (string, optional) A string to be displayed as the label text above this
+                        box in the UI, often composed of the class label, class name, and/or scores
+
+            class_labels: (dictionary, optional) A map of integer class labels to their readable class names
+
+        key: (string)
+            The readable name or id for this set of bounding boxes (e.g. predictions, ground_truth)
+
+    Examples:
+        ### Log bounding boxes for a single image
+        <!--yeadoc-test:boundingbox-2d-->
+        ```python
+        import numpy as np
+        import wandb
+
+        wandb.init()
+        image = np.random.randint(low=0, high=256, size=(200, 300, 3))
+
+        class_labels = {0: "person", 1: "car", 2: "road", 3: "building"}
+
+        img = wandb.Image(
+            image,
+            boxes={
+                "predictions": {
+                    "box_data": [
+                        {
+                            # one box expressed in the default relative/fractional domain
+                            "position": {"minX": 0.1, "maxX": 0.2, "minY": 0.3, "maxY": 0.4},
+                            "class_id": 1,
+                            "box_caption": class_labels[1],
+                            "scores": {"acc": 0.2, "loss": 1.2},
+                        },
+                        {
+                            # another box expressed in the pixel domain
+                            "position": {"middle": [150, 20], "width": 68, "height": 112},
+                            "domain": "pixel",
+                            "class_id": 3,
+                            "box_caption": "a building",
+                            "scores": {"acc": 0.5, "loss": 0.7},
+                        },
+                        # Log as many boxes an as needed
+                    ],
+                    "class_labels": class_labels,
+                }
+            },
+        )
+
+        wandb.log({"driving_scene": img})
+        ```
+
+        ### Log a bounding box overlay to a Table
+        <!--yeadoc-test:bb2d-image-with-labels-->
+        ```python
+        import numpy as np
+        import wandb
+
+        wandb.init()
+        image = np.random.randint(low=0, high=256, size=(200, 300, 3))
+
+        class_labels = {0: "person", 1: "car", 2: "road", 3: "building"}
+
+        class_set = wandb.Classes(
+            [
+                {"name": "person", "id": 0},
+                {"name": "car", "id": 1},
+                {"name": "road", "id": 2},
+                {"name": "building", "id": 3},
+            ]
+        )
+
+        img = wandb.Image(
+            image,
+            boxes={
+                "predictions": {
+                    "box_data": [
+                        {
+                            # one box expressed in the default relative/fractional domain
+                            "position": {"minX": 0.1, "maxX": 0.2, "minY": 0.3, "maxY": 0.4},
+                            "class_id": 1,
+                            "box_caption": class_labels[1],
+                            "scores": {"acc": 0.2, "loss": 1.2},
+                        },
+                        {
+                            # another box expressed in the pixel domain
+                            "position": {"middle": [150, 20], "width": 68, "height": 112},
+                            "domain": "pixel",
+                            "class_id": 3,
+                            "box_caption": "a building",
+                            "scores": {"acc": 0.5, "loss": 0.7},
+                        },
+                        # Log as many boxes an as needed
+                    ],
+                    "class_labels": class_labels,
+                }
+            },
+            classes=class_set,
+        )
+
+        table = wandb.Table(columns=["image"])
+        table.add_data(img)
+        wandb.log({"driving_scene": table})
+        ```
+    """
+
+    _log_type = "bounding-boxes"
+    # TODO: when the change is made to have this produce a dict with a _type, define
+    # it here as _log_type, associate it in to_json
+
+    def __init__(self, val: dict, key: str) -> None:
+        """Initialize a BoundingBoxes object.
+
+        The input dictionary `val` should contain the keys:
+            box_data: a list of dictionaries, each of which describes a bounding box.
+            class_labels: (optional) A map of integer class labels to their readable
+                class names.
+
+        Each bounding box dictionary should contain the following keys:
+            position: (dictionary) the position and size of the bounding box.
+            domain: (string) One of two options for the bounding box coordinate domain.
+            class_id: (integer) The class label id for this box.
+            scores: (dictionary of string to number, optional) A mapping of named fields
+                to numerical values (float or int).
+            box_caption: (optional) The label text, often composed of the class label,
+                class name, and/or scores.
+
+        The position dictionary should be in one of two formats:
+            {"minX", "minY", "maxX", "maxY"}: (dictionary) A set of coordinates defining
+                the upper and lower bounds of the box (the bottom left and top right
+                corners).
+            {"middle", "width", "height"}: (dictionary) A set of coordinates defining
+                the center and dimensions of the box, with "middle" as a list [x, y] for
+                the center point and "width" and "height" as numbers.
+        Note that boxes need not all use the same format.
+
+        Args:
+            val: (dictionary) A dictionary containing the bounding box data.
+            key: (string) The readable name or id for this set of bounding boxes (e.g.
+                predictions, ground_truth)
+        """
+        super().__init__(val)
+        self._val = val["box_data"]
+        self._key = key
+        # Add default class mapping
+        if "class_labels" not in val:
+            np = util.get_module(
+                "numpy", required="Bounding box support requires numpy"
+            )
+            classes = (
+                np.unique(list(box["class_id"] for box in val["box_data"]))
+                .astype(np.int32)
+                .tolist()
+            )
+            class_labels = {c: "class_" + str(c) for c in classes}
+            self._class_labels = class_labels
+        else:
+            self._class_labels = val["class_labels"]
+
+    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)
+        run._add_singleton(
+            "bounding_box/class_labels",
+            str(key) + "_wandb_delimeter_" + self._key,
+            self._class_labels,
+        )
+
+    @classmethod
+    def type_name(cls) -> str:
+        return "boxes2D"
+
+    def validate(self, val: dict) -> bool:
+        # 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 string"
+                    )
+
+        boxes = val["box_data"]
+        if not isinstance(boxes, list):
+            raise TypeError("Boxes must be a list")
+
+        for box in boxes:
+            # Required arguments
+            error_str = (
+                "Each box must contain a position with: middle, width, and height or \
+                    \nminX, maxX, minY, maxY."
+            )
+            if "position" not in box:
+                raise TypeError(error_str)
+            else:
+                valid = False
+                if (
+                    "middle" in box["position"]
+                    and len(box["position"]["middle"]) == 2
+                    and has_num(box["position"], "width")
+                    and has_num(box["position"], "height")
+                ):
+                    valid = True
+                elif (
+                    has_num(box["position"], "minX")
+                    and has_num(box["position"], "maxX")
+                    and has_num(box["position"], "minY")
+                    and has_num(box["position"], "maxY")
+                ):
+                    valid = True
+
+                if not valid:
+                    raise TypeError(error_str)
+
+            # Optional arguments
+            if ("scores" in box) and not isinstance(box["scores"], dict):
+                raise TypeError("Box scores must be a dictionary")
+            elif "scores" in box:
+                for k, v in list(box["scores"].items()):
+                    if not isinstance(k, str):
+                        raise TypeError("A score key must be a string")
+                    if not isinstance(v, numbers.Number):
+                        raise TypeError("A score value must be a number")
+
+            if ("class_id" in box) and not isinstance(box["class_id"], int):
+                raise TypeError("A box's class_id must be an integer")
+
+            # Optional
+            if ("box_caption" in box) and not isinstance(box["box_caption"], str):
+                raise TypeError("A box's caption must be a string")
+        return True
+
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+        from wandb.sdk.wandb_run import Run
+
+        if isinstance(run_or_artifact, Run):
+            return super().to_json(run_or_artifact)
+        elif isinstance(run_or_artifact, wandb.Artifact):
+            # TODO (tim): I would like to log out a proper dictionary representing this object, but don't
+            # want to mess with the visualizations that are currently available in the UI. This really should output
+            # an object with a _type key. Will need to push this change to the UI first to ensure backwards compat
+            return self._val
+        else:
+            raise ValueError("to_json accepts wandb_run.Run or wandb.Artifact")
+
+    @classmethod
+    def from_json(
+        cls: Type["BoundingBoxes2D"], json_obj: dict, source_artifact: "Artifact"
+    ) -> "BoundingBoxes2D":
+        return cls({"box_data": json_obj}, "")