KFrimps commited on
Commit
87cbda7
·
verified ·
1 Parent(s): 9c37b60

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +38 -0
app.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import ViltProcessor, ViltForQuestionAnswering
3
+ from PIL import Image
4
+ import torch
5
+
6
+ # Load the processor and model
7
+ processor = ViltProcessor.from_pretrained("MariaK/vilt_finetuned_200")
8
+ model = ViltForQuestionAnswering.from_pretrained("MariaK/vilt_finetuned_200")
9
+ device = "cuda" if torch.cuda.is_available() else "cpu"
10
+ model.to(device)
11
+
12
+ def predict(image, question):
13
+ # prepare inputs
14
+ inputs = processor(image, question, return_tensors="pt").to(device)
15
+
16
+ # forward pass
17
+ with torch.no_grad():
18
+ outputs = model(**inputs)
19
+
20
+ logits = outputs.logits
21
+ idx = logits.argmax(-1).item()
22
+ predicted_answer = model.config.id2label[idx]
23
+ return predicted_answer
24
+
25
+ # Create the Gradio interface
26
+ iface = gr.Interface(
27
+ fn=predict,
28
+ inputs=[
29
+ gr.Image(type="pil"),
30
+ gr.Textbox(lines=1, placeholder="Enter your question here..."),
31
+ ],
32
+ outputs="text",
33
+ title="Visual Question Answering with Fine-tuned Vilt",
34
+ description="Upload an image and ask a question about it!",
35
+ )
36
+
37
+ # Launch the interface
38
+ iface.launch(share=True) # Set share=True to share the space