| import tensorflow as tf | |
| from tensorflow import keras | |
| from .transformer import TransformerEncoder, TransformerDecoderLayer, get_look_ahead_mask, get_padding_mask | |
| class TFPort(keras.Model): | |
| def __init__(self, crop_size, embed_dim, encoder_layers, decoder_layers, input_seq_length, | |
| target_seq_length, hidden_dim, num_heads, input_vocab_size, dropout_rate=0.1): | |
| super().__init__() | |
| self.image_encoder = keras.applications.InceptionV3( | |
| include_top=False, | |
| weights='imagenet', | |
| input_shape=crop_size + (3,), | |
| ) | |
| self.image_encoder.trainable = False | |
| self.conv = keras.layers.Conv2D(embed_dim, 1) | |
| self.ingredient_encoder = TransformerEncoder(encoder_layers, input_seq_length, embed_dim, hidden_dim, | |
| num_heads, input_vocab_size, dropout_rate=dropout_rate) | |
| self.portion_embedding = PortionEmbedding(target_seq_length, embed_dim) | |
| self.dropout = keras.layers.Dropout(dropout_rate) | |
| self.decoder_layers = [ | |
| TransformerDecoderLayer(embed_dim, hidden_dim, num_heads, dropout_rate=dropout_rate) | |
| for _ in range(decoder_layers) | |
| ] | |
| self.linear = keras.layers.Dense(1, activation="relu") | |
| def call(self, inputs, training=False): | |
| image, ingredients, targets = inputs | |
| padding_mask = get_padding_mask(ingredients) | |
| encoded_img = self.image_encoder(image, training=False) | |
| encoded_img = self.conv(encoded_img, training=training) | |
| encoded_img = tf.reshape(encoded_img, (tf.shape(encoded_img)[0], -1, tf.shape(encoded_img)[3])) | |
| encoded_ingr = self.ingredient_encoder(ingredients, padding_mask, training=training) | |
| encoder_outputs = tf.concat([encoded_img, encoded_ingr], axis=1) | |
| img_mask = tf.ones((tf.shape(encoded_img)[0], 1, tf.shape(encoded_img)[1]), dtype=tf.int32) | |
| padding_mask = tf.concat([img_mask, padding_mask], axis=2) | |
| look_ahead_mask = get_look_ahead_mask(targets) | |
| x = self.portion_embedding(targets) | |
| x = self.dropout(x, training=training) | |
| for i in range(len(self.decoder_layers)): | |
| x = self.decoder_layers[i](x, encoder_outputs, look_ahead_mask, training=training, | |
| padding_mask=padding_mask) | |
| x = self.linear(x) | |
| return tf.squeeze(x) | |
| class PortionEmbedding(keras.layers.Layer): | |
| def __init__(self, seq_length, embed_dim): | |
| super().__init__() | |
| self.linear = keras.layers.Dense(embed_dim) | |
| self.position_embeddings = keras.layers.Embedding( | |
| input_dim=seq_length, output_dim=embed_dim | |
| ) | |
| def call(self, inputs): | |
| positions = tf.range(start=0, limit=tf.shape(inputs)[-1], delta=1) | |
| embedded_portions = self.linear(inputs[:, :, tf.newaxis]) | |
| embedded_positions = self.position_embeddings(positions) | |
| return embedded_portions + embedded_positions | |