|
import torch |
|
import numpy as np |
|
from PIL import Image |
|
import streamlit as st |
|
from torchvision.transforms import v2 |
|
from transformers import GenerationConfig |
|
from transformers import GPT2TokenizerFast |
|
from transformers import ViTImageProcessor |
|
from transformers import VisionEncoderDecoderModel |
|
|
|
|
|
st.set_page_config( |
|
layout="centered", |
|
page_title="Generate Caption", |
|
initial_sidebar_state="collapsed", |
|
) |
|
|
|
|
|
if all(key not in st.session_state.keys() for key in ("generate", "image")): |
|
st.session_state["generate"] = False |
|
st.session_state["image"] = None |
|
|
|
|
|
@st.cache_resource(show_spinner="Loading Resources...") |
|
def loadResources(): |
|
encoder = 'microsoft/swin-base-patch4-window7-224-in22k' |
|
decoder = 'gpt2' |
|
|
|
model = VisionEncoderDecoderModel.from_encoder_decoder_pretrained( |
|
encoder, decoder |
|
) |
|
|
|
processor = ViTImageProcessor.from_pretrained(encoder) |
|
tokenizer = GPT2TokenizerFast.from_pretrained(decoder) |
|
|
|
if 'gpt2' in decoder: |
|
tokenizer.pad_token = tokenizer.eos_token |
|
model.config.eos_token_id = tokenizer.eos_token_id |
|
model.config.pad_token_id = tokenizer.pad_token_id |
|
model.config.decoder_start_token_id = tokenizer.bos_token_id |
|
else: |
|
model.config.decoder_start_token_id = tokenizer.cls_token_id |
|
model.config.pad_token_id = tokenizer.pad_token_id |
|
|
|
model = torch.load("generator_model.pkl", map_location=torch.device("cpu")) |
|
model.eval() |
|
return processor, tokenizer, model |
|
|
|
|
|
@st.cache_data |
|
def preprocess_image(_processor, image): |
|
transforms = v2.Compose([ |
|
v2.Resize(size=(224,224)), |
|
v2.ToDtype(torch.float32, scale = True), |
|
]) |
|
image = transforms(image) |
|
img = _processor(image, return_tensors = 'pt') |
|
return img |
|
|
|
|
|
@st.cache_data |
|
def get_caption(_processor, _tokenizer, _model, image): |
|
image = preprocess_image(_processor, image) |
|
output = _model.generate( |
|
**image, |
|
generation_config = GenerationConfig( |
|
pad_token_id = _tokenizer.pad_token_id |
|
) |
|
) |
|
|
|
caption = _tokenizer.batch_decode( |
|
output, |
|
skip_special_tokens = True |
|
) |
|
|
|
caption = " ".join([item[0].upper()+item[1:] for item in caption[0].split(" ")]) |
|
|
|
return caption |
|
|
|
|
|
def DisplayInteractionElements(): |
|
st.markdown('<div style="display: flex; justify-content: center;"><p style="font-size: 40px; font-weight: bold;">π Caption Generator π</p></div>', unsafe_allow_html=True) |
|
st.file_uploader(accept_multiple_files=False, label='Upload an Image', type=['jpg', 'jpeg', 'png'], key="image_uploader") |
|
|
|
if st.session_state['image_uploader']: |
|
image = st.session_state['image_uploader'] |
|
im_file = Image.open(image).convert("RGB") |
|
im_file = np.array(im_file) |
|
|
|
st.session_state['image'] = im_file |
|
|
|
col1, col2, col3 = st.columns(3) |
|
|
|
col2.image(image=image, caption='Uploaded Image') |
|
|
|
st.button(label='Generate Caption', use_container_width=True, type='primary', on_click=generateCaption) |
|
|
|
|
|
def generateCaption(): |
|
st.session_state['generate'] = True |
|
|
|
def main(): |
|
|
|
DisplayInteractionElements() |
|
|
|
processor, tokenizer, model = loadResources() |
|
|
|
if not st.session_state['image_uploader']: |
|
st.session_state['generate'] = False |
|
|
|
if st.session_state['generate'] and st.session_state['image_uploader']: |
|
caption = get_caption(processor, tokenizer, model, st.session_state['image']) |
|
st.markdown(f'<div style="display: flex; justify-content: center;"><p style="font-size: 35px; font-weight: bold; color: blue;">{caption}</p></div>', unsafe_allow_html = True) |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
|
|
|
|
|
|
|
|
|
|