Skip to content

AutoML Edge (Vertex AI) classification models

AutoML Edge models trained using Google Cloud's Vertex AI offering can now be easily deployed on MD.ai. The following steps outline the process for deploying image classification models -

  • Export your model as a TF Saved Model by following steps for the Container format mentioned here. This will download a saved_model.pb file to the Google cloud storage bucket of your choice.
  • Download this saved_model.pb file from your bucket.
  • Download the sample code for deploying AutoML edge models on MD.ai from here - vertex-ai-mdai-template.zip
  • For binary and multiclass models, you simply need to replace the saved_model.pb file in this sample zip file with your own exported model file.
  • Optionally, in the .mdai/config.yaml file change the device_type to gpu if required to run inference on a GPU.
  • For multilabel or more complex models, make changes in the .mdai/mdai-deploy.py file as required depending on your model definition.
  • Finally, follow our documentation for steps to deploy models into a project.

Technical considerations

AutoML edge models do not support DICOM images and only work on JPEG/PNG inputs. Note that if you are using an Upload dataset, we convert all images to DICOM and so your model might fail to work if you do not add steps to convert the loaded DICOM image to JPEG/PNG before sending it as an input to the TF saved model.

If you are using Google Cloud Storage or Amazon S3 datasets then the model will work natively with the help of our template. For our viewer, we still create DICOM tags for such images and the model does need to map the outputs to the specific DICOM IDs that we create, but as shown below the mdai_deploy.py file takes care of this mapping -

import os
from io import BytesIO
import numpy as np
import tensorflow as tf

class MDAIModel:
    def __init__(self):
        # Load vertex AI saved model
        modelpath = os.path.dirname(os.path.dirname(__file__))
        model = tf.saved_model.load(modelpath)
        self.model = model.signatures["serving_default"]

    def predict(self, data):
        input_files = data["files"]
        outputs = []

        for file in input_files:
            # Ignore images that are not JPEG/PNG
            if file["content_type"] not in ["image/jpeg", "image/png"]:
                continue

            # DICOM UID mapping
            uids = file["content_uids"]

            # Load JPEG/PNG image
            image = BytesIO(file["content"]).getvalue()

            # Run inference
            prediction = self.model(
                image_bytes=tf.convert_to_tensor([image]),
                key=tf.convert_to_tensor([uids["SOPInstanceUID"]]),
            )
            probabilities = prediction["scores"].numpy()[0]

            # For binary and multiclass classification, we return the index
            # having the maximum probability. This can be changed to a list
            # of indices in case of multilabel classification.
            class_index = int(np.argmax(probabilities))

            outputs.append(
                {
                    "type": "ANNOTATION",
                    "study_uid": uids["StudyInstanceUID"],
                    "series_uid": uids["SeriesInstanceUID"],
                    "instance_uid": uids["SOPInstanceUID"],
                    "class_index": class_index,
                    "probability": [
                        {"class_index": i, "probability": float(j)}
                        for i, j in enumerate(probabilities)
                    ],
                }
            )

        return outputs

You might need to change certain parts of this file such as how the predicted class indices are chosen depending on your model requirements.