File size: 2,730 Bytes
a7ce59e |
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 |
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 = from_pretrained_keras(model_id)
# 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.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
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, axis=-1)
self.single_output_unit = (
self.model.output_shape[1] == 1
) # if there are two classes
if self.single_output_unit:
score = predictions[0][0]
labels = [
{"label": str(self.id2label["1"]), "score": float(score)},
{"label": str(self.id2label["0"]), "score": float(1 - score)},
]
else:
labels = [
{
"label": str(self.id2label[str(i)]),
"mask": base64.b64encode(predictions[0][i]),
"score": float(score),
}
for i, score in enumerate(predictions[0])
]
return sorted(labels, key=lambda tup: tup["score"], reverse=True)[: self.top_k]
|