chrisjcc commited on
Commit
949fb1d
·
verified ·
1 Parent(s): 0fa5bd0

Add building an image captioning app

Browse files
Files changed (1) hide show
  1. app.py +39 -0
app.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import io
3
+ import IPython.display
4
+ from PIL import Image
5
+ import base64
6
+
7
+ from transformers import pipeline
8
+ import gradio as gr
9
+
10
+ hf_api_key = os.environ['HF_API_KEY']
11
+
12
+ get_completion = pipeline("ner", model="Salesforce/blip-image-captioning-base")
13
+
14
+ def image_to_base64_str(pil_image):
15
+ byte_arr = io.BytesIO()
16
+ pil_image.save(byte_arr, format='PNG')
17
+ byte_arr = byte_arr.getvalue()
18
+
19
+ return str(base64.b64encode(byte_arr).decode('utf-8'))
20
+
21
+ def captioner(image):
22
+ base64_image = image_to_base64_str(image)
23
+ result = get_completion(base64_image)
24
+
25
+ return result[0]['generated_text']
26
+
27
+ gr.close_all()
28
+ demo = gr.Interface(fn=captioner,
29
+ inputs=[gr.Image(label="Upload image", type="pil")],
30
+ outputs=[gr.Textbox(label="Caption")],
31
+ title="Image Captioning with BLIP",
32
+ description="Caption any image using the BLIP model",
33
+ allow_flagging="never",
34
+ examples=["christmas_dog.jpeg", "bird_flight.jpeg", "cow.jpeg"])
35
+
36
+ demo.launch(
37
+ share=True,
38
+ #server_port=int(os.environ['PORT1'])
39
+ )