import torch import clip import PIL.Image from PIL import Image import skimage.io as io import streamlit as st from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup from transformers import GPT2TokenizerFast, ViTImageProcessor, VisionEncoderDecoderModel from model import generate2,ClipCaptionModel from engine import inference model_trained = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning") model_trained.load_state_dict(torch.load('model_trained.pth',map_location=torch.device('cpu')),strict=False) image_processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning") tokenizer = GPT2TokenizerFast.from_pretrained("nlpconnect/vit-gpt2-image-captioning") def show_n_generate(img, model, greedy = True): image = Image.open(img) pixel_values = image_processor(image, return_tensors ="pt").pixel_values if greedy: generated_ids = model.generate(pixel_values, max_new_tokens = 30) else: generated_ids = model.generate( pixel_values, do_sample=True, max_new_tokens = 30, top_k=5) generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] return generated_text device = "cpu" clip_model, preprocess = clip.load("ViT-B/32", device=device, jit=False) tokenizer = GPT2Tokenizer.from_pretrained("gpt2") prefix_length = 10 model = ClipCaptionModel(prefix_length) model.load_state_dict(torch.load('model.h5',map_location=torch.device('cpu')),strict=False) model = model.eval() coco_model = ClipCaptionModel(prefix_length) coco_model.load_state_dict(torch.load('COCO_model.h5',map_location=torch.device('cpu')),strict=False) model = model.eval() def ui(): st.markdown("# Image Captioning") # st.markdown("## Done By- Vageesh and Rushil") uploaded_file = st.file_uploader("Upload an Image", type=['png', 'jpeg', 'jpg']) if uploaded_file is not None: image = io.imread(uploaded_file) pil_image = PIL.Image.fromarray(image) image = preprocess(pil_image).unsqueeze(0).to(device) option = st.selectbox('Please select the Model',('Clip Captioning','Attention Decoder','VIT+GPT2')) if option=='Clip Captioning': with torch.no_grad(): prefix = clip_model.encode_image(image).to(device, dtype=torch.float32) prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1) generated_text_prefix = generate2(model, tokenizer, embed=prefix_embed) st.image(uploaded_file, width = 500, channels = 'RGB') st.markdown("**PREDICTION:** " + generated_text_prefix) elif option=='Attention Decoder': out = inference(uploaded_file) st.image(uploaded_file, width = 500, channels = 'RGB') st.markdown("**PREDICTION:** " + out) # elif option=='VIT+GPT2': # out=show_n_generate(uploaded_file, greedy = False, model = model_trained) # st.image(uploaded_file, width = 500, channels = 'RGB') # st.markdown("**PREDICTION:** " + out) if __name__ == '__main__': ui()