burman/Image-Caption

1import torch2import re3import gradio as gr4from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel5 6device = 'cpu'7encoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"8decoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"9model_checkpoint = "nlpconnect/vit-gpt2-image-captioning"10feature_extractor = ViTFeatureExtractor.from_pretrained(encoder_checkpoint)11tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)12model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)13 14 15def greet(name):16    return "Hello " + name + "!!"17 18 19def predict(image, max_length=64, num_beams=4):20    image = image.convert('RGB')21    image = feature_extractor(22        image, return_tensors="pt").pixel_values.to(device)23 24    def clean_text(x): return x.replace('<|endoftext|>', '').split('\n')[0]25    caption_ids = model.generate(image, max_length=max_length)[0]26    caption_text = clean_text(tokenizer.decode(caption_ids))27    return caption_text28 29 30input = gr.inputs.Image(label="Upload your Image", type='pil', optional=True)31output = gr.outputs.Textbox(type="auto", label="Captions")32examples = [f"example{i}.jpg" for i in range(1, 7)]33 34description = "Image captioning application made using transformers"35title = "Image Captioning 🖼️"36 37article = "Created By : Shreyas Dixit "38 39interface = gr.Interface(40    fn=predict,41    inputs=input,42    theme="grass",43    outputs=output,44    examples=examples,45    title=title,46    description=description,47    article=article,48    show_api=True49)50interface.launch()51