wandb 0.21.2__py3-none-macosx_12_0_arm64.whl

wandb/sdk/data_types/utils.py

@@ -0,0 +1,260 @@
+import datetime
+import logging
+import os
+from decimal import Decimal
+from typing import TYPE_CHECKING, Optional, Sequence, Union, cast
+
+import wandb
+from wandb import util
+
+from ..internal import incremental_table_util
+from .base_types.media import BatchableMedia, Media
+from .base_types.wb_value import WBValue
+from .image import _server_accepts_image_filenames
+from .plotly import Plotly
+
+if TYPE_CHECKING:  # pragma: no cover
+    import matplotlib  # type: ignore
+    import pandas as pd
+    import plotly  # type: ignore
+
+    from ..wandb_run import Run as LocalRun
+
+    ValToJsonType = Union[
+        dict,
+        "WBValue",
+        Sequence["WBValue"],
+        "plotly.Figure",
+        "matplotlib.artist.Artist",
+        "pd.DataFrame",
+        object,
+    ]
+
+
+def history_dict_to_json(
+    run: Optional["LocalRun"],
+    payload: dict,
+    step: Optional[int] = None,
+    ignore_copy_err: Optional[bool] = None,
+) -> dict:
+    # Converts a History row dict's elements so they're friendly for JSON serialization.
+
+    if step is None:
+        # We should be at the top level of the History row; assume this key is set.
+        step = payload["_step"]
+
+    # We use list here because we were still seeing cases of RuntimeError dict changed size
+    for key in list(payload):
+        val = payload[key]
+        if isinstance(val, dict):
+            payload[key] = history_dict_to_json(
+                run, val, step=step, ignore_copy_err=ignore_copy_err
+            )
+        else:
+            payload[key] = val_to_json(
+                run, key, val, namespace=step, ignore_copy_err=ignore_copy_err
+            )
+
+    return payload
+
+
+# TODO: refine this
+def val_to_json(
+    run: Optional["LocalRun"],
+    key: str,
+    val: "ValToJsonType",
+    namespace: Optional[Union[str, int]] = None,
+    ignore_copy_err: Optional[bool] = None,
+) -> Union[Sequence, dict]:
+    # Converts a wandb datatype to its JSON representation.
+    if namespace is None:
+        raise ValueError(
+            "val_to_json must be called with a namespace(a step number, or 'summary') argument"
+        )
+
+    converted = val
+
+    if isinstance(val, (int, float, str, bool)):
+        # These are already JSON-serializable,
+        # no need to do the expensive checks below.
+        return converted  # type: ignore[return-value]
+
+    typename = util.get_full_typename(val)
+
+    if util.is_pandas_data_frame(val):
+        val = wandb.Table(dataframe=val)
+
+    elif util.is_matplotlib_typename(typename) or util.is_plotly_typename(typename):
+        val = Plotly.make_plot_media(val)
+    elif isinstance(val, (list, tuple, range)) and all(
+        isinstance(v, WBValue) for v in val
+    ):
+        assert run
+        # This check will break down if Image/Audio/... have child classes.
+        if (
+            len(val)
+            and isinstance(val[0], BatchableMedia)
+            and all(isinstance(v, type(val[0])) for v in val)
+        ):
+            if TYPE_CHECKING:
+                val = cast(Sequence["BatchableMedia"], val)
+
+            items = _prune_max_seq(val)
+
+            if _server_accepts_image_filenames(run):
+                for item in items:
+                    item.bind_to_run(
+                        run=run,
+                        key=key,
+                        step=namespace,
+                        ignore_copy_err=ignore_copy_err,
+                    )
+            else:
+                for i, item in enumerate(items):
+                    item.bind_to_run(
+                        run=run,
+                        key=key,
+                        step=namespace,
+                        id_=i,
+                        ignore_copy_err=ignore_copy_err,
+                    )
+                if run._attach_id and run._init_pid != os.getpid():
+                    wandb.termwarn(
+                        f"Attempting to log a sequence of {items[0].__class__.__name__} objects from multiple processes might result in data loss. Please upgrade your wandb server",
+                        repeat=False,
+                    )
+
+            return items[0].seq_to_json(items, run, key, namespace)
+        else:
+            # TODO(adrian): Good idea to pass on the same key here? Maybe include
+            # the array index?
+            # There is a bug here: if this array contains two arrays of the same type of
+            # anonymous media objects, their eventual names will collide.
+            # This used to happen. The frontend doesn't handle heterogeneous arrays
+            # raise ValueError(
+            #    "Mixed media types in the same list aren't supported")
+            return [
+                val_to_json(
+                    run, key, v, namespace=namespace, ignore_copy_err=ignore_copy_err
+                )
+                for v in val
+            ]
+
+    if isinstance(val, WBValue):
+        assert run
+        if isinstance(val, Media) and not val.is_bound():
+            if hasattr(val, "_log_type") and val._log_type in [
+                "table",
+                "partitioned-table",
+                "joined-table",
+            ]:
+                _log_table_artifact(val, key, run)
+
+            # Partitioned tables and joined tables do not support being bound to runs.
+            if not (
+                hasattr(val, "_log_type")
+                and val._log_type in ["partitioned-table", "joined-table"]
+            ):
+                val.bind_to_run(run, key, namespace)
+
+        res = val.to_json(run)
+
+        if isinstance(val, wandb.Table) and val.log_mode == "INCREMENTAL":
+            # Set the _last_logged_idx AFTER the Table has been logged and
+            # bound to the run.
+            val._last_logged_idx = len(val.data) - 1
+        return res
+
+    return converted  # type: ignore
+
+
+def _log_table_artifact(val: "Media", key: str, run: "LocalRun") -> None:
+    """Log a table to the run based on the table type and logging mode.
+
+    Creates and logs a `run_table` type for Table, PartitionedTable, and
+    JoinedTable values. For tables with log_mode="INCREMENTAL", creates and
+    logs an incremental artifact of type `wandb-run-incremental-table.`
+
+    Args:
+        val: A wbvalue with log_type "table", "partitioned-table",
+            or "joined-table."
+        key: The key used to log val.
+        run: The LocalRun used to log val.
+    """
+    from wandb.sdk.artifacts._internal_artifact import InternalArtifact
+
+    if isinstance(val, wandb.Table) and val.log_mode == "INCREMENTAL":
+        if (
+            run.resumed
+            and val._previous_increments_paths is None
+            and val._increment_num is None
+        ):
+            val._load_incremental_table_state_from_resumed_run(run, key)
+        else:
+            val._set_incremental_table_run_target(run)
+        art = incremental_table_util.init_artifact(run, key)
+        entry_name = incremental_table_util.get_entry_name(val, key)
+    else:
+        art = InternalArtifact(f"run-{run.id}-{key}", "run_table")
+        entry_name = key
+
+    art.add(val, entry_name)
+    run.log_artifact(art)
+
+
+def _prune_max_seq(seq: Sequence["BatchableMedia"]) -> Sequence["BatchableMedia"]:
+    # If media type has a max respect it
+    items = seq
+    if hasattr(seq[0], "MAX_ITEMS") and seq[0].MAX_ITEMS < len(seq):
+        logging.warning(
+            f"Only {seq[0].MAX_ITEMS} {seq[0].__class__.__name__} will be uploaded."
+        )
+        items = seq[: seq[0].MAX_ITEMS]
+    return items
+
+
+def _json_helper(val, artifact):
+    if isinstance(val, WBValue):
+        return val.to_json(artifact)
+    elif val.__class__ is dict:
+        res = {}
+        for key in val:
+            res[key] = _json_helper(val[key], artifact)
+        return res
+
+    if hasattr(val, "tolist"):
+        py_val = val.tolist()
+        if val.__class__.__name__ == "datetime64" and isinstance(py_val, int):
+            # when numpy datetime64 .tolist() returns an int, it is nanoseconds.
+            # need to convert to milliseconds
+            return _json_helper(py_val / int(1e6), artifact)
+        return _json_helper(py_val, artifact)
+    elif hasattr(val, "item"):
+        return _json_helper(val.item(), artifact)
+
+    if isinstance(val, datetime.datetime):
+        if val.tzinfo is None:
+            val = datetime.datetime(
+                val.year,
+                val.month,
+                val.day,
+                val.hour,
+                val.minute,
+                val.second,
+                val.microsecond,
+                tzinfo=datetime.timezone.utc,
+            )
+        return int(val.timestamp() * 1000)
+    elif isinstance(val, datetime.date):
+        return int(
+            datetime.datetime(
+                val.year, val.month, val.day, tzinfo=datetime.timezone.utc
+            ).timestamp()
+            * 1000
+        )
+    elif isinstance(val, (list, tuple)):
+        return [_json_helper(i, artifact) for i in val]
+    elif isinstance(val, Decimal):
+        return float(val)
+    else:
+        return util.json_friendly(val)[0]

wandb/sdk/data_types/video.py

@@ -0,0 +1,303 @@
+import functools
+import logging
+import os
+import pathlib
+from io import BytesIO
+from typing import TYPE_CHECKING, Any, Literal, Optional, Sequence, Type, Union
+
+import wandb
+from wandb import env, util
+from wandb.sdk import wandb_setup
+from wandb.sdk.lib import filesystem, printer, printer_asyncio, 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
+
+    import numpy as np
+
+    from wandb.sdk.artifacts.artifact import Artifact
+
+    from ..wandb_run import Run as LocalRun
+
+
+def _should_print_spinner() -> bool:
+    settings = wandb_setup.singleton().settings_if_loaded
+    if settings and (settings.quiet or settings.silent):
+        return False
+
+    return not env.is_quiet() and not env.is_silent()
+
+
+# This helper function is a workaround for the issue discussed here:
+# https://github.com/wandb/wandb/issues/3472
+#
+# Essentially, the issue is that moviepy's write_gif function fails to close
+# the open write / file descriptor returned from `imageio.save`. The following
+# function is a simplified copy of the function in the moviepy source code.
+# See https://github.com/Zulko/moviepy/blob/7e3e8bb1b739eb6d1c0784b0cb2594b587b93b39/moviepy/video/io/gif_writers.py#L428
+#
+# Except, we close the writer!
+def write_gif_with_image_io(
+    clip: Any, filename: str, fps: Optional[int] = None
+) -> None:
+    from packaging.version import parse
+
+    imageio = util.get_module(
+        "imageio",
+        required='wandb.Video requires imageio when passing raw data. Install with "pip install wandb[media]"',
+    )
+
+    if parse(imageio.__version__) < parse("2.28.1"):
+        raise ValueError(
+            "imageio version 2.28.1 or higher is required to encode gifs. "
+            "Please upgrade imageio with `pip install imageio>=2.28.1`"
+        )
+
+    writer = imageio.save(
+        filename,
+        quantizer=0,
+        palettesize=256,
+        loop=0,
+        duration=1000 / clip.fps,
+    )
+
+    for frame in clip.iter_frames(fps=fps, dtype="uint8"):
+        writer.append_data(frame)
+
+    writer.close()
+
+
+class Video(BatchableMedia):
+    """A class for logging videos to W&B."""
+
+    _log_type = "video-file"
+    EXTS = ("gif", "mp4", "webm", "ogg")
+    _width: Optional[int]
+    _height: Optional[int]
+
+    def __init__(
+        self,
+        data_or_path: Union[str, pathlib.Path, "np.ndarray", "TextIO", "BytesIO"],
+        caption: Optional[str] = None,
+        fps: Optional[int] = None,
+        format: Optional[Literal["gif", "mp4", "webm", "ogg"]] = None,
+    ):
+        """Initialize a W&B Video object.
+
+        Args:
+            data_or_path: Video can be initialized with a path to a file or an io object.
+                Video can be initialized with a numpy tensor. The numpy tensor
+                must be either 4 dimensional or 5 dimensional.
+                The dimensions should be (number of frames, channel, height, width) or
+                (batch, number of frames, channel, height, width)
+                The format parameter must be specified with the format argument
+                when initializing with a numpy array
+                or io object.
+            caption: Caption associated with the video for display.
+            fps: The frame rate to use when encoding raw video frames.
+                Default value is 4.
+                This parameter has no effect when data_or_path is a string, or bytes.
+            format: Format of video, necessary if initializing with a numpy array
+                or io object. This parameter will be used to determine the format
+                to use when encoding the video data. Accepted values are "gif",
+                "mp4", "webm", or "ogg".
+                If no value is provided, the default format will be "gif".
+
+        Examples:
+        Log a numpy array as a video
+
+        ```python
+        import numpy as np
+        import wandb
+
+        with wandb.init() as run:
+            # axes are (number of frames, channel, height, width)
+            frames = np.random.randint(
+                low=0, high=256, size=(10, 3, 100, 100), dtype=np.uint8
+            )
+            run.log({"video": wandb.Video(frames, format="mp4", fps=4)})
+        ```
+        """
+        super().__init__(caption=caption)
+
+        if format is None:
+            wandb.termwarn(
+                "`format` argument was not provided, defaulting to `gif`. "
+                "This parameter will be required in v0.20.0, "
+                "please specify the format explicitly."
+            )
+        self._format = format or "gif"
+        self._width = None
+        self._height = None
+        self._channels = None
+        if self._format not in Video.EXTS:
+            raise ValueError(
+                "wandb.Video accepts {} formats".format(", ".join(Video.EXTS))
+            )
+
+        if isinstance(data_or_path, (BytesIO, str)) and fps:
+            msg = (
+                "`fps` argument does not affect the frame rate of the video "
+                "when providing a file path or raw bytes."
+            )
+            wandb.termwarn(msg)
+
+        if isinstance(data_or_path, BytesIO):
+            filename = os.path.join(
+                MEDIA_TMP.name, runid.generate_id() + "." + self._format
+            )
+            with open(filename, "wb") as f:
+                f.write(data_or_path.read())
+            self._set_file(filename, is_tmp=True)
+        elif isinstance(data_or_path, (str, pathlib.Path)):
+            data_or_path = str(data_or_path)
+
+            _, ext = os.path.splitext(data_or_path)
+            ext = ext[1:].lower()
+            if ext not in Video.EXTS:
+                raise ValueError(
+                    "wandb.Video accepts {} formats".format(", ".join(Video.EXTS))
+                )
+            self._set_file(data_or_path, is_tmp=False)
+            # ffprobe -v error -select_streams v:0 -show_entries stream=width,height -of csv=p=0 data_or_path
+        else:
+            if hasattr(data_or_path, "numpy"):  # TF data eager tensors
+                self.data = data_or_path.numpy()
+            elif util.is_numpy_array(data_or_path):
+                self.data = data_or_path
+            else:
+                raise ValueError(
+                    "wandb.Video accepts a file path or numpy like data as input"
+                )
+            fps = fps or 4
+
+            if _should_print_spinner():
+                printer_asyncio.run_async_with_spinner(
+                    printer.new_printer(),
+                    "Encoding video...",
+                    functools.partial(self.encode, fps=fps),
+                )
+            else:
+                self.encode(fps=fps)
+
+    def encode(self, fps: int = 4) -> None:
+        """Encode the video data to a file.
+
+        <!-- lazydoc-ignore: internal -->
+        """
+        # import ImageSequenceClip from the appropriate MoviePy module
+        mpy = util.get_module(
+            "moviepy.video.io.ImageSequenceClip",
+            required='wandb.Video requires moviepy when passing raw data. Install with "pip install wandb[media]"',
+        )
+
+        tensor = self._prepare_video(self.data)
+        _, self._height, self._width, self._channels = tensor.shape  # type: ignore
+
+        # encode sequence of images into gif string
+        clip = mpy.ImageSequenceClip(list(tensor), fps=fps)
+
+        filename = os.path.join(
+            MEDIA_TMP.name, runid.generate_id() + "." + self._format
+        )
+
+        if self._format == "gif":
+            write_gif_with_image_io(clip, filename)
+        else:
+            clip.write_videofile(filename, logger=None)
+
+        self._set_file(filename, is_tmp=True)
+
+    @classmethod
+    def get_media_subdir(cls: Type["Video"]) -> str:
+        """Get media subdirectory for video files.
+
+        <!-- lazydoc-ignore-classmethod: internal -->
+        """
+        return os.path.join("media", "videos")
+
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+        """Returns the JSON representation expected by the backend.
+
+        <!-- lazydoc-ignore: internal -->
+        """
+        json_dict = super().to_json(run_or_artifact)
+        json_dict["_type"] = self._log_type
+
+        if self._width is not None:
+            json_dict["width"] = self._width
+        if self._height is not None:
+            json_dict["height"] = self._height
+
+        return json_dict
+
+    def _prepare_video(self, video: "np.ndarray") -> "np.ndarray":
+        """This logic was mostly taken from tensorboardX."""
+        np = util.get_module(
+            "numpy",
+            required='wandb.Video requires numpy when passing raw data. To get it, run "pip install numpy".',
+        )
+        if video.ndim < 4:
+            raise ValueError(
+                "Video must be at least 4 dimensions: time, channels, height, width"
+            )
+        if video.ndim == 4:
+            video = video.reshape(1, *video.shape)
+        b, t, c, h, w = video.shape
+
+        if video.dtype != np.uint8:
+            logging.warning("Converting video data to uint8")
+            video = video.astype(np.uint8)
+
+        def is_power2(num: int) -> bool:
+            return num != 0 and ((num & (num - 1)) == 0)
+
+        # pad to nearest power of 2, all at once
+        if not is_power2(video.shape[0]):
+            len_addition = int(2 ** video.shape[0].bit_length() - video.shape[0])
+            video = np.concatenate(
+                (video, np.zeros(shape=(len_addition, t, c, h, w))), axis=0
+            )
+
+        n_rows = 2 ** ((b.bit_length() - 1) // 2)
+        n_cols = video.shape[0] // n_rows
+
+        video = video.reshape(n_rows, n_cols, t, c, h, w)
+        video = np.transpose(video, axes=(2, 0, 4, 1, 5, 3))
+        video = video.reshape(t, n_rows * h, n_cols * w, c)
+        return video
+
+    @classmethod
+    def seq_to_json(
+        cls: Type["Video"],
+        seq: Sequence["BatchableMedia"],
+        run: "LocalRun",
+        key: str,
+        step: Union[int, str],
+    ) -> dict:
+        """Convert a sequence of Video objects to a JSON representation.
+
+        <!-- lazydoc-ignore-classmethod: internal -->
+        """
+        base_path = os.path.join(run.dir, cls.get_media_subdir())
+        filesystem.mkdir_exists_ok(base_path)
+
+        meta = {
+            "_type": "videos",
+            "count": len(seq),
+            "videos": [v.to_json(run) for v in seq],
+            "captions": Video.captions(seq),
+        }
+        return meta
+
+
+class _VideoFileType(_dtypes.Type):
+    name = "video-file"
+    types = [Video]
+
+
+_dtypes.TypeRegistry.add(_VideoFileType)