Create pipeline.py

pipeline.py

@@ -0,0 +1,54 @@
+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.word_to_index = tf.keras.layers.StringLookup(
+          mask_token="",
+          vocabulary=self.tokenizer.get_vocabulary())
+        
+        self.index_to_word = tf.keras.layers.StringLookup(
+          mask_token="",
+          vocabulary=self.tokenizer.get_vocabulary(),
+          invert=True)
+
+    def load_image(img):
+        #img = tf.io.read_file(image_path)
+        img = tf.io.decode_jpeg(img, channels=3)
+        img = tf.image.resize(img, IMAGE_SHAPE[:-1])
+        return img
+
+    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
+        """
+        image = load_image(inputs)
+        initial = self.word_to_index([['[START]']]) # (batch, sequence)
+        img_features = self.feature_extractor(image[tf.newaxis, ...])
+        temperature = 0
+        tokens = initial # (batch, sequence)
+        for n in range(50):
+          preds = self((img_features, tokens)).numpy()  # (batch, sequence, vocab)
+          preds = preds[:,-1, :]  #(batch, vocab)
+          if temperature==0:
+              next = tf.argmax(preds, axis=-1)[:, tf.newaxis]  # (batch, 1)
+          else:
+              next = tf.random.categorical(preds/temperature, num_samples=1)  # (batch, 1)
+          tokens = tf.concat([tokens, next], axis=1) # (batch, sequence)
+
+          if next[0] == self.word_to_index('[END]'):
+            break
+        words = self.index_to_word(tokens[0, 1:-1])
+        result = tf.strings.reduce_join(words, axis=-1, separator=' ')
+        return result.numpy().decode()