Update app.py

app.py

@@ -1,15 +1,37 @@
-from transformers import AutoProcessor, AutoModelForCausalLM
-import requests
 from PIL import Image
+from transformers import VisionEncoderDecoderModel , ViTFeatureExtractor , PreTrainedTokenizerFast
+import gradio as gr
 
-processor = AutoProcessor.from_pretrained("microsoft/git-base-coco")
-model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-coco")
+model = VisionEncoderDecoderModel.from_pretrained("ydshieh/vit-gpt2-coco-en")
+vit_feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch32-224-in21k")
+tokenizer = PreTrainedTokenizerFast.from_pretrained("distilgpt2")
 
-url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-image = Image.open(requests.get(url, stream=True).raw)
 
-pixel_values = processor(images=image, return_tensors="pt").pixel_values
+def caption_images(image):
+  pixel_values = vit_feature_extractor(images=image,return_tensors="pt").pixel_values
+  encoder_outputs = model.generate(pixel_values.to('cpu'),num_beams=5)
+  generated_sentence = tokenizer.batch_decode(encoder_outputs,skip_special_tokens=True)
 
-generated_ids = model.generate(pixel_values=pixel_values, max_length=50)
-generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-print(generated_caption)
+  return (generated_sentence[0].strip())
+
+
+inputs = [
+    gr.components.Image(type='pil',label='Original Image')
+]
+
+outputs = [
+    gr.components.Textbox(label='Caption')
+]
+
+title = "Simple Image captioning Application"
+description = "Upload an image to see the caption generated"
+example =['/content/messi.jpg']
+
+gr.Interface(
+    caption_images,
+    inputs,
+    outputs,
+    title=title,
+    description = description,
+    examples = example,
+).launch(debug=True)