| from PIL import Image |
| from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, PreTrainedTokenizerFast |
| import requests |
|
|
| model = VisionEncoderDecoderModel.from_pretrained("Zayn/vit2distilgpt2") |
| vit_feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k") |
| tokenizer = PreTrainedTokenizerFast.from_pretrained("distilgpt2") |
|
|
| def vit2distilgpt2(img): |
| pixel_values = vit_feature_extractor(images=img, return_tensors="pt").pixel_values |
| encoder_outputs = model.generate(pixel_values.to('cpu'),num_beams=5) |
| generated_sentences = tokenizer.batch_decode(encoder_outputs, skip_special_tokens =True) |
|
|
| return(generated_sentences[0].split('.')[0]) |
|
|
| import gradio as gr |
|
|
| inputs = [ |
| gr.inputs.Image(type="pil", label = "Original Image") |
| ] |
|
|
| outputs = [ |
| gr.outputs.Textbox(label = 'Caption') |
| ] |
| title = "Image Captioning using Vision Transformer and GPT-2" |
| description = "Developed by Zayn" |
| article = "< a href='https://huggingface.co/Zayn/vit2distilgpt2'>Hugging Face AI Community</a>" |
| examples = [ |
| ["car.jpg"] |
| ] |
| gr.Interface( |
| vit2distilgpt2, |
| inputs, |
| outputs, |
| title = title, |
| description = description, |
| article = article, |
| examples = examples, |
| theme = "huggingface", |
| ).launch(debug=True,enable_queue=True) |
