CoolFace
Apppublic

Zayn/Image_Captioning_Using_Vision_Transformer_and_GPT-2

sourceHugging Facemitupdated 2y agoView on Hugging Face
0likes
app.py40 linesDownload Raw Back to root
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)