File size: 2,482 Bytes
a7ce59e eb496ad a7ce59e eb496ad a7ce59e 6cb57f7 a7ce59e 6cb57f7 a7ce59e 6cb57f7 a7ce59e 6cb57f7 a7ce59e eb496ad 6cb57f7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import json
from typing import Any, Dict, List
import tensorflow as tf
from tensorflow import keras
from app.pipelines import Pipeline
from huggingface_hub import from_pretrained_keras, hf_hub_download
from PIL import Image
import base64
MODEL_FILENAME = "saved_model.pb"
CONFIG_FILENAME = "config.json"
class PreTrainedPipeline(Pipeline):
def __init__(self, model_id: str):
# Reload Keras SavedModel
self.model = keras.models.load_model('./model.h5')
# Number of labels
self.num_labels = self.model.output_shape[1]
# Config is required to know the mapping to label.
#config_file = hf_hub_download(model_id, filename=CONFIG_FILENAME)
#with open(config_file) as config:
# config = json.load(config)
self.num_labels = 3
self.id2label = config.get(
"id2label", {str(i): f"LABEL_{i}" for i in range(self.num_labels)}
)
def __call__(self, inputs: "Image.Image") -> List[Dict[str, Any]]:
"""
Args:
inputs (:obj:`PIL.Image`):
The raw image representation as PIL.
No transformation made whatsoever from the input. Make all necessary transformations here.
Return:
A :obj:`list`:. The list contains items that are dicts should be liked {"label": "XXX" (str), mask: "base64 encoding of the mask" (str), "score": float}
It is preferred if the returned list is in decreasing `score` order
"""
# Resize image to expected size
expected_input_size = self.model.input_shape
with Image.open(inputs) as im:
inputs = np.array(im)
if expected_input_size[-1] == 1:
inputs = inputs.convert("L")
target_size = (expected_input_size[1], expected_input_size[2])
img = tf.image.resize(inputs, target_size)
img_array = tf.keras.preprocessing.image.img_to_array(img)
img_array = img_array[tf.newaxis, ...]
predictions = self.model.predict(img_array)
self.single_output_unit = (
self.model.output_shape[1] == 1
) # if there are two classes
labels = []
for i in enumerate(predictions):
labels.append({
"label": str(i[0]),
"mask": base64.b64encode(i[1]),
"score": 1.0,
})
return sorted(labels)
|