Add model card

README.md

@@ -0,0 +1,96 @@
+---
+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.
+
+# 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)
+
+```
+
+
+