File size: 7,258 Bytes
79a66df
 
 
 
 
 
 
 
 
 
 
 
 
0832cb7
79a66df
 
 
 
6cc1e81
 
79a66df
09f2656
79a66df
 
 
 
 
 
6cc1e81
79a66df
 
 
 
 
 
 
 
 
09f2656
79a66df
 
6cc1e81
79a66df
 
 
 
 
 
 
 
 
 
 
09f2656
79a66df
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
from typing import Dict, List, Any
from PIL import Image
from tfing import TFIng
from tfport import TFPort, get_look_ahead_mask, get_padding_mask

import os
import json
import tensorflow as tf
import numpy as np


class PreTrainedPipeline():
    def __init__(self, path=""):
        embed_dim = 256
        num_layers = 3
        seq_length = 20
        hidden_dim = 1024
        num_heads = 8
        self.crop_size = (224, 224)
        self.img_size = 256
        self.nutr_names = ('energy', 'fat', 'protein', 'carbs')
        with open(os.path.join(path, "ingredients_metadata.json"), encoding='UTF-8') as f:
            self.ingredients = json.load(f)
        self.ing_names = {ing['name']: int(ing_id) for ing_id, ing in self.ingredients.items()}
        self.vocab_size = len(self.ingredients) + 3
        self.seq_length = seq_length

        self.tfing = TFIng(
            self.crop_size, 
            embed_dim, 
            num_layers, 
            seq_length, 
            hidden_dim, 
            num_heads, 
            self.vocab_size
        )
        self.tfing.compile()
        self.tfing((tf.zeros((1, 224, 224, 3)), tf.zeros((1, seq_length))))
        self.tfing.load_weights(os.path.join(path, 'tfing.h5'))

        self.tfport = TFPort(
            self.crop_size, 
            embed_dim, 
            num_layers, 
            num_layers, 
            seq_length, 
            seq_length, 
            hidden_dim, 
            num_heads, 
            self.vocab_size
        )
        self.tfport.compile()
        self.tfport((tf.zeros((1, 224, 224, 3)), tf.zeros((1, seq_length)), tf.zeros((1, seq_length))))
        self.tfport.load_weights(os.path.join(path, 'tfport.h5'))

    def __call__(self, inputs: "Image.Image") -> List[Dict[str, Any]]:
        image = tf.keras.preprocessing.image.img_to_array(inputs)
        height = tf.shape(image)[0]
        width = tf.shape(image)[1]
        if width > height:
            image = tf.image.resize(image, (self.img_size, int(float(self.img_size * width) / float(height))))
        else:
            image = tf.image.resize(image, (int(float(self.img_size * height) / float(width)), self.img_size))

        image = tf.keras.applications.inception_v3.preprocess_input(image)
        image = tf.keras.layers.CenterCrop(*self.crop_size)(image)
        prediction = self.predict(image)
        return [
            {
                "label": prediction['ingredients'][i], 
                "score": prediction['portions'][i]
            } 
            for i in range(len(prediction['ingredients']))
        ]

    def encode_image(self, image):
            encoder_out = self.tfing.encoder(image)
            encoder_out = self.tfing.conv(encoder_out)
            encoder_out = tf.reshape(
                encoder_out, 
                (tf.shape(encoder_out)[0], -1, tf.shape(encoder_out)[3])
            )
            return encoder_out

    def encode_ingredients(self, ingredients, padding_mask):
        return self.tfport.ingredient_encoder(ingredients, padding_mask)

    def decode_ingredients(self, encoded_img, decoder_in):
        decoder_outputs = self.tfing.decoder(decoder_in, encoded_img)
        output = self.tfing.linear(decoder_outputs)
        return output + self.tfing.get_replacement_mask(decoder_in)

    def decode_portions(self, encoded_img, encoded_ingr, decoder_in, padding_mask):
        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(decoder_in)
        
        x = self.tfport.portion_embedding(decoder_in)
        for i in range(len(self.tfport.decoder_layers)):
            x = self.tfport.decoder_layers[i](x, encoder_outputs, look_ahead_mask, padding_mask=padding_mask)
        x = self.tfport.linear(x)
        return tf.squeeze(x)

    def predict_ingredients(self, encoded_img, known_ing=None):
        predicted = np.zeros((1, self.seq_length + 1), dtype=int)
        predicted[0, 0] = self.vocab_size - 2
        start_index = 0
        if known_ing:
            predicted[0, 1:len(known_ing) + 1] = known_ing
            start_index = len(known_ing)
        for i in range(start_index, self.seq_length):
            decoded = self.decode_ingredients(encoded_img, predicted[:, :-1])
            next_token = int(np.argmax(decoded[0, i]))
            predicted[0, i + 1] = next_token
            if next_token == self.vocab_size - 1:
                return predicted[0, 1:]
            if i == self.seq_length - 1:
                predicted[0, i + 1] = self.vocab_size - 1
                return predicted[0, 1:]

    def predict_portions(self, encoded_image, ingredients):
        predicted = np.zeros((1, self.seq_length + 1), dtype=float)
        predicted[0, 0] = -1
        padding_mask = get_padding_mask(ingredients)
        encoded_ingr = self.encode_ingredients(ingredients, padding_mask)
        for i in range(self.seq_length):
            if ingredients[0, i] == self.vocab_size - 1:
                return predicted[0, 1:]
            next_proportion = float(
                self.decode_portions(
                    encoded_image, 
                    encoded_ingr, 
                    predicted[:, :-1], 
                    padding_mask
                )[i]
            )
            predicted[0, i + 1] = next_proportion
        return predicted[0, 1:]

    def process_ingredients(self, ingredients):
        processed = []
        for ingredient in ingredients.split('\n'):
            stripped = ingredient.strip()
            if stripped == '.':
                return processed, True
            if stripped in self.ing_names:
                processed.append(self.ing_names[stripped])
        return processed, False

    def predict(self, image, known_ing=None):
        encoded_image = self.encode_image(image[tf.newaxis, :])
        known_ing, skip_ing = self.process_ingredients(known_ing)\
            if known_ing else (None, False)
        if not skip_ing:
            ingredients = self.predict_ingredients(encoded_image, known_ing=known_ing)
        else:
            ingredients = known_ing[:self.seq_length - 1]
            ingredients.append(self.vocab_size - 1)
            ingredients = np.pad(ingredients, (0, self.seq_length - len(ingredients)))
        readable_ingredients = [
            self.ingredients[str(token)]['name'] for token in ingredients 
            if token != 0 and token != self.vocab_size - 1
        ]
        portions = self.predict_portions(encoded_image, ingredients[tf.newaxis, :])\
            if len(readable_ingredients) > 1 else [100]
        portions_slice = portions[:len(readable_ingredients)]
        scale = 100 / sum(portions_slice)
        return {
            'ingredients': readable_ingredients,
            'portions': [portion * scale for portion in portions_slice],
            'nutrition': {
                name: sum(
                    self.ingredients[str(ingredients[i])][name] * portions[i] / 100 
                    for i in range(len(readable_ingredients))
                ) for name in self.nutr_names
            }
        }