Add VQA

app_caption.py

@@ -18,7 +18,7 @@ def create_demo():
             model_name = gr.Dropdown(label='Model', choices=['Prismer-Base, Prismer-Large'], value='Prismer-Base')
             run_button = gr.Button('Run')
         with gr.Column(scale=1.5):
-            caption = gr.Text(label='Caption')
+            caption = gr.Text(label='Model Prediction')
             with gr.Row():
                 depth = gr.Image(label='Depth')
                 edge = gr.Image(label='Edge')

app_vqa.py

@@ -16,9 +16,10 @@ def create_demo():
         with gr.Column():
             image = gr.Image(label='Input', type='filepath')
             model_name = gr.Dropdown(label='Model', choices=['Prismer-Base', 'Prismer-Large'], value='Prismer-Base')
+            question = gr.Text(label='Question')
             run_button = gr.Button('Run')
         with gr.Column(scale=1.5):
-            caption = gr.Text(label='Caption')
+            answer = gr.Text(label='Model Prediction')
             with gr.Row():
                 depth = gr.Image(label='Depth')
                 edge = gr.Image(label='Edge')
@@ -28,8 +29,8 @@ def create_demo():
                 object_detection = gr.Image(label='Object Detection')
                 ocr = gr.Image(label='OCR Detection')
 
-    inputs = [image, model_name]
-    outputs = [caption, depth, edge, normals, segmentation, object_detection, ocr]
+    inputs = [image, model_name, question]
+    outputs = [answer, depth, edge, normals, segmentation, object_detection, ocr]
 
     # paths = sorted(pathlib.Path('prismer/images').glob('*'))
     # examples = [[path.as_posix(), 'prismer_base'] for path in paths]
@@ -44,9 +45,9 @@ def create_demo():
     gr.Examples(examples=examples,
                 inputs=inputs,
                 outputs=outputs,
-                fn=model.run_caption)
+                fn=model.run_vqa_model)
 
-    run_button.click(fn=model.run_caption, inputs=inputs, outputs=outputs)
+    run_button.click(fn=model.run_vqa_model, inputs=inputs, outputs=outputs)
 
 
 if __name__ == '__main__':

prismer_model.py

@@ -16,7 +16,9 @@ submodule_dir = repo_dir / 'prismer'
 sys.path.insert(0, submodule_dir.as_posix())
 
 from dataset import create_dataset, create_loader
+from dataset.utils import pre_question
 from model.prismer_caption import PrismerCaption
+from model.prismer_vqa import PrismerVQA
 
 
 def download_models() -> None:
@@ -73,6 +75,11 @@ class Model:
         if exp_name == self.exp_name:
             return
 
+        if self.exp_name == 'Prismer-Base':
+            model_name = 'prismer_base'
+        elif self.exp_name == 'Prismer-Large':
+            model_name = 'prismer_large'
+
         if self.mode == 'caption':
             config = {
                 'dataset': 'demo',
@@ -80,13 +87,12 @@ class Model:
                 'label_path': 'prismer/helpers/labels',
                 'experts': ['depth', 'normal', 'seg_coco', 'edge', 'obj_detection', 'ocr_detection'],
                 'image_resolution': 480,
-                'prismer_model': 'prismer_base' if self.exp_name == 'Prismer-Base' else 'prismer_large',
+                'prismer_model': model_name,
                 'freeze': 'freeze_vision',
-                'prefix': 'A picture of',
+                'prefix': '',
             }
-
             model = PrismerCaption(config)
-            state_dict = torch.load(f'prismer/logging/caption_{exp_name}/pytorch_model.bin', map_location='cuda:0')
+            state_dict = torch.load(f'prismer/logging/pretrain_{model_name}/pytorch_model.bin', map_location='cuda:0')
 
         elif self.mode == 'vqa':
             config = {
@@ -95,13 +101,12 @@ class Model:
                 'label_path': 'prismer/helpers/labels',
                 'experts': ['depth', 'normal', 'seg_coco', 'edge', 'obj_detection', 'ocr_detection'],
                 'image_resolution': 480,
-                'prismer_model': 'prismer_base' if self.exp_name == 'Prismer-Base' else 'prismer_large',
+                'prismer_model': model_name,
                 'freeze': 'freeze_vision',
-                'prefix': 'A picture of',
             }
 
-            model = PrismerCaption(config)
-            state_dict = torch.load(f'prismer/logging/caption_{exp_name}/pytorch_model.bin', map_location='cuda:0')
+            model = PrismerVQA(config)
+            state_dict = torch.load(f'prismer/logging/vqa_{model_name}/pytorch_model.bin', map_location='cuda:0')
 
         model.load_state_dict(state_dict)
         model.eval()
@@ -131,14 +136,21 @@ class Model:
         return caption, *out_paths
 
     @torch.inference_mode()
-    def run_vqa_model(self, exp_name: str) -> str:
+    def run_vqa_model(self, exp_name: str, question: str) -> str:
         self.set_model(exp_name)
-        _, test_dataset = create_dataset('vqa', self.config)
+        _, test_dataset = create_dataset('caption', self.config)
         test_loader = create_loader(test_dataset, batch_size=1, num_workers=4, train=False)
         experts, _ = next(iter(test_loader))
-        captions = self.model(experts, train=False, prefix=self.config['prefix'])
-        captions = self.tokenizer(captions, max_length=30, padding='max_length', return_tensors='pt').input_ids
-        caption = captions.to(experts['rgb'].device)[0]
-        caption = self.tokenizer.decode(caption, skip_special_tokens=True)
-        caption = caption.capitalize() + '.'
-        return caption
+        question = pre_question(question)
+        answer = self.model(experts, question, train=False, inference='generate')
+        answer = self.tokenizer(answer, max_length=30, padding='max_length', return_tensors='pt').input_ids
+        answer = answer.to(experts['rgb'].device)[0]
+        answer = self.tokenizer.decode(answer, skip_special_tokens=True)
+        answer = answer.capitalize() + '.'
+        return answer
+
+    def run_vqa(self, image_path: str, model_name: str, question: str) -> tuple[str | None, ...]:
+        out_paths = run_experts(image_path)
+        answer = self.run_vqa_model(model_name, question)
+        label_prettify(image_path, out_paths)
+        return answer, *out_paths