chore: import upstream snapshot with attribution

examples/flower_classifier/image_pyfunc.py

@@ -0,0 +1,193 @@
+"""
+Example of a custom python function implementing image classifier with image preprocessing embedded
+in the model.
+"""
+
+import base64
+import importlib.metadata
+import os
+from io import BytesIO
+from typing import Any
+
+import keras
+import numpy as np
+import pandas as pd
+import PIL
+import tensorflow as tf
+import yaml
+from PIL import Image
+
+import mlflow
+from mlflow.utils import PYTHON_VERSION
+from mlflow.utils.file_utils import TempDir
+
+
+def decode_and_resize_image(raw_bytes, size):
+    """
+    Read, decode and resize raw image bytes (e.g. raw content of a jpeg file).
+
+    Args:
+        raw_bytes: Image bits, e.g. jpeg image.
+        size: Requested output dimensions.
+
+    Returns:
+        Multidimensional numpy array representing the resized image.
+    """
+    return np.asarray(Image.open(BytesIO(raw_bytes)).resize(size), dtype=np.float32)
+
+
+class KerasImageClassifierPyfunc:
+    """
+    Image classification model with embedded pre-processing.
+
+    This class is essentially an MLflow custom python function wrapper around a Keras model.
+    The wrapper provides image preprocessing so that the model can be applied to images directly.
+    The input to the model is base64 encoded image binary data (e.g. contents of a jpeg file).
+    The output is the predicted class label, predicted class id followed by probabilities for each
+    class.
+
+    The model declares current local versions of Keras, Tensorlow and pillow as dependencies in its
+    conda environment file.
+    """
+
+    def __init__(self, graph, session, model, image_dims, domain):
+        self._graph = graph
+        self._session = session
+        self._model = model
+        self._image_dims = image_dims
+        self._domain = domain
+        probs_names = [f"p({x})" for x in domain]
+        self._column_names = ["predicted_label", "predicted_label_id"] + probs_names
+
+    def predict(
+        self,
+        input,
+        params: dict[str, Any] | None = None,
+    ):
+        """
+        Generate predictions for the data.
+
+        Args:
+            input: pandas.DataFrame with one column containing images to be scored. The image
+                column must contain base64 encoded binary content of the image files. The image
+                format must be supported by PIL (e.g. jpeg or png).
+            params: Additional parameters to pass to the model for inference.
+
+        Returns:
+            pandas.DataFrame containing predictions with the following schema:
+                Predicted class: string,
+                Predicted class index: int,
+                Probability(class==0): float,
+                ...,
+                Probability(class==N): float,
+        """
+
+        # decode image bytes from base64 encoding
+        def decode_img(x):
+            return pd.Series(base64.decodebytes(bytearray(x[0], encoding="utf8")))
+
+        images = input.apply(axis=1, func=decode_img)
+        probs = self._predict_images(images)
+        m, n = probs.shape
+        label_idx = np.argmax(probs, axis=1)
+        labels = np.array([self._domain[i] for i in label_idx], dtype=str).reshape(m, 1)
+        output_data = np.concatenate((labels, label_idx.reshape(m, 1), probs), axis=1)
+        res = pd.DataFrame(columns=self._column_names, data=output_data)
+        res.index = input.index
+        return res
+
+    def _predict_images(self, images):
+        """
+        Generate predictions for input images.
+
+        Args:
+            images: Binary image data.
+
+        Returns:
+            Predicted probabilities for each class.
+        """
+
+        def preprocess_f(z):
+            return decode_and_resize_image(z, self._image_dims[:2])
+
+        x = np.array(images[images.columns[0]].apply(preprocess_f).tolist())
+        with self._graph.as_default():
+            with self._session.as_default():
+                return self._model.predict(x)
+
+
+def log_model(keras_model, signature, artifact_path, image_dims, domain):
+    """
+    Log a KerasImageClassifierPyfunc model as an MLflow artifact for the current run.
+
+    Args:
+        keras_model: Keras model to be saved.
+        signature: Model signature.
+        artifact_path: Run-relative artifact path this model is to be saved to.
+        image_dims: Image dimensions the Keras model expects.
+        domain: Labels for the classes this model can predict.
+    """
+
+    with TempDir() as tmp:
+        data_path = tmp.path("image_model")
+        os.mkdir(data_path)
+        conf = {"image_dims": "/".join(map(str, image_dims)), "domain": "/".join(map(str, domain))}
+        with open(os.path.join(data_path, "conf.yaml"), "w") as f:
+            yaml.safe_dump(conf, stream=f)
+        keras_path = os.path.join(data_path, "keras_model")
+        mlflow.tensorflow.save_model(model=keras_model, path=keras_path)
+        conda_env = tmp.path("conda_env.yaml")
+        with open(conda_env, "w") as f:
+            f.write(
+                conda_env_template.format(
+                    python_version=PYTHON_VERSION,
+                    keras_version=keras.__version__,
+                    tf_name=tf.__name__,  # can have optional -gpu suffix
+                    tf_version=tf.__version__,
+                    pip_version=importlib.metadata.version("pip"),
+                    pillow_version=PIL.__version__,
+                )
+            )
+
+        mlflow.pyfunc.log_model(
+            name=artifact_path,
+            signature=signature,
+            loader_module=__name__,
+            code_paths=[__file__],
+            data_path=data_path,
+            conda_env=conda_env,
+        )
+
+
+def _load_pyfunc(path):
+    """
+    Load the KerasImageClassifierPyfunc model.
+    """
+    with open(os.path.join(path, "conf.yaml")) as f:
+        conf = yaml.safe_load(f)
+    keras_model_path = os.path.join(path, "keras_model")
+    domain = conf["domain"].split("/")
+    image_dims = np.array([int(x) for x in conf["image_dims"].split("/")], dtype=np.int32)
+    # NOTE: TensorFlow based models depend on global state (Graph and Session) given by the context.
+    # To make sure we score the model in the same session as we loaded it in, we create a new
+    # session and a new graph here and store them with the model.
+    with tf.Graph().as_default() as g:
+        with tf.Session().as_default() as sess:
+            keras.backend.set_session(sess)
+            keras_model = mlflow.tensorflow.load_model(keras_model_path)
+    return KerasImageClassifierPyfunc(g, sess, keras_model, image_dims, domain=domain)
+
+
+conda_env_template = """
+name: flower_classifier
+channels:
+  - conda-forge
+dependencies:
+  - python=={python_version}
+  - pip=={pip_version}
+  - pip:
+    - mlflow>=1.6
+    - pillow=={pillow_version}
+    - keras=={keras_version}
+    - {tf_name}=={tf_version}
+"""