Fix ocr

app.py

@@ -16,6 +16,7 @@ if os.getenv('SYSTEM') == 'spaces':
                     dirs_exist_ok=True)
 
 from app_caption import create_demo as create_demo_caption
+from app_vqa import create_demo as create_demo_vqa
 from prismer_model import build_deformable_conv, download_models
 
 
@@ -40,5 +41,7 @@ with gr.Blocks() as demo:
     with gr.Tabs():
         with gr.TabItem('Zero-shot Image Captioning'):
             create_demo_caption()
+        with gr.TabItem('Visual Question Answering'):
+            create_demo_vqa()
 
 demo.queue(api_open=False).launch()

app_vqa.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import os
+import pathlib
+import gradio as gr
+
+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')
+            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')
+            with gr.Row():
+                depth = gr.Image(label='Depth')
+                edge = gr.Image(label='Edge')
+                normals = gr.Image(label='Normals')
+            with gr.Row():
+                segmentation = gr.Image(label='Segmentation')
+                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]
+
+    # paths = sorted(pathlib.Path('prismer/images').glob('*'))
+    # examples = [[path.as_posix(), 'prismer_base'] for path in paths]
+    # gr.Examples(examples=examples,
+    #             inputs=inputs,
+    #             outputs=outputs,
+    #             fn=model.run_caption,
+    #             cache_examples=os.getenv('SYSTEM') == 'spaces')
+
+    paths = sorted(pathlib.Path('prismer/images').glob('*'))
+    examples = [[path.as_posix(), 'Prismer-Base'] for path in paths]
+    gr.Examples(examples=examples,
+                inputs=inputs,
+                outputs=outputs,
+                fn=model.run_caption)
+
+    run_button.click(fn=model.run_caption, inputs=inputs, outputs=outputs)
+
+
+if __name__ == '__main__':
+    demo = create_demo()
+    demo.queue().launch()

label_prettify.py

@@ -1,5 +1,6 @@
 import os
 import json
+import random
 import torch
 import matplotlib.pyplot as plt
 import matplotlib
@@ -65,7 +66,8 @@ def seg_prettify(rgb_path, file_name):
 
     for i in np.unique(seg_labels):
         obj_idx_all = np.where(seg_labels == i)
-        x, y = obj_idx_all[1].mean(), obj_idx_all[0].mean()
+        obj_idx = random.randint(0, len(obj_idx_all[0]))
+        x, y = obj_idx_all[1][obj_idx], obj_idx_all[0][obj_idx]
         obj_name = coco_label_map[int(i * 255)]
         obj_name = obj_name.split(',')[0]
         if islight(seg_map[int(y), int(x)]):
@@ -105,8 +107,9 @@ def ocr_detection_prettify(rgb_path, file_name):
 
         x, y = rgb.shape[1] / 2, rgb.shape[0] / 2
         plt.text(x, y, 'No text detected', c='black', horizontalalignment='center', verticalalignment='center', clip_on=True)
-
         plt.axis('off')
+
+        os.makedirs(os.path.dirname(file_name), exist_ok=True)
         plt.savefig(file_name, bbox_inches='tight', transparent=True, pad_inches=0)
         plt.close()
 

prismer_model.py

@@ -24,10 +24,10 @@ def download_models() -> None:
         subprocess.run(shlex.split('python download_checkpoints.py --download_experts=True'), cwd='prismer')
 
     model_names = [
-        # 'vqa_prismer_base',
-        # 'vqa_prismer_large',
-        'caption_prismer_base',
-        'caption_prismer_large',
+        'vqa_prismer_base',
+        'vqa_prismer_large',
+        'pretrain_prismer_base',
+        'pretrain_prismer_large',
     ]
     for model_name in model_names:
         if pathlib.Path(f'prismer/logging/{model_name}').exists():
@@ -126,6 +126,19 @@ class Model:
 
     def run_caption(self, image_path: str, model_name: str) -> tuple[str | None, ...]:
         out_paths = run_experts(image_path)
-        # caption = self.run_caption_model(model_name)
+        caption = self.run_caption_model(model_name)
         label_prettify(image_path, out_paths)
-        return None, *out_paths
+        return caption, *out_paths
+
+    @torch.inference_mode()
+    def run_vqa_model(self, exp_name: str) -> str:
+        self.set_model(exp_name)
+        _, test_dataset = create_dataset('vqa', 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