| from typing import Dict | |
| from PIL import Image | |
| import numpy as np | |
| import os | |
| import json | |
| import tensorflow as tf | |
| from tensorflow import keras | |
| class PreTrainedPipeline(): | |
| def __init__(self, path=""): | |
| self.model = keras.saving.load_model("./") | |
| with open(os.path.join(path, "config.json")) as config: | |
| config = json.load(config) | |
| self.id2label = config["id2label"] | |
| def __call__(self, inputs: "Image.Image")-> Dict[str, str]: | |
| """ | |
| 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", "score": 0.82} | |
| It is preferred if the returned list is in decreasing `score` order | |
| """ | |
| img = keras.preprocessing.image.load_img(input, target_size=(224, 224)) | |
| x = keras.preprocessing.image.img_to_array(img) | |
| x = np.expand_dims(x, axis=0) | |
| x = keras.applications.vgg16.preprocess_input(x) | |
| prediction = self.model.predict(x) | |
| return { 'label': "detected", 'score': "dragon" if prediction[0][0] >= 0.99 else "not-dragon" } | |