Compare commits
10 Commits
45bde45ae6
...
5bebf1e9bb
Author | SHA1 | Date |
---|---|---|
|
5bebf1e9bb | |
|
d273135c52 | |
|
100350d7e6 | |
|
65636df6de | |
|
395ab07079 | |
|
b0acf67ee3 | |
|
ea737d2858 | |
|
973da61a93 | |
|
f6d2d95250 | |
|
462822078f |
84
README.md
84
README.md
|
@ -1,7 +1,11 @@
|
||||||
---
|
---
|
||||||
tags:
|
tags:
|
||||||
- image-classification
|
- image-to-text
|
||||||
library_name: generic
|
widget:
|
||||||
|
- src: https://huggingface.co/datasets/mishig/sample_images/resolve/main/football-match.jpg
|
||||||
|
example_title: Football Match
|
||||||
|
- src: https://huggingface.co/datasets/mishig/sample_images/resolve/main/dog-cat.jpg
|
||||||
|
example_title: Dog & Cat
|
||||||
---
|
---
|
||||||
|
|
||||||
## Example
|
## Example
|
||||||
|
@ -12,35 +16,89 @@ as a proof-of-concept for the 🤗 FlaxVisionEncoderDecoder Framework.
|
||||||
|
|
||||||
The model can be used as follows:
|
The model can be used as follows:
|
||||||
|
|
||||||
|
**In PyTorch**
|
||||||
```python
|
```python
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import requests
|
||||||
|
from PIL import Image
|
||||||
|
from transformers import ViTFeatureExtractor, AutoTokenizer, VisionEncoderDecoderModel
|
||||||
|
|
||||||
|
|
||||||
|
loc = "ydshieh/vit-gpt2-coco-en"
|
||||||
|
|
||||||
|
feature_extractor = ViTFeatureExtractor.from_pretrained(loc)
|
||||||
|
tokenizer = AutoTokenizer.from_pretrained(loc)
|
||||||
|
model = VisionEncoderDecoderModel.from_pretrained(loc)
|
||||||
|
model.eval()
|
||||||
|
|
||||||
|
|
||||||
|
def predict(image):
|
||||||
|
|
||||||
|
pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
output_ids = model.generate(pixel_values, max_length=16, num_beams=4, return_dict_in_generate=True).sequences
|
||||||
|
|
||||||
|
preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
||||||
|
preds = [pred.strip() for pred in preds]
|
||||||
|
|
||||||
|
return preds
|
||||||
|
|
||||||
|
|
||||||
|
# We will verify our results on an image of cute cats
|
||||||
|
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
|
||||||
|
with Image.open(requests.get(url, stream=True).raw) as image:
|
||||||
|
preds = predict(image)
|
||||||
|
|
||||||
|
print(preds)
|
||||||
|
# should produce
|
||||||
|
# ['a cat laying on top of a couch next to another cat']
|
||||||
|
|
||||||
|
```
|
||||||
|
|
||||||
|
**In Flax**
|
||||||
|
```python
|
||||||
|
|
||||||
|
import jax
|
||||||
import requests
|
import requests
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
from transformers import ViTFeatureExtractor, AutoTokenizer, FlaxVisionEncoderDecoderModel
|
from transformers import ViTFeatureExtractor, AutoTokenizer, FlaxVisionEncoderDecoderModel
|
||||||
|
|
||||||
|
|
||||||
loc = "ydshieh/vit-gpt2-coco-en"
|
loc = "ydshieh/vit-gpt2-coco-en"
|
||||||
|
|
||||||
feature_extractor = ViTFeatureExtractor.from_pretrained(loc)
|
feature_extractor = ViTFeatureExtractor.from_pretrained(loc)
|
||||||
tokenizer = AutoTokenizer.from_pretrained(loc)
|
tokenizer = AutoTokenizer.from_pretrained(loc)
|
||||||
model = FlaxVisionEncoderDecoderModel.from_pretrained(loc)
|
model = FlaxVisionEncoderDecoderModel.from_pretrained(loc)
|
||||||
|
|
||||||
# We will verify our results on an image of cute cats
|
gen_kwargs = {"max_length": 16, "num_beams": 4}
|
||||||
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
|
|
||||||
with Image.open(requests.get(url, stream=True).raw) as img:
|
|
||||||
pixel_values = feature_extractor(images=img, return_tensors="np").pixel_values
|
|
||||||
|
|
||||||
def generate_step(pixel_values):
|
|
||||||
|
|
||||||
output_ids = model.generate(pixel_values, max_length=16, num_beams=4).sequences
|
# This takes sometime when compiling the first time, but the subsequent inference will be much faster
|
||||||
|
@jax.jit
|
||||||
|
def generate(pixel_values):
|
||||||
|
output_ids = model.generate(pixel_values, **gen_kwargs).sequences
|
||||||
|
return output_ids
|
||||||
|
|
||||||
|
|
||||||
|
def predict(image):
|
||||||
|
|
||||||
|
pixel_values = feature_extractor(images=image, return_tensors="np").pixel_values
|
||||||
|
output_ids = generate(pixel_values)
|
||||||
preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
||||||
preds = [pred.strip() for pred in preds]
|
preds = [pred.strip() for pred in preds]
|
||||||
|
|
||||||
return preds
|
return preds
|
||||||
|
|
||||||
preds = generate_step(pixel_values)
|
|
||||||
|
# We will verify our results on an image of cute cats
|
||||||
|
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
|
||||||
|
with Image.open(requests.get(url, stream=True).raw) as image:
|
||||||
|
preds = predict(image)
|
||||||
|
|
||||||
print(preds)
|
print(preds)
|
||||||
|
|
||||||
# should produce
|
# should produce
|
||||||
# ['a cat laying on top of a couch next to another cat']
|
# ['a cat laying on top of a couch next to another cat']
|
||||||
|
|
||||||
```
|
```
|
14
pipeline.py
14
pipeline.py
|
@ -20,7 +20,7 @@ class PreTrainedPipeline():
|
||||||
max_length = 16
|
max_length = 16
|
||||||
num_beams = 4
|
num_beams = 4
|
||||||
# self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
|
# self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
|
||||||
self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "return_dict_in_generate": True}
|
self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "return_dict_in_generate": True, "output_scores": True}
|
||||||
|
|
||||||
self.model.to("cpu")
|
self.model.to("cpu")
|
||||||
self.model.eval()
|
self.model.eval()
|
||||||
|
@ -29,8 +29,12 @@ class PreTrainedPipeline():
|
||||||
def _generate(pixel_values):
|
def _generate(pixel_values):
|
||||||
|
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
output_ids = self.model.generate(pixel_values, **self.gen_kwargs).sequences
|
|
||||||
return output_ids
|
outputs = self.model.generate(pixel_values, **self.gen_kwargs)
|
||||||
|
output_ids = outputs.sequences
|
||||||
|
sequences_scores = outputs.sequences_scores
|
||||||
|
|
||||||
|
return output_ids, sequences_scores
|
||||||
|
|
||||||
self.generate = _generate
|
self.generate = _generate
|
||||||
|
|
||||||
|
@ -49,10 +53,10 @@ class PreTrainedPipeline():
|
||||||
# pixel_values = self.feature_extractor(images=inputs, return_tensors="np").pixel_values
|
# pixel_values = self.feature_extractor(images=inputs, return_tensors="np").pixel_values
|
||||||
pixel_values = self.feature_extractor(images=inputs, return_tensors="pt").pixel_values
|
pixel_values = self.feature_extractor(images=inputs, return_tensors="pt").pixel_values
|
||||||
|
|
||||||
output_ids = self.generate(pixel_values)
|
output_ids, sequences_scores = self.generate(pixel_values)
|
||||||
preds = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
preds = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
||||||
preds = [pred.strip() for pred in preds]
|
preds = [pred.strip() for pred in preds]
|
||||||
|
|
||||||
preds = [{"label": preds[0], "score": 1.0}]
|
preds = [{"label": preds[0], "score": float(sequences_scores[0])}]
|
||||||
|
|
||||||
return preds
|
return preds
|
||||||
|
|
|
@ -1,6 +1,7 @@
|
||||||
{
|
{
|
||||||
"do_normalize": true,
|
"do_normalize": true,
|
||||||
"do_resize": true,
|
"do_resize": true,
|
||||||
|
"feature_extractor_type": "ViTFeatureExtractor",
|
||||||
"image_mean": [
|
"image_mean": [
|
||||||
0.5,
|
0.5,
|
||||||
0.5,
|
0.5,
|
||||||
|
|
Binary file not shown.
Loading…
Reference in New Issue