truefoundry 0.3.3__py3-none-any.whl → 0.4.0.dev1__py3-none-any.whl

truefoundry/ml/log_types/artifacts/utils.py

@@ -0,0 +1,161 @@
+import json
+import logging
+import os
+import posixpath
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, Optional, Sequence, Tuple, Union
+
+from truefoundry.ml.exceptions import MlFoundryException
+from truefoundry.ml.log_types.artifacts.constants import DESCRIPTION_MAX_LENGTH
+
+logger = logging.getLogger("truefoundry.ml")
+
+
+def _copy_tree(src_path, dest_path, symlinks=False, ignore_dangling_symlinks=False):
+    os.makedirs(dest_path, exist_ok=True)
+    for item in os.listdir(src_path):
+        src = os.path.join(src_path, item)
+        dest = os.path.join(dest_path, item)
+        if os.path.isdir(src):
+            _copy_tree(
+                src,
+                dest,
+                symlinks=symlinks,
+                ignore_dangling_symlinks=ignore_dangling_symlinks,
+            )
+        else:
+            shutil.copy2(src, dest, follow_symlinks=True)
+
+
+def is_destination_path_dirlike(dest_path) -> bool:
+    if not dest_path:
+        return True
+
+    if dest_path.endswith(os.sep) or dest_path.endswith(posixpath.sep):
+        return True
+
+    if os.path.exists(dest_path) and os.path.isdir(dest_path):
+        return True
+
+    return False
+
+
+def _copy_additional_files(
+    root_dir: str,
+    files_dir: str,  # relative to root dir e.g. "files/"
+    model_dir: Optional[str],  # relative to files_dir e.g "model/"
+    additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]],
+    ignore_model_dir_dest_conflict: bool = False,
+):
+    """
+
+    File copying examples:
+        # non ambiguous
+        # a.txt -> /tmp/                            result /tmp/a.txt
+        # a.txt -> /tmp/a/                          result /tmp/a/a.txt
+        # a.txt -> /tmp/a/b/c/d.txt                 result /tmp/a/b/c/d.txt
+        # .gitignore -> /tmp/.gitignore             result /tmp/.gitignore
+
+        # ambiguous but destination directory exists
+        # a.txt -> /tmp                             result /tmp/a.txt
+        # a.txt -> /tmp/a (and /tmp/a/ exists)      result /tmp/a/a.txt
+
+        # ambiguous - when the destination can't be reliably distinguished as a directory
+        # a -> /tmp/a                                result /tmp/a
+        # a -> /tmp/b                                result /tmp/b
+        # a -> /tmp/a.txt                            result /tmp/a.txt
+        # .gitignore -> /tmp/.gitinclude             result /tmp/.gitinclude
+        # a.txt -> /tmp/a                            result /tmp/a
+    """
+    for src_path, dest_path in additional_files:
+        src_path = str(src_path)
+        if not os.path.exists(src_path):
+            raise MlFoundryException(
+                f"Source path {src_path!r} in `additional_files` does not exist."
+            )
+        dest_path = dest_path or ""
+        normalized_path = os.path.normpath(dest_path)
+        if dest_path.endswith(os.sep) or dest_path.endswith(posixpath.sep):
+            normalized_path += os.sep
+        dest_path = normalized_path.lstrip(os.sep)
+
+        if (
+            model_dir
+            and dest_path.startswith(model_dir)
+            and not ignore_model_dir_dest_conflict
+        ):
+            logger.warning(
+                f"Destination path {dest_path!r} in `additional_files` conflicts with "
+                f"reserved path {model_dir!r}/ which is being used to store the model. "
+                f"This might cause errors"
+            )
+
+        files_abs_dir = os.path.join(root_dir, files_dir)
+        dest_abs_path = os.path.join(files_abs_dir, dest_path)
+
+        if os.path.isfile(src_path):
+            _src = src_path
+            if is_destination_path_dirlike(dest_abs_path):
+                os.makedirs(dest_abs_path, exist_ok=True)
+                _dst = os.path.relpath(
+                    os.path.join(dest_abs_path, os.path.basename(_src)), files_abs_dir
+                )
+            else:
+                os.makedirs(os.path.dirname(dest_abs_path), exist_ok=True)
+                _dst = os.path.relpath(dest_abs_path, files_abs_dir)
+            logger.info(f"Adding file {_src} as /{_dst}")
+            shutil.copy2(src_path, dest_abs_path, follow_symlinks=True)
+        elif os.path.isdir(src_path):
+            os.makedirs(dest_abs_path, exist_ok=True)
+            _src = src_path.rstrip("/")
+            _dst = os.path.relpath(dest_abs_path, files_abs_dir).rstrip("/")
+            logger.info(f"Adding contents of {_src}/ to /{_dst}/")
+            _copy_tree(
+                src_path=src_path,
+                dest_path=dest_abs_path,
+                symlinks=True,
+                ignore_dangling_symlinks=False,
+            )
+
+
+def _validate_description(description: Optional[str]):
+    if description is not None:
+        if not isinstance(description, str):
+            raise MlFoundryException(
+                "`description` must be either `None` or type `str`"
+            )
+        if len(description) > DESCRIPTION_MAX_LENGTH:
+            raise MlFoundryException(
+                f"`description` cannot be longer than {DESCRIPTION_MAX_LENGTH} characters"
+            )
+
+
+def _validate_artifact_metadata(metadata: Dict[str, Any]):
+    if not isinstance(metadata, dict):
+        raise MlFoundryException("`metadata` must be json serializable dict")
+    try:
+        json.dumps(metadata)
+    except ValueError as ve:
+        raise MlFoundryException("`metadata` must be json serializable dict") from ve
+
+
+def calculate_local_directory_size(
+    directory: tempfile.TemporaryDirectory,  # type: ignore[type-arg]
+):
+    """
+    Tells about the size of the artifact
+
+    Args:
+        directory (str): directory path
+
+    Returns:
+        total size of the artifact
+    """
+    total_size = 0
+    for path, _dirs, files in os.walk(directory.name):
+        for f in files:
+            file_path = os.path.join(path, f)
+            total_size += os.stat(file_path).st_size
+    return total_size

truefoundry/ml/log_types/image/__init__.py

@@ -0,0 +1,3 @@
+from truefoundry.ml.log_types.image.image import Image
+
+__all__ = ["Image"]

truefoundry/ml/log_types/image/constants.py

@@ -0,0 +1,8 @@
+import re
+
+IMAGE_METADATA_FILE_NAME = "image_metadata.json"
+DEFAULT_IMAGE_FORMAT = "png"
+MISSING_PILLOW_PACKAGE_MESSAGE = (
+    "We need PIL package to save image.\nTo install, run `pip install pillow`"
+)
+IMAGE_KEY_REGEX = re.compile(r"^[a-zA-Z0-9-_]+$")

truefoundry/ml/log_types/image/image.py

@@ -0,0 +1,358 @@
+import json
+import os
+import posixpath
+import shutil
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+import numpy as np
+
+from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
+    ArtifactType,
+)
+from truefoundry.ml.exceptions import MlFoundryException
+from truefoundry.ml.log_types.artifacts.artifact import (
+    ArtifactVersionInternalMetadata,
+    _log_artifact_version_helper,
+)
+from truefoundry.ml.log_types.artifacts.constants import (
+    FILES_DIR,
+    INTERNAL_METADATA_PATH,
+)
+from truefoundry.ml.log_types.image.constants import (
+    DEFAULT_IMAGE_FORMAT,
+    IMAGE_KEY_REGEX,
+    IMAGE_METADATA_FILE_NAME,
+    MISSING_PILLOW_PACKAGE_MESSAGE,
+)
+from truefoundry.ml.log_types.image.image_normalizer import normalize_image
+from truefoundry.ml.log_types.image.types import BoundingBoxGroups, ClassGroups
+from truefoundry.ml.log_types.pydantic_base import PydanticBase
+from truefoundry.ml.run_utils import get_module
+
+if TYPE_CHECKING:
+    # noinspection PyUnresolvedReferences
+    import PIL.Image
+    from numpy.typing import NDArray
+
+    from truefoundry.ml.mlfoundry_run import MlFoundryRun
+
+ImageLikeNDArray = Union[
+    "NDArray[np.int8]",
+    "NDArray[np.uint8]",
+    "NDArray[np.float32]",
+    "NDArray[np.float64]",
+    "NDArray[np.bool_]",
+]
+DataOrPathType = Union[str, Path, ImageLikeNDArray, "PIL.Image.Image"]
+ClassGroupsType = Dict[str, Union[str, List[str]]]
+BBoxGrouptype = Dict[str, List[Dict[str, Any]]]
+
+
+def validate_key_name(key: str):
+    if not key or not IMAGE_KEY_REGEX.match(key):
+        raise MlFoundryException(
+            f"Invalid run image key: {key} should only contain alphanumeric, hyphen or underscore"
+        )
+
+
+class ImageVersionInternalMetadata(ArtifactVersionInternalMetadata):
+    image_file: str
+    image_metadata_file: str
+
+
+class Image:
+    def __init__(
+        self,
+        data_or_path: DataOrPathType,
+        caption: Optional[str] = None,
+        class_groups: Optional[ClassGroupsType] = None,
+        bbox_groups: Optional[BBoxGrouptype] = None,
+    ):
+        """Represent and log image using this class in `mlfoundry`.
+
+        You can initialize `truefoundry.ml.Image` by either by using a local path
+        or you can use a numpy array / PIL.Image object.
+
+        If you are using numpy array, we only support the following data types,
+        - bool
+        - integer [0 - 255]
+        - unsigned integer [0 - 255]
+        - float [0.0 - 1.0]
+
+        Any out of range value will be clipped.
+
+        As for array shape/dim, we follow the following structures,
+        - H x W (Grayscale)
+        - H x W x 1 (Grayscale)
+        - H x W x 3 (an RGB channel order is assumed)
+        - H x W x 4 (an RGBA channel order is assumed)
+
+        `PIL` package is required to log images. To install the `PIL` package,
+        run `pip install pillow`.
+
+        We can also log class names and bounding boxes associated with the image.
+        Class names and bounding boxes should be always grouped under `actuals` or
+        `predictions`. For example, if we have an image where the ground truth class
+        is "cat" and predicted class is "dog", we can represent it like,
+
+        ```python
+        from truefoundry.ml import Image
+
+        Image(
+            data_or_path=imarray,
+            class_groups={"actuals": "dog", "predictions": "cat"}
+        )
+        ```
+
+        You can define a bounding box using the following dictionary structure,
+        ```python
+        {
+            "position": {"min_x": 15, "min_y": 5, "max_x": 20, "max_y": 30}, # required, defines the position of the bounding box
+                                                                             # (min_x, min_y) defines the top left and
+                                                                             # (max_x, max_y) defines the bottom right corner of the box.
+            "class_name": "dog", # required, the class name of the bounding box
+            "caption": "dog", # optional, the caption of the bounding box.
+                              # If not passed, the class name is set as caption.
+        }
+        ```
+
+        Args:
+            data_or_path (Union[str, Path, "numpy.ndarray", "PIL.Image.Image"]):
+                Either the local path or the image object (Numpy array or PIL Image).
+            caption (Optional[str], optional): A string caption or label for the image.
+            class_groups (Optional[Dict[str, Union[str, List[str]]]], optional):
+                Class names associated with the image. Expects class name(s) grouped by
+                `predictions` or `actuals`.
+            bbox_groups (Optional[Dict[str, List[Dict]]], optional): Bounding boxes
+                associated with the image. Expects bounding boxes grouped by `predictions`
+                or `actuals`.
+
+        Examples:
+        ### Logging images with caption and class names
+
+        ```python
+        from truefoundry.ml import get_client, Image
+        import numpy as np
+
+        client = get_client()
+        run = client.create_run(
+            ml_repo="my-classification-project",
+        )
+
+        imarray = np.random.randint(low=0, high=256, size=(100, 100, 3))
+
+        images_to_log = {
+            "logged-image-array": Image(
+                data_or_path=imarray,
+                caption="testing image logging",
+                class_groups={"actuals": "dog", "predictions": "cat"},
+            ),
+        }
+        run.log_images(images_to_log, step=1)
+
+        run.end()
+        ```
+
+        ### Logging images for a multi-label classification problem
+
+        ```python
+        from truefoundry.ml import Image
+
+        images_to_log = {
+            "logged-image-array": Image(
+                data_or_path=imarray,
+                caption="testing image logging",
+                class_groups={"actuals": ["dog", "human"], "predictions": ["cat", "human"]},
+            ),
+        }
+
+        run.log_images(images_to_log, step=1)
+        ```
+
+        ### Logging images with bounding boxes
+
+        ```python
+        from truefoundry.ml import Image
+
+        images_to_log = {
+            "logged-image-array": Image(
+                data_or_path=imarray,
+                caption="testing image logging",
+                bbox_groups={
+                    "predictions": [
+                        {
+                            "position": {"min_x": 5, "min_y": 5, "max_x": 20, "max_y": 30},
+                            "class_name": "cat",
+                        }
+                    ],
+                    "actuals": [
+                        {
+                            "position": {"min_x": 15, "min_y": 5, "max_x": 20, "max_y": 30},
+                            "class_name": "dog",
+                            "caption": "dog",
+                        }
+                    ],
+                },
+            ),
+        }
+
+        run.log_images(images_to_log, step=1)
+        ```
+        """
+
+        self._caption: Optional[str] = None
+        self._class_groups: Optional[ClassGroups] = None
+        self._bbox_groups: Optional[BoundingBoxGroups] = None
+
+        self._image: Optional["PIL.Image.Image"] = None
+        self._image_artifact_path = None
+        self._local_image_path: Optional[str] = None
+
+        self._init_image(data_or_path)
+        self._init_metadata(
+            caption=caption, class_groups=class_groups, bbox_groups=bbox_groups
+        )
+
+    @property
+    def image(self) -> "PIL.Image.Image":
+        if self._image is None:
+            raise MlFoundryException("Image is not initialized")
+        return self._image
+
+    @property
+    def caption(self) -> Optional[str]:
+        return self._caption
+
+    @property
+    def class_groups(self) -> Optional[ClassGroups]:
+        return self._class_groups
+
+    @property
+    def bbox_groups(self) -> Optional[BoundingBoxGroups]:
+        return self._bbox_groups
+
+    def save_image_to_local_dir(self, local_dir) -> str:
+        if self._local_image_path is None:
+            image_file_name = self._serialize_and_save_image_as_artifact(local_dir)
+        else:
+            image_file_name = self._copy_local_image_as_artifact(local_dir)
+
+        return image_file_name
+
+    def _serialize_and_save_image_as_artifact(self, local_dir: str) -> str:
+        file_name = f"image.{DEFAULT_IMAGE_FORMAT}"
+        local_path = os.path.join(local_dir, file_name)
+        self.image.save(local_path)
+
+        return file_name
+
+    def _copy_local_image_as_artifact(self, local_dir: str) -> str:
+        assert self._local_image_path is not None
+        file_name = os.path.basename(self._local_image_path)
+        new_local_image_path = os.path.join(local_dir, file_name)
+        shutil.copy2(self._local_image_path, new_local_image_path)
+        return file_name
+
+    def _init_image(self, data_or_path: DataOrPathType):
+        pil_image_module = get_module(
+            module_name="PIL.Image",
+            required=True,
+            error_message=MISSING_PILLOW_PACKAGE_MESSAGE,
+        )
+        if isinstance(data_or_path, (str, Path)):
+            self._local_image_path = os.path.abspath(data_or_path)
+            with pil_image_module.open(data_or_path) as image:
+                image.load()
+                self._image = image
+        else:
+            self._image = normalize_image(data_or_path)
+
+    def _init_metadata(
+        self,
+        caption: Optional[str],
+        class_groups: Optional[ClassGroupsType],
+        bbox_groups: Optional[Dict[str, Any]],
+    ):
+        if caption is not None:
+            self._caption = str(caption)
+
+        if class_groups:
+            self._class_groups = ClassGroups.parse_obj(class_groups)
+
+        if bbox_groups:
+            self._bbox_groups = BoundingBoxGroups.parse_obj(bbox_groups)
+
+    def _to_dict(self) -> Dict[str, Any]:
+        dict_ = {}
+        dict_["caption"] = self._caption
+        dict_["class_groups"] = (
+            self.class_groups.to_dict() if self.class_groups is not None else None
+        )
+        dict_["bbox_groups"] = (
+            self.bbox_groups.to_dict() if self.bbox_groups is not None else None
+        )
+
+        return dict_
+
+    def save(
+        self,
+        run: "MlFoundryRun",
+        key: str,
+        step: int = 0,
+    ):
+        validate_key_name(key)
+
+        # creating temp dir, files directory, .truefoundry directory
+        temp_dir = tempfile.TemporaryDirectory(prefix="truefoundry-")
+
+        local_internal_metadata_path = os.path.join(
+            temp_dir.name, INTERNAL_METADATA_PATH
+        )
+        os.makedirs(os.path.dirname(local_internal_metadata_path), exist_ok=True)
+
+        local_files_dir = os.path.join(temp_dir.name, FILES_DIR)
+        os.makedirs(local_files_dir, exist_ok=True)
+
+        try:
+            # saving the image file
+            image_file_name = self.save_image_to_local_dir(local_files_dir)
+
+            # saving the image_metadata
+            local_image_metadata_file_path = os.path.join(
+                local_files_dir, IMAGE_METADATA_FILE_NAME
+            )
+            with open(local_image_metadata_file_path, "w") as fp:
+                fp.write(ImageRunLogType(value=self._to_dict()).json())
+
+            # creating and saving the internal_metadata file
+            internal_metadata = ImageVersionInternalMetadata(
+                files_dir=FILES_DIR,
+                # joining with posixpath for avoiding backslash in case logged from windows machine
+                image_file=posixpath.join(FILES_DIR, image_file_name),
+                image_metadata_file=posixpath.join(FILES_DIR, IMAGE_METADATA_FILE_NAME),
+            )
+            with open(local_internal_metadata_path, "w") as f:
+                json.dump(internal_metadata.dict(), f)
+
+        except Exception as e:
+            temp_dir.cleanup()
+            raise MlFoundryException("Failed to log Image") from e
+
+        _log_artifact_version_helper(
+            run=run,
+            name=key,
+            artifact_type=ArtifactType.IMAGE,
+            artifact_dir=temp_dir,
+            internal_metadata=internal_metadata,
+            step=step,
+        )
+
+
+class ImageRunLogType(PydanticBase):
+    value: Dict[str, Any]
+
+    @staticmethod
+    def get_log_type() -> str:
+        return "image"

truefoundry/ml/log_types/image/image_normalizer.py

@@ -0,0 +1,101 @@
+# The logic is copied from
+# https://github.com/truefoundry/mlfoundry-server/blob/c2334ace8f35322cd5e50c275b8df08327688c01/mlflow/tracking/client.py#L1188
+from typing import TYPE_CHECKING, Union
+
+import numpy as np
+
+from truefoundry.ml.exceptions import MlFoundryException
+from truefoundry.ml.log_types.image.constants import MISSING_PILLOW_PACKAGE_MESSAGE
+from truefoundry.ml.logger import logger
+from truefoundry.ml.run_utils import get_module
+
+if TYPE_CHECKING:
+    import PIL.Image
+    from numpy.typing import NDArray
+
+ImageLikeNDArray = Union[
+    "NDArray[np.int8]",
+    "NDArray[np.uint8]",
+    "NDArray[np.float32]",
+    "NDArray[np.float64]",
+    "NDArray[np.bool_]",
+]
+
+
+def _normalize_to_uint8(x):
+    # Based on: https://github.com/matplotlib/matplotlib/blob/06567e021f21be046b6d6dcf00380c1cb9adaf3c/lib/matplotlib/image.py#L684
+
+    is_int = np.issubdtype(x.dtype, np.integer)
+    low = 0
+    high = 255 if is_int else 1
+    if x.min() < low or x.max() > high:
+        msg = (
+            "Out-of-range values are detected. "
+            "Clipping array (dtype: '{}') to [{}, {}]".format(x.dtype, low, high)
+        )
+        logger.warning(msg)
+        x = np.clip(x, low, high)
+
+    # float or bool
+    if not is_int:
+        x = x * 255
+
+    return x.astype(np.uint8)
+
+
+def _convert_numpy_to_pil_image(image: ImageLikeNDArray) -> "PIL.Image.Image":
+    pil_image_module = get_module(
+        module_name="PIL.Image",
+        required=True,
+        error_message=MISSING_PILLOW_PACKAGE_MESSAGE,
+    )
+    valid_data_types = {
+        "b": "bool",
+        "i": "signed integer",
+        "u": "unsigned integer",
+        "f": "floating",
+    }
+
+    if image.dtype.kind not in valid_data_types.keys():
+        raise TypeError(
+            "Invalid array data type: '{}'. Must be one of {}".format(
+                image.dtype, list(valid_data_types.values())
+            )
+        )
+
+    if image.ndim not in [2, 3]:
+        raise ValueError(
+            "`image` must be a 2D or 3D array but got a {}D array".format(image.ndim)
+        )
+
+    if (image.ndim == 3) and (image.shape[2] not in [1, 3, 4]):
+        raise ValueError(
+            "Invalid channel length: {}. Must be one of [1, 3, 4]".format(
+                image.shape[2]
+            )
+        )
+
+    # squeeze a 3D grayscale image since `Image.fromarray` doesn't accept it.
+    if image.ndim == 3 and image.shape[2] == 1:
+        image = image[:, :, 0]
+
+    image = _normalize_to_uint8(image)
+
+    return pil_image_module.fromarray(image)
+
+
+def normalize_image(
+    image: Union["PIL.Image.Image", ImageLikeNDArray],
+) -> "PIL.Image.Image":
+    pil_image_module = get_module(
+        module_name="PIL.Image",
+        required=True,
+        error_message=MISSING_PILLOW_PACKAGE_MESSAGE,
+    )
+    if isinstance(image, pil_image_module.Image):
+        return image
+    if isinstance(image, np.ndarray):
+        return _convert_numpy_to_pil_image(image)
+    raise MlFoundryException(
+        f"image should be of type PIL Image/np.ndarray got type {type(image)}"
+    )

truefoundry/ml/log_types/image/types.py

@@ -0,0 +1,68 @@
+import enum
+from typing import Dict, List, Optional
+
+from truefoundry.pydantic_v1 import BaseModel, NonEmptyStr, root_validator
+
+
+class Group(enum.Enum):
+    ACTUALS = "actuals"
+    PREDICTIONS = "predictions"
+
+
+class Position(BaseModel):
+    min_x: float
+    min_y: float
+    max_x: float
+    max_y: float
+
+    class Config:
+        allow_mutation = False
+
+
+class BoundingBox(BaseModel):
+    position: Position
+    class_name: NonEmptyStr
+    caption: Optional[str] = None
+
+    class Config:
+        allow_mutation = False
+
+    @root_validator(pre=True)
+    def set_caption_if_not_passed(cls, values):
+        if not values.get("caption"):
+            values["caption"] = values.get("class_name")
+
+        return values
+
+
+class BoundingBoxGroups(BaseModel):
+    __root__: Dict[Group, List[BoundingBox]]
+
+    class Config:
+        allow_mutation = False
+        use_enum_values = True
+
+    def to_dict(self):
+        return self.dict()["__root__"]
+
+
+class ClassGroups(BaseModel):
+    __root__: Dict[Group, List[NonEmptyStr]]
+
+    class Config:
+        allow_mutation = False
+        use_enum_values = True
+
+    @root_validator(pre=True)
+    def transform_class_to_classes(cls, values):
+        values = values["__root__"]
+        processed_values = {}
+        for group, classes in values.items():
+            if isinstance(classes, str):
+                processed_values[group] = [classes]
+            else:
+                processed_values[group] = classes
+        return {"__root__": processed_values}
+
+    def to_dict(self):
+        return self.dict()["__root__"]