Update app.py

app.py

@@ -6,56 +6,67 @@ Automatically generated by Colaboratory.
 Original file is located at
     https://colab.research.google.com/drive/1Uvn7yZCyrMpOYNPb7K0G45tQZJVx8LyX
 """
-!pip install torch
-
-from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
-import gradio as gr
-import torch
 from PIL import Image
+from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, PreTrainedTokenizerFast
+import requests
 
-model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-feature_extractor = ViTFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+model = VisionEncoderDecoderModel.from_pretrained("sachin/vit2distilgpt2")
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+vit_feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
 
+tokenizer = PreTrainedTokenizerFast.from_pretrained("distilgpt2")
 
+# url = 'https://d2gp644kobdlm6.cloudfront.net/wp-content/uploads/2016/06/bigstock-Shocked-and-surprised-boy-on-t-113798588-300x212.jpg'
 
-max_length = 16
-num_beams = 4
-gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
+# with Image.open(requests.get(url, stream=True).raw) as img:
+#     pixel_values = vit_feature_extractor(images=img, return_tensors="pt").pixel_values
 
-def predict_step(image):
-#   images = []
-#   for image_path in image_paths:
-#     i_image = Image.open(image_path)
-#     if i_image.mode != "RGB":
-#       i_image = i_image.convert(mode="RGB")
+#encoder_outputs = model.generate(pixel_values.to('cpu'),num_beams=5)
 
-#     images.append(i_image)
+#generated_sentences = tokenizer.batch_decode(encoder_outputs, skip_special_tokens=True)
 
-  pixel_values = feature_extractor(images = image, return_tensors = "pt").pixel_values
-  pixel_values = pixel_values.to(device)
+#generated_sentences
 
-  output_ids = model.generate(pixel_values, **gen_kwargs)
+#naive text processing 
+#generated_sentences[0].split('.')[0]
 
-  preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
-  preds = [pred.strip() for pred in preds]
-  return preds
+# inference function
 
-inputs = [ gr.inputs.Image(type = 'pil', label = 'Original Image')]
-outputs = [ gr.outputs.Textbox(label = 'Caption')]
-title = 'Image Captioning using ViT + GPT2'
-description = 'ViT and GPT2 are used here to generate Image Caption for the user uploaded image.'
-article = " <a href=' https://huggingface.co/sachin/vit2distilgpt2 '>Model Repository on Hugging Face Model Hub</a>" 
+def vit2distilgpt2(img):
+  pixel_values = vit_feature_extractor(images=img, return_tensors="pt").pixel_values
+  encoder_outputs = generated_ids = model.generate(pixel_values.to('cpu'),num_beams=5)
+  generated_sentences = tokenizer.batch_decode(encoder_outputs, skip_special_tokens=True)
 
-gr.Interface(
-    predict_step,
-    inputs, outputs, 
-    title = title,
-    description = description, 
-    article = article,
-    theme = 'huggingface'
-).launch(debug = True, enable_queue = True)
+  return(generated_sentences[0].split('.')[0])
+
+#!wget https://media.glamour.com/photos/5f171c4fd35176eaedb36823/master/w_2560%2Cc_limit/bike.jpg
+
+import gradio as gr
 
+inputs = [
+    gr.inputs.Image(type="pil", label="Original Image")
+]
+
+outputs = [
+    gr.outputs.Textbox(label = 'Caption')
+]
+
+title = "Image Captioning using ViT + GPT2"
+description = "ViT and GPT2 are used to generate Image Caption for the uploaded image. COCO Dataset was used for training. This image captioning model might have some biases that we couldn't figure during our stress testing, so if you find any bias (gender, race and so on) please use `Flag` button to flag the image with bias"
+article = " <a href='https://huggingface.co/sachin/vit2distilgpt2'>Model Repo on Hugging Face Model Hub</a>"
+examples = [
+    ["people-walking-street-pedestrian-crossing-traffic-light-city.jpeg"],
+    ["elonmusk.jpeg"]
+
+]
+
+gr.Interface(
+    vit2distilgpt2,
+    inputs,
+    outputs,
+    title=title,
+    description=description,
+    article=article,
+    examples=examples,
+    theme="huggingface",
+).launch(debug=True, enable_queue=True)