Update README.md
Browse files
README.md
CHANGED
@@ -1,3 +1,45 @@
|
|
1 |
---
|
2 |
license: apache-2.0
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
3 |
---
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
---
|
2 |
license: apache-2.0
|
3 |
+
datasets:
|
4 |
+
- atasoglu/flickr8k-dataset
|
5 |
+
language:
|
6 |
+
- en
|
7 |
+
metrics:
|
8 |
+
- rouge
|
9 |
+
pipeline_tag: image-to-text
|
10 |
+
tags:
|
11 |
+
- image
|
12 |
+
- vision
|
13 |
---
|
14 |
+
|
15 |
+
Vision Encoder Decoder (ViT + BERT) model that fine-tuned on [flickr8k-dataset](https://huggingface.co/datasets/atasoglu/flickr8k-dataset) for image-to-text task.
|
16 |
+
|
17 |
+
Example:
|
18 |
+
|
19 |
+
```py
|
20 |
+
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, BertTokenizer
|
21 |
+
import torch
|
22 |
+
from PIL import Image
|
23 |
+
|
24 |
+
# load models
|
25 |
+
feature_extractor = ViTImageProcessor.from_pretrained("atasoglu/vit-bert-flickr8k")
|
26 |
+
tokenizer = BertTokenizer.from_pretrained("atasoglu/vit-bert-flickr8k")
|
27 |
+
model = VisionEncoderDecoderModel.from_pretrained("atasoglu/vit-bert-flickr8k")
|
28 |
+
|
29 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
30 |
+
model.to(device)
|
31 |
+
|
32 |
+
# load image
|
33 |
+
img = Image.open("example.jpg")
|
34 |
+
|
35 |
+
# encode (extracting features)
|
36 |
+
pixel_values = feature_extractor(images=[img], return_tensors="pt").pixel_values
|
37 |
+
pixel_values = pixel_values.to(device)
|
38 |
+
|
39 |
+
# generate caption
|
40 |
+
output_ids = model.generate(pixel_values)
|
41 |
+
|
42 |
+
# decode
|
43 |
+
preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
44 |
+
print(preds)
|
45 |
+
```
|