|
--- |
|
language: en |
|
tags: |
|
- image-classification |
|
- image-captioning |
|
|
|
--- |
|
|
|
# Poster2Plot |
|
|
|
An image captioning model to generate movie/t.v show plot from poster. It generates decent plots but is no way perfect. We are still working on improving the model. |
|
|
|
## Live demo on Hugging Face Spaces: https://huggingface.co/spaces/deepklarity/poster2plot |
|
|
|
# Model Details |
|
|
|
The base model uses a Vision Transformer (ViT) model as an image encoder and GPT-2 as a decoder. |
|
|
|
We used the following models: |
|
|
|
* Encoder: [google/vit-base-patch16-224-in21k](https://huggingface.co/google/vit-base-patch16-224-in21k) |
|
* Decoder: [gpt2](https://huggingface.co/gpt2) |
|
|
|
# Datasets |
|
|
|
Publicly available IMDb datasets were used to train the model. |
|
|
|
# How to use |
|
|
|
## In PyTorch |
|
|
|
```python |
|
import torch |
|
import re |
|
import requests |
|
from PIL import Image |
|
from transformers import AutoTokenizer, AutoFeatureExtractor, VisionEncoderDecoderModel |
|
|
|
# Pattern to ignore all the text after 2 or more full stops |
|
regex_pattern = "[.]{2,}" |
|
|
|
|
|
def post_process(text): |
|
try: |
|
text = text.strip() |
|
text = re.split(regex_pattern, text)[0] |
|
except Exception as e: |
|
print(e) |
|
pass |
|
return text |
|
|
|
|
|
def predict(image, max_length=64, num_beams=4): |
|
pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values |
|
pixel_values = pixel_values.to(device) |
|
|
|
with torch.no_grad(): |
|
output_ids = model.generate( |
|
pixel_values, |
|
max_length=max_length, |
|
num_beams=num_beams, |
|
return_dict_in_generate=True, |
|
).sequences |
|
|
|
preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True) |
|
pred = post_process(preds[0]) |
|
|
|
return pred |
|
|
|
|
|
model_name_or_path = "deepklarity/poster2plot" |
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
|
|
|
# Load model. |
|
|
|
model = VisionEncoderDecoderModel.from_pretrained(model_name_or_path) |
|
model.to(device) |
|
print("Loaded model") |
|
|
|
feature_extractor = AutoFeatureExtractor.from_pretrained(model.encoder.name_or_path) |
|
print("Loaded feature_extractor") |
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model.decoder.name_or_path, use_fast=True) |
|
if model.decoder.name_or_path == "gpt2": |
|
tokenizer.pad_token = tokenizer.eos_token |
|
|
|
print("Loaded tokenizer") |
|
|
|
url = "https://upload.wikimedia.org/wikipedia/en/2/26/Moana_Teaser_Poster.jpg" |
|
with Image.open(requests.get(url, stream=True).raw) as image: |
|
pred = predict(image) |
|
|
|
print(pred) |
|
|
|
``` |
|
|
|
|
|
|
|
|
