Iblouse commited on
Commit
973104b
1 Parent(s): ce36534

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +24 -6
app.py CHANGED
@@ -2,15 +2,33 @@ import os
2
  import gradio as gr
3
  from transformers import pipeline
4
 
5
- pipe = pipeline("image-to-text",
6
- model="./models/Salesforce/blip-image-captioning-base")
 
 
 
 
 
7
 
8
  def launch(input):
 
 
 
 
 
 
 
 
 
9
  out = pipe(input)
10
  return out[0]['generated_text']
11
 
12
- iface = gr.Interface(launch,
13
- inputs=gr.Image(type='pil'),
14
- outputs="text")
 
 
 
15
 
16
- iface.launch()
 
 
2
  import gradio as gr
3
  from transformers import pipeline
4
 
5
+ # Initialize the pipeline with the image-to-text model
6
+ model_path = "Salesforce/blip-image-captioning-base"
7
+ if not os.path.exists(model_path):
8
+ raise FileNotFoundError(f"Model path {model_path} does not exist. Please provide a valid path.")
9
+
10
+ # Initialize the image-to-text pipeline with the specified model
11
+ pipe = pipeline("image-to-text", model=model_path)
12
 
13
  def launch(input):
14
+ """
15
+ Function to generate image caption.
16
+
17
+ Args:
18
+ input (PIL.Image): Input image for captioning.
19
+
20
+ Returns:
21
+ str: Generated caption for the input image.
22
+ """
23
  out = pipe(input)
24
  return out[0]['generated_text']
25
 
26
+ # Create a Gradio interface for the image-to-text pipeline
27
+ iface = gr.Interface(
28
+ fn=launch, # Function to generate captions
29
+ inputs=gr.Image(type='pil'), # Input type: Image (PIL format)
30
+ outputs="text" # Output type: Text
31
+ )
32
 
33
+ # Launch the Gradio interface
34
+ iface.launch()