rogman commited on
Commit
808e48c
1 Parent(s): 2fa6d13

Create new file

Browse files
Files changed (1) hide show
  1. app.py +48 -0
app.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import gradio as gr
3
+ import torch
4
+ import PIL
5
+
6
+ from flamingo_mini import FlamingoConfig, FlamingoModel, FlamingoProcessor
7
+
8
+
9
+
10
+ EXAMPLES_DIR = 'examples'
11
+ DEFAULT_PROMPT = "<image>"
12
+
13
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
14
+
15
+ model = FlamingoModel.from_pretrained('dhansmair/flamingo-mini')
16
+ model.to(device)
17
+ model.eval()
18
+
19
+ processor = FlamingoProcessor(model.config, load_vision_processor=True)
20
+
21
+ # setup some example images
22
+ examples = []
23
+ if os.path.isdir(EXAMPLES_DIR):
24
+ for file in os.listdir(EXAMPLES_DIR):
25
+ path = EXAMPLES_DIR + "/" + file
26
+ examples.append([path, DEFAULT_PROMPT])
27
+
28
+
29
+ def predict_caption(image, prompt):
30
+ assert isinstance(prompt, str)
31
+
32
+ features = processor.extract_features(image).to(device)
33
+ caption = model.generate_captions(processor,
34
+ visual_features=features,
35
+ prompt=prompt)
36
+
37
+ if isinstance(caption, list):
38
+ caption = caption[0]
39
+
40
+ return caption
41
+
42
+
43
+ iface = gr.Interface(fn=predict_caption,
44
+ inputs=[gr.Image(type="pil"), gr.Textbox(value=DEFAULT_PROMPT, label="Prompt")],
45
+ examples=examples,
46
+ outputs="text")
47
+
48
+ iface.launch()