| import tensorflow as tf |
| import numpy as np |
| from tensorflow import keras |
| import os |
| from typing import Dict, List, Any |
| import pickle |
| from PIL import Image |
| class PreTrainedPipeline(): |
| def __init__(self, path=""): |
| |
| |
| self.decoder = keras.models.load_model(os.path.join(path, "decoder")) |
| self.encoder = keras.models.load_model(os.path.join(path, "encoder")) |
| |
| image_model = tf.keras.applications.InceptionV3(include_top=False, |
| weights='imagenet') |
| new_input = image_model.input |
| hidden_layer = image_model.layers[-1].output |
|
|
| self.image_features_extract_model = tf.keras.Model(new_input, hidden_layer) |
| |
| with open('./tokenizer.pickle', 'rb') as handle: |
| self.tokenizer = pickle.load(handle) |
|
|
|
|
|
|
| def load_image(img): |
| img = tf.io.decode_jpeg(img, channels=3) |
| img = tf.image.resize(img, (299, 299)) |
| img = tf.keras.applications.inception_v3.preprocess_input(img) |
| return img, image_path |
|
|
| 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", "score": 0.82} |
| It is preferred if the returned list is in decreasing `score` order |
| """ |
|
|
| hidden = tf.zeros((1, 512)) |
| max_length = 46 |
| temp_input = tf.expand_dims(load_image(image)[0], 0) |
| img_tensor_val = self.image_features_extract_model(temp_input) |
| img_tensor_val = tf.reshape(img_tensor_val, (img_tensor_val.shape[0], |
| -1, |
| img_tensor_val.shape[3])) |
|
|
| features = self.encoder(img_tensor_val) |
|
|
| dec_input = tf.expand_dims([self.tokenizer.word_index['<start>']], 0) |
| result = [] |
|
|
| for i in range(max_length): |
| predictions, hidden, attention_weights = self.decoder(dec_input, |
| features, |
| hidden) |
|
|
| predicted_id = tf.random.categorical(predictions, 1)[0][0].numpy() |
| result.append(self.tokenizer.index_word[predicted_id]) |
|
|
| if self.tokenizer.index_word[predicted_id] == '<end>': |
| return result |
|
|
| dec_input = tf.expand_dims([predicted_id], 0) |
| resp = "" |
| for i in result: |
| if i!=1: |
| resp = resp + " " + result[i] |
| else: |
| resp += result[i] |
| return resp |
|
|
