Spaces:
Sleeping
Sleeping
File size: 2,430 Bytes
57231ec 08ec34d 7e847dc 57231ec 08ec34d 57231ec 7e847dc 57231ec 08ec34d 7e847dc 57231ec 08ec34d 57231ec 08ec34d 57231ec 08ec34d 57231ec |
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 |
import streamlit as st
from PIL import Image
from transformers import GPT2TokenizerFast, ViTImageProcessor, VisionEncoderDecoderModel,RobertaTokenizerFast, VisionEncoderDecoderModel
#import torch
#from transformers import BlipProcessor, BlipForConditionalGeneration
def set_page_config():
st.set_page_config(
page_title='Caption an Cartoon Image',
page_icon=':camera:',
layout='wide',
)
def initialize_model():
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# load a fine-tuned image captioning model and corresponding tokenizer and image processor
model = VisionEncoderDecoderModel.from_pretrained("sourabhbargi11/caption4").to(device)
tokenizer = GPT2TokenizerFast.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
image_processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
return image_processor, model,tokenizer, device
def upload_image():
return st.sidebar.file_uploader("Upload an image (we aren't storing anything)", type=["jpg", "jpeg", "png"])
def image_preprocess(image):
image = image.resize((max_width, height))
if image.mode == "L":
image = image.convert("RGB")
return image
def generate_caption(processor, model, device, image):
#inputs = processor(image, return_tensors='pt').to(device)
#out = model.generate(**inputs, max_new_tokens=20)
#caption = processor.decode(out[0], skip_special_tokens=True)
caption="im here "
return caption
def main():
set_page_config()
st.header("Caption an Image :camera:")
uploaded_image = upload_image()
if uploaded_image is not None:
image = Image.open(uploaded_image)
image = image_preprocess(image)
st.image(image, caption='Your image')
with st.sidebar:
st.divider()
if st.sidebar.button('Generate Caption'):
with st.spinner('Generating caption...'):
image_processor, model,tokenizer, device = initialize_model()
caption = generate_caption(image_processor, model, device, image)
#caption="im here man"
st.header("Caption:")
st.markdown(f'**{caption}**')
if __name__ == '__main__':
main()
st.markdown("""
---
You are looking at partial finetuned model , please JUDGE ME!!! """)
|