Add VQA

label_prettify.py

@@ -87,7 +87,7 @@ def ocr_detection_prettify(rgb_path, file_name):
         ocr_labels_dict = torch.load(file_name.replace('.png', '.pt'))
 
         plt.imshow(rgb)
-        plt.imshow((1 - ocr_labels) < 1, cmap='gray', alpha=0.8)
+        plt.imshow(ocr_labels, cmap='gray', alpha=0.8)
 
         for i in np.unique(ocr_labels)[:-1]:
             text_idx_all = np.where(ocr_labels == i)

prismer_model.py

@@ -75,11 +75,13 @@ class Model:
         if exp_name == self.exp_name:
             return
 
+        # remap model name
         if self.exp_name == 'Prismer-Base':
-            model_name = 'prismer_base'
+            self.exp_name = 'prismer_base'
         elif self.exp_name == 'Prismer-Large':
-            model_name = 'prismer_large'
+            self.exp_name = 'prismer_large'
 
+        # load checkpoints
         if self.mode == 'caption':
             config = {
                 'dataset': 'demo',
@@ -87,12 +89,12 @@ class Model:
                 'label_path': 'prismer/helpers/labels',
                 'experts': ['depth', 'normal', 'seg_coco', 'edge', 'obj_detection', 'ocr_detection'],
                 'image_resolution': 480,
-                'prismer_model': model_name,
+                'prismer_model': self.exp_name,
                 'freeze': 'freeze_vision',
                 'prefix': '',
             }
             model = PrismerCaption(config)
-            state_dict = torch.load(f'prismer/logging/pretrain_{model_name}/pytorch_model.bin', map_location='cuda:0')
+            state_dict = torch.load(f'prismer/logging/pretrain_{self.exp_name}/pytorch_model.bin', map_location='cuda:0')
 
         elif self.mode == 'vqa':
             config = {
@@ -101,12 +103,12 @@ class Model:
                 'label_path': 'prismer/helpers/labels',
                 'experts': ['depth', 'normal', 'seg_coco', 'edge', 'obj_detection', 'ocr_detection'],
                 'image_resolution': 480,
-                'prismer_model': model_name,
+                'prismer_model': self.exp_name,
                 'freeze': 'freeze_vision',
             }
 
             model = PrismerVQA(config)
-            state_dict = torch.load(f'prismer/logging/vqa_{model_name}/pytorch_model.bin', map_location='cuda:0')
+            state_dict = torch.load(f'prismer/logging/vqa_{self.exp_name}/pytorch_model.bin', map_location='cuda:0')
 
         model.load_state_dict(state_dict)
         model.eval()