peaceAsh commited on
Commit
bee5682
·
verified ·
1 Parent(s): 57a3d98

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +38 -0
app.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from PIL import Image
2
+ from transformers import VisionEncoderDecoderModel , ViTFeatureExtractor , PreTrainedTokenizerFast
3
+ import gradio as gr
4
+
5
+ model = VisionEncoderDecoderModel.from_pretrained("ydshieh/vit-gpt2-coco-en")
6
+ vit_feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch32-224-in21k")
7
+ tokenizer = PreTrainedTokenizerFast.from_pretrained("distilgpt2")
8
+
9
+
10
+ def caption_images(image):
11
+ pixel_values = vit_feature_extractor(images=image,return_tensors="pt").pixel_values
12
+ encoder_outputs = model.generate(pixel_values.to('cpu'),num_beams=5)
13
+ generated_sentence = tokenizer.batch_decode(encoder_outputs,skip_special_tokens=True)
14
+
15
+ return (generated_sentence[0].strip())
16
+
17
+
18
+ inputs = [
19
+ gr.components.Image(type='pil',label='Original Image')
20
+ ]
21
+
22
+ outputs = [
23
+ gr.components.Textbox(label='Caption')
24
+ ]
25
+
26
+ title = "Simple Image captioning Application"
27
+ description = "Upload an image to see the caption generated"
28
+ example =['/content/messi.jpg']
29
+
30
+ gr.Interface(
31
+ caption_images,
32
+ inputs,
33
+ outputs,
34
+ title=title,
35
+ description = description,
36
+ examples = example,
37
+ ).launch(debug=True)
38
+