example

.gitignore

@@ -3,3 +3,4 @@
 *.yaml
 *ipynb_checkpoints
 __pycache__
+*.json

model.py

@@ -0,0 +1,89 @@
+from huggingface_hub import hf_hub_url, cached_download
+from PIL import Image
+import os
+import json
+import glob
+import random
+from typing import Any, Dict, List
+import torch
+import torchvision
+
+import wordsegment as ws
+
+from virtex.config import Config
+from virtex.factories import TokenizerFactory, PretrainingModelFactory
+from virtex.utils.checkpointing import CheckpointManager
+
+CONFIG_PATH = "config.yaml"
+MODEL_PATH = "checkpoint_last5.pth"
+VALID_SUBREDDITS_PATH = "subreddit_list.json"
+SAMPLES_PATH = "./samples/*.jpg"
+
+class ImageLoader():
+    def __init__(self):
+        self.transformer = torchvision.transforms.Compose([torchvision.transforms.Resize(256),
+                               torchvision.transforms.CenterCrop(224),
+                               torchvision.transforms.ToTensor()])
+    def load(self, im_path, prompt = ""):
+        im = torch.FloatTensor(self.transformer(Image.open(im_path))).unsqueeze(0)
+        return {"image": im, "decode_prompt": prompt}
+    def transform(self, image, prompt = ""):
+        im = torch.FloatTensor(self.transformer(image)).unsqueeze(0)
+        return {"image": im, "decode_prompt": prompt}
+
+class VirTexModel():
+    def __init__(self):
+        self.config = Config(CONFIG_PATH)
+        ws.load()
+        self.device = 'cpu'
+        self.tokenizer = TokenizerFactory.from_config(self.config)
+        self.model = PretrainingModelFactory.from_config(self.config).to(self.device)
+        CheckpointManager(model=self.model).load("./checkpoint_last5.pth")
+        self.model.eval()
+        self.valid_subs = json.load(open(VALID_SUBREDDITS_PATH))
+        
+    def predict(self, image_dict, sub_prompt = None, prompt = ""):
+        if sub_prompt is None:
+            subreddit_tokens = torch.tensor([self.model.sos_index], device=self.device).long()
+        else:
+            subreddit_tokens = torch.tensor([self.tokenizer.token_to_id(sub_prompt)], device=self.device).long()
+        predictions: List[Dict[str, Any]] = []
+        
+        is_valid_subreddit = False
+        subreddit, rest_of_caption = "", ""
+        while not is_valid_subreddit:
+            
+            with torch.no_grad():
+                caption = self.model(image_dict)["predictions"][0].tolist()
+                
+            if self.tokenizer.token_to_id("[SEP]") in caption:
+                sep_index = caption.index(self.tokenizer.token_to_id("[SEP]"))
+                caption[sep_index] = self.tokenizer.token_to_id("://")
+            
+            caption = self.tokenizer.decode(caption)
+            
+            if "://" in caption:
+                subreddit, rest_of_caption = caption.split("://")
+                subreddit = "".join(subreddit.split())
+                rest_of_caption = rest_of_caption.strip()
+            else:
+                subreddit, rest_of_caption = "", caption
+            
+            is_valid_subreddit = True if sub_prompt is not None else subreddit in self.valid_subs
+            
+            
+        return subreddit, rest_of_caption
+
+def download_files():
+    #download model files
+    download_files = [CONFIG_PATH, MODEL_PATH, VALID_SUBREDDITS_PATH]
+    for f in download_files:
+        fp = cached_download(hf_hub_url("zamborg/redcaps", filename=f))
+        os.system(f"cp {fp} ./{f}")
+
+def get_samples():
+    return glob.glob(SAMPLES_PATH)
+
+def get_rand_img(samples):
+    return samples[random.randint(0,len(samples)-1)]
+

requirements.txt

@@ -0,0 +1,19 @@
+albumentations>=0.5.0
+Cython>=0.25
+ftfy==5.8
+future==0.18.0
+huggingface-hub==0.1.2
+lmdb==0.97
+loguru==0.3.2
+mypy_extensions==0.4.1
+lvis==0.5.3
+numpy>=1.17
+opencv-python==4.1.2.30
+scikit-learn==0.21.3
+sentencepiece>=0.1.90
+torch==1.7.0
+torchvision==0.8
+tqdm>=4.50.0
+wordsegment==1.3.1
+git+git://github.com/facebookresearch/fvcore.git#egg=fvcore
+git+git://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI