Update with md5sum and half precision inference

app.py

@@ -3,18 +3,8 @@
 from __future__ import annotations
 
 import os
-import shutil
-import subprocess
-
 import gradio as gr
 
-if os.getenv('SYSTEM') == 'spaces':
-    with open('patch') as f:
-        subprocess.run('patch -p1'.split(), cwd='prismer', stdin=f)
-    shutil.copytree('prismer/helpers/images',
-                    'prismer/images',
-                    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
@@ -36,7 +26,7 @@ if (SPACE_ID := os.getenv('SPACE_ID')) is not None:
     description += f'For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. <a href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a>'
 
 
-with gr.Blocks() as demo:
+with gr.Blocks(theme='sudeepshouche/minimalist') as demo:
     gr.Markdown(description)
     with gr.Tabs():
         with gr.TabItem('Zero-shot Image Captioning'):

app_vqa.py

@@ -35,7 +35,7 @@ def create_demo() -> gr.Blocks:
         paths = sorted(pathlib.Path('prismer/images').glob('*'))
         ex_questions = ['What is the man on the left doing?',
                         'What is this person doing?',
-                        'How many cows in this image?',
+                        'How many cows are in this image?',
                         'What is the type of animal in this image?',
                         'What toy is it?']
         examples = [[path.as_posix(), 'Prismer-Base', ex_questions[i]] for i, path in enumerate(paths)]

label_prettify.py

@@ -5,6 +5,7 @@ import torch
 import matplotlib.pyplot as plt
 import matplotlib
 import numpy as np
+import shutil
 
 from prismer.utils import create_ade20k_label_colormap
 
@@ -23,101 +24,109 @@ def islight(rgb):
 
 
 def depth_prettify(file_path):
-    depth = plt.imread(file_path)
-    plt.imsave(file_path, depth, cmap='rainbow')
+    pretty_path = file_path.replace('.png', '_p.png')
+    if not os.path.exists(pretty_path):
+        depth = plt.imread(file_path)
+        plt.imsave(pretty_path, depth, cmap='rainbow')
 
 
 def obj_detection_prettify(rgb_path, path_name):
-    rgb = plt.imread(rgb_path)
-    obj_labels = plt.imread(path_name)
-    obj_labels_dict = json.load(open(path_name.replace('.png', '.json')))
+    pretty_path = path_name.replace('.png', '_p.png')
+    if not os.path.exists(pretty_path):
+        rgb = plt.imread(rgb_path)
+        obj_labels = plt.imread(path_name)
+        obj_labels_dict = json.load(open(path_name.replace('.png', '.json')))
 
-    plt.imshow(rgb)
+        plt.imshow(rgb)
 
-    if len(np.unique(obj_labels)) == 1:
-        plt.axis('off')
-        plt.savefig(path_name, bbox_inches='tight', transparent=True, pad_inches=0)
-        plt.close()
-    else:
-        num_objs = np.unique(obj_labels)[:-1].max()
-        plt.imshow(obj_labels, cmap='terrain', vmax=num_objs + 1 / 255., alpha=0.8)
-        cmap = matplotlib.colormaps.get_cmap('terrain')
-        for i in np.unique(obj_labels)[:-1]:
-            obj_idx_all = np.where(obj_labels == i)
-            x, y = obj_idx_all[1].mean(), obj_idx_all[0].mean()
-            obj_name = obj_label_map[obj_labels_dict[str(int(i * 255))]]
-            obj_name = obj_name.split(',')[0]
-            if islight([c*255 for c in cmap(i / num_objs)[:3]]):
-                plt.text(x, y, obj_name, c='black', horizontalalignment='center', verticalalignment='center', clip_on=True)
-            else:
-                plt.text(x, y, obj_name, c='white', horizontalalignment='center', verticalalignment='center', clip_on=True)
-
-        plt.axis('off')
-        plt.savefig(path_name, bbox_inches='tight', transparent=True, pad_inches=0)
-        plt.close()
+        if len(np.unique(obj_labels)) == 1:
+            plt.axis('off')
+            plt.savefig(path_name, bbox_inches='tight', transparent=True, pad_inches=0)
+            plt.close()
+        else:
+            num_objs = np.unique(obj_labels)[:-1].max()
+            plt.imshow(obj_labels, cmap='terrain', vmax=num_objs + 1 / 255., alpha=0.8)
+            cmap = matplotlib.colormaps.get_cmap('terrain')
+            for i in np.unique(obj_labels)[:-1]:
+                obj_idx_all = np.where(obj_labels == i)
+                x, y = obj_idx_all[1].mean(), obj_idx_all[0].mean()
+                obj_name = obj_label_map[obj_labels_dict[str(int(i * 255))]]
+                obj_name = obj_name.split(',')[0]
+                if islight([c*255 for c in cmap(i / num_objs)[:3]]):
+                    plt.text(x, y, obj_name, c='black', horizontalalignment='center', verticalalignment='center', clip_on=True)
+                else:
+                    plt.text(x, y, obj_name, c='white', horizontalalignment='center', verticalalignment='center', clip_on=True)
+
+            plt.axis('off')
+            plt.savefig(pretty_path, bbox_inches='tight', transparent=True, pad_inches=0)
+            plt.close()
 
 
 def seg_prettify(rgb_path, file_name):
-    rgb = plt.imread(rgb_path)
-    seg_labels = plt.imread(file_name)
+    pretty_path = file_name.replace('.png', '_p.png')
+    if not os.path.exists(pretty_path):
+        rgb = plt.imread(rgb_path)
+        seg_labels = plt.imread(file_name)
 
-    plt.imshow(rgb)
+        plt.imshow(rgb)
 
-    seg_map = np.zeros(list(seg_labels.shape) + [3], dtype=np.int16)
-    for i in np.unique(seg_labels):
-        seg_map[seg_labels == i] = ade_color[int(i * 255)]
+        seg_map = np.zeros(list(seg_labels.shape) + [3], dtype=np.int16)
+        for i in np.unique(seg_labels):
+            seg_map[seg_labels == i] = ade_color[int(i * 255)]
 
-    plt.imshow(seg_map, alpha=0.8)
+        plt.imshow(seg_map, alpha=0.8)
 
-    for i in np.unique(seg_labels):
-        obj_idx_all = np.where(seg_labels == i)
-        if len(obj_idx_all[0]) > 20:  # only plot the label with its number of labelled pixel more than 20
-            obj_idx = random.randint(0,  len(obj_idx_all[0]) - 1)
-            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)]):
-                plt.text(x, y, obj_name, c='black', horizontalalignment='center', verticalalignment='center', clip_on=True)
-            else:
-                plt.text(x, y, obj_name, c='white', horizontalalignment='center', verticalalignment='center', clip_on=True)
+        for i in np.unique(seg_labels):
+            obj_idx_all = np.where(seg_labels == i)
+            if len(obj_idx_all[0]) > 20:  # only plot the label with its number of labelled pixel more than 20
+                obj_idx = random.randint(0,  len(obj_idx_all[0]) - 1)
+                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)]):
+                    plt.text(x, y, obj_name, c='black', horizontalalignment='center', verticalalignment='center', clip_on=True)
+                else:
+                    plt.text(x, y, obj_name, c='white', horizontalalignment='center', verticalalignment='center', clip_on=True)
 
-    plt.axis('off')
-    plt.savefig(file_name, bbox_inches='tight', transparent=True, pad_inches=0)
-    plt.close()
+        plt.axis('off')
+        plt.savefig(pretty_path, bbox_inches='tight', transparent=True, pad_inches=0)
+        plt.close()
 
 
 def ocr_detection_prettify(rgb_path, file_name):
-    if os.path.exists(file_name):
-        rgb = plt.imread(rgb_path)
-        ocr_labels = plt.imread(file_name)
-        ocr_labels_dict = torch.load(file_name.replace('.png', '.pt'))
+    pretty_path = file_name.replace('.png', '_p.png')
+    if not os.path.exists(pretty_path):
+        if os.path.exists(file_name):
+            rgb = plt.imread(rgb_path)
+            ocr_labels = plt.imread(file_name)
+            ocr_labels_dict = torch.load(file_name.replace('.png', '.pt'))
 
-        plt.imshow(rgb)
-        plt.imshow(ocr_labels, cmap='gray', alpha=0.8)
+            plt.imshow(rgb)
+            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)
-            x, y = text_idx_all[1].mean(), text_idx_all[0].mean()
-            text = ocr_labels_dict[int(i * 255)]['text']
-            plt.text(x, y, text, c='white', horizontalalignment='center', verticalalignment='center', clip_on=True)
+            for i in np.unique(ocr_labels)[:-1]:
+                text_idx_all = np.where(ocr_labels == i)
+                x, y = text_idx_all[1].mean(), text_idx_all[0].mean()
+                text = ocr_labels_dict[int(i * 255)]['text']
+                plt.text(x, y, text, c='white', horizontalalignment='center', verticalalignment='center', clip_on=True)
 
-        plt.axis('off')
-        plt.savefig(file_name, bbox_inches='tight', transparent=True, pad_inches=0)
-        plt.close()
-    else:
-        rgb = plt.imread(rgb_path)
-        ocr_labels = np.ones_like(rgb, dtype=np.float32())
+            plt.axis('off')
+            plt.savefig(pretty_path, bbox_inches='tight', transparent=True, pad_inches=0)
+            plt.close()
+        else:
+            rgb = plt.imread(rgb_path)
+            ocr_labels = np.ones_like(rgb, dtype=np.float32())
 
-        plt.imshow(rgb)
-        plt.imshow(ocr_labels, cmap='gray', alpha=0.8)
+            plt.imshow(rgb)
+            plt.imshow(ocr_labels, cmap='gray', alpha=0.8)
 
-        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')
+            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()
+            os.makedirs(os.path.dirname(file_name), exist_ok=True)
+            plt.savefig(pretty_path, bbox_inches='tight', transparent=True, pad_inches=0)
+            plt.close()
 
 
 def label_prettify(rgb_path, expert_paths):
@@ -130,4 +139,7 @@ def label_prettify(rgb_path, expert_paths):
             ocr_detection_prettify(rgb_path, expert_path)
         elif 'obj' in expert_path:
             obj_detection_prettify(rgb_path, expert_path)
-
+        else:
+            pretty_path = expert_path.replace('.png', '_p.png')
+            if not os.path.exists(pretty_path):
+                shutil.copyfile(expert_path, pretty_path)

prismer/configs/experts.yaml

@@ -1,2 +1,3 @@
-data_path: 'helpers'
-save_path: 'helpers/labels'
+data_path: helpers
+im_name: 87dfaeb4978ce05aa7be5e5b4cc1273a
+save_path: helpers/labels

prismer/dataset/caption_dataset.py

@@ -32,10 +32,7 @@ class Caption(Dataset):
             elif self.dataset == 'nocaps':
                 self.data_list = json.load(open(os.path.join(self.data_path, 'nocaps_val.json'), 'r'))
             elif self.dataset == 'demo':
-                data_folders = glob.glob(f'{self.data_path}/*/')
-                self.data_list = [{'image': data} for f in data_folders for data in glob.glob(f + '*.jpg')]
-                self.data_list += [{'image': data} for f in data_folders for data in glob.glob(f + '*.png')]
-                self.data_list += [{'image': data} for f in data_folders for data in glob.glob(f + '*.jpeg')]
+                self.data_list = [{'image': f'helpers/images/{config["im_name"]}.jpg'}]
 
     def __len__(self):
         return len(self.data_list)
@@ -50,10 +47,11 @@ class Caption(Dataset):
         elif self.dataset == 'demo':
             img_path_split = self.data_list[index]['image'].split('/')
             img_name = img_path_split[-2] + '/' + img_path_split[-1]
-            image, labels, labels_info = get_expert_labels('', self.label_path, img_name, 'helpers', self.experts)
+            image, labels, labels_info = get_expert_labels('prismer', self.label_path, img_name, 'helpers', self.experts)
 
         experts = self.transform(image, labels)
         experts = post_label_process(experts, labels_info)
+        experts['rgb'] = experts['rgb'].half()
 
         if self.train:
             caption = pre_caption(self.prefix + ' ' + self.data_list[index]['caption'], max_words=30)

prismer/dataset/utils.py

@@ -5,6 +5,7 @@
 # https://github.com/NVlabs/prismer/blob/main/LICENSE
 
 import os
+import pathlib
 import re
 import json
 import torch
@@ -14,10 +15,12 @@ import torchvision.transforms as transforms
 import torchvision.transforms.functional as transforms_f
 from dataset.randaugment import RandAugment
 
-COCO_FEATURES = torch.load('dataset/coco_features.pt')['features']
-ADE_FEATURES = torch.load('dataset/ade_features.pt')['features']
-DETECTION_FEATURES = torch.load('dataset/detection_features.pt')['features']
-BACKGROUND_FEATURES = torch.load('dataset/background_features.pt')
+cur_dir = pathlib.Path(__file__).parent
+
+COCO_FEATURES = torch.load(cur_dir / 'coco_features.pt')['features']
+ADE_FEATURES = torch.load(cur_dir / 'ade_features.pt')['features']
+DETECTION_FEATURES = torch.load(cur_dir / 'detection_features.pt')['features']
+BACKGROUND_FEATURES = torch.load(cur_dir / 'background_features.pt')
 
 
 class Transform:
@@ -119,7 +122,8 @@ def post_label_process(inputs, labels_info):
     for exp in inputs:
         if exp in ['depth', 'normal', 'edge']:  # remap to -1 to 1 range
             inputs[exp] = 2 * (inputs[exp] - inputs[exp].min()) / (inputs[exp].max() - inputs[exp].min() + eps) - 1
-        
+            inputs[exp] = inputs[exp].half()
+
         elif exp == 'seg_coco':  # in-paint with CLIP features
             text_emb = torch.empty([64, *inputs[exp].shape[1:]])
             for l in inputs[exp].unique():
@@ -127,7 +131,7 @@ def post_label_process(inputs, labels_info):
                     text_emb[:, (inputs[exp][0] == l)] = BACKGROUND_FEATURES.unsqueeze(-1)
                 else:
                     text_emb[:, (inputs[exp][0] == l)] = COCO_FEATURES[l].unsqueeze(-1)
-            inputs[exp] = text_emb
+            inputs[exp] = text_emb.half()
 
         elif exp == 'seg_ade':  # in-paint with CLIP features
             text_emb = torch.empty([64, *inputs[exp].shape[1:]])
@@ -136,7 +140,7 @@ def post_label_process(inputs, labels_info):
                     text_emb[:, (inputs[exp][0] == l)] = BACKGROUND_FEATURES.unsqueeze(-1)
                 else:
                     text_emb[:, (inputs[exp][0] == l)] = ADE_FEATURES[l].unsqueeze(-1)
-            inputs[exp] = text_emb
+            inputs[exp] = text_emb.half()
 
         elif exp == 'obj_detection':  # in-paint with CLIP features
             text_emb = torch.empty([64, *inputs[exp].shape[1:]])
@@ -146,7 +150,7 @@ def post_label_process(inputs, labels_info):
                     text_emb[:, (inputs[exp][0] == l)] = BACKGROUND_FEATURES.unsqueeze(-1)
                 else:
                     text_emb[:, (inputs[exp][0] == l)] = DETECTION_FEATURES[label_map[str(l.item())]].unsqueeze(-1)
-            inputs[exp] = {'label': text_emb, 'instance': inputs[exp]}
+            inputs[exp] = {'label': text_emb.half(), 'instance': inputs[exp].half()}
 
         elif exp == 'ocr_detection':  # in-paint with CLIP features
             text_emb = torch.empty([64, *inputs[exp].shape[1:]])
@@ -156,7 +160,7 @@ def post_label_process(inputs, labels_info):
                     text_emb[:, (inputs[exp][0] == l)] = BACKGROUND_FEATURES.unsqueeze(-1)
                 else:
                     text_emb[:, (inputs[exp][0] == l)] = label_map[l.item()]['features'].unsqueeze(-1)
-            inputs[exp] = text_emb
+            inputs[exp] = text_emb.half()
     return inputs
 
 

prismer/experts/depth/generate_dataset.py

@@ -14,12 +14,10 @@ ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 
 class Dataset(Dataset):
-    def __init__(self, data_path, transform):
-        self.data_path = data_path
+    def __init__(self, config, transform):
+        self.data_path = config['data_path']
         self.transform = transform
-        data_folders = glob.glob(f'{data_path}/*/')
-        self.data_list = [data for f in data_folders for data in glob.glob(f + '*.JPEG')]
-        self.data_list += [data for f in data_folders for data in glob.glob(f + '*.jpg')]
+        self.data_list = [f'helpers/images/{config["im_name"]}.jpg']
 
     def __len__(self):
         return len(self.data_list)
@@ -29,4 +27,4 @@ class Dataset(Dataset):
         image = Image.open(image_path).convert('RGB')
         img_size = [image.size[0], image.size[1]]
         image = self.transform(image)
-        return image, image_path, img_size
+        return image.half(), image_path, img_size

prismer/experts/edge/generate_dataset.py

@@ -14,12 +14,10 @@ ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 
 class Dataset(Dataset):
-    def __init__(self, data_path, transform):
-        self.data_path = data_path
+    def __init__(self, config, transform):
+        self.data_path = config['data_path']
         self.transform = transform
-        data_folders = glob.glob(f'{data_path}/*/')
-        self.data_list = [data for f in data_folders for data in glob.glob(f + '*.JPEG')]
-        self.data_list += [data for f in data_folders for data in glob.glob(f + '*.jpg')]
+        self.data_list = [f'helpers/images/{config["im_name"]}.jpg']
 
     def __len__(self):
         return len(self.data_list)
@@ -29,4 +27,4 @@ class Dataset(Dataset):
         image = Image.open(image_path).convert('RGB')
         img_size = [image.size[0], image.size[1]]
         image = self.transform(image)
-        return torch.flip(image, dims=(0, )) * 255., image_path, img_size
+        return torch.flip(image.half(), dims=(0, )) * 255., image_path, img_size

prismer/experts/generate_depth.py

@@ -21,11 +21,10 @@ model, transform = load_expert_model(task='depth')
 accelerator = Accelerator(mixed_precision='fp16')
 
 config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
-data_path = config['data_path']
 save_path = os.path.join(config['save_path'], 'depth')
 
 batch_size = 64
-dataset = Dataset(data_path, transform)
+dataset = Dataset(config, transform)
 data_loader = torch.utils.data.DataLoader(
     dataset=dataset,
     batch_size=batch_size,

prismer/experts/generate_edge.py

@@ -23,11 +23,10 @@ model, transform = load_expert_model(task='edge')
 accelerator = Accelerator(mixed_precision='fp16')
 
 config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
-data_path = config['data_path']
 save_path = os.path.join(config['save_path'], 'edge')
 
 batch_size = 64
-dataset = Dataset(data_path, transform)
+dataset = Dataset(config, transform)
 data_loader = torch.utils.data.DataLoader(
     dataset=dataset,
     batch_size=batch_size,

prismer/experts/generate_normal.py

@@ -23,11 +23,10 @@ model, transform = load_expert_model(task='normal')
 accelerator = Accelerator(mixed_precision='fp16')
 
 config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
-data_path = config['data_path']
 save_path = os.path.join(config['save_path'], 'normal')
 
 batch_size = 64
-dataset = CustomDataset(data_path, transform)
+dataset = CustomDataset(config, transform)
 data_loader = torch.utils.data.DataLoader(
     dataset=dataset,
     batch_size=batch_size,

prismer/experts/generate_objdet.py

@@ -26,9 +26,8 @@ config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
 data_path = config['data_path']
 save_path = config['save_path']
 
-depth_path = os.path.join(save_path, 'depth', data_path.split('/')[-1])
 batch_size = 32
-dataset = Dataset(data_path, depth_path, transform)
+dataset = Dataset(config, transform)
 data_loader = torch.utils.data.DataLoader(
     dataset=dataset,
     batch_size=batch_size,

prismer/experts/generate_ocrdet.py

@@ -27,11 +27,10 @@ accelerator = Accelerator(mixed_precision='fp16')
 pca_clip = pk.load(open('dataset/clip_pca.pkl', 'rb'))
 
 config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
-data_path = config['data_path']
 save_path = os.path.join(config['save_path'], 'ocr_detection')
 
 batch_size = 32
-dataset = Dataset(data_path, transform)
+dataset = Dataset(config, transform)
 data_loader = torch.utils.data.DataLoader(
     dataset=dataset,
     batch_size=batch_size,

prismer/experts/generate_segmentation.py

@@ -21,11 +21,10 @@ model, transform = load_expert_model(task='seg_coco')
 accelerator = Accelerator(mixed_precision='fp16')
 
 config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
-data_path = config['data_path']
 save_path = os.path.join(config['save_path'], 'seg_coco')
 
 batch_size = 4
-dataset = Dataset(data_path, transform)
+dataset = Dataset(config, transform)
 data_loader = torch.utils.data.DataLoader(
     dataset=dataset,
     batch_size=batch_size,

prismer/experts/model_bank.py

@@ -131,6 +131,8 @@ def load_expert_model(task=None):
         model = None
         transform = None
 
+    if 'seg' not in task:
+        model = model.half()
     model.eval()
     return model, transform
 

prismer/experts/normal/generate_dataset.py

@@ -14,12 +14,10 @@ ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 
 class CustomDataset(Dataset):
-    def __init__(self, data_path, transform):
-        self.data_path = data_path
+    def __init__(self, config, transform):
+        self.data_path = config['data_path']
         self.transform = transform
-        data_folders = glob.glob(f'{data_path}/*/')
-        self.data_list = [data for f in data_folders for data in glob.glob(f + '*.JPEG')]
-        self.data_list += [data for f in data_folders for data in glob.glob(f + '*.jpg')]
+        self.data_list = [f'helpers/images/{config["im_name"]}.jpg']
 
     def __len__(self):
         return len(self.data_list)
@@ -29,6 +27,6 @@ class CustomDataset(Dataset):
         image = Image.open(image_path).convert('RGB')
         img_size = [image.size[0], image.size[1]]
         image = self.transform(image)
-        return image, image_path, img_size
+        return image.half(), image_path, img_size
 
 

prismer/experts/obj_detection/generate_dataset.py

@@ -5,6 +5,7 @@
 # https://github.com/NVlabs/prismer/blob/main/LICENSE
 
 import glob
+import os
 import torch
 
 from torch.utils.data import Dataset
@@ -15,13 +16,11 @@ ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 
 class Dataset(Dataset):
-    def __init__(self, data_path, depth_path, transform):
-        self.data_path = data_path
-        self.depth_path = depth_path
+    def __init__(self, config, transform):
+        self.data_path = config['data_path']
+        self.depth_path = os.path.join(config['save_path'], 'depth', self.data_path.split('/')[-1])
         self.transform = transform
-        data_folders = glob.glob(f'{data_path}/*/')
-        self.data_list = [data for f in data_folders for data in glob.glob(f + '*.JPEG')]
-        self.data_list += [data for f in data_folders for data in glob.glob(f + '*.jpg')]
+        self.data_list = [f'helpers/images/{config["im_name"]}.jpg']
 
     def __len__(self):
         return len(self.data_list)
@@ -43,7 +42,7 @@ class Dataset(Dataset):
         depth = self.transform(depth)
         depth = torch.tensor(np.array(depth)).float() / 255.
         img_size = image.shape
-        return {"image": image, "height": img_size[1], "width": img_size[2],
+        return {"image": image.half(), "height": img_size[1], "width": img_size[2],
                 "true_height": true_img_size[0], "true_width": true_img_size[1],
                 'image_path': image_path, 'depth': depth}
 

prismer/experts/ocr_detection/generate_dataset.py

@@ -14,12 +14,10 @@ ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 
 class Dataset(Dataset):
-    def __init__(self, data_path, transform):
-        self.data_path = data_path
+    def __init__(self, config, transform):
+        self.data_path = config['data_path']
         self.transform = transform
-        data_folders = glob.glob(f'{data_path}/*/')
-        self.data_list = [data for f in data_folders for data in glob.glob(f + '*.JPEG')]
-        self.data_list += [data for f in data_folders for data in glob.glob(f + '*.jpg')]
+        self.data_list = [f'helpers/images/{config["im_name"]}.jpg']
 
     def __len__(self):
         return len(self.data_list)
@@ -30,7 +28,7 @@ class Dataset(Dataset):
 
         image, scale_w, scale_h, original_w, original_h = resize(original_image)
         image = self.transform(image)
-        return image, image_path, scale_w, scale_h, original_w, original_h
+        return image.half(), image_path, scale_w, scale_h, original_w, original_h
 
 
 def resize(im):

prismer/experts/segmentation/generate_dataset.py

@@ -16,12 +16,10 @@ ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 
 class Dataset(Dataset):
-    def __init__(self, data_path, transform):
-        self.data_path = data_path
+    def __init__(self, config, transform):
+        self.data_path = config['data_path']
         self.transform = transform
-        data_folders = glob.glob(f'{data_path}/*/')
-        self.data_list = [data for f in data_folders for data in glob.glob(f + '*.JPEG')]
-        self.data_list += [data for f in data_folders for data in glob.glob(f + '*.jpg')]
+        self.data_list = [f'helpers/images/{config["im_name"]}.jpg']
 
     def __len__(self):
         return len(self.data_list)

prismer/model/prismer.py

@@ -5,6 +5,7 @@
 # https://github.com/NVlabs/prismer/blob/main/LICENSE
 
 import json
+import pathlib
 import torch.nn as nn
 
 from model.modules.vit import load_encoder
@@ -12,6 +13,9 @@ from model.modules.roberta import load_decoder
 from transformers import RobertaTokenizer, RobertaConfig
 
 
+cur_dir = pathlib.Path(__file__).parent
+
+
 class Prismer(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -26,7 +30,7 @@ class Prismer(nn.Module):
             elif exp in ['obj_detection', 'ocr_detection']:
                 self.experts[exp] = 64
 
-        prismer_config = json.load(open('configs/prismer.json', 'r'))[config['prismer_model']]
+        prismer_config = json.load(open(f'{cur_dir.parent}/configs/prismer.json', 'r'))[config['prismer_model']]
         roberta_config = RobertaConfig.from_dict(prismer_config['roberta_model'])
 
         self.tokenizer = RobertaTokenizer.from_pretrained(prismer_config['roberta_model']['model_name'])
@@ -35,7 +39,7 @@ class Prismer(nn.Module):
 
         self.prepare_to_train(config['freeze'])
         self.ignored_modules = self.get_ignored_modules(config['freeze'])
-    
+
     def prepare_to_train(self, mode='none'):
         for name, params in self.named_parameters():
             if mode == 'freeze_lang':

prismer_model.py

@@ -7,6 +7,12 @@ import shlex
 import shutil
 import subprocess
 import sys
+import hashlib
+from typing import Tuple
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
 
 import cv2
 import torch
@@ -55,27 +61,43 @@ def run_expert(expert_name: str):
                    check=True)
 
 
-def run_experts(image_path: str) -> tuple[str | None, ...]:
-    helper_dir = submodule_dir / 'helpers'
-    shutil.rmtree(helper_dir, ignore_errors=True)
-    image_dir = helper_dir / 'images'
-    image_dir.mkdir(parents=True, exist_ok=True)
-    out_path = image_dir / 'image.jpg'
-    cv2.imwrite(out_path.as_posix(), cv2.imread(image_path))
+def compute_md5(image_path: str) -> str:
+    with open(image_path, 'rb') as f:
+        s = f.read()
+    return hashlib.md5(s).hexdigest()
 
-    # expert_names = ['edge', 'normal', 'objdet', 'ocrdet', 'segmentation']
-    # run_expert('depth')
-    # with concurrent.futures.ProcessPoolExecutor() as executor:
-    #     executor.map(run_expert, expert_names)
-
-    # no parallelization just to be safe
-    expert_names = ['depth', 'edge', 'normal', 'objdet', 'ocrdet', 'segmentation']
-    for exp in expert_names:
-        run_expert(exp)
 
+def run_experts(image_path: str) -> Tuple[str, Tuple[str, ...]]:
+    im_name = compute_md5(image_path)
+    out_path = submodule_dir / 'helpers' / 'images' / f'{im_name}.jpg'
     keys = ['depth', 'edge', 'normal', 'seg_coco', 'obj_detection', 'ocr_detection']
-    results = [pathlib.Path('prismer/helpers/labels') / key / 'helpers/images/image.png' for key in keys]
-    return tuple(path.as_posix() for path in results)
+    results = [pathlib.Path('prismer/helpers/labels') / key / f'helpers/images/{im_name}.png' for key in keys]
+    results_pretty = [pathlib.Path('prismer/helpers/labels') / key / f'helpers/images/{im_name}_p.png' for key in keys]
+    out_paths = tuple(path.as_posix() for path in results)
+    pretty_paths = tuple(path.as_posix() for path in results_pretty)
+
+    config = yaml.load(open('prismer/configs/experts.yaml', 'r'), Loader=yaml.Loader)
+    config['im_name'] = im_name
+    with open('prismer/configs/experts.yaml', 'w') as yaml_file:
+        yaml.dump(config, yaml_file, default_flow_style=False)
+
+    if not os.path.exists(out_paths[0]):
+        cv2.imwrite(out_path.as_posix(), cv2.imread(image_path))
+
+        # paralleled inference
+        expert_names = ['edge', 'normal', 'objdet', 'ocrdet', 'segmentation']
+        run_expert('depth')
+        with concurrent.futures.ProcessPoolExecutor() as executor:
+            executor.map(run_expert, expert_names)
+        executor.shutdown(wait=True)
+
+        # no parallelization just to be safe
+        # expert_names = ['depth', 'edge', 'normal', 'objdet', 'ocrdet', 'segmentation']
+        # for exp in expert_names:
+        #     run_expert(exp)
+
+    label_prettify(image_path, out_paths)
+    return im_name, pretty_paths
 
 
 class Model:
@@ -126,20 +148,28 @@ class Model:
                                                                                       len(model.expert_encoder.positional_embedding))
 
         model.load_state_dict(state_dict)
+        model = model.half()
         model.eval()
 
         self.config = config
-        self.model = model
+        self.model = model.to('cuda:0')
         self.tokenizer = model.tokenizer
         self.exp_name = exp_name
         self.mode = mode
 
     @torch.inference_mode()
-    def run_caption_model(self, exp_name: str) -> str:
+    def run_caption_model(self, exp_name: str, im_name: str) -> str:
         self.set_model(exp_name, 'caption')
+        self.config['im_name'] = im_name
         _, 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))
+        for exp in experts:
+            if exp == 'obj_detection':
+                experts[exp]['label'] = experts['obj_detection']['label'].to('cuda:0')
+                experts[exp]['instance'] = experts['obj_detection']['instance'].to('cuda:0')
+            else:
+                experts[exp] = experts[exp].to('cuda:0')
         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]
@@ -148,17 +178,23 @@ class Model:
         return caption
 
     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)
-        label_prettify(image_path, out_paths)
-        return caption, *out_paths
+        im_name, pretty_paths = run_experts(image_path)
+        caption = self.run_caption_model(model_name, im_name)
+        return caption, *pretty_paths
 
     @torch.inference_mode()
-    def run_vqa_model(self, exp_name: str, question: str) -> str:
+    def run_vqa_model(self, exp_name: str, im_name: str, question: str) -> str:
         self.set_model(exp_name, 'vqa')
+        self.config['im_name'] = im_name
         _, 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))
+        for exp in experts:
+            if exp == 'obj_detection':
+                experts[exp]['label'] = experts['obj_detection']['label'].to('cuda:0')
+                experts[exp]['instance'] = experts['obj_detection']['instance'].to('cuda:0')
+            else:
+                experts[exp] = experts[exp].to('cuda:0')
         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
@@ -168,7 +204,6 @@ class Model:
         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
+        im_name, pretty_paths = run_experts(image_path)
+        answer = self.run_vqa_model(model_name, im_name, question)
+        return answer, *pretty_paths

requirements.txt

@@ -6,7 +6,7 @@ fire==0.5.0
 geffnet==1.0.2
 git+https://github.com/facebookresearch/detectron2.git@5aeb252b194b93dc2879b4ac34bc51a31b5aee13
 git+https://github.com/openai/CLIP.git@a9b1bf5
-gradio==3.20.1
+gradio==3.24.1
 huggingface-hub==0.12.1
 opencv-python-headless==4.7.0.72
 pyclipper==1.3.0.post4