Spaces:
Running
Running
File size: 2,703 Bytes
d1de1d0 |
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 |
import pickle
import tensorflow as tf
import pandas as pd
import numpy as np
MAX_LENGTH = 40
BATCH_SIZE = 32
BUFFER_SIZE = 1000
EMBEDDING_DIM = 512
UNITS = 512
# LOADING DATA
vocab = pickle.load(open('vocabulary/vocab_coco.file', 'rb'))
tokenizer = tf.keras.layers.TextVectorization(
standardize = None,
output_sequence_length = MAX_LENGTH,
vocabulary = vocab
)
idx2word = tf.keras.layers.StringLookup(
mask_token = "",
vocabulary = tokenizer.get_vocabulary(),
invert = True
)
def load_image_from_path(img_path):
img = tf.io.read_file(img_path)
img = tf.io.decode_jpeg(img, channels=3)
img = tf.keras.layers.Resizing(299, 299)(img)
img = tf.keras.applications.inception_v3.preprocess_input(img)
return img
def generate_caption(img, caption_model, add_noise=False):
if isinstance(img, str):
img = load_image_from_path(img)
if add_noise == True:
noise = tf.random.normal(img.shape)*0.1
img = (img + noise)
img = (img - tf.reduce_min(img))/(tf.reduce_max(img) - tf.reduce_min(img))
img = tf.expand_dims(img, axis=0)
img_embed = caption_model.cnn_model(img)
img_encoded = caption_model.encoder(img_embed, training=False)
y_inp = '[start]'
for i in range(MAX_LENGTH-1):
tokenized = tokenizer([y_inp])[:, :-1]
mask = tf.cast(tokenized != 0, tf.int32)
pred = caption_model.decoder(
tokenized, img_encoded, training=False, mask=mask)
pred_idx = np.argmax(pred[0, i, :])
pred_word = idx2word(pred_idx).numpy().decode('utf-8')
if pred_word == '[end]':
break
y_inp += ' ' + pred_word
y_inp = y_inp.replace('[start] ', '')
return y_inp
def get_caption_model():
encoder = TransformerEncoderLayer(EMBEDDING_DIM, 1)
decoder = TransformerDecoderLayer(EMBEDDING_DIM, UNITS, 8)
cnn_model = CNN_Encoder()
caption_model = ImageCaptioningModel(
cnn_model=cnn_model, encoder=encoder, decoder=decoder, image_aug=None,
)
def call_fn(batch, training):
return batch
caption_model.call = call_fn
sample_x, sample_y = tf.random.normal((1, 299, 299, 3)), tf.zeros((1, 40))
caption_model((sample_x, sample_y))
sample_img_embed = caption_model.cnn_model(sample_x)
sample_enc_out = caption_model.encoder(sample_img_embed, training=False)
caption_model.decoder(sample_y, sample_enc_out, training=False)
try:
caption_model.load_weights('models/trained_coco_weights.h5')
except FileNotFoundError:
caption_model.load_weights('image-caption-generator/models/trained_coco_weights.h5')
return caption_model |