NicolasVana/image-captioning

1import json2import os, shutil3import random4import streamlit as st5import os6from pathlib import Path7import numpy as np8 9from PIL import Image10import tensorflow as tf11from tensorflow.keras.applications.inception_v3 import preprocess_input12from tensorflow.keras.preprocessing import image13from tensorflow.keras.applications.inception_v3 import InceptionV314from tensorflow.keras.models import Model15from tensorflow.keras.preprocessing.sequence import pad_sequences16 17 18 19root = Path(os.getcwd())20aux_pre = root / 'Inception' / 'PretrainedInceptionLSTM'21aux_re = root / 'Inception' / 'RetrainedInceptionLSTM'22 23model_re_path = root / 'Inception' / 'RetrainedInceptionLSTM' / 'Model'24model_inception_path = root / 'Inception' / 'RetrainedInceptionFeatureExtraction' / 'Model'25model_pre_path = root / 'Inception' / 'PretrainedInceptionLSTM' / 'Model'26 27# Must create28 29def get_pretrained_inceptionV3():30    model = InceptionV3(weights='imagenet')31    model2 = Model(model.input, model.layers[-2].output)32    return model233 34def fetch_auxiliary_files(type):35    if type == 'Pretrained Inception':36        word2Index = np.load(aux_pre / "word2Index.npy", allow_pickle=True).item()37        index2Word = np.load(aux_pre / "index2Word.npy", allow_pickle=True).item()38        variable_params = np.load(aux_pre / "variable_params.npy", allow_pickle=True).item()39        return word2Index, index2Word, variable_params40    if type == 'Retrained Inception':41        word2Index = np.load(aux_re / "word2Index.npy", allow_pickle=True).item()42        index2Word = np.load(aux_re / "index2Word.npy", allow_pickle=True).item()43        variable_params = np.load(aux_re / "variable_params.npy", allow_pickle=True).item()44        return word2Index, index2Word, variable_params45 46@st.cache(allow_output_mutation=True, show_spinner=False)47def fetch_model(type):48    with st.spinner(text="Fetching Model"):49        if type == 'Pretrained Inception':50            model_pre = tf.keras.models.load_model(model_pre_path)51            model_inc = get_pretrained_inceptionV3()52            return model_inc, model_pre53        if type == 'Retrained Inception':54            model_re = tf.keras.models.load_model(model_re_path)55            model_inc = tf.keras.models.load_model(model_inception_path)56            return model_inc, model_re57 58def preprocess_image_inception(image):59    if image.mode != "RGB":60        image = image.convert(mode="RGB")61 62    x = np.array(image)63    x = np.expand_dims(x, axis = 0)64    x = preprocess_input(x)65    x = x.reshape(1, 299, 299, 3)66 67    return x68 69def extract_features(model, image):70    features = model.predict(image, verbose = 0)71    return features72 73def generate_caption(model, features, max_len, word2Index, index2Word, beam_index = 3):74    caption = beam_search(model, features, max_len, word2Index, index2Word, beam_index)75    return caption76 77def beam_search(model, features, max_len, word2Index, index2Word, beam_index):78    start = [word2Index["startseq"]]79    start_word = [[start, 1]]80 81    final_preds = []82    live_seqs = beam_index83    features = np.tile(features, (beam_index,1))84    count = 085    while len(start_word) > 0:86        #print(count)87        count+=188        temp = []89        padded_seqs = []90        #Get padded seqs for each of the starting seqs so far, misnamed as start_word91        for s in start_word:92            par_caps = pad_sequences([s[0]], maxlen=max_len, padding='post')93            padded_seqs.append(par_caps)94 95        #Formatting input so that it can be used for a prediction96        padded_seqs = np.array(padded_seqs).reshape(len(start_word), max_len)97 98        preds = model.predict([features[:len(start_word)],padded_seqs], verbose=0)99 100        #Getting the best branches for each of the start seqs that we had101        for index, pred in enumerate(preds):102            word_preds = np.argsort(pred)[-live_seqs:]103            for w in word_preds:104                next_cap, prob = start_word[index][0][:], start_word[index][1]105                next_cap.append(w)106                prob *= pred[w]107                temp.append([next_cap, prob])108 109        start_word = temp110        # Sorting according to the probabilities111        start_word = sorted(start_word, reverse=False, key=lambda l: l[1])112        # Getting the top words from all branches113        start_word = start_word[-live_seqs:]114 115        for pair in start_word:116            if index2Word[pair[0][-1]] == 'endseq':117                final_preds.append([pair[0][:-1], pair[1]])118                start_word = start_word[:-1]119                live_seqs -= 1120            if len(pair[0]) == max_len:121                final_preds.append(pair)122                start_word = start_word[:-1]123                live_seqs -= 1124 125    # Between all the finished sequences (either max len or predicted endseq), decide which is best126    max_prob = 0127    for index, pred in enumerate(final_preds):128        if pred[1] > max_prob:129            best_index = index130            max_prob = pred[1]131 132    # Convert to readable text133    final_pred = final_preds[best_index]134    final_caption = [index2Word[i] for i in final_pred[0]]135    final_caption = ' '.join(final_caption[1:])136    return final_caption137 138# # create target model directory139# model_dir = './models/'140# os.makedirs(model_dir, exist_ok=True)141#142# files_to_download = [143#     "config.json",144#     "flax_model.msgpack",145#     "merges.txt",146#     "special_tokens_map.json",147#     "tokenizer.json",148#     "tokenizer_config.json",149#     "vocab.json",150#     "preprocessor_config.json",151# ]152 153def _compile():154 155    image_path = 'samples/ROCO_00929.jpg'156    image = Image.open(image_path)157    #predict(image)158    image.close()159 160 161_compile()162 163 164sample_dir = './samples/'165sample_image_ids = tuple(["None"] + [int(f.replace('ROCO_', '').replace('.jpg', '')) for f in os.listdir(sample_dir) if f.startswith('ROCO_')])166 167with open(os.path.join(sample_dir, "Roco-img-ids.json"), "r", encoding="UTF-8") as fp:168    roco_image_ids = json.load(fp)169 170 171def get_random_image_id():172 173    image_id = random.sample(roco_image_ids, k=1)[0]174    return image_id175