generate_captions.py (4990B)
1 import numpy as np 2 import os, sys 3 import torch 4 from fairseq import utils, tasks 5 from fairseq import checkpoint_utils 6 import PIL 7 import torchvision 8 import json 9 import pathlib 10 import pandas as pd 11 import tqdm 12 13 import mmsrl 14 sys.path.append(mmsrl.__path__[0] + "/../OFA") 15 16 from utils.eval_utils import eval_step 17 from tasks.mm_tasks.caption import CaptionTask 18 from models.ofa import OFAModel 19 20 21 class ImageCaption(): 22 def __init__(self, images_path): 23 # Register caption task 24 tasks.register_task('caption',CaptionTask) 25 self.images_path = images_path 26 27 # turn on cuda if GPU is available 28 self.use_cuda = torch.cuda.is_available() 29 # use fp16 only when GPU is available 30 self.use_fp16 = False 31 32 # Load pretrained ckpt & config 33 overrides={"bpe_dir":"OFA/utils/BPE", "eval_cider":False, "beam":5, "max_len_b":16, "no_repeat_ngram_size":3, "seed":7} 34 self.models, self.cfg, self.task = checkpoint_utils.load_model_ensemble_and_task( 35 utils.split_paths('OFA/checkpoints/caption_large_best_clean.pt'), 36 arg_overrides=overrides 37 ) 38 39 # Move models to GPU 40 for model in self.models: 41 model.eval() 42 if self.use_fp16: 43 model.half() 44 if self.use_cuda and not self.cfg.distributed_training.pipeline_model_parallel: 45 model.cuda() 46 model.prepare_for_inference_(self.cfg) 47 48 # Initialize generator 49 self.generator = self.task.build_generator(self.models, self.cfg.generation) 50 51 # Text preprocess 52 self.bos_item = torch.LongTensor([self.task.src_dict.bos()]) 53 self.eos_item = torch.LongTensor([self.task.src_dict.eos()]) 54 self.pad_idx = self.task.src_dict.pad() 55 56 mean = [0.5, 0.5, 0.5] 57 std = [0.5, 0.5, 0.5] 58 59 self.patch_resize_transform = torchvision.transforms.Compose([ 60 lambda image: image.convert("RGB"), 61 torchvision.transforms.Resize((self.cfg.task.patch_image_size, self.cfg.task.patch_image_size), interpolation=PIL.Image.BICUBIC), 62 torchvision.transforms.ToTensor(), 63 torchvision.transforms.Normalize(mean=mean, std=std), 64 ]) 65 66 def encode_text(self, text, length=None, append_bos=False, append_eos=False): 67 s = self.task.tgt_dict.encode_line( 68 line=self.task.bpe.encode(text), 69 add_if_not_exist=False, 70 append_eos=False 71 ).long() 72 if length is not None: 73 s = s[:length] 74 if append_bos: 75 s = torch.cat([self.bos_item, s]) 76 if append_eos: 77 s = torch.cat([s, self.eos_item]) 78 return s 79 80 def construct_sample(self, data, length=None, ocr=False): 81 image = PIL.Image.open(self.images_path / data['image']) 82 patch_image = self.patch_resize_transform(image).unsqueeze(0) 83 patch_mask = torch.tensor([True]) 84 src_text = self.encode_text(" what does the image describe?", append_bos=True, append_eos=True) 85 if ocr: 86 ocr_text = data['OCR'].replace('\n', ' ') 87 ocr_encoded = self.encode_text(ocr_text, length=length, append_bos=True, append_eos=True) 88 src_text = torch.cat((ocr_encoded, src_text)) 89 90 src_length = torch.LongTensor([src_text.ne(self.pad_idx).long().sum()]) 91 sample = { 92 "id":np.array(['42']), 93 "net_input": { 94 "src_tokens": src_text.unsqueeze(0), 95 "src_lengths": src_length, 96 "patch_images": patch_image, 97 "patch_masks": patch_mask 98 } 99 } 100 return sample 101 102 def generate(self, data, length=None, ocr = False): 103 104 # Construct input sample & preprocess for GPU if cuda available 105 sample = self.construct_sample(data, length, ocr) 106 sample = utils.move_to_cuda(sample) if self.use_cuda else sample 107 sample = utils.apply_to_sample(apply_half, sample) if self.use_fp16 else sample 108 109 with torch.no_grad(): 110 result, scores = eval_step(self.task, self.generator, self.models, sample) 111 112 return result[0]['caption'] 113 114 115 if __name__=='__main__': 116 117 data_path = pathlib.Path(os.environ['DATA_PATH']) 118 images_path = data_path / "mmsrl" / "all" / 'images' 119 caption_generator = ImageCaption(images_path) 120 captions_path = data_path / "mmsrl" / "all" / "captions.csv" 121 with open(captions_path, 'w', buffering=1) as f: 122 for split in ['val', 'unseen_test', 'train']: 123 annot_file = data_path / "mmsrl" / "all" / 'annotations' / f"{split}.jsonl" 124 with annot_file.open('r') as json_file: 125 dataset = list(map(json.loads, json_file)) 126 #print("split: ", len(dataset), "samples") 127 for data in tqdm.tqdm(dataset, desc=f"{split}", leave=False): 128 caption = caption_generator.generate(data) 129 f.write(data['image'] + '\t' + caption + '\n')