| import json | |
| from typing import Any, Dict, List | |
| import tensorflow as tf | |
| from tensorflow import keras | |
| from huggingface_hub import from_pretrained_keras, hf_hub_download | |
| import base64 | |
| import io | |
| import numpy as np | |
| from PIL import Image | |
| class PreTrainedPipeline(): | |
| def __init__(self, model_id: str): | |
| self.model = keras.models.load_model("./tf_model.h5") | |
| def __call__(self, inputs: "Image.Image")-> List[Dict[str, Any]]: | |
| with Image.open(inputs) as img: | |
| img = np.array(img) | |
| im = tf.image.resize(img, (128, 128)) | |
| im = tf.cast(im, tf.float32) / 255.0 | |
| pred_mask = model.predict(im[tf.newaxis, ...]) | |
| pred_mask_arg = tf.argmax(pred_mask, axis=-1) | |
| labels = [] | |
| binary_masks = {} | |
| mask_codes = {} | |
| for cls in range(pred_mask.shape[-1]): | |
| binary_masks[f"mask_{cls}"] = np.zeros(shape = (pred_mask.shape[1], pred_mask.shape[2])) | |
| for row in range(pred_mask_arg[0][1].get_shape().as_list()[0]): | |
| for col in range(pred_mask_arg[0][2].get_shape().as_list()[0]): | |
| if pred_mask_arg[0][row][col] == cls: | |
| binary_masks[f"mask_{cls}"][row][col] = 1 | |
| else: | |
| binary_masks[f"mask_{cls}"][row][col] = 0 | |
| mask = binary_masks[f"mask_{cls}"] | |
| mask *= 255 | |
| img = Image.fromarray(mask.astype(np.int8), mode="L") | |
| with io.BytesIO() as out: | |
| img.save(out, format="PNG") | |
| png_string = out.getvalue() | |
| mask = base64.b64encode(png_string).decode("utf-8") | |
| mask_codes[f"mask_{cls}"] = mask | |
| labels.append({ | |
| "label": f"LABEL_{cls}", | |
| "mask": mask_codes[f"mask_{cls}"], | |
| "score": 1.0, | |
| }) | |
| return labels |