CoolFace
Apppublic

ipvikas/ImageProcessing

sourceHugging Faceupdated 1y agoView on Hugging Face
3likes
image_summary.py70 linesDownload Raw Back to root
1import os2import gradio as gr3from PIL import Image4import requests5 6# from transformers import ViTFeatureExtractor7# feature_extractor = ViTFeatureExtractor()8from transformers import ViTImageProcessor9feature_extractor = ViTImageProcessor()10 11 12# or, to load one that corresponds to a checkpoint on the hub:13# feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")14feature_extractor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")15 16 17 18from transformers import VisionEncoderDecoderModel19# initialize a vit-bert from a pretrained ViT and a pretrained BERT model. Note that the cross-attention layers will be randomly initialized20model = VisionEncoderDecoderModel.from_encoder_decoder_pretrained(21    "google/vit-base-patch16-224-in21k", "bert-base-uncased"22)23# saving model after fine-tuning24model.save_pretrained("./vit-bert")25# load fine-tuned model26model = VisionEncoderDecoderModel.from_pretrained("./vit-bert")27 28#####################29from transformers import AutoTokenizer30repo_name = "ydshieh/vit-gpt2-coco-en"31# feature_extractor = ViTFeatureExtractor.from_pretrained(repo_name)32feature_extractor = ViTImageProcessor.from_pretrained(repo_name)33 34 35tokenizer = AutoTokenizer.from_pretrained(repo_name)36model = VisionEncoderDecoderModel.from_pretrained(repo_name)37 38def get_quote(image):39    40    ##############41    pixel_values = feature_extractor(image, return_tensors="pt").pixel_values42    # autoregressively generate text (using beam search or other decoding strategy)43    generated_ids = model.generate(pixel_values, max_length=16, num_beams=4, return_dict_in_generate=True)44    45    ################46    # decode into text47    preds = tokenizer.batch_decode(generated_ids[0], skip_special_tokens=True)48    preds = [pred.strip() for pred in preds]49    return preds50 51#1: Text to Speech52title = "Sentence, listing all the items present in the image file"53 54description = "Summary of items in the Image"55examples=[["english.png"],["Parag_Letter_j.jpg"]]56 57 58 59image_summary_demo = gr.Interface(fn=get_quote, 60                                  inputs=gr.Image(type="pil"), 61                                  outputs=['text'],62                                  title = title, 63                                  description = "Upload an image file and get text from it" ,                                64                                  cache_examples=False, 65                                  examples=examples, 66                                  enable_queue=True)67 68# if __name__ == "__main__":69    70#     demo.launch(debug=True, cache_examples=True)