Zayn/Image_Captioning_Using_Vision_Transformer_and_GPT-2
0
1from PIL import Image2from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, PreTrainedTokenizerFast3import requests4 5model = VisionEncoderDecoderModel.from_pretrained("Zayn/vit2distilgpt2")6vit_feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")7tokenizer = PreTrainedTokenizerFast.from_pretrained("distilgpt2")8 9def vit2distilgpt2(img):10 pixel_values = vit_feature_extractor(images=img, return_tensors="pt").pixel_values11 encoder_outputs = model.generate(pixel_values.to('cpu'),num_beams=5)12 generated_sentences = tokenizer.batch_decode(encoder_outputs, skip_special_tokens =True)13 14 return(generated_sentences[0].split('.')[0])15 16import gradio as gr17 18inputs = [19 gr.inputs.Image(type="pil", label = "Original Image")20]21 22outputs = [23 gr.outputs.Textbox(label = 'Caption')24] 25title = "Image Captioning using Vision Transformer and GPT-2"26description = "Developed by Zayn"27article = "< a href='https://huggingface.co/Zayn/vit2distilgpt2'>Hugging Face AI Community</a>" 28examples = [29 ["car.jpg"]30]31gr.Interface(32 vit2distilgpt2,33 inputs,34 outputs,35 title = title,36 description = description,37 article = article,38 examples = examples,39 theme = "huggingface", 40).launch(debug=True,enable_queue=True)