Repository: mlflow/mlflow

Issue:
[FR] Add APIs for loading text, JSON, and images to mlflow.artifacts
### Willingness to contribute

No. I cannot contribute this feature at this time.

### Proposal Summary

## MLflow Roadmap Item

This is an MLflow Roadmap item that has been prioritized by the MLflow maintainers. We’ve identified this feature as a highly requested addition to the MLflow package based on community feedback. 
We're seeking a community contribution for the implementation of this feature and will enthusiastically support the development and review of a submitted PR for this.

## Feature scope

This roadmap feature’s complexity is classified as:

- [X] `good-first-issue`: This feature is limited in complexity and effort required to implement.
- [ ] `simple`: This feature does not require a large amount of effort to implement and / or is clear enough to not need a design discussion with maintainers.
- [ ] `involved`: This feature will require a substantial amount of development effort but does not require an agreed-upon design from the maintainers. The feedback given during the PR phase may be involved and necessitate multiple iterations before approval. (Please bear with us as we collaborate with you to make a great contribution)
- [ ] `design-recommended`: This is a substantial feature that should have a design document approved prior to working on an implementation (to save your time, not ours). After agreeing to work on this feature, a maintainer will be assigned to support you throughout the development process.

## Proposal Summary

Introduce the following APIs to MLflow:

1. `mlflow.artifacts.load_text()`: Takes a URI referring to a text file as a parameter, downloads the file, and loads the file's contents as a string that is returned to the caller.

2. `mlflow.artifacts.load_json()`: Takes a URI referring to a JSON file as a parameter, downloads the file, and loads the file's contents as a dictionary that is returned to the caller via JSON parsing, throwing if the file cannot be parsed as JSON.

3. `mlflow.artifacts.load_image()`: Takes a URI referring to an image file as a parameter, downloads the file, and loads the file's contents as an image using Pillow's `Image.load()` function (https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.load), throwing if the file cannot be loaded as an image. The Pillow image object should be returned to the caller; if the function is called in a Jupyter notebook environment, it should also be displayed automatically.

### Motivation

> #### What is the use case for this feature?

Provide users a way to easily load common MLflow artifacts as Python objects.

> #### Why is this use case valuable to support for MLflow users in general?

Many users log text, JSON, and images, and they need to load them later for further analysis / processing.

> #### Why is it currently difficult to achieve this use case?

MLflow's `mlflow.artifacts.download_artifacts()` API (https://mlflow.org/docs/latest/python_api/mlflow.artifacts.html#mlflow.artifacts.download_artifacts) only downloads file contents to disk; it is up to users to parse them.

### Details

_No response_

### What component(s) does this bug affect?

- [X] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
@dbczumar I would like to work on this :) Please assign it to me.
@subramaniam02 Fantastic! I've assigned this to you. Looking forward to your pull request! Let me know if you have any questions.

Resolving patch:
diff --git a/mlflow/artifacts/__init__.py b/mlflow/artifacts/__init__.py
--- a/mlflow/artifacts/__init__.py
+++ b/mlflow/artifacts/__init__.py
@@ -1,11 +1,13 @@
 """
 APIs for interacting with artifacts in MLflow
 """
+import json
 import pathlib
+import tempfile
 from typing import Optional
 
 from mlflow.exceptions import MlflowException
-from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, BAD_REQUEST
 from mlflow.tracking import _get_store
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri, get_artifact_repository
 
@@ -57,3 +59,108 @@ def download_artifacts(
     artifact_repo = get_artifact_repository(artifact_uri)
     artifact_location = artifact_repo.download_artifacts(artifact_path, dst_path=dst_path)
     return artifact_location
+
+
+def load_text(artifact_uri: str) -> str:
+    """
+    Loads the artifact contents as a string.
+
+    :param artifact_uri: artifact location
+    :return: str
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+        with mlflow.start_run() as run:
+            artifact_uri = run.info.artifact_uri
+            mlflow.log_text("This is a sentence", "file.txt")
+            file_content = mlflow.artifacts.load_text(artifact_uri + "/file.txt")
+            print(file_content)
+
+    .. code-block:: text
+        :caption: Output
+
+        This is a sentence
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        local_artifact = download_artifacts(artifact_uri, dst_path=tmpdir)
+        with open(local_artifact) as local_artifact_fd:
+            try:
+                return str(local_artifact_fd.read())
+            except Exception:
+                raise MlflowException("Unable to form a str object from file content", BAD_REQUEST)
+
+
+def load_dict(artifact_uri: str) -> dict:
+    """
+    Loads the artifact contents as a dictionary.
+
+    :param artifact_uri: artifact location
+    :return: dict
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+        with mlflow.start_run() as run:
+            artifact_uri = run.info.artifact_uri
+            mlflow.log_dict({"mlflow-version": "0.28", "n_cores": "10"}, "config.json")
+            config_json = mlflow.artifacts.load_dict(artifact_uri + "/config.json")
+            print(config_json)
+
+    .. code-block:: text
+        :caption: Output
+
+        {'mlflow-version': '0.28', 'n_cores': '10'}
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        local_artifact = download_artifacts(artifact_uri, dst_path=tmpdir)
+        with open(local_artifact) as local_artifact_fd:
+            try:
+                return json.load(local_artifact_fd)
+            except json.JSONDecodeError:
+                raise MlflowException("Unable to form a JSON object from file content", BAD_REQUEST)
+
+
+def load_image(artifact_uri: str):
+    """
+    Loads artifact contents as a ``PIL.Image.Image`` object
+
+    :param artifact_uri: artifact location
+    :return: PIL.Image
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+        from PIL import Image
+        with mlflow.start_run() as run:
+            image = Image.new("RGB", (100, 100))
+            artifact_uri = run.info.artifact_uri
+            mlflow.log_image(image, "image.png")
+            image  = mlflow.artifacts.load_image(artifact_uri + "/image.png")
+            print(image)
+
+    .. code-block:: text
+        :caption: Output
+
+        <PIL.PngImagePlugin.PngImageFile image mode=RGB size=100x100 at 0x11D2FA3D0>
+    """
+    try:
+        from PIL import Image
+    except ImportError as exc:
+        raise ImportError(
+            "`load_image` requires Pillow. Please install it via: pip install Pillow"
+        ) from exc
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        local_artifact = download_artifacts(artifact_uri, dst_path=tmpdir)
+        try:
+            image_obj = Image.open(local_artifact)
+            image_obj.load()
+            return image_obj
+        except Exception:
+            raise MlflowException(
+                "Unable to form a PIL Image object from file content", BAD_REQUEST
+            )


Test patch:
diff --git a/tests/artifacts/test_artifacts.py b/tests/artifacts/test_artifacts.py
--- a/tests/artifacts/test_artifacts.py
+++ b/tests/artifacts/test_artifacts.py
@@ -4,6 +4,9 @@
 
 import mlflow
 from mlflow.exceptions import MlflowException
+from collections import namedtuple
+
+Artifact = namedtuple("Artifact", ["uri", "content"])
 
 
 @pytest.fixture()
@@ -67,3 +70,69 @@ def test_download_artifacts_throws_for_invalid_arguments():
 
     with pytest.raises(MlflowException, match="`artifact_path` cannot be specified"):
         mlflow.artifacts.download_artifacts(artifact_path="path", artifact_uri="uri")
+
+
+@pytest.fixture()
+def run_with_text_artifact():
+    artifact_path = "test/file.txt"
+    artifact_content = "This is a sentence"
+    with mlflow.start_run() as run:
+        mlflow.log_text(artifact_content, artifact_path)
+
+    artifact_uri = str(pathlib.PurePosixPath(run.info.artifact_uri) / artifact_path)
+    return Artifact(artifact_uri, artifact_content)
+
+
+@pytest.fixture()
+def run_with_json_artifact():
+    artifact_path = "test/config.json"
+    artifact_content = {"mlflow-version": "0.28", "n_cores": "10"}
+    with mlflow.start_run() as run:
+        mlflow.log_dict(artifact_content, artifact_path)
+
+    artifact_uri = str(pathlib.PurePosixPath(run.info.artifact_uri) / artifact_path)
+    return Artifact(artifact_uri, artifact_content)
+
+
+@pytest.fixture()
+def run_with_image_artifact():
+    from PIL import Image
+
+    artifact_path = "test/image.png"
+    image = Image.new("RGB", (100, 100))
+    with mlflow.start_run() as run:
+        mlflow.log_image(image, artifact_path)
+
+    artifact_uri = str(pathlib.PurePosixPath(run.info.artifact_uri) / artifact_path)
+    return Artifact(artifact_uri, image)
+
+
+def test_load_text(run_with_text_artifact):
+    artifact = run_with_text_artifact
+    assert mlflow.artifacts.load_text(artifact.uri) == artifact.content
+
+
+def test_load_dict(run_with_json_artifact):
+    artifact = run_with_json_artifact
+    assert mlflow.artifacts.load_dict(artifact.uri) == artifact.content
+
+
+def test_load_json_invalid_json(run_with_text_artifact):
+    artifact = run_with_text_artifact
+    with pytest.raises(mlflow.exceptions.MlflowException, match="Unable to form a JSON object"):
+        mlflow.artifacts.load_dict(artifact.uri)
+
+
+def test_load_image(run_with_image_artifact):
+    from PIL import Image
+
+    artifact = run_with_image_artifact
+    assert isinstance(mlflow.artifacts.load_image(artifact.uri), Image.Image)
+
+
+def test_load_image_invalid_image(run_with_text_artifact):
+    artifact = run_with_text_artifact
+    with pytest.raises(
+        mlflow.exceptions.MlflowException, match="Unable to form a PIL Image object"
+    ):
+        mlflow.artifacts.load_image(artifact.uri)