Ashish Soni commited on
Commit
ae789f7
·
verified ·
1 Parent(s): f356bda

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -2
app.py CHANGED
@@ -1,8 +1,30 @@
 
 
1
  import gradio as gr
 
 
2
  from PIL import Image
 
3
  import spaces
 
4
  from transformers import pipeline
5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
  # Initialize Model
7
  get_completion = pipeline("image-to-text",model="Salesforce/blip-image-captioning-base", device=0)
8
 
@@ -17,7 +39,9 @@ def captioner(input: Image.Image) -> str:
17
  Returns:
18
  str: The generated caption text.
19
  """
20
- output = get_completion(input)
 
 
21
  return output[0]['generated_text']
22
 
23
  ####### GRADIO APP #######
@@ -42,6 +66,7 @@ with demo:
42
  interface = gr.Interface(fn=captioner,
43
  inputs=[gr.Image(label="Upload image", type="pil")],
44
  outputs=[gr.Textbox(label="Caption")],
45
- allow_flagging="never")
 
46
 
47
  demo.launch()
 
1
+ import base64
2
+
3
  import gradio as gr
4
+ import io
5
+
6
  from PIL import Image
7
+
8
  import spaces
9
+
10
  from transformers import pipeline
11
 
12
+ def image_to_base64_str(pil_image: Image.Image) -> str:
13
+ """
14
+ Converts a PIL image to a base64 encoded string.
15
+
16
+ Args:
17
+ pil_image (Image.Image): The PIL image to be converted.
18
+
19
+ Returns:
20
+ str: The base64 encoded string representation of the image.
21
+ """
22
+ byte_arr = io.BytesIO()
23
+ pil_image.save(byte_arr, format='PNG')
24
+ byte_arr = byte_arr.getvalue()
25
+ return str(base64.b64encode(byte_arr).decode('utf-8'))
26
+
27
+
28
  # Initialize Model
29
  get_completion = pipeline("image-to-text",model="Salesforce/blip-image-captioning-base", device=0)
30
 
 
39
  Returns:
40
  str: The generated caption text.
41
  """
42
+
43
+ base64_image = image_to_base64_str(image)
44
+ output = get_completion(base64_image)
45
  return output[0]['generated_text']
46
 
47
  ####### GRADIO APP #######
 
66
  interface = gr.Interface(fn=captioner,
67
  inputs=[gr.Image(label="Upload image", type="pil")],
68
  outputs=[gr.Textbox(label="Caption")],
69
+ allow_flagging="never",
70
+ examples=["christmas_dog.jpeg", "bird_flight.jpeg", "cow.jpeg"]))
71
 
72
  demo.launch()