| import ast |
| import json |
| import os |
|
|
| import numpy as np |
| import pandas as pd |
| from autogluon.multimodal import MultiModalPredictor |
| from constants import ( |
| ALLOWED_INPUT_FORMATS, |
| ALLOWED_OUTPUT_FORMATS, |
| BRACKET_FORMATTER, |
| COMMA_DELIMITER, |
| IMAGE_COLUMN_NAME, |
| JSON_FORMAT, |
| LABELS, |
| NUM_GPU, |
| PREDICTED_LABEL, |
| PROBABILITIES, |
| PROBABILITY, |
| SAGEMAKER_INFERENCE_OUTPUT, |
| ) |
| from utils import infer_type_and_cast_value |
|
|
| INFERENCE_OUTPUT = ( |
| infer_type_and_cast_value(os.getenv(SAGEMAKER_INFERENCE_OUTPUT)) |
| if SAGEMAKER_INFERENCE_OUTPUT in os.environ |
| else [PREDICTED_LABEL] |
| ) |
| NUM_GPUS = infer_type_and_cast_value(os.getenv(NUM_GPU)) |
|
|
|
|
| def generate_single_csv_line_inference_selection(data): |
| """Generate a single csv line response. |
| |
| :param data: list of output generated from the model |
| :return: csv line for the predictions |
| """ |
| contents: str |
| for single_prediction in data: |
| contents = ( |
| BRACKET_FORMATTER.format(single_prediction) |
| if isinstance(single_prediction, list) |
| else str(single_prediction) |
| ) |
| return contents |
|
|
|
|
| def model_fn(model_dir): |
| """Load model from previously saved artifact. |
| |
| :param model_dir: local path to the model directory |
| :return: loaded model |
| """ |
| predictor = MultiModalPredictor.load(model_dir) |
| if NUM_GPUS is not None: |
| predictor._config.env.num_gpus = NUM_GPUS |
|
|
| return predictor |
|
|
|
|
| def convert_to_json_compatible_type(value): |
| """Convert the input value to a JSON compatible type. |
| |
| :param value: input value |
| :return: JSON compatible value |
| """ |
| string_value = "{}".format(value) |
| try: |
| return ast.literal_eval(string_value) |
| except Exception: |
| return string_value |
|
|
|
|
| def transform_fn(model, request_body, input_content_type, output_content_type): |
| """Transform function for serving inference requests. |
| |
| If INFERENCE_OUTPUT is provided, then the predictions are generated in the requested format and concatenated in the |
| same order. Otherwise, prediction_labels are generated by default. |
| |
| :param model: loaded model |
| :param request_body: request body |
| :param input_content_type: content type of the input |
| :param output_content_type: content type of the response |
| :return: prediction response |
| """ |
| if input_content_type.lower() not in ALLOWED_INPUT_FORMATS: |
| raise Exception( |
| f"{input_content_type} input content type not supported. Supported formats are {ALLOWED_INPUT_FORMATS}" |
| ) |
|
|
| if output_content_type.lower() not in ALLOWED_OUTPUT_FORMATS: |
| raise Exception( |
| f"{output_content_type} output content type not supported. Supported formats are {ALLOWED_OUTPUT_FORMATS}" |
| ) |
|
|
| data = pd.DataFrame({IMAGE_COLUMN_NAME: [request_body]}) |
| result_dict = dict() |
|
|
| result = [] |
| inference_output_list = ( |
| INFERENCE_OUTPUT if isinstance(INFERENCE_OUTPUT, list) else [INFERENCE_OUTPUT] |
| ) |
| for output_type in inference_output_list: |
| if output_type == PREDICTED_LABEL: |
| prediction = model.predict(data) |
| result_dict[PREDICTED_LABEL] = convert_to_json_compatible_type(prediction.squeeze()) |
| elif output_type == PROBABILITIES: |
| predict_probs = model.predict_proba(data) |
| prediction = predict_probs.to_numpy() |
| result_dict[PROBABILITIES] = predict_probs.squeeze().tolist() |
| elif output_type == LABELS: |
| labels = model.class_labels |
| prediction = np.array([labels]).astype("str") |
| result_dict[LABELS] = labels.tolist() |
| else: |
| predict_probabilities = model.predict_proba(data).to_numpy() |
| prediction = np.max(predict_probabilities, axis=1) |
| result_dict[PROBABILITY] = prediction.squeeze().tolist() |
| result.append(generate_single_csv_line_inference_selection(prediction.tolist())) |
| response = COMMA_DELIMITER.join(result) |
|
|
| if output_content_type == JSON_FORMAT: |
| response = json.dumps(result_dict) |
|
|
| return response, output_content_type |
|
|
