Add VQA

app_caption.py

@@ -28,7 +28,7 @@ def create_demo():
                 object_detection = gr.Image(label='Object Detection')
                 ocr = gr.Image(label='OCR Detection')
 
-    inputs = [image, model_name]
+    inputs = [image, model_name, 'caption']
     outputs = [caption, depth, edge, normals, segmentation, object_detection, ocr]
 
     # paths = sorted(pathlib.Path('prismer/images').glob('*'))

app_vqa.py

@@ -11,7 +11,6 @@ from prismer_model import Model
 
 def create_demo():
     model = Model()
-    model.mode = 'vqa'
     with gr.Row():
         with gr.Column():
             image = gr.Image(label='Input', type='filepath')
@@ -29,7 +28,7 @@ def create_demo():
                 object_detection = gr.Image(label='Object Detection')
                 ocr = gr.Image(label='OCR Detection')
 
-    inputs = [image, model_name, question]
+    inputs = [image, model_name, 'vqa', question]
     outputs = [answer, depth, edge, normals, segmentation, object_detection, ocr]
 
     # paths = sorted(pathlib.Path('prismer/images').glob('*'))

prismer_model.py

@@ -68,20 +68,16 @@ class Model:
         self.config = None
         self.model = None
         self.tokenizer = None
+        self.model_name = ''
         self.exp_name = ''
         self.mode = ''
 
-    def set_model(self, exp_name: str) -> None:
-        if exp_name == self.exp_name:
+    def set_model(self, exp_name: str, mode: str) -> None:
+        if exp_name == self.exp_name and mode == self.mode:
             return
 
-        # remap model name
-        if self.exp_name == 'Prismer-Base':
-            self.exp_name = 'prismer_base'
-        elif self.exp_name == 'Prismer-Large':
-            self.exp_name = 'prismer_large'
-
         # load checkpoints
+        model_name = exp_name.lower().replace('-', '_')
         if self.mode == 'caption':
             config = {
                 'dataset': 'demo',
@@ -89,12 +85,12 @@ class Model:
                 'label_path': 'prismer/helpers/labels',
                 'experts': ['depth', 'normal', 'seg_coco', 'edge', 'obj_detection', 'ocr_detection'],
                 'image_resolution': 480,
-                'prismer_model': self.exp_name,
+                'prismer_model': model_name,
                 'freeze': 'freeze_vision',
                 'prefix': '',
             }
             model = PrismerCaption(config)
-            state_dict = torch.load(f'prismer/logging/pretrain_{self.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 = {
@@ -103,12 +99,12 @@ class Model:
                 'label_path': 'prismer/helpers/labels',
                 'experts': ['depth', 'normal', 'seg_coco', 'edge', 'obj_detection', 'ocr_detection'],
                 'image_resolution': 480,
-                'prismer_model': self.exp_name,
+                'prismer_model': model_name,
                 'freeze': 'freeze_vision',
             }
 
             model = PrismerVQA(config)
-            state_dict = torch.load(f'prismer/logging/vqa_{self.exp_name}/pytorch_model.bin', map_location='cuda:0')
+            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()
@@ -117,10 +113,11 @@ class Model:
         self.model = model
         self.tokenizer = model.tokenizer
         self.exp_name = exp_name
+        self.mode = mode
 
     @torch.inference_mode()
-    def run_caption_model(self, exp_name: str) -> str:
-        self.set_model(exp_name)
+    def run_caption_model(self, exp_name: str, mode: str) -> str:
+        self.set_model(exp_name, mode)
         _, 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))
@@ -131,15 +128,15 @@ class Model:
         caption = caption.capitalize() + '.'
         return caption
 
-    def run_caption(self, image_path: str, model_name: str) -> tuple[str | None, ...]:
+    def run_caption(self, image_path: str, model_name: str, mode: str) -> tuple[str | None, ...]:
         out_paths = run_experts(image_path)
-        caption = self.run_caption_model(model_name)
+        caption = self.run_caption_model(model_name, mode)
         label_prettify(image_path, out_paths)
         return caption, *out_paths
 
     @torch.inference_mode()
-    def run_vqa_model(self, exp_name: str, question: str) -> str:
-        self.set_model(exp_name)
+    def run_vqa_model(self, exp_name: str, mode: str, question: str) -> str:
+        self.set_model(exp_name, mode)
         _, 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))
@@ -151,8 +148,8 @@ class Model:
         answer = answer.capitalize() + '.'
         return answer
 
-    def run_vqa(self, image_path: str, model_name: str, question: str) -> tuple[str | None, ...]:
+    def run_vqa(self, image_path: str, model_name: str, mode: str, question: str) -> tuple[str | None, ...]:
         out_paths = run_experts(image_path)
-        answer = self.run_vqa_model(model_name, question)
+        answer = self.run_vqa_model(model_name, mode, question)
         label_prettify(image_path, out_paths)
         return answer, *out_paths