ravi-vc commited on
Commit
f1fec60
·
verified ·
1 Parent(s): 5b624fc

first commit

Browse files
Files changed (1) hide show
  1. app.py +33 -0
app.py CHANGED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
3
+ from PIL import Image
4
+ import torch
5
+
6
+ # Load ViT-GPT2 (Apache-2.0 licensed, safe to use)
7
+ model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
8
+ feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
9
+ tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
10
+
11
+ device = "cuda" if torch.cuda.is_available() else "cpu"
12
+ model.to(device)
13
+
14
+ def caption_image(image):
15
+ # Convert image to tensor
16
+ pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
17
+
18
+ # Generate caption
19
+ output_ids = model.generate(pixel_values, max_length=50, num_beams=4)
20
+ caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
21
+ return caption
22
+
23
+ # Build Gradio app
24
+ demo = gr.Interface(
25
+ fn=caption_image,
26
+ inputs=gr.Image(type="pil"),
27
+ outputs="text",
28
+ title="Chart Analyzer",
29
+ description="Upload a chart/visualization image and get a description of it."
30
+ )
31
+
32
+ if __name__ == "__main__":
33
+ demo.launch()