Test 3 experts

app.py

@@ -3,8 +3,14 @@ import gradio as gr
 
 from gradio_caption import create_demo as create_caption
 from gradio_vqa import create_demo as create_vqa
+from prismer_model import build_deformable_conv, download_models
 
 
+# Prepare Prismer checkpoints
+download_models()
+build_deformable_conv()
+
+# Official Demo here
 description = """
 # Prismer 
 The official demo for **Prismer: A Vision-Language Model with An Ensemble of Experts**.
@@ -22,6 +28,5 @@ with gr.Blocks() as demo:
     with gr.Tab("Visual Question Answering"):
         create_vqa()
 
+demo.queue(api_open=False).launch()
 
-if __name__ == '__main__':
-    demo.queue().launch()

gradio_caption.py

@@ -1,32 +1,44 @@
-import gradio as gr
-from PIL import Image
-import tempfile
+from __future__ import annotations
 
+import os
+import pathlib
+import gradio as gr
 
-def predict_depth(model, image):
-    depth = model.infer_pil(image)
-    return depth
+from prismer_model import run_experts
 
 
 def create_demo():
     with gr.Row():
-        with gr.Column(scale=1):
-            model_type = gr.Dropdown(["Prismer-Base", "Prismer-Large"], label="Model Size", value="Prismer-Base")
-            rgb = gr.Image(label="Input Image", type='pil')
-            submit = gr.Button("Submit")
-        with gr.Column(scale=2):
-            pred = gr.Textbox(label="Model Prediction")
+        with gr.Column():
+            image = gr.Image(label='Input', type='filepath')
+            model_name = gr.Dropdown(label='Model', choices=['prismer_base'], 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")
+                depth = gr.Image(label='Depth')
+                edge = gr.Image(label='Edge')
+                normals = gr.Image(label='Normals')
             with gr.Row():
-                seg = gr.Image(label="Segmentation")
-                obj_det = gr.Image(label="Object Detection")
-                ocr_det = gr.Image(label="OCR Detection")
+                segmentation = gr.Image(label='Segmentation')
+                object_detection = gr.Image(label='Object Detection')
+                ocr = gr.Image(label='OCR Detection')
+
+    inputs = [image, model_name]
+    outputs = [depth, edge, normals]
+
+    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=run_experts,
+                cache_examples=os.getenv('SYSTEM') == 'spaces')
+
+    run_button.click(fn=run_experts, inputs=inputs, outputs=outputs)
+
 
-    def on_submit(im, model_type):
-        return pred, depth, edge, normals, seg, obj_det, ocr_det
+if __name__ == '__main__':
+    demo = create_demo()
+    demo.queue().launch()
 
-    submit.click(on_submit, inputs=[rgb, model_type], outputs=[pred, depth, edge, normals, seg, obj_det, ocr_det])
-    examples = gr.Examples(examples=["examples/1.png"], inputs=[rgb])

prismer/dataset/__init__.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+from torch.utils.data import DataLoader
+
+from dataset.pretrain_dataset import Pretrain
+from dataset.vqa_dataset import VQA
+from dataset.caption_dataset import Caption
+from dataset.classification_dataset import Classification
+
+
+def create_dataset(dataset, config):
+    if dataset == 'pretrain':
+        dataset = Pretrain(config)
+        return dataset
+
+    elif dataset == 'vqa':
+        train_dataset = VQA(config, train=True)
+        test_dataset = VQA(config, train=False)
+        return train_dataset, test_dataset
+
+    elif dataset == 'caption':
+        train_dataset = Caption(config, train=True)
+        test_dataset = Caption(config, train=False)
+        return train_dataset, test_dataset
+    
+    elif dataset == 'classification':
+        train_dataset = Classification(config, train=True)
+        test_dataset = Classification(config, train=False)
+        return train_dataset, test_dataset
+
+
+def create_loader(dataset, batch_size, num_workers, train, collate_fn=None):
+    data_loader = DataLoader(dataset,
+                             batch_size=batch_size,
+                             num_workers=num_workers,
+                             collate_fn=collate_fn,
+                             shuffle=True if train else False,
+                             drop_last=True if train else False)
+    return data_loader

prismer/dataset/caption_dataset.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import glob
+
+from torch.utils.data import Dataset
+from dataset.utils import *
+from PIL import ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+class Caption(Dataset):
+    def __init__(self, config, train=True):
+        self.data_path = config['data_path']
+        self.label_path = config['label_path']
+        self.experts = config['experts']
+        self.prefix = config['prefix']
+        self.dataset = config['dataset']
+        self.transform = Transform(resize_resolution=config['image_resolution'], scale_size=[0.5, 1.0], train=train)
+        self.train = train
+
+        if train:
+            self.data_list = []
+            if self.dataset in ['coco', 'nocaps']:
+                self.data_list += json.load(open(os.path.join(self.data_path, 'coco_karpathy_train.json'), 'r'))
+        else:
+            if self.dataset == 'coco':
+                self.data_list = json.load(open(os.path.join(self.data_path, 'coco_karpathy_test.json'), 'r'))
+            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')]
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        data = self.data_list[index]
+
+        if self.dataset == 'coco':
+            image, labels, labels_info = get_expert_labels(self.data_path, self.label_path, data['image'], 'vqav2', self.experts)
+        elif self.dataset == 'nocaps':
+            image, labels, labels_info = get_expert_labels(self.data_path, self.label_path, data['image'], 'nocaps', self.experts)
+        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)
+
+        experts = self.transform(image, labels)
+        experts = post_label_process(experts, labels_info)
+
+        if self.train:
+            caption = pre_caption(self.prefix + ' ' + self.data_list[index]['caption'], max_words=30)
+            return experts, caption
+        else:
+            return experts, index
+

prismer/dataset/coco_features.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccf18221afe8dddef3ffb9daad31d5c7a92cdc2f2f434d77cbeb48031bc75756
+size 36651

prismer/dataset/detection_features.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c38ba9352b2a9f832b14fdc19ac407527ffeaa2903958a73f6eb649f78119c76
+size 198443

prismer/dataset/utils.py

@@ -0,0 +1,188 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import os
+import re
+import json
+import torch
+import PIL.Image as Image
+import numpy as np
+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')
+
+
+class Transform:
+    def __init__(self, resize_resolution=384, scale_size=[0.5, 1.0], train=False):
+        self.resize_size = [resize_resolution, resize_resolution]
+        self.scale_size = scale_size
+        self.train = train
+        self.randaugment = RandAugment(2, 5)
+
+    def __call__(self, image, labels):
+        if self.train:
+            # random resize crop
+            i, j, h, w = transforms.RandomResizedCrop.get_params(img=image, scale=self.scale_size, ratio=[3. / 4, 4. / 3])
+            image = transforms_f.crop(image, i, j, h, w)
+            if labels is not None:
+                for exp in labels:
+                    labels[exp] = transforms_f.crop(labels[exp], i, j, h, w)
+
+        # resize to the defined shape
+        image = transforms_f.resize(image, self.resize_size, transforms_f.InterpolationMode.BICUBIC)
+        if labels is not None:
+            for exp in labels:
+                labels[exp] = transforms_f.resize(labels[exp], [224, 224], transforms_f.InterpolationMode.NEAREST)
+
+        if self.train:
+            # random flipping
+            if torch.rand(1) > 0.5:
+                image = transforms_f.hflip(image)
+                if labels is not None:
+                    for exp in labels:
+                        labels[exp] = transforms_f.hflip(labels[exp])
+
+            # random augmentation
+            image, labels = self.randaugment(image, labels)
+
+        # transform to tensor
+        image = transforms_f.to_tensor(image)
+        if labels is not None:
+            for exp in labels:
+                if exp in ['depth', 'normal', 'edge']:
+                    labels[exp] = transforms_f.to_tensor(labels[exp])
+                else:
+                    labels[exp] = (transforms_f.to_tensor(labels[exp]) * 255).long()
+
+        # apply normalisation:
+        image = transforms_f.normalize(image, mean=[0.48145466, 0.4578275, 0.40821073],
+                                       std=[0.26862954, 0.26130258, 0.27577711])
+        if labels is not None:
+            return {'rgb': image, **labels}
+        else:
+            return{'rgb': image}
+
+
+def get_expert_labels(data_path, label_path, image_path, dataset, experts):
+    image_full_path = os.path.join(data_path, dataset, image_path)
+    image = Image.open(image_full_path).convert('RGB')
+    if experts != 'none':
+        labels = {}
+        labels_info = {}
+        ps = image_path.split('.')[-1]
+        for exp in experts:
+            if exp in ['seg_coco', 'seg_ade', 'edge', 'depth']:
+                label_full_path = os.path.join(label_path, exp, dataset, image_path.replace(f'.{ps}', '.png'))
+                if os.stat(label_full_path).st_size > 0:
+                    labels[exp] = Image.open(label_full_path).convert('L')
+                else:
+                    labels[exp] = Image.fromarray(np.zeros([image.size[1], image.size[0]])).convert('L')
+            elif exp == 'normal':
+                label_full_path = os.path.join(label_path, exp, dataset, image_path.replace(f'.{ps}', '.png'))
+                if os.stat(label_full_path).st_size > 0:
+                    labels[exp] = Image.open(label_full_path).convert('RGB')
+                else:
+                    labels[exp] = Image.fromarray(np.zeros([image.size[1], image.size[0], 3])).convert('RGB')
+            elif exp == 'obj_detection':
+                label_full_path = os.path.join(label_path, exp, dataset, image_path.replace(f'.{ps}', '.png'))
+                if os.stat(label_full_path).st_size > 0:
+                    labels[exp] = Image.open(label_full_path).convert('L')
+                else:
+                    labels[exp] = Image.fromarray(255 * np.ones([image.size[1], image.size[0]])).convert('L')
+                label_info_path = os.path.join(label_path, exp, dataset, image_path.replace(f'.{ps}', '.json'))
+                labels_info[exp] = json.load(open(label_info_path, 'r'))
+            elif exp == 'ocr_detection':
+                label_full_path = os.path.join(label_path, exp, dataset, image_path.replace(f'.{ps}', '.png'))
+                label_info_path = os.path.join(label_path, exp, dataset, image_path.replace(f'.{ps}', '.pt'))
+                if os.path.exists(label_info_path):
+                    labels[exp] = Image.open(label_full_path).convert('L')
+                    labels_info[exp] = torch.load(label_info_path)
+                else:
+                    labels[exp] = Image.fromarray(255 * np.ones([image.size[1], image.size[0]])).convert('L')
+                    labels_info[exp] = None
+
+    else:
+        labels, labels_info = None, None
+    return image, labels, labels_info
+
+
+def post_label_process(inputs, labels_info):
+    eps = 1e-6
+    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
+        
+        elif exp == 'seg_coco':  # in-paint with CLIP features
+            text_emb = torch.empty([64, *inputs[exp].shape[1:]])
+            for l in inputs[exp].unique():
+                if l == 255:
+                    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
+
+        elif exp == 'seg_ade':  # in-paint with CLIP features
+            text_emb = torch.empty([64, *inputs[exp].shape[1:]])
+            for l in inputs[exp].unique():
+                if l == 255:
+                    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
+
+        elif exp == 'obj_detection':  # in-paint with CLIP features
+            text_emb = torch.empty([64, *inputs[exp].shape[1:]])
+            label_map = labels_info[exp]
+            for l in inputs[exp].unique():
+                if l == 255:
+                    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]}
+
+        elif exp == 'ocr_detection':  # in-paint with CLIP features
+            text_emb = torch.empty([64, *inputs[exp].shape[1:]])
+            label_map = labels_info[exp]
+            for l in inputs[exp].unique():
+                if l == 255:
+                    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
+    return inputs
+
+
+def pre_caption(caption, max_words=50):
+    caption = re.sub(r"([.!\"()*#:;~])", ' ', caption.capitalize())  # remove special characters
+    caption = re.sub(r"\s{2,}", ' ', caption)  # remove two white spaces
+
+    caption = caption.rstrip('\n')  # remove \num_ans_per_q symbol
+    caption = caption.strip(' ')    # remove leading and trailing white spaces
+
+    # truncate caption to the max words
+    caption_words = caption.split(' ')
+    if len(caption_words) > max_words:
+        caption = ' '.join(caption_words[:max_words])
+    return caption
+
+
+def pre_question(question, max_words=50):
+    question = re.sub(r"([.!\"()*#:;~])", ' ', question.capitalize())  # remove special characters
+    question = question.strip()
+
+    # truncate question
+    question_words = question.split(' ')
+    if len(question_words) > max_words:
+        question = ' '.join(question_words[:max_words])
+    if question[-1] != '?':
+        question += '?'
+    return question
+

prismer/dataset/vqa_dataset.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+from torch.utils.data import Dataset
+from dataset.utils import *
+
+
+class VQA(Dataset):
+    def __init__(self, config, train=True):
+        self.data_path = config['data_path']
+        self.label_path = config['label_path']
+        self.experts = config['experts']
+        self.transform = Transform(resize_resolution=config['image_resolution'], scale_size=[0.5, 1.0], train=train)
+        self.train = train
+
+        if train:
+            self.data_list = []
+            if 'vqav2' in config['datasets']:
+                self.data_list += json.load(open(os.path.join(self.data_path, 'vqav2_train_val.json'), 'r'))
+            if 'vg' in config['datasets']:
+                self.data_list += json.load(open(os.path.join(self.data_path, 'vg_qa.json'), 'r'))
+        else:
+            self.data_list = json.load(open(os.path.join(self.data_path, 'vqav2_test.json'), 'r'))
+            self.answer_list = json.load(open(os.path.join(self.data_path, 'answer_list.json'), 'r'))
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        data = self.data_list[index]
+
+        if data['dataset'] == 'vqa':
+            image, labels, labels_info = get_expert_labels(self.data_path, self.label_path, data['image'], 'vqav2', self.experts)
+        elif data['dataset'] == 'vg':
+            image, labels, labels_info = get_expert_labels(self.data_path, self.label_path, data['image'], 'vg', self.experts)
+
+        experts = self.transform(image, labels)
+        experts = post_label_process(experts, labels_info)
+
+        if self.train:
+            question = pre_question(data['question'], max_words=30)
+            answers = data['answer']
+            weights = torch.tensor(data['weight']) if data['dataset'] != 'vg' else torch.tensor(0.2)
+            return experts, question, answers, weights
+        else:
+            question = pre_question(data['question'], max_words=30)
+            question_id = data['question_id']
+            return experts, index, question, question_id

prismer/demo.py

@@ -0,0 +1,77 @@
+import os
+import argparse
+import torch
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+
+
+from model.prismer_caption import PrismerCaption
+from dataset import create_dataset, create_loader
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--mode', default='')
+parser.add_argument('--port', default='')
+
+parser.add_argument('--exp_name', default='', type=str)
+args = parser.parse_args()
+
+# load config
+config = yaml.load(open('configs/caption.yaml', 'r'), Loader=yaml.Loader)['demo']
+
+# generate expert labels
+if len(config['experts']) > 0:
+    script_name = f'python experts/generate_depth.py'
+    os.system(script_name)
+    print('***** Generated Depth *****')
+
+    script_name = f'python experts/generate_edge.py'
+    os.system(script_name)
+    print('***** Generated Edge *****')
+
+    script_name = f'python experts/generate_normal.py'
+    os.system(script_name)
+    print('***** Generated Surface Normals *****')
+
+    script_name = f'python experts/generate_objdet.py'
+    os.system(script_name)
+    print('***** Generated Object Detection Labels *****')
+
+    script_name = f'python experts/generate_ocrdet.py'
+    os.system(script_name)
+    print('***** Generated OCR Detection Labels *****')
+
+    script_name = f'python experts/generate_segmentation.py'
+    os.system(script_name)
+    print('***** Generated Segmentation Labels *****')
+
+# load datasets
+_, test_dataset = create_dataset('caption', config)
+test_loader = create_loader(test_dataset, batch_size=1, num_workers=4, train=False)
+
+# load pre-trained model
+model = PrismerCaption(config)
+state_dict = torch.load(f'logging/caption_{args.exp_name}/pytorch_model.bin', map_location='cuda:0')
+model.load_state_dict(state_dict)
+tokenizer = model.tokenizer
+
+# inference
+model.eval()
+with torch.no_grad():
+    for step, (experts, data_ids) in enumerate(tqdm(test_loader)):
+        captions = model(experts, train=False, prefix=config['prefix'])
+
+        captions = tokenizer(captions, max_length=30, padding='max_length', return_tensors='pt').input_ids
+        caption = captions.to(experts['rgb'].device)[0]
+
+        caption = tokenizer.decode(caption, skip_special_tokens=True)
+        caption = caption.capitalize() + '.'
+
+        # save caption
+        save_path = test_loader.dataset.data_list[data_ids[0]]['image'].replace('jpg', 'txt')
+        with open(save_path, 'w') as f:
+            f.write(caption)
+
+print('All Done.')

prismer/demo_vis.py

@@ -0,0 +1,161 @@
+import glob
+import os
+import json
+import torch
+import random
+import matplotlib.pyplot as plt
+import numpy as np
+
+from utils import create_ade20k_label_colormap
+
+obj_label_map = torch.load('dataset/detection_features.pt')['labels']
+coco_label_map = torch.load('dataset/coco_features.pt')['labels']
+ade_color = create_ade20k_label_colormap()
+
+file_path = 'helpers/images'
+expert_path = 'helpers/labels'
+plt.ioff()
+
+
+def get_label_path(file_name, expert_name, with_suffix=False):
+    file_suffix = '.png' if not with_suffix else '_.png'
+    label_name = ''.join(file_name.split('.')[:-1] + [file_suffix])
+    label_path = os.path.join(expert_path, expert_name, label_name)
+    return label_path
+
+
+def depth_prettify(file_name):
+    label_path = get_label_path(file_name, 'depth')
+    save_path = get_label_path(file_name, 'depth', True)
+    depth = plt.imread(label_path)
+    plt.imsave(save_path, depth, cmap='rainbow')
+
+
+def obj_detection_prettify(file_name):
+    label_path = get_label_path(file_name, 'obj_detection')
+    save_path = get_label_path(file_name, 'obj_detection', True)
+
+    rgb = plt.imread(file_name)
+    obj_labels = plt.imread(label_path)
+    obj_labels_dict = json.load(open(label_path.replace('.png', '.json')))
+
+    plt.imshow(rgb)
+
+    num_objs = np.unique(obj_labels)[:-1].max()
+    plt.imshow(obj_labels, cmap='terrain', vmax=num_objs + 1 / 255., alpha=0.5)
+
+    for i in np.unique(obj_labels)[:-1]:
+        obj_idx_all = np.where(obj_labels == i)
+        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 = obj_label_map[obj_labels_dict[str(int(i * 255))]]
+        plt.text(x, y, obj_name, c='white', horizontalalignment='center', verticalalignment='center')
+
+    plt.axis('off')
+    plt.savefig(save_path, bbox_inches='tight', transparent=True, pad_inches=0)
+    plt.close()
+
+
+def seg_prettify(file_name):
+    label_path = get_label_path(file_name, 'seg_coco')
+    save_path = get_label_path(file_name, 'seg_coco', True)
+
+    rgb = plt.imread(file_name)
+    seg_labels = plt.imread(label_path)
+
+    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)]
+
+    plt.imshow(seg_map, alpha=0.5)
+
+    for i in np.unique(seg_labels):
+        obj_idx_all = np.where(seg_labels == i)
+        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)]
+        plt.text(x, y, obj_name, c='white', horizontalalignment='center', verticalalignment='center')
+
+    plt.axis('off')
+    plt.savefig(save_path, bbox_inches='tight', transparent=True, pad_inches=0)
+    plt.close()
+
+
+def ocr_detection_prettify(file_name):
+    label_path = get_label_path(file_name, 'ocr_detection')
+    save_path = get_label_path(file_name, 'ocr_detection', True)
+
+    if os.path.exists(label_path):
+        rgb = plt.imread(file_name)
+        ocr_labels = plt.imread(label_path)
+        ocr_labels_dict = torch.load(label_path.replace('.png', '.pt'))
+
+        plt.imshow(rgb)
+        plt.imshow((1 - ocr_labels) < 1, 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')
+
+        plt.axis('off')
+        plt.savefig(save_path, bbox_inches='tight', transparent=True, pad_inches=0)
+        plt.close()
+    else:
+        rgb = plt.imread(file_name)
+        ocr_labels = np.ones_like(rgb, dtype=np.float32())
+
+        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')
+
+        plt.axis('off')
+        plt.savefig(save_path, bbox_inches='tight', transparent=True, pad_inches=0)
+        plt.close()
+
+
+im_list = glob.glob(file_path + '/*.jpg') + glob.glob(file_path + '/*.png') + glob.glob(file_path + '/*.jpeg')
+
+# prettify labels first:
+for i in range(len(im_list)):
+    depth_prettify(im_list[i])
+    seg_prettify(im_list[i])
+    ocr_detection_prettify(im_list[i])
+    obj_detection_prettify(im_list[i])
+
+pretty = {'depth': True, 'normal': False, 'edge': False,
+          'obj_detection': True, 'ocr_detection': True, 'seg_coco': True}
+
+# plot expert labels
+for im_path in im_list:
+    fig, axs = plt.subplots(1, 7, figsize=(20, 4))
+    rgb = plt.imread(im_path)
+    axs[0].imshow(rgb)
+    axs[0].axis('off')
+    axs[0].set_title('RGB')
+
+    for j in range(6):
+        label_name = list(pretty.keys())[j]
+        label_path = get_label_path(im_path, label_name, with_suffix=pretty[label_name])
+        label = plt.imread(label_path)
+        if label_name != 'edge':
+            axs[j + 1].imshow(label)
+        else:
+            axs[j + 1].imshow(label, cmap='gray')
+
+        axs[j + 1].axis('off')
+        axs[j + 1].set_title(label_name)
+
+    caption_path = ''.join(im_path.split('.')[:-1] + ['.txt'])
+    with open(caption_path) as f:
+        caption = f.readlines()[0]
+
+    plt.suptitle(caption)
+    plt.tight_layout()
+
+plt.show()

prismer/download_checkpoints.py

@@ -0,0 +1,124 @@
+from huggingface_hub import hf_hub_download, hf_hub_url, get_hf_file_metadata
+from huggingface_hub.utils import disable_progress_bars
+from pathlib import Path
+from rich.progress import Progress
+from fire import Fire
+from typing import Union, List
+
+_EXPERTS = [
+    "10_model.pth",
+    "Unified_learned_OCIM_RS200_6x+2x.pth",
+    "dpt_hybrid-midas-501f0c75.pt",
+    "icdar2015_hourglass88.pth",
+    "model_final_e0c58e.pkl",
+    "model_final_f07440.pkl",
+    "scannet.pt",
+]
+
+_MODELS = [
+    "vqa_prismer_base",
+    "vqa_prismer_large",
+    "vqa_prismerz_base",
+    "vqa_prismerz_large",
+    "caption_prismerz_base",
+    "caption_prismerz_large",
+    "caption_prismer_base",
+    "caption_prismer_large",
+    "pretrain_prismer_base",
+    "pretrain_prismer_large",
+    "pretrain_prismerz_base",
+    "pretrain_prismerz_large",
+]
+
+_REPO_ID = "lorenmt/prismer"
+
+
+def download_checkpoints(
+        download_experts: bool = False,
+        download_models: Union[bool, List] = False,
+        hide_tqdm: bool = False,
+        force_redownload: bool = False,
+):
+    if hide_tqdm:
+        disable_progress_bars()
+    # Convert to list and check for invalid names
+    download_experts = _EXPERTS if download_experts else []
+    if download_models:
+        # only download single model
+        if isinstance(download_models, str):
+            download_models = [download_models]
+
+        assert all([m in _MODELS for m in download_models]), f"Invalid model name. Must be one of {_MODELS}"
+        download_models = _MODELS if isinstance(download_models, bool) else download_models
+    else:
+        download_models = []
+
+    # Check if files already exist
+    if not force_redownload:
+        download_experts = [e for e in download_experts if not Path(f"./experts/expert_weights/{e}").exists()]
+        download_models = [m for m in download_models if not Path(f"{m}/pytorch_model.bin").exists()]
+    
+    assert download_experts or download_models, "Nothing to download."
+
+    with Progress() as progress:
+        # Calculate total download size
+        progress.print("[blue]Calculating download size...")
+        total_size = 0
+        for expert in download_experts:
+            url = hf_hub_url(
+                filename=expert,
+                repo_id=_REPO_ID,
+                subfolder="expert_weights"
+            )
+            total_size += get_hf_file_metadata(url).size
+
+        for model in download_models:
+            url = hf_hub_url(
+                filename=f"pytorch_model.bin",
+                repo_id=_REPO_ID,
+                subfolder=model
+            )
+            total_size += get_hf_file_metadata(url).size
+        progress.print(f"[blue]Total download size: {total_size / 1e9:.2f} GB")
+
+        # Download files
+        total_files = len(download_experts) + len(download_models)
+        total_task = progress.add_task(f"[green]Downloading files", total=total_files)
+        if download_experts:
+            expert_task = progress.add_task(
+                f"[green]Downloading experts...", total=len(download_experts)
+                )
+            out_folder = Path("experts/expert_weights")
+            out_folder.mkdir(parents=True, exist_ok=True)
+            for expert in download_experts:
+                path = Path(hf_hub_download(
+                    filename=expert,
+                    repo_id=_REPO_ID,
+                    subfolder="expert_weights"
+                ))
+                path.resolve().rename(out_folder/path.name)
+                path.unlink()
+                progress.advance(expert_task)
+                progress.advance(total_task)
+
+        if download_models:
+            model_task = progress.add_task(
+                f"[green]Downloading models...", total=len(download_models)
+                )
+            for model in download_models:
+                path = Path(hf_hub_download(
+                    filename=f"pytorch_model.bin",
+                    repo_id=_REPO_ID,
+                    subfolder=model
+                ))
+                out_folder = Path("./logging")/model
+                out_folder.mkdir(parents=True, exist_ok=True)
+                path.resolve().rename(out_folder/"pytorch_model.bin")
+                path.unlink()
+                progress.advance(model_task)
+                progress.advance(total_task)
+        progress.print("[green]Done!")
+
+
+if __name__ == "__main__":
+    Fire(download_checkpoints)

prismer/experts/depth/base_model.py

@@ -0,0 +1,16 @@
+import torch
+
+
+class BaseModel(torch.nn.Module):
+    def load(self, path):
+        """Load model from file.
+
+        Args:
+            path (str): file path
+        """
+        parameters = torch.load(path, map_location=torch.device("cpu"))
+
+        if "optimizer" in parameters:
+            parameters = parameters["model"]
+
+        self.load_state_dict(parameters)

prismer/experts/depth/blocks.py

@@ -0,0 +1,383 @@
+import torch
+import torch.nn as nn
+
+from .vit import (
+    _make_pretrained_vitb_rn50_384,
+    _make_pretrained_vitl16_384,
+    _make_pretrained_vitb16_384,
+    forward_vit,
+)
+
+
+def _make_encoder(
+    backbone,
+    features,
+    use_pretrained,
+    groups=1,
+    expand=False,
+    exportable=True,
+    hooks=None,
+    use_vit_only=False,
+    use_readout="ignore",
+    enable_attention_hooks=False,
+):
+    if backbone == "vitl16_384":
+        pretrained = _make_pretrained_vitl16_384(
+            use_pretrained,
+            hooks=hooks,
+            use_readout=use_readout,
+            enable_attention_hooks=enable_attention_hooks,
+        )
+        scratch = _make_scratch(
+            [256, 512, 1024, 1024], features, groups=groups, expand=expand
+        )  # ViT-L/16 - 85.0% Top1 (backbone)
+    elif backbone == "vitb_rn50_384":
+        pretrained = _make_pretrained_vitb_rn50_384(
+            use_pretrained,
+            hooks=hooks,
+            use_vit_only=use_vit_only,
+            use_readout=use_readout,
+            enable_attention_hooks=enable_attention_hooks,
+        )
+        scratch = _make_scratch(
+            [256, 512, 768, 768], features, groups=groups, expand=expand
+        )  # ViT-H/16 - 85.0% Top1 (backbone)
+    elif backbone == "vitb16_384":
+        pretrained = _make_pretrained_vitb16_384(
+            use_pretrained,
+            hooks=hooks,
+            use_readout=use_readout,
+            enable_attention_hooks=enable_attention_hooks,
+        )
+        scratch = _make_scratch(
+            [96, 192, 384, 768], features, groups=groups, expand=expand
+        )  # ViT-B/16 - 84.6% Top1 (backbone)
+    elif backbone == "resnext101_wsl":
+        pretrained = _make_pretrained_resnext101_wsl(use_pretrained)
+        scratch = _make_scratch(
+            [256, 512, 1024, 2048], features, groups=groups, expand=expand
+        )  # efficientnet_lite3
+    else:
+        print(f"Backbone '{backbone}' not implemented")
+        assert False
+
+    return pretrained, scratch
+
+
+def _make_scratch(in_shape, out_shape, groups=1, expand=False):
+    scratch = nn.Module()
+
+    out_shape1 = out_shape
+    out_shape2 = out_shape
+    out_shape3 = out_shape
+    out_shape4 = out_shape
+    if expand == True:
+        out_shape1 = out_shape
+        out_shape2 = out_shape * 2
+        out_shape3 = out_shape * 4
+        out_shape4 = out_shape * 8
+
+    scratch.layer1_rn = nn.Conv2d(
+        in_shape[0],
+        out_shape1,
+        kernel_size=3,
+        stride=1,
+        padding=1,
+        bias=False,
+        groups=groups,
+    )
+    scratch.layer2_rn = nn.Conv2d(
+        in_shape[1],
+        out_shape2,
+        kernel_size=3,
+        stride=1,
+        padding=1,
+        bias=False,
+        groups=groups,
+    )
+    scratch.layer3_rn = nn.Conv2d(
+        in_shape[2],
+        out_shape3,
+        kernel_size=3,
+        stride=1,
+        padding=1,
+        bias=False,
+        groups=groups,
+    )
+    scratch.layer4_rn = nn.Conv2d(
+        in_shape[3],
+        out_shape4,
+        kernel_size=3,
+        stride=1,
+        padding=1,
+        bias=False,
+        groups=groups,
+    )
+
+    return scratch
+
+
+def _make_resnet_backbone(resnet):
+    pretrained = nn.Module()
+    pretrained.layer1 = nn.Sequential(
+        resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1
+    )
+
+    pretrained.layer2 = resnet.layer2
+    pretrained.layer3 = resnet.layer3
+    pretrained.layer4 = resnet.layer4
+
+    return pretrained
+
+
+def _make_pretrained_resnext101_wsl(use_pretrained):
+    resnet = torch.hub.load("facebookresearch/WSL-Images", "resnext101_32x8d_wsl")
+    return _make_resnet_backbone(resnet)
+
+
+class Interpolate(nn.Module):
+    """Interpolation module."""
+
+    def __init__(self, scale_factor, mode, align_corners=False):
+        """Init.
+
+        Args:
+            scale_factor (float): scaling
+            mode (str): interpolation mode
+        """
+        super(Interpolate, self).__init__()
+
+        self.interp = nn.functional.interpolate
+        self.scale_factor = scale_factor
+        self.mode = mode
+        self.align_corners = align_corners
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (tensor): input
+
+        Returns:
+            tensor: interpolated data_list
+        """
+
+        x = self.interp(
+            x,
+            scale_factor=self.scale_factor,
+            mode=self.mode,
+            align_corners=self.align_corners,
+        )
+
+        return x
+
+
+class ResidualConvUnit(nn.Module):
+    """Residual convolution module."""
+
+    def __init__(self, features):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(
+            features, features, kernel_size=3, stride=1, padding=1, bias=True
+        )
+
+        self.conv2 = nn.Conv2d(
+            features, features, kernel_size=3, stride=1, padding=1, bias=True
+        )
+
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (tensor): input
+
+        Returns:
+            tensor: output
+        """
+        out = self.relu(x)
+        out = self.conv1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+
+        return out + x
+
+
+class FeatureFusionBlock(nn.Module):
+    """Feature fusion block."""
+
+    def __init__(self, features):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super(FeatureFusionBlock, self).__init__()
+
+        self.resConfUnit1 = ResidualConvUnit(features)
+        self.resConfUnit2 = ResidualConvUnit(features)
+
+    def forward(self, *xs):
+        """Forward pass.
+
+        Returns:
+            tensor: output
+        """
+        output = xs[0]
+
+        if len(xs) == 2:
+            output += self.resConfUnit1(xs[1])
+
+        output = self.resConfUnit2(output)
+
+        output = nn.functional.interpolate(
+            output, scale_factor=2, mode="bilinear", align_corners=True
+        )
+
+        return output
+
+
+class ResidualConvUnit_custom(nn.Module):
+    """Residual convolution module."""
+
+    def __init__(self, features, activation, bn):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super().__init__()
+
+        self.bn = bn
+
+        self.groups = 1
+
+        self.conv1 = nn.Conv2d(
+            features,
+            features,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=not self.bn,
+            groups=self.groups,
+        )
+
+        self.conv2 = nn.Conv2d(
+            features,
+            features,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=not self.bn,
+            groups=self.groups,
+        )
+
+        if self.bn == True:
+            self.bn1 = nn.BatchNorm2d(features)
+            self.bn2 = nn.BatchNorm2d(features)
+
+        self.activation = activation
+
+        self.skip_add = nn.quantized.FloatFunctional()
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (tensor): input
+
+        Returns:
+            tensor: output
+        """
+
+        out = self.activation(x)
+        out = self.conv1(out)
+        if self.bn == True:
+            out = self.bn1(out)
+
+        out = self.activation(out)
+        out = self.conv2(out)
+        if self.bn == True:
+            out = self.bn2(out)
+
+        if self.groups > 1:
+            out = self.conv_merge(out)
+
+        return self.skip_add.add(out, x)
+
+        # return out + x
+
+
+class FeatureFusionBlock_custom(nn.Module):
+    """Feature fusion block."""
+
+    def __init__(
+        self,
+        features,
+        activation,
+        deconv=False,
+        bn=False,
+        expand=False,
+        align_corners=True,
+    ):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super(FeatureFusionBlock_custom, self).__init__()
+
+        self.deconv = deconv
+        self.align_corners = align_corners
+
+        self.groups = 1
+
+        self.expand = expand
+        out_features = features
+        if self.expand == True:
+            out_features = features // 2
+
+        self.out_conv = nn.Conv2d(
+            features,
+            out_features,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+            groups=1,
+        )
+
+        self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn)
+        self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn)
+
+        self.skip_add = nn.quantized.FloatFunctional()
+
+    def forward(self, *xs):
+        """Forward pass.
+
+        Returns:
+            tensor: output
+        """
+        output = xs[0]
+
+        if len(xs) == 2:
+            res = self.resConfUnit1(xs[1])
+            output = self.skip_add.add(output, res)
+            # output += res
+
+        output = self.resConfUnit2(output)
+
+        output = nn.functional.interpolate(
+            output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
+        )
+
+        output = self.out_conv(output)
+
+        return output

prismer/experts/depth/generate_dataset.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import glob
+
+from torch.utils.data import Dataset
+from PIL import ImageFile
+from dataset.utils import *
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+class Dataset(Dataset):
+    def __init__(self, data_path, transform):
+        self.data_path = 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')]
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        image_path = self.data_list[index]
+        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

prismer/experts/depth/models.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .base_model import BaseModel
+from .blocks import (
+    FeatureFusionBlock,
+    FeatureFusionBlock_custom,
+    Interpolate,
+    _make_encoder,
+    forward_vit,
+)
+
+
+def _make_fusion_block(features, use_bn):
+    return FeatureFusionBlock_custom(
+        features,
+        nn.ReLU(False),
+        deconv=False,
+        bn=use_bn,
+        expand=False,
+        align_corners=True,
+    )
+
+
+class DPT(BaseModel):
+    def __init__(
+        self,
+        head,
+        features=256,
+        backbone="vitb_rn50_384",
+        readout="project",
+        channels_last=False,
+        use_bn=False,
+        enable_attention_hooks=False,
+    ):
+
+        super(DPT, self).__init__()
+
+        self.channels_last = channels_last
+
+        hooks = {
+            "vitb_rn50_384": [0, 1, 8, 11],
+            "vitb16_384": [2, 5, 8, 11],
+            "vitl16_384": [5, 11, 17, 23],
+        }
+
+        # Instantiate backbone and reassemble blocks
+        self.pretrained, self.scratch = _make_encoder(
+            backbone,
+            features,
+            False,  # Set to true of you want to train from scratch, uses ImageNet weights
+            groups=1,
+            expand=False,
+            exportable=False,
+            hooks=hooks[backbone],
+            use_readout=readout,
+            enable_attention_hooks=enable_attention_hooks,
+        )
+
+        self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
+
+        self.scratch.output_conv = head
+
+    def forward(self, x):
+        if self.channels_last == True:
+            x.contiguous(memory_format=torch.channels_last)
+
+        layer_1, layer_2, layer_3, layer_4 = forward_vit(self.pretrained, x)
+
+        layer_1_rn = self.scratch.layer1_rn(layer_1)
+        layer_2_rn = self.scratch.layer2_rn(layer_2)
+        layer_3_rn = self.scratch.layer3_rn(layer_3)
+        layer_4_rn = self.scratch.layer4_rn(layer_4)
+
+        path_4 = self.scratch.refinenet4(layer_4_rn)
+        path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
+        path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
+        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
+
+        out = self.scratch.output_conv(path_1)
+
+        return out
+
+
+class DPTDepthModel(DPT):
+    def __init__(
+        self, path=None, non_negative=True, scale=1.0, shift=0.0, invert=False, **kwargs
+    ):
+        features = kwargs["features"] if "features" in kwargs else 256
+
+        self.scale = scale
+        self.shift = shift
+        self.invert = invert
+
+        head = nn.Sequential(
+            nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),
+            Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
+            nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(True),
+            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
+            nn.ReLU(True) if non_negative else nn.Identity(),
+            nn.Identity(),
+        )
+
+        super().__init__(head, **kwargs)
+
+        if path is not None:
+            self.load(path)
+
+    def forward(self, x):
+        inv_depth = super().forward(x).squeeze(dim=1)
+
+        if self.invert:
+            depth = self.scale * inv_depth + self.shift
+            depth[depth < 1e-8] = 1e-8
+            depth = 1.0 / depth
+            return depth
+        else:
+            return inv_depth
+

prismer/experts/depth/vit.py

@@ -0,0 +1,576 @@
+import torch
+import torch.nn as nn
+import timm
+import types
+import math
+import torch.nn.functional as F
+
+
+activations = {}
+
+
+def get_activation(name):
+    def hook(model, input, output):
+        activations[name] = output
+
+    return hook
+
+
+attention = {}
+
+
+def get_attention(name):
+    def hook(module, input, output):
+        x = input[0]
+        B, N, C = x.shape
+        qkv = (
+            module.qkv(x)
+            .reshape(B, N, 3, module.num_heads, C // module.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * module.scale
+
+        attn = attn.softmax(dim=-1)  # [:,:,1,1:]
+        attention[name] = attn
+
+    return hook
+
+
+def get_mean_attention_map(attn, token, shape):
+    attn = attn[:, :, token, 1:]
+    attn = attn.unflatten(2, torch.Size([shape[2] // 16, shape[3] // 16])).float()
+    attn = torch.nn.functional.interpolate(
+        attn, size=shape[2:], mode="bicubic", align_corners=False
+    ).squeeze(0)
+
+    all_attn = torch.mean(attn, 0)
+
+    return all_attn
+
+
+class Slice(nn.Module):
+    def __init__(self, start_index=1):
+        super(Slice, self).__init__()
+        self.start_index = start_index
+
+    def forward(self, x):
+        return x[:, self.start_index :]
+
+
+class AddReadout(nn.Module):
+    def __init__(self, start_index=1):
+        super(AddReadout, self).__init__()
+        self.start_index = start_index
+
+    def forward(self, x):
+        if self.start_index == 2:
+            readout = (x[:, 0] + x[:, 1]) / 2
+        else:
+            readout = x[:, 0]
+        return x[:, self.start_index :] + readout.unsqueeze(1)
+
+
+class ProjectReadout(nn.Module):
+    def __init__(self, in_features, start_index=1):
+        super(ProjectReadout, self).__init__()
+        self.start_index = start_index
+
+        self.project = nn.Sequential(nn.Linear(2 * in_features, in_features), nn.GELU())
+
+    def forward(self, x):
+        readout = x[:, 0].unsqueeze(1).expand_as(x[:, self.start_index :])
+        features = torch.cat((x[:, self.start_index :], readout), -1)
+
+        return self.project(features)
+
+
+class Transpose(nn.Module):
+    def __init__(self, dim0, dim1):
+        super(Transpose, self).__init__()
+        self.dim0 = dim0
+        self.dim1 = dim1
+
+    def forward(self, x):
+        x = x.transpose(self.dim0, self.dim1)
+        return x
+
+
+def forward_vit(pretrained, x):
+    b, c, h, w = x.shape
+
+    glob = pretrained.model.forward_flex(x)
+
+    layer_1 = pretrained.activations["1"]
+    layer_2 = pretrained.activations["2"]
+    layer_3 = pretrained.activations["3"]
+    layer_4 = pretrained.activations["4"]
+
+    layer_1 = pretrained.act_postprocess1[0:2](layer_1)
+    layer_2 = pretrained.act_postprocess2[0:2](layer_2)
+    layer_3 = pretrained.act_postprocess3[0:2](layer_3)
+    layer_4 = pretrained.act_postprocess4[0:2](layer_4)
+
+    unflatten = nn.Sequential(
+        nn.Unflatten(
+            2,
+            torch.Size(
+                [
+                    h // pretrained.model.patch_size[1],
+                    w // pretrained.model.patch_size[0],
+                ]
+            ),
+        )
+    )
+
+    if layer_1.ndim == 3:
+        layer_1 = unflatten(layer_1)
+    if layer_2.ndim == 3:
+        layer_2 = unflatten(layer_2)
+    if layer_3.ndim == 3:
+        layer_3 = unflatten(layer_3)
+    if layer_4.ndim == 3:
+        layer_4 = unflatten(layer_4)
+
+    layer_1 = pretrained.act_postprocess1[3 : len(pretrained.act_postprocess1)](layer_1)
+    layer_2 = pretrained.act_postprocess2[3 : len(pretrained.act_postprocess2)](layer_2)
+    layer_3 = pretrained.act_postprocess3[3 : len(pretrained.act_postprocess3)](layer_3)
+    layer_4 = pretrained.act_postprocess4[3 : len(pretrained.act_postprocess4)](layer_4)
+
+    return layer_1, layer_2, layer_3, layer_4
+
+
+def _resize_pos_embed(self, posemb, gs_h, gs_w):
+    posemb_tok, posemb_grid = (
+        posemb[:, : self.start_index],
+        posemb[0, self.start_index :],
+    )
+
+    gs_old = int(math.sqrt(len(posemb_grid)))
+
+    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
+    posemb_grid = F.interpolate(posemb_grid, size=(gs_h, gs_w), mode="bilinear")
+    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_h * gs_w, -1)
+
+    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
+
+    return posemb
+
+
+def forward_flex(self, x):
+    b, c, h, w = x.shape
+
+    pos_embed = self._resize_pos_embed(
+        self.pos_embed, h // self.patch_size[1], w // self.patch_size[0]
+    )
+
+    B = x.shape[0]
+
+    if hasattr(self.patch_embed, "backbone"):
+        x = self.patch_embed.backbone(x)
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+
+    x = self.patch_embed.proj(x).flatten(2).transpose(1, 2)
+
+    if getattr(self, "dist_token", None) is not None:
+        cls_tokens = self.cls_token.expand(
+            B, -1, -1
+        )  # stole cls_tokens impl from Phil Wang, thanks
+        dist_token = self.dist_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, dist_token, x), dim=1)
+    else:
+        cls_tokens = self.cls_token.expand(
+            B, -1, -1
+        )  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+
+    x = x + pos_embed
+    x = self.pos_drop(x)
+
+    for blk in self.blocks:
+        x = blk(x)
+
+    x = self.norm(x)
+
+    return x
+
+
+def get_readout_oper(vit_features, features, use_readout, start_index=1):
+    if use_readout == "ignore":
+        readout_oper = [Slice(start_index)] * len(features)
+    elif use_readout == "add":
+        readout_oper = [AddReadout(start_index)] * len(features)
+    elif use_readout == "project":
+        readout_oper = [
+            ProjectReadout(vit_features, start_index) for out_feat in features
+        ]
+    else:
+        assert (
+            False
+        ), "wrong operation for readout token, use_readout can be 'ignore', 'add', or 'project'"
+
+    return readout_oper
+
+
+def _make_vit_b16_backbone(
+    model,
+    features=[96, 192, 384, 768],
+    size=[384, 384],
+    hooks=[2, 5, 8, 11],
+    vit_features=768,
+    use_readout="ignore",
+    start_index=1,
+    enable_attention_hooks=False,
+):
+    pretrained = nn.Module()
+
+    pretrained.model = model
+    pretrained.model.blocks[hooks[0]].register_forward_hook(get_activation("1"))
+    pretrained.model.blocks[hooks[1]].register_forward_hook(get_activation("2"))
+    pretrained.model.blocks[hooks[2]].register_forward_hook(get_activation("3"))
+    pretrained.model.blocks[hooks[3]].register_forward_hook(get_activation("4"))
+
+    pretrained.activations = activations
+
+    if enable_attention_hooks:
+        pretrained.model.blocks[hooks[0]].attn.register_forward_hook(
+            get_attention("attn_1")
+        )
+        pretrained.model.blocks[hooks[1]].attn.register_forward_hook(
+            get_attention("attn_2")
+        )
+        pretrained.model.blocks[hooks[2]].attn.register_forward_hook(
+            get_attention("attn_3")
+        )
+        pretrained.model.blocks[hooks[3]].attn.register_forward_hook(
+            get_attention("attn_4")
+        )
+        pretrained.attention = attention
+
+    readout_oper = get_readout_oper(vit_features, features, use_readout, start_index)
+
+    # 32, 48, 136, 384
+    pretrained.act_postprocess1 = nn.Sequential(
+        readout_oper[0],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[0],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+        nn.ConvTranspose2d(
+            in_channels=features[0],
+            out_channels=features[0],
+            kernel_size=4,
+            stride=4,
+            padding=0,
+            bias=True,
+            dilation=1,
+            groups=1,
+        ),
+    )
+
+    pretrained.act_postprocess2 = nn.Sequential(
+        readout_oper[1],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[1],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+        nn.ConvTranspose2d(
+            in_channels=features[1],
+            out_channels=features[1],
+            kernel_size=2,
+            stride=2,
+            padding=0,
+            bias=True,
+            dilation=1,
+            groups=1,
+        ),
+    )
+
+    pretrained.act_postprocess3 = nn.Sequential(
+        readout_oper[2],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[2],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+    )
+
+    pretrained.act_postprocess4 = nn.Sequential(
+        readout_oper[3],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[3],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+        nn.Conv2d(
+            in_channels=features[3],
+            out_channels=features[3],
+            kernel_size=3,
+            stride=2,
+            padding=1,
+        ),
+    )
+
+    pretrained.model.start_index = start_index
+    pretrained.model.patch_size = [16, 16]
+
+    # We inject this function into the VisionTransformer instances so that
+    # we can use it with interpolated position embeddings without modifying the library source.
+    pretrained.model.forward_flex = types.MethodType(forward_flex, pretrained.model)
+    pretrained.model._resize_pos_embed = types.MethodType(
+        _resize_pos_embed, pretrained.model
+    )
+
+    return pretrained
+
+
+def _make_vit_b_rn50_backbone(
+    model,
+    features=[256, 512, 768, 768],
+    size=[384, 384],
+    hooks=[0, 1, 8, 11],
+    vit_features=768,
+    use_vit_only=False,
+    use_readout="ignore",
+    start_index=1,
+    enable_attention_hooks=False,
+):
+    pretrained = nn.Module()
+
+    pretrained.model = model
+
+    if use_vit_only == True:
+        pretrained.model.blocks[hooks[0]].register_forward_hook(get_activation("1"))
+        pretrained.model.blocks[hooks[1]].register_forward_hook(get_activation("2"))
+    else:
+        pretrained.model.patch_embed.backbone.stages[0].register_forward_hook(
+            get_activation("1")
+        )
+        pretrained.model.patch_embed.backbone.stages[1].register_forward_hook(
+            get_activation("2")
+        )
+
+    pretrained.model.blocks[hooks[2]].register_forward_hook(get_activation("3"))
+    pretrained.model.blocks[hooks[3]].register_forward_hook(get_activation("4"))
+
+    if enable_attention_hooks:
+        pretrained.model.blocks[2].attn.register_forward_hook(get_attention("attn_1"))
+        pretrained.model.blocks[5].attn.register_forward_hook(get_attention("attn_2"))
+        pretrained.model.blocks[8].attn.register_forward_hook(get_attention("attn_3"))
+        pretrained.model.blocks[11].attn.register_forward_hook(get_attention("attn_4"))
+        pretrained.attention = attention
+
+    pretrained.activations = activations
+
+    readout_oper = get_readout_oper(vit_features, features, use_readout, start_index)
+
+    if use_vit_only == True:
+        pretrained.act_postprocess1 = nn.Sequential(
+            readout_oper[0],
+            Transpose(1, 2),
+            nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+            nn.Conv2d(
+                in_channels=vit_features,
+                out_channels=features[0],
+                kernel_size=1,
+                stride=1,
+                padding=0,
+            ),
+            nn.ConvTranspose2d(
+                in_channels=features[0],
+                out_channels=features[0],
+                kernel_size=4,
+                stride=4,
+                padding=0,
+                bias=True,
+                dilation=1,
+                groups=1,
+            ),
+        )
+
+        pretrained.act_postprocess2 = nn.Sequential(
+            readout_oper[1],
+            Transpose(1, 2),
+            nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+            nn.Conv2d(
+                in_channels=vit_features,
+                out_channels=features[1],
+                kernel_size=1,
+                stride=1,
+                padding=0,
+            ),
+            nn.ConvTranspose2d(
+                in_channels=features[1],
+                out_channels=features[1],
+                kernel_size=2,
+                stride=2,
+                padding=0,
+                bias=True,
+                dilation=1,
+                groups=1,
+            ),
+        )
+    else:
+        pretrained.act_postprocess1 = nn.Sequential(
+            nn.Identity(), nn.Identity(), nn.Identity()
+        )
+        pretrained.act_postprocess2 = nn.Sequential(
+            nn.Identity(), nn.Identity(), nn.Identity()
+        )
+
+    pretrained.act_postprocess3 = nn.Sequential(
+        readout_oper[2],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[2],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+    )
+
+    pretrained.act_postprocess4 = nn.Sequential(
+        readout_oper[3],
+        Transpose(1, 2),
+        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
+        nn.Conv2d(
+            in_channels=vit_features,
+            out_channels=features[3],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        ),
+        nn.Conv2d(
+            in_channels=features[3],
+            out_channels=features[3],
+            kernel_size=3,
+            stride=2,
+            padding=1,
+        ),
+    )
+
+    pretrained.model.start_index = start_index
+    pretrained.model.patch_size = [16, 16]
+
+    # We inject this function into the VisionTransformer instances so that
+    # we can use it with interpolated position embeddings without modifying the library source.
+    pretrained.model.forward_flex = types.MethodType(forward_flex, pretrained.model)
+
+    # We inject this function into the VisionTransformer instances so that
+    # we can use it with interpolated position embeddings without modifying the library source.
+    pretrained.model._resize_pos_embed = types.MethodType(
+        _resize_pos_embed, pretrained.model
+    )
+
+    return pretrained
+
+
+def _make_pretrained_vitb_rn50_384(
+    pretrained,
+    use_readout="ignore",
+    hooks=None,
+    use_vit_only=False,
+    enable_attention_hooks=False,
+):
+    model = timm.create_model("vit_base_resnet50_384", pretrained=pretrained)
+
+    hooks = [0, 1, 8, 11] if hooks == None else hooks
+    return _make_vit_b_rn50_backbone(
+        model,
+        features=[256, 512, 768, 768],
+        size=[384, 384],
+        hooks=hooks,
+        use_vit_only=use_vit_only,
+        use_readout=use_readout,
+        enable_attention_hooks=enable_attention_hooks,
+    )
+
+
+def _make_pretrained_vitl16_384(
+    pretrained, use_readout="ignore", hooks=None, enable_attention_hooks=False
+):
+    model = timm.create_model("vit_large_patch16_384", pretrained=pretrained)
+
+    hooks = [5, 11, 17, 23] if hooks == None else hooks
+    return _make_vit_b16_backbone(
+        model,
+        features=[256, 512, 1024, 1024],
+        hooks=hooks,
+        vit_features=1024,
+        use_readout=use_readout,
+        enable_attention_hooks=enable_attention_hooks,
+    )
+
+
+def _make_pretrained_vitb16_384(
+    pretrained, use_readout="ignore", hooks=None, enable_attention_hooks=False
+):
+    model = timm.create_model("vit_base_patch16_384", pretrained=pretrained)
+
+    hooks = [2, 5, 8, 11] if hooks == None else hooks
+    return _make_vit_b16_backbone(
+        model,
+        features=[96, 192, 384, 768],
+        hooks=hooks,
+        use_readout=use_readout,
+        enable_attention_hooks=enable_attention_hooks,
+    )
+
+
+def _make_pretrained_deitb16_384(
+    pretrained, use_readout="ignore", hooks=None, enable_attention_hooks=False
+):
+    model = timm.create_model("vit_deit_base_patch16_384", pretrained=pretrained)
+
+    hooks = [2, 5, 8, 11] if hooks == None else hooks
+    return _make_vit_b16_backbone(
+        model,
+        features=[96, 192, 384, 768],
+        hooks=hooks,
+        use_readout=use_readout,
+        enable_attention_hooks=enable_attention_hooks,
+    )
+
+
+def _make_pretrained_deitb16_distil_384(
+    pretrained, use_readout="ignore", hooks=None, enable_attention_hooks=False
+):
+    model = timm.create_model(
+        "vit_deit_base_distilled_patch16_384", pretrained=pretrained
+    )
+
+    hooks = [2, 5, 8, 11] if hooks == None else hooks
+    return _make_vit_b16_backbone(
+        model,
+        features=[96, 192, 384, 768],
+        hooks=hooks,
+        use_readout=use_readout,
+        start_index=2,
+        enable_attention_hooks=enable_attention_hooks,
+    )

prismer/experts/edge/generate_dataset.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import glob
+
+from torch.utils.data import Dataset
+from dataset.utils import *
+from PIL import ImageFile
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+class Dataset(Dataset):
+    def __init__(self, data_path, transform):
+        self.data_path = 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')]
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        image_path = self.data_list[index]
+        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

prismer/experts/edge/images.py

@@ -0,0 +1,50 @@
+import cv2
+import numpy as np
+import torch
+
+
+def image_normalization(img, img_min=0, img_max=255,
+                        epsilon=1e-12):
+    """This is a typical image normalization function
+    where the minimum and maximum of the image is needed
+    source: https://en.wikipedia.org/wiki/Normalization_(image_processing)
+
+    :param img: an image could be gray scale or color
+    :param img_min:  for default is 0
+    :param img_max: for default is 255
+
+    :return: a normalized image, if max is 255 the dtype is uint8
+    """
+
+    img = np.float32(img)
+    # whenever an inconsistent image
+    img = (img - np.min(img)) * (img_max - img_min) / \
+        ((np.max(img) - np.min(img)) + epsilon) + img_min
+    return img
+
+
+def fuse_edge(pred):
+    edge_maps = []
+    for i in pred:
+        tmp = torch.sigmoid(i).cpu().detach().numpy()
+        edge_maps.append(tmp)
+    tensor = np.array(edge_maps)
+
+    fuses = []
+    for idx in range(tensor.shape[1]):
+        tmp = tensor[:, idx, ...]
+        tmp = np.squeeze(tmp)
+
+        # Iterate our all 7 NN outputs for a particular image
+        for i in range(tmp.shape[0]):
+            tmp_img = tmp[i]
+            tmp_img = np.uint8(image_normalization(tmp_img))
+            tmp_img = cv2.bitwise_not(tmp_img)
+
+            if i == 6:
+                fuse = tmp_img
+                fuse = fuse.astype(np.uint8)
+        fuses.append(fuse)
+    return fuses
+
+

prismer/experts/edge/model.py

@@ -0,0 +1,286 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def weight_init(m):
+    if isinstance(m, (nn.Conv2d,)):
+        # torch.nn.init.xavier_uniform_(m.weight, gain=1.0)
+        torch.nn.init.xavier_normal_(m.weight, gain=1.0)
+        # torch.nn.init.normal_(m.weight, mean=0.0, std=0.01)
+        if m.weight.data.shape[1] == torch.Size([1]):
+            torch.nn.init.normal_(m.weight, mean=0.0)
+
+        if m.bias is not None:
+            torch.nn.init.zeros_(m.bias)
+
+    # for fusion layer
+    if isinstance(m, (nn.ConvTranspose2d,)):
+        # torch.nn.init.xavier_uniform_(m.weight, gain=1.0)
+        torch.nn.init.xavier_normal_(m.weight, gain=1.0)
+        # torch.nn.init.normal_(m.weight, mean=0.0, std=0.01)
+
+        if m.weight.data.shape[1] == torch.Size([1]):
+            torch.nn.init.normal_(m.weight, std=0.1)
+        if m.bias is not None:
+            torch.nn.init.zeros_(m.bias)
+
+
+class CoFusion(nn.Module):
+
+    def __init__(self, in_ch, out_ch):
+        super(CoFusion, self).__init__()
+        self.conv1 = nn.Conv2d(in_ch, 64, kernel_size=3,
+                               stride=1, padding=1)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3,
+                               stride=1, padding=1)
+        self.conv3 = nn.Conv2d(64, out_ch, kernel_size=3,
+                               stride=1, padding=1)
+        self.relu = nn.ReLU()
+
+        self.norm_layer1 = nn.GroupNorm(4, 64)
+        self.norm_layer2 = nn.GroupNorm(4, 64)
+
+    def forward(self, x):
+        # fusecat = torch.cat(x, dim=1)
+        attn = self.relu(self.norm_layer1(self.conv1(x)))
+        attn = self.relu(self.norm_layer2(self.conv2(attn)))
+        attn = F.softmax(self.conv3(attn), dim=1)
+
+        # return ((fusecat * attn).sum(1)).unsqueeze(1)
+        return ((x * attn).sum(1)).unsqueeze(1)
+
+class _DenseLayer(nn.Sequential):
+    def __init__(self, input_features, out_features):
+        super(_DenseLayer, self).__init__()
+
+        # self.add_module('relu2', nn.ReLU(inplace=True)),
+        self.add_module('conv1', nn.Conv2d(input_features, out_features,
+                                           kernel_size=3, stride=1, padding=2, bias=True)),
+        self.add_module('norm1', nn.BatchNorm2d(out_features)),
+        self.add_module('relu1', nn.ReLU(inplace=True)),
+        self.add_module('conv2', nn.Conv2d(out_features, out_features,
+                                           kernel_size=3, stride=1, bias=True)),
+        self.add_module('norm2', nn.BatchNorm2d(out_features))
+
+    def forward(self, x):
+        x1, x2 = x
+
+        new_features = super(_DenseLayer, self).forward(F.relu(x1))  # F.relu()
+        # if new_features.shape[-1]!=x2.shape[-1]:
+        #     new_features =F.interpolate(new_features,size=(x2.shape[2],x2.shape[-1]), mode='bicubic',
+        #                                 align_corners=False)
+        return 0.5 * (new_features + x2), x2
+
+
+class _DenseBlock(nn.Sequential):
+    def __init__(self, num_layers, input_features, out_features):
+        super(_DenseBlock, self).__init__()
+        for i in range(num_layers):
+            layer = _DenseLayer(input_features, out_features)
+            self.add_module('denselayer%d' % (i + 1), layer)
+            input_features = out_features
+
+
+class UpConvBlock(nn.Module):
+    def __init__(self, in_features, up_scale):
+        super(UpConvBlock, self).__init__()
+        self.up_factor = 2
+        self.constant_features = 16
+
+        layers = self.make_deconv_layers(in_features, up_scale)
+        assert layers is not None, layers
+        self.features = nn.Sequential(*layers)
+
+    def make_deconv_layers(self, in_features, up_scale):
+        layers = []
+        all_pads=[0,0,1,3,7]
+        for i in range(up_scale):
+            kernel_size = 2 ** up_scale
+            pad = all_pads[up_scale]  # kernel_size-1
+            out_features = self.compute_out_features(i, up_scale)
+            layers.append(nn.Conv2d(in_features, out_features, 1))
+            layers.append(nn.ReLU(inplace=True))
+            layers.append(nn.ConvTranspose2d(
+                out_features, out_features, kernel_size, stride=2, padding=pad))
+            in_features = out_features
+        return layers
+
+    def compute_out_features(self, idx, up_scale):
+        return 1 if idx == up_scale - 1 else self.constant_features
+
+    def forward(self, x):
+        return self.features(x)
+
+
+class SingleConvBlock(nn.Module):
+    def __init__(self, in_features, out_features, stride,
+                 use_bs=True
+                 ):
+        super(SingleConvBlock, self).__init__()
+        self.use_bn = use_bs
+        self.conv = nn.Conv2d(in_features, out_features, 1, stride=stride,
+                              bias=True)
+        self.bn = nn.BatchNorm2d(out_features)
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.use_bn:
+            x = self.bn(x)
+        return x
+
+
+class DoubleConvBlock(nn.Module):
+    def __init__(self, in_features, mid_features,
+                 out_features=None,
+                 stride=1,
+                 use_act=True):
+        super(DoubleConvBlock, self).__init__()
+
+        self.use_act = use_act
+        if out_features is None:
+            out_features = mid_features
+        self.conv1 = nn.Conv2d(in_features, mid_features,
+                               3, padding=1, stride=stride)
+        self.bn1 = nn.BatchNorm2d(mid_features)
+        self.conv2 = nn.Conv2d(mid_features, out_features, 3, padding=1)
+        self.bn2 = nn.BatchNorm2d(out_features)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        if self.use_act:
+            x = self.relu(x)
+        return x
+
+
+class DexiNed(nn.Module):
+    """ Definition of the DXtrem network. """
+
+    def __init__(self):
+        super(DexiNed, self).__init__()
+        self.block_1 = DoubleConvBlock(3, 32, 64, stride=2,)
+        self.block_2 = DoubleConvBlock(64, 128, use_act=False)
+        self.dblock_3 = _DenseBlock(2, 128, 256) # [128,256,100,100]
+        self.dblock_4 = _DenseBlock(3, 256, 512)
+        self.dblock_5 = _DenseBlock(3, 512, 512)
+        self.dblock_6 = _DenseBlock(3, 512, 256)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        # left skip connections, figure in Journal
+        self.side_1 = SingleConvBlock(64, 128, 2)
+        self.side_2 = SingleConvBlock(128, 256, 2)
+        self.side_3 = SingleConvBlock(256, 512, 2)
+        self.side_4 = SingleConvBlock(512, 512, 1)
+        self.side_5 = SingleConvBlock(512, 256, 1) # Sory I forget to comment this line :(
+
+        # right skip connections, figure in Journal paper
+        self.pre_dense_2 = SingleConvBlock(128, 256, 2)
+        self.pre_dense_3 = SingleConvBlock(128, 256, 1)
+        self.pre_dense_4 = SingleConvBlock(256, 512, 1)
+        self.pre_dense_5 = SingleConvBlock(512, 512, 1)
+        self.pre_dense_6 = SingleConvBlock(512, 256, 1)
+
+
+        self.up_block_1 = UpConvBlock(64, 1)
+        self.up_block_2 = UpConvBlock(128, 1)
+        self.up_block_3 = UpConvBlock(256, 2)
+        self.up_block_4 = UpConvBlock(512, 3)
+        self.up_block_5 = UpConvBlock(512, 4)
+        self.up_block_6 = UpConvBlock(256, 4)
+        self.block_cat = SingleConvBlock(6, 1, stride=1, use_bs=False) # hed fusion method
+        # self.block_cat = CoFusion(6,6)# cats fusion method
+
+
+        self.apply(weight_init)
+
+    def slice(self, tensor, slice_shape):
+        t_shape = tensor.shape
+        height, width = slice_shape
+        if t_shape[-1]!=slice_shape[-1]:
+            new_tensor = F.interpolate(
+                tensor, size=(height, width), mode='bicubic',align_corners=False)
+        else:
+            new_tensor=tensor
+        # tensor[..., :height, :width]
+        return new_tensor
+
+    def forward(self, x):
+        assert x.ndim == 4, x.shape
+
+        # Block 1
+        block_1 = self.block_1(x)
+        block_1_side = self.side_1(block_1)
+
+        # Block 2
+        block_2 = self.block_2(block_1)
+        block_2_down = self.maxpool(block_2)
+        block_2_add = block_2_down + block_1_side
+        block_2_side = self.side_2(block_2_add)
+
+        # Block 3
+        block_3_pre_dense = self.pre_dense_3(block_2_down)
+        block_3, _ = self.dblock_3([block_2_add, block_3_pre_dense])
+        block_3_down = self.maxpool(block_3) # [128,256,50,50]
+        block_3_add = block_3_down + block_2_side
+        block_3_side = self.side_3(block_3_add)
+
+        # Block 4
+        block_2_resize_half = self.pre_dense_2(block_2_down)
+        block_4_pre_dense = self.pre_dense_4(block_3_down+block_2_resize_half)
+        block_4, _ = self.dblock_4([block_3_add, block_4_pre_dense])
+        block_4_down = self.maxpool(block_4)
+        block_4_add = block_4_down + block_3_side
+        block_4_side = self.side_4(block_4_add)
+
+        # Block 5
+        block_5_pre_dense = self.pre_dense_5(
+            block_4_down) #block_5_pre_dense_512 +block_4_down
+        block_5, _ = self.dblock_5([block_4_add, block_5_pre_dense])
+        block_5_add = block_5 + block_4_side
+
+        # Block 6
+        block_6_pre_dense = self.pre_dense_6(block_5)
+        block_6, _ = self.dblock_6([block_5_add, block_6_pre_dense])
+
+        # upsampling blocks
+        out_1 = self.up_block_1(block_1)
+        out_2 = self.up_block_2(block_2)
+        out_3 = self.up_block_3(block_3)
+        out_4 = self.up_block_4(block_4)
+        out_5 = self.up_block_5(block_5)
+        out_6 = self.up_block_6(block_6)
+        results = [out_1, out_2, out_3, out_4, out_5, out_6]
+
+        # concatenate multiscale outputs
+        block_cat = torch.cat(results, dim=1)  # Bx6xHxW
+        block_cat = self.block_cat(block_cat)  # Bx1xHxW
+
+        # return results
+        results.append(block_cat)
+        return results
+
+
+if __name__ == '__main__':
+    batch_size = 8
+    img_height = 352
+    img_width = 352
+
+    # device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = "cpu"
+    input = torch.rand(batch_size, 3, img_height, img_width).to(device)
+    # target = torch.rand(batch_size, 1, img_height, img_width).to(device)
+    print(f"input shape: {input.shape}")
+    model = DexiNed().to(device)
+    output = model(input)
+    print(f"output shapes: {[t.shape for t in output]}")
+
+    # for i in range(20000):
+    #     print(i)
+    #     output = model(input)
+    #     loss = nn.MSELoss()(output[-1], target)
+    #     loss.backward()

prismer/experts/generate_depth.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import torch
+import os
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
+from experts.model_bank import load_expert_model
+from experts.depth.generate_dataset import Dataset
+import PIL.Image as Image
+from accelerate import Accelerator
+from tqdm import tqdm
+
+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)
+data_loader = torch.utils.data.DataLoader(
+    dataset=dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=4,
+    pin_memory=True
+)
+
+model, data_loader = accelerator.prepare(model, data_loader)
+
+with torch.no_grad():
+    for i, (test_data, img_path, img_size) in enumerate(tqdm(data_loader)):
+        test_pred = model(test_data)
+
+        for k in range(len(test_pred)):
+            img_path_split = img_path[k].split('/')
+            ps = img_path[k].split('.')[-1]
+            im_save_path = os.path.join(save_path, img_path_split[-3], img_path_split[-2])
+            os.makedirs(im_save_path, exist_ok=True)
+
+            im_size = img_size[0][k].item(), img_size[1][k].item()
+            depth = test_pred[k]
+            depth = (depth - depth.min()) / (depth.max() - depth.min())
+            depth = torch.nn.functional.interpolate(depth.unsqueeze(0).unsqueeze(1), size=(im_size[1], im_size[0]), mode='bilinear', align_corners=True)
+            depth_im = Image.fromarray(255 * depth[0, 0].detach().cpu().numpy()).convert('L')
+            depth_im.save(os.path.join(im_save_path, img_path_split[-1].replace(f'.{ps}', '.png')))
+
+

prismer/experts/generate_edge.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import torch
+import os
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+
+from experts.model_bank import load_expert_model
+from experts.edge.generate_dataset import Dataset
+from experts.edge.images import fuse_edge
+import PIL.Image as Image
+from accelerate import Accelerator
+from tqdm import tqdm
+
+
+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)
+data_loader = torch.utils.data.DataLoader(
+    dataset=dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=4,
+    pin_memory=True
+)
+
+model, data_loader = accelerator.prepare(model, data_loader)
+
+with torch.no_grad():
+    for i, (test_data, img_path, img_size) in enumerate(tqdm(data_loader)):
+        test_pred = model(test_data)
+        fuses = fuse_edge(test_pred)
+        for k in range(len(fuses)):
+            edge = fuses[k]
+            img_path_split = img_path[k].split('/')
+            ps = img_path[k].split('.')[-1]
+            im_save_path = os.path.join(save_path, img_path_split[-3], img_path_split[-2])
+            os.makedirs(im_save_path, exist_ok=True)
+
+            im_size = img_size[0][k].item(), img_size[1][k].item()
+            edge = Image.fromarray(edge).convert('L')
+            edge = edge.resize((im_size[0], im_size[1]), resample=Image.Resampling.BILINEAR)
+            edge.save(os.path.join(im_save_path, img_path_split[-1].replace(f'.{ps}', '.png')))
+
+

prismer/experts/generate_normal.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import torch
+import os
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
+from experts.model_bank import load_expert_model
+from experts.normal.generate_dataset import CustomDataset
+import PIL.Image as Image
+from accelerate import Accelerator
+from tqdm import tqdm
+import numpy as np
+
+
+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)
+data_loader = torch.utils.data.DataLoader(
+    dataset=dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=4,
+    pin_memory=True
+)
+
+model, data_loader = accelerator.prepare(model, data_loader)
+
+with torch.no_grad():
+    for i, (test_data, img_path, img_size) in enumerate(tqdm(data_loader)):
+        test_pred = model(test_data)
+        pred_norm = test_pred[0][-1][:, :3]
+        for k in range(len(pred_norm)):
+            img_path_split = img_path[k].split('/')
+            ps = img_path[k].split('.')[-1]
+            im_save_path = os.path.join(save_path, img_path_split[-3], img_path_split[-2])
+            os.makedirs(im_save_path, exist_ok=True)
+
+            im_size = img_size[0][k].item(), img_size[1][k].item()
+            norm = pred_norm[k]
+            norm = ((norm + 1) * 0.5).clip(0, 1)
+            norm = torch.nn.functional.interpolate(norm.unsqueeze(0), size=(im_size[1], im_size[0]), mode='bilinear', align_corners=True)
+            norm_im = Image.fromarray((norm[0] * 255).permute(1, 2, 0).detach().cpu().numpy().astype(np.uint8)).convert('RGB')
+            norm_im.save(os.path.join(im_save_path, img_path_split[-1].replace(f'.{ps}', '.png')))
+
+

prismer/experts/generate_objdet.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import torch
+import os
+import json
+import copy
+import PIL.Image as Image
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
+from experts.model_bank import load_expert_model
+from experts.obj_detection.generate_dataset import Dataset, collate_fn
+from accelerate import Accelerator
+from tqdm import tqdm
+
+model, transform = load_expert_model(task='obj_detection')
+accelerator = Accelerator(mixed_precision='fp16')
+
+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)
+data_loader = torch.utils.data.DataLoader(
+    dataset=dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=4,
+    pin_memory=True,
+    collate_fn=collate_fn,
+)
+
+model, data_loader = accelerator.prepare(model, data_loader)
+
+
+def get_mask_labels(depth, instance_boxes, instance_id):
+    obj_masks = []
+    obj_ids = []
+    for i in range(len(instance_boxes)):
+        is_duplicate = False
+        mask = torch.zeros_like(depth)
+        x1, y1, x2, y2 = instance_boxes[i][0].item(), instance_boxes[i][1].item(), \
+                         instance_boxes[i][2].item(), instance_boxes[i][3].item()
+        mask[int(y1):int(y2), int(x1):int(x2)] = 1
+        for j in range(len(obj_masks)):
+            if ((mask + obj_masks[j]) == 2).sum() / ((mask + obj_masks[j]) > 0).sum() > 0.95:
+                is_duplicate = True
+                break
+        if not is_duplicate:
+            obj_masks.append(mask)
+            obj_ids.append(instance_id[i])
+
+    obj_masked_modified = copy.deepcopy(obj_masks[:])
+    for i in range(len(obj_masks) - 1):
+        mask1 = obj_masks[i]
+        mask1_ = obj_masked_modified[i]
+        for j in range(i + 1, len(obj_masks)):
+            mask2 = obj_masks[j]
+            mask2_ = obj_masked_modified[j]
+            # case 1: if they don't intersect we don't touch them
+            if ((mask1 + mask2) == 2).sum() == 0:
+                continue
+            # case 2: the entire object 1 is inside of object 2, we say object 1 is in front of object 2:
+            elif (((mask1 + mask2) == 2).float() - mask1).sum() == 0:
+                mask2_ -= mask1_
+            # case 3: the entire object 2 is inside of object 1, we say object 2 is in front of object 1:
+            elif (((mask1 + mask2) == 2).float() - mask2).sum() == 0:
+                mask1_ -= mask2_
+            # case 4: use depth to check object order:
+            else:
+                # object 1 is closer
+                if (depth * mask1).sum() / mask1.sum() > (depth * mask2).sum() / mask2.sum():
+                    mask2_ -= ((mask1 + mask2) == 2).float()
+                # object 2 is closer
+                if (depth * mask1).sum() / mask1.sum() < (depth * mask2).sum() / mask2.sum():
+                    mask1_ -= ((mask1 + mask2) == 2).float()
+
+    final_mask = torch.ones_like(depth) * 255
+    instance_labels = {}
+    for i in range(len(obj_masked_modified)):
+        final_mask = final_mask.masked_fill(obj_masked_modified[i] > 0, i)
+        instance_labels[i] = obj_ids[i].item()
+    return final_mask, instance_labels
+
+
+with torch.no_grad():
+    for i, test_data in enumerate(tqdm(data_loader)):
+        test_pred = model(test_data)
+        for k in range(len(test_pred)):
+            instance_boxes = test_pred[k]['instances'].get_fields()['pred_boxes'].tensor
+            instance_id = test_pred[k]['instances'].get_fields()['pred_classes']
+            depth = test_data[k]['depth']
+
+            final_mask, instance_labels = get_mask_labels(depth, instance_boxes, instance_id)
+
+            img_path_split = test_data[k]['image_path'].split('/')
+            im_save_path = os.path.join(save_path, 'obj_detection', img_path_split[-3], img_path_split[-2])
+            ps = test_data[k]['image_path'].split('.')[-1]
+            os.makedirs(im_save_path, exist_ok=True)
+
+            height, width = test_data[k]['true_height'], test_data[k]['true_width']
+            final_mask = Image.fromarray(final_mask.cpu().numpy()).convert('L')
+            final_mask = final_mask.resize((height, width), resample=Image.Resampling.NEAREST)
+            final_mask.save(os.path.join(im_save_path, img_path_split[-1].replace(f'.{ps}', '.png')))
+
+            with open(os.path.join(im_save_path, img_path_split[-1].replace(f'.{ps}', '.json')), 'w') as fp:
+                json.dump(instance_labels, fp)

prismer/experts/generate_ocrdet.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import torch
+import os
+import PIL.Image as Image
+import numpy as np
+import cv2
+import clip
+import pickle as pk
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+
+from experts.model_bank import load_expert_model
+from experts.ocr_detection.generate_dataset import Dataset
+from accelerate import Accelerator
+from tqdm import tqdm
+
+
+model, transform = load_expert_model(task='ocr_detection')
+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)
+data_loader = torch.utils.data.DataLoader(
+    dataset=dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=4,
+    pin_memory=True,
+)
+
+clip_model, _ = clip.load("ViT-L/14", device=accelerator.device)
+model, data_loader = accelerator.prepare(model, data_loader)
+
+
+def get_label(w, h, word_instances):
+    word_lists = []
+    final_mask = np.ones([h, w],  dtype=np.uint8) * 255
+    counter = 0
+    for word_instance in word_instances[::-1]:
+        mask = np.zeros([h ,w])
+        mask = cv2.fillPoly(mask, [np.int32(word_instance.word_bbox.reshape(-1, 2))], 1)
+        text = word_instance.text.lower()
+        final_mask[mask > 0] = counter
+        word_lists.append(text)
+        counter += 1
+    return final_mask, word_lists
+
+
+with torch.no_grad():
+    for i, (test_data, image_path, scale_w, scale_h, original_w, original_h) in enumerate(tqdm(data_loader)):
+        word_instance_lists = model(test_data, scale_w, scale_h, original_w, original_h)
+        for k in range(len(word_instance_lists)):
+            word_instance = word_instance_lists[k]
+            if len(word_instance) == 0:
+                continue
+            else:
+                final_mask, word_lists = get_label(original_w[k], original_h[k], word_instance)
+
+            final_mask = Image.fromarray(final_mask)
+            img_path_split = image_path[k].split('/')
+            ps = image_path[k].split('.')[-1]
+            im_save_path = os.path.join(save_path, img_path_split[-3], img_path_split[-2])
+            os.makedirs(im_save_path, exist_ok=True)
+
+            final_mask.save(os.path.join(im_save_path, img_path_split[-1].replace(f'.{ps}', '.png')))
+
+            if len(word_lists) > 0:
+                word_embed = clip.tokenize(word_lists).to(accelerator.device)
+                word_features = pca_clip.transform(clip_model.encode_text(word_embed).float().cpu())
+                word_lists = {j: {'features': torch.from_numpy(word_features[j]).float(),
+                                  'text': word_lists[j]} for j in range(len(word_lists))}
+                torch.save(word_lists, os.path.join(im_save_path, img_path_split[-1].replace(f'.{ps}', '.pt')))
+
+

prismer/experts/generate_segmentation.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import torch
+import os
+import PIL.Image as Image
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
+from experts.model_bank import load_expert_model
+from experts.segmentation.generate_dataset import Dataset, collate_fn
+from accelerate import Accelerator
+from tqdm import tqdm
+
+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)
+data_loader = torch.utils.data.DataLoader(
+    dataset=dataset,
+    batch_size=batch_size,
+    shuffle=False,
+    num_workers=4,
+    pin_memory=True,
+    collate_fn=collate_fn,
+)
+
+
+model, data_loader = accelerator.prepare(model, data_loader)
+
+with torch.no_grad():
+    for i, test_data in enumerate(tqdm(data_loader)):
+        test_pred = model(test_data)
+
+        for k in range(len(test_pred)):
+            pred = test_pred[k]['sem_seg']
+            labels = torch.argmax(pred, dim=0)
+
+            img_path_split = test_data[k]['image_path'].split('/')
+            ps = test_data[k]['image_path'].split('.')[-1]
+            im_save_path = os.path.join(save_path, img_path_split[-3], img_path_split[-2])
+            os.makedirs(im_save_path, exist_ok=True)
+
+            seg = Image.fromarray(labels.float().detach().cpu().numpy()).convert('L')
+            seg.save(os.path.join(im_save_path, img_path_split[-1].replace(f'.{ps}', '.png')))
+

prismer/experts/model_bank.py

@@ -0,0 +1,139 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import torch
+import torchvision.transforms as transforms
+
+
+def load_expert_model(task=None):
+    if task == 'depth':
+        # DPT model is a standard pytorch model class
+        from experts.depth.models import DPTDepthModel
+
+        model = DPTDepthModel(path='experts/expert_weights/dpt_hybrid-midas-501f0c75.pt',
+                              backbone="vitb_rn50_384",
+                              non_negative=True,
+                              enable_attention_hooks=False)
+        transform = transforms.Compose([
+            transforms.Resize([480, 480]),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=0.5, std=0.5)]
+        )
+
+    elif task == 'seg_coco':
+        # Mask2Former is wrapped in detection2,
+        # the model takes input in the format of: {"image": image (BGR), "height": height, "width": width}
+        import argparse
+        from detectron2.engine.defaults import DefaultPredictor
+        from experts.segmentation.utils import setup_cfg
+
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--mode", default="client")
+        parser.add_argument("--port", default=2)
+        args = parser.parse_args()
+
+        args.config_file = 'experts/segmentation/configs/coco/panoptic-segmentation/swin/maskformer2_swin_large_IN21k_384_bs16_100ep.yaml'
+        args.opts = ['MODEL.WEIGHTS', 'experts/expert_weights/model_final_f07440.pkl']
+        cfg = setup_cfg(args)
+        model = DefaultPredictor(cfg).model
+        transform = transforms.Compose([
+            transforms.Resize(size=479, max_size=480)
+        ])
+
+    elif task == 'seg_ade':
+        # Mask2Former is wrapped in detection2,
+        # the model takes input in the format of: {"image": image (BGR), "height": height, "width": width}
+        import argparse
+        from detectron2.engine.defaults import DefaultPredictor
+        from experts.segmentation.utils import setup_cfg
+
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--mode", default="client")
+        parser.add_argument("--port", default=2)
+        args = parser.parse_args()
+
+        args.config_file = 'experts/segmentation/configs/ade20k/panoptic-segmentation/swin/maskformer2_swin_large_IN21k_384_bs16_160k.yaml'
+        args.opts = ['MODEL.WEIGHTS', 'experts/expert_weights/model_final_e0c58e.pkl']
+        cfg = setup_cfg(args)
+        model = DefaultPredictor(cfg).model
+        transform = transforms.Compose([
+            transforms.Resize(size=479, max_size=480)
+        ])
+
+    elif task == 'obj_detection':
+        # UniDet is wrapped in detection2,
+        # the model takes input in the format of: {"image": image (BGR), "height": height, "width": width}
+        import argparse
+        from detectron2.engine.defaults import DefaultPredictor
+        from experts.obj_detection.utils import setup_cfg
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--mode", default="client")
+        parser.add_argument("--port", default=2)
+        parser.add_argument("--confidence-threshold", type=float, default=0.5)
+        args = parser.parse_args()
+
+        args.config_file = 'experts/obj_detection/configs/Unified_learned_OCIM_RS200_6x+2x.yaml'
+        args.opts = ['MODEL.WEIGHTS', 'experts/expert_weights/Unified_learned_OCIM_RS200_6x+2x.pth']
+
+        cfg = setup_cfg(args)
+        model = DefaultPredictor(cfg).model
+        transform = transforms.Compose([
+            transforms.Resize(size=479, max_size=480)
+        ])
+
+    elif task == 'ocr_detection':
+        from experts.ocr_detection.charnet.modeling.model import CharNet
+        model = CharNet()
+        model.load_state_dict(torch.load('experts/expert_weights/icdar2015_hourglass88.pth'))
+        transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+
+    elif task == 'normal':
+        # NLL-AngMF model is a standard pytorch model class
+        import argparse
+        from experts.normal.models.NNET import NNET
+        from experts.normal.utils import utils
+
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--mode", default="client")
+        parser.add_argument("--port", default=2)
+        parser.add_argument('--architecture', default='BN', type=str, help='{BN, GN}')
+        parser.add_argument("--pretrained", default='scannet', type=str, help="{nyu, scannet}")
+        parser.add_argument('--sampling_ratio', type=float, default=0.4)
+        parser.add_argument('--importance_ratio', type=float, default=0.7)
+        args = parser.parse_args()
+        model = NNET(args)
+        model = utils.load_checkpoint('experts/expert_weights/scannet.pt', model)
+
+        transform = transforms.Compose([
+            transforms.Resize([480, 480]),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+
+    elif task == 'edge':
+        # NLL-AngMF model is a standard pytorch model class
+        from experts.edge.model import DexiNed
+        model = DexiNed()
+        model.load_state_dict(torch.load('experts/expert_weights/10_model.pth', map_location='cpu'))
+        transform = transforms.Compose([
+            transforms.Resize([480, 480]),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[1.0, 1.0, 1.0])
+        ])
+    else:
+        print('Task not supported')
+        model = None
+        transform = None
+
+    model.eval()
+    return model, transform
+
+
+
+

prismer/experts/normal/generate_dataset.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import glob
+
+from torch.utils.data import Dataset
+from PIL import Image
+from PIL import ImageFile
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+class CustomDataset(Dataset):
+    def __init__(self, data_path, transform):
+        self.data_path = 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')]
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        image_path = self.data_list[index]
+        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
+
+

prismer/experts/normal/models/NNET.py

@@ -0,0 +1,22 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from experts.normal.models.submodules.encoder import Encoder
+from experts.normal.models.submodules.decoder import Decoder
+
+
+class NNET(nn.Module):
+    def __init__(self, args):
+        super(NNET, self).__init__()
+        self.encoder = Encoder()
+        self.decoder = Decoder(args)
+
+    def get_1x_lr_params(self):  # lr/10 learning rate
+        return self.encoder.parameters()
+
+    def get_10x_lr_params(self):  # lr learning rate
+        return self.decoder.parameters()
+
+    def forward(self, img, **kwargs):
+        return self.decoder(self.encoder(img), **kwargs)

prismer/experts/normal/models/baseline.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from experts.normal.models.submodules.submodules import UpSampleBN, norm_normalize
+
+
+# This is the baseline encoder-decoder we used in the ablation study
+class NNET(nn.Module):
+    def __init__(self, args=None):
+        super(NNET, self).__init__()
+        self.encoder = Encoder()
+        self.decoder = Decoder(num_classes=4)
+
+    def forward(self, x, **kwargs):
+        out = self.decoder(self.encoder(x), **kwargs)
+
+        # Bilinearly upsample the output to match the input resolution
+        up_out = F.interpolate(out, size=[x.size(2), x.size(3)], mode='bilinear', align_corners=False)
+        
+        # L2-normalize the first three channels / ensure positive value for concentration parameters (kappa)
+        up_out = norm_normalize(up_out) 
+        return up_out
+
+    def get_1x_lr_params(self):  # lr/10 learning rate
+        return self.encoder.parameters()
+
+    def get_10x_lr_params(self):  # lr learning rate
+        modules = [self.decoder]
+        for m in modules:
+            yield from m.parameters()
+
+
+# Encoder
+class Encoder(nn.Module):
+    def __init__(self):
+        super(Encoder, self).__init__()
+
+        basemodel_name = 'tf_efficientnet_b5_ap'
+        basemodel = torch.hub.load('rwightman/gen-efficientnet-pytorch', basemodel_name, pretrained=True)
+
+        # Remove last layer
+        basemodel.global_pool = nn.Identity()
+        basemodel.classifier = nn.Identity()
+
+        self.original_model = basemodel
+
+    def forward(self, x):
+        features = [x]
+        for k, v in self.original_model._modules.items():
+            if (k == 'blocks'):
+                for ki, vi in v._modules.items():
+                    features.append(vi(features[-1]))
+            else:
+                features.append(v(features[-1]))
+        return features
+
+
+# Decoder (no pixel-wise MLP, no uncertainty-guided sampling)
+class Decoder(nn.Module):
+    def __init__(self, num_classes=4):
+        super(Decoder, self).__init__()
+        self.conv2 = nn.Conv2d(2048, 2048, kernel_size=1, stride=1, padding=0)
+        self.up1 = UpSampleBN(skip_input=2048 + 176, output_features=1024)
+        self.up2 = UpSampleBN(skip_input=1024 + 64, output_features=512)
+        self.up3 = UpSampleBN(skip_input=512 + 40, output_features=256)
+        self.up4 = UpSampleBN(skip_input=256 + 24, output_features=128)
+        self.conv3 = nn.Conv2d(128, num_classes, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, features):
+        x_block0, x_block1, x_block2, x_block3, x_block4 = features[4], features[5], features[6], features[8], features[11]
+        x_d0 = self.conv2(x_block4)
+        x_d1 = self.up1(x_d0, x_block3)
+        x_d2 = self.up2(x_d1, x_block2)
+        x_d3 = self.up3(x_d2, x_block1)
+        x_d4 = self.up4(x_d3, x_block0)
+        out = self.conv3(x_d4)
+        return out
+
+
+if __name__ == '__main__':
+    model = Baseline()
+    x = torch.rand(2, 3, 480, 640)
+    out = model(x)
+    print(out.shape)

prismer/experts/normal/models/submodules/decoder.py

@@ -0,0 +1,202 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from experts.normal.models.submodules.submodules import UpSampleBN, UpSampleGN, norm_normalize, sample_points
+
+
+class Decoder(nn.Module):
+    def __init__(self, args):
+        super(Decoder, self).__init__()
+
+        # hyper-parameter for sampling
+        self.sampling_ratio = args.sampling_ratio
+        self.importance_ratio = args.importance_ratio
+
+        # feature-map
+        self.conv2 = nn.Conv2d(2048, 2048, kernel_size=1, stride=1, padding=0)
+        if args.architecture == 'BN':
+            self.up1 = UpSampleBN(skip_input=2048 + 176, output_features=1024)
+            self.up2 = UpSampleBN(skip_input=1024 + 64, output_features=512)
+            self.up3 = UpSampleBN(skip_input=512 + 40, output_features=256)
+            self.up4 = UpSampleBN(skip_input=256 + 24, output_features=128)
+
+        elif args.architecture == 'GN':
+            self.up1 = UpSampleGN(skip_input=2048 + 176, output_features=1024)
+            self.up2 = UpSampleGN(skip_input=1024 + 64, output_features=512)
+            self.up3 = UpSampleGN(skip_input=512 + 40, output_features=256)
+            self.up4 = UpSampleGN(skip_input=256 + 24, output_features=128)
+
+        else:
+            raise Exception('invalid architecture')
+
+        # produces 1/8 res output
+        self.out_conv_res8 = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
+
+        # produces 1/4 res output
+        self.out_conv_res4 = nn.Sequential(
+            nn.Conv1d(512 + 4, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 4, kernel_size=1),
+        )
+
+        # produces 1/2 res output
+        self.out_conv_res2 = nn.Sequential(
+            nn.Conv1d(256 + 4, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 4, kernel_size=1),
+        )
+
+        # produces 1/1 res output
+        self.out_conv_res1 = nn.Sequential(
+            nn.Conv1d(128 + 4, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 128, kernel_size=1), nn.ReLU(),
+            nn.Conv1d(128, 4, kernel_size=1),
+        )
+
+    def forward(self, features, gt_norm_mask=None, mode='test'):
+        x_block0, x_block1, x_block2, x_block3, x_block4 = features[4], features[5], features[6], features[8], features[11]
+
+        # generate feature-map
+
+        x_d0 = self.conv2(x_block4)                     # x_d0 : [2, 2048, 15, 20]      1/32 res
+        x_d1 = self.up1(x_d0, x_block3)                 # x_d1 : [2, 1024, 30, 40]      1/16 res
+        x_d2 = self.up2(x_d1, x_block2)                 # x_d2 : [2, 512, 60, 80]       1/8 res
+        x_d3 = self.up3(x_d2, x_block1)                 # x_d3: [2, 256, 120, 160]      1/4 res
+        x_d4 = self.up4(x_d3, x_block0)                 # x_d4: [2, 128, 240, 320]      1/2 res
+
+        # 1/8 res output
+        out_res8 = self.out_conv_res8(x_d2)             # out_res8: [2, 4, 60, 80]      1/8 res output
+        out_res8 = norm_normalize(out_res8)             # out_res8: [2, 4, 60, 80]      1/8 res output
+
+        ################################################################################################################
+        # out_res4
+        ################################################################################################################
+
+        if mode == 'train':
+            # upsampling ... out_res8: [2, 4, 60, 80] -> out_res8_res4: [2, 4, 120, 160]
+            out_res8_res4 = F.interpolate(out_res8, scale_factor=2, mode='bilinear', align_corners=True)
+            B, _, H, W = out_res8_res4.shape
+
+            # samples: [B, 1, N, 2]
+            point_coords_res4, rows_int, cols_int = sample_points(out_res8_res4.detach(), gt_norm_mask,
+                                                                  sampling_ratio=self.sampling_ratio,
+                                                                  beta=self.importance_ratio)
+
+            # output (needed for evaluation / visualization)
+            out_res4 = out_res8_res4
+
+            # grid_sample feature-map
+            feat_res4 = F.grid_sample(x_d2, point_coords_res4, mode='bilinear', align_corners=True)  # (B, 512, 1, N)
+            init_pred = F.grid_sample(out_res8, point_coords_res4, mode='bilinear', align_corners=True)  # (B, 4, 1, N)
+            feat_res4 = torch.cat([feat_res4, init_pred], dim=1)  # (B, 512+4, 1, N)
+
+            # prediction (needed to compute loss)
+            samples_pred_res4 = self.out_conv_res4(feat_res4[:, :, 0, :])  # (B, 4, N)
+            samples_pred_res4 = norm_normalize(samples_pred_res4)  # (B, 4, N) - normalized
+
+            for i in range(B):
+                out_res4[i, :, rows_int[i, :], cols_int[i, :]] = samples_pred_res4[i, :, :]
+
+        else:
+            # grid_sample feature-map
+            feat_map = F.interpolate(x_d2, scale_factor=2, mode='bilinear', align_corners=True)
+            init_pred = F.interpolate(out_res8, scale_factor=2, mode='bilinear', align_corners=True)
+            feat_map = torch.cat([feat_map, init_pred], dim=1)  # (B, 512+4, H, W)
+            B, _, H, W = feat_map.shape
+
+            # try all pixels
+            out_res4 = self.out_conv_res4(feat_map.view(B, 512 + 4, -1))  # (B, 4, N)
+            out_res4 = norm_normalize(out_res4)  # (B, 4, N) - normalized
+            out_res4 = out_res4.view(B, 4, H, W)
+            samples_pred_res4 = point_coords_res4 = None
+
+        ################################################################################################################
+        # out_res2
+        ################################################################################################################
+
+        if mode == 'train':
+
+            # upsampling ... out_res4: [2, 4, 120, 160] -> out_res4_res2: [2, 4, 240, 320]
+            out_res4_res2 = F.interpolate(out_res4, scale_factor=2, mode='bilinear', align_corners=True)
+            B, _, H, W = out_res4_res2.shape
+
+            # samples: [B, 1, N, 2]
+            point_coords_res2, rows_int, cols_int = sample_points(out_res4_res2.detach(), gt_norm_mask,
+                                                                  sampling_ratio=self.sampling_ratio,
+                                                                  beta=self.importance_ratio)
+
+            # output (needed for evaluation / visualization)
+            out_res2 = out_res4_res2
+
+            # grid_sample feature-map
+            feat_res2 = F.grid_sample(x_d3, point_coords_res2, mode='bilinear', align_corners=True)  # (B, 256, 1, N)
+            init_pred = F.grid_sample(out_res4, point_coords_res2, mode='bilinear', align_corners=True)  # (B, 4, 1, N)
+            feat_res2 = torch.cat([feat_res2, init_pred], dim=1)  # (B, 256+4, 1, N)
+
+            # prediction (needed to compute loss)
+            samples_pred_res2 = self.out_conv_res2(feat_res2[:, :, 0, :])  # (B, 4, N)
+            samples_pred_res2 = norm_normalize(samples_pred_res2)  # (B, 4, N) - normalized
+
+            for i in range(B):
+                out_res2[i, :, rows_int[i, :], cols_int[i, :]] = samples_pred_res2[i, :, :]
+
+        else:
+            # grid_sample feature-map
+            feat_map = F.interpolate(x_d3, scale_factor=2, mode='bilinear', align_corners=True)
+            init_pred = F.interpolate(out_res4, scale_factor=2, mode='bilinear', align_corners=True)
+            feat_map = torch.cat([feat_map, init_pred], dim=1)  # (B, 512+4, H, W)
+            B, _, H, W = feat_map.shape
+
+            out_res2 = self.out_conv_res2(feat_map.view(B, 256 + 4, -1))  # (B, 4, N)
+            out_res2 = norm_normalize(out_res2)  # (B, 4, N) - normalized
+            out_res2 = out_res2.view(B, 4, H, W)
+            samples_pred_res2 = point_coords_res2 = None
+
+        ################################################################################################################
+        # out_res1
+        ################################################################################################################
+
+        if mode == 'train':
+            # upsampling ... out_res4: [2, 4, 120, 160] -> out_res4_res2: [2, 4, 240, 320]
+            out_res2_res1 = F.interpolate(out_res2, scale_factor=2, mode='bilinear', align_corners=True)
+            B, _, H, W = out_res2_res1.shape
+
+            # samples: [B, 1, N, 2]
+            point_coords_res1, rows_int, cols_int = sample_points(out_res2_res1.detach(), gt_norm_mask,
+                                                                  sampling_ratio=self.sampling_ratio,
+                                                                  beta=self.importance_ratio)
+
+            # output (needed for evaluation / visualization)
+            out_res1 = out_res2_res1
+
+            # grid_sample feature-map
+            feat_res1 = F.grid_sample(x_d4, point_coords_res1, mode='bilinear', align_corners=True)  # (B, 128, 1, N)
+            init_pred = F.grid_sample(out_res2, point_coords_res1, mode='bilinear', align_corners=True)  # (B, 4, 1, N)
+            feat_res1 = torch.cat([feat_res1, init_pred], dim=1)  # (B, 128+4, 1, N)
+
+            # prediction (needed to compute loss)
+            samples_pred_res1 = self.out_conv_res1(feat_res1[:, :, 0, :])  # (B, 4, N)
+            samples_pred_res1 = norm_normalize(samples_pred_res1)  # (B, 4, N) - normalized
+
+            for i in range(B):
+                out_res1[i, :, rows_int[i, :], cols_int[i, :]] = samples_pred_res1[i, :, :]
+
+        else:
+            # grid_sample feature-map
+            feat_map = F.interpolate(x_d4, scale_factor=2, mode='bilinear', align_corners=True)
+            init_pred = F.interpolate(out_res2, scale_factor=2, mode='bilinear', align_corners=True)
+            feat_map = torch.cat([feat_map, init_pred], dim=1)  # (B, 512+4, H, W)
+            B, _, H, W = feat_map.shape
+
+            out_res1 = self.out_conv_res1(feat_map.view(B, 128 + 4, -1))  # (B, 4, N)
+            out_res1 = norm_normalize(out_res1)  # (B, 4, N) - normalized
+            out_res1 = out_res1.view(B, 4, H, W)
+            samples_pred_res1 = point_coords_res1 = None
+
+        return [out_res8, out_res4, out_res2, out_res1], \
+               [out_res8, samples_pred_res4, samples_pred_res2, samples_pred_res1], \
+               [None, point_coords_res4, point_coords_res2, point_coords_res1]
+

prismer/experts/normal/models/submodules/encoder.py

@@ -0,0 +1,32 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Encoder(nn.Module):
+    def __init__(self):
+        super(Encoder, self).__init__()
+
+        basemodel_name = 'tf_efficientnet_b5_ap'
+        print('Loading base model ()...'.format(basemodel_name), end='')
+        basemodel = torch.hub.load('rwightman/gen-efficientnet-pytorch', basemodel_name, pretrained=True)
+        print('Done.')
+
+        # Remove last layer
+        print('Removing last two layers (global_pool & classifier).')
+        basemodel.global_pool = nn.Identity()
+        basemodel.classifier = nn.Identity()
+
+        self.original_model = basemodel
+
+    def forward(self, x):
+        features = [x]
+        for k, v in self.original_model._modules.items():
+            if (k == 'blocks'):
+                for ki, vi in v._modules.items():
+                    features.append(vi(features[-1]))
+            else:
+                features.append(v(features[-1]))
+        return features
+
+

prismer/experts/normal/models/submodules/submodules.py

@@ -0,0 +1,140 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+########################################################################################################################
+
+
+# Upsample + BatchNorm
+class UpSampleBN(nn.Module):
+    def __init__(self, skip_input, output_features):
+        super(UpSampleBN, self).__init__()
+
+        self._net = nn.Sequential(nn.Conv2d(skip_input, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.BatchNorm2d(output_features),
+                                  nn.LeakyReLU(),
+                                  nn.Conv2d(output_features, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.BatchNorm2d(output_features),
+                                  nn.LeakyReLU())
+
+    def forward(self, x, concat_with):
+        up_x = F.interpolate(x, size=[concat_with.size(2), concat_with.size(3)], mode='bilinear', align_corners=True)
+        f = torch.cat([up_x, concat_with], dim=1)
+        return self._net(f)
+
+
+# Upsample + GroupNorm + Weight Standardization
+class UpSampleGN(nn.Module):
+    def __init__(self, skip_input, output_features):
+        super(UpSampleGN, self).__init__()
+
+        self._net = nn.Sequential(Conv2d(skip_input, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.GroupNorm(8, output_features),
+                                  nn.LeakyReLU(),
+                                  Conv2d(output_features, output_features, kernel_size=3, stride=1, padding=1),
+                                  nn.GroupNorm(8, output_features),
+                                  nn.LeakyReLU())
+
+    def forward(self, x, concat_with):
+        up_x = F.interpolate(x, size=[concat_with.size(2), concat_with.size(3)], mode='bilinear', align_corners=True)
+        f = torch.cat([up_x, concat_with], dim=1)
+        return self._net(f)
+
+
+# Conv2d with weight standardization
+class Conv2d(nn.Conv2d):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+                 padding=0, dilation=1, groups=1, bias=True):
+        super(Conv2d, self).__init__(in_channels, out_channels, kernel_size, stride,
+                 padding, dilation, groups, bias)
+
+    def forward(self, x):
+        weight = self.weight
+        weight_mean = weight.mean(dim=1, keepdim=True).mean(dim=2,
+                                  keepdim=True).mean(dim=3, keepdim=True)
+        weight = weight - weight_mean
+        std = weight.view(weight.size(0), -1).std(dim=1).view(-1, 1, 1, 1) + 1e-5
+        weight = weight / std.expand_as(weight)
+        return F.conv2d(x, weight, self.bias, self.stride,
+                        self.padding, self.dilation, self.groups)
+
+
+# normalize
+def norm_normalize(norm_out):
+    min_kappa = 0.01
+    norm_x, norm_y, norm_z, kappa = torch.split(norm_out, 1, dim=1)
+    norm = torch.sqrt(norm_x ** 2.0 + norm_y ** 2.0 + norm_z ** 2.0) + 1e-10
+    kappa = F.elu(kappa) + 1.0 + min_kappa
+    final_out = torch.cat([norm_x / norm, norm_y / norm, norm_z / norm, kappa], dim=1)
+    return final_out
+
+
+# uncertainty-guided sampling (only used during training)
+@torch.no_grad()
+def sample_points(init_normal, gt_norm_mask, sampling_ratio, beta):
+    device = init_normal.device
+    B, _, H, W = init_normal.shape
+    N = int(sampling_ratio * H * W)
+    beta = beta
+
+    # uncertainty map
+    uncertainty_map = -1 * init_normal[:, 3, :, :]  # B, H, W
+
+    # gt_invalid_mask (B, H, W)
+    if gt_norm_mask is not None:
+        gt_invalid_mask = F.interpolate(gt_norm_mask.float(), size=[H, W], mode='nearest')
+        gt_invalid_mask = gt_invalid_mask[:, 0, :, :] < 0.5
+        uncertainty_map[gt_invalid_mask] = -1e4
+
+    # (B, H*W)
+    _, idx = uncertainty_map.view(B, -1).sort(1, descending=True)
+
+    # importance sampling
+    if int(beta * N) > 0:
+        importance = idx[:, :int(beta * N)]    # B, beta*N
+
+        # remaining
+        remaining = idx[:, int(beta * N):]     # B, H*W - beta*N
+
+        # coverage
+        num_coverage = N - int(beta * N)
+
+        if num_coverage <= 0:
+            samples = importance
+        else:
+            coverage_list = []
+            for i in range(B):
+                idx_c = torch.randperm(remaining.size()[1])    # shuffles "H*W - beta*N"
+                coverage_list.append(remaining[i, :][idx_c[:num_coverage]].view(1, -1))     # 1, N-beta*N
+            coverage = torch.cat(coverage_list, dim=0)                                      # B, N-beta*N
+            samples = torch.cat((importance, coverage), dim=1)                              # B, N
+
+    else:
+        # remaining
+        remaining = idx[:, :]  # B, H*W
+
+        # coverage
+        num_coverage = N
+
+        coverage_list = []
+        for i in range(B):
+            idx_c = torch.randperm(remaining.size()[1])  # shuffles "H*W - beta*N"
+            coverage_list.append(remaining[i, :][idx_c[:num_coverage]].view(1, -1))  # 1, N-beta*N
+        coverage = torch.cat(coverage_list, dim=0)  # B, N-beta*N
+        samples = coverage
+
+    # point coordinates
+    rows_int = samples // W         # 0 for first row, H-1 for last row
+    rows_float = rows_int / float(H-1)         # 0 to 1.0
+    rows_float = (rows_float * 2.0) - 1.0       # -1.0 to 1.0
+
+    cols_int = samples % W          # 0 for first column, W-1 for last column
+    cols_float = cols_int / float(W-1)         # 0 to 1.0
+    cols_float = (cols_float * 2.0) - 1.0       # -1.0 to 1.0
+
+    point_coords = torch.zeros(B, 1, N, 2)
+    point_coords[:, 0, :, 0] = cols_float             # x coord
+    point_coords[:, 0, :, 1] = rows_float             # y coord
+    point_coords = point_coords.to(device)
+    return point_coords, rows_int, cols_int

prismer/experts/normal/utils/losses.py

@@ -0,0 +1,178 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+
+
+# compute loss
+class compute_loss(nn.Module):
+    def __init__(self, args):
+        """args.loss_fn can be one of following:
+            - L1            - L1 loss (no uncertainty)
+            - L2            - L2 loss (no uncertainty)
+            - AL            - Angular loss (no uncertainty)
+            - NLL_vMF       - NLL of vonMF distribution
+            - NLL_ours      - NLL of Angular vonMF distribution
+            - UG_NLL_vMF    - NLL of vonMF distribution (+ pixel-wise MLP + uncertainty-guided sampling)
+            - UG_NLL_ours   - NLL of Angular vonMF distribution (+ pixel-wise MLP + uncertainty-guided sampling)
+        """
+        super(compute_loss, self).__init__()
+        self.loss_type = args.loss_fn
+        if self.loss_type in ['L1', 'L2', 'AL', 'NLL_vMF', 'NLL_ours']:
+            self.loss_fn = self.forward_R
+        elif self.loss_type in ['UG_NLL_vMF', 'UG_NLL_ours']:
+            self.loss_fn = self.forward_UG
+        else:
+            raise Exception('invalid loss type')
+
+    def forward(self, *args):
+        return self.loss_fn(*args)
+
+    def forward_R(self, norm_out, gt_norm, gt_norm_mask):
+        pred_norm, pred_kappa = norm_out[:, 0:3, :, :], norm_out[:, 3:, :, :]
+
+        if self.loss_type == 'L1':
+            l1 = torch.sum(torch.abs(gt_norm - pred_norm), dim=1, keepdim=True)
+            loss = torch.mean(l1[gt_norm_mask])
+
+        elif self.loss_type == 'L2':
+            l2 = torch.sum(torch.square(gt_norm - pred_norm), dim=1, keepdim=True)
+            loss = torch.mean(l2[gt_norm_mask])
+
+        elif self.loss_type == 'AL':
+            dot = torch.cosine_similarity(pred_norm, gt_norm, dim=1)
+
+            valid_mask = gt_norm_mask[:, 0, :, :].float() \
+                         * (dot.detach() < 0.999).float() \
+                         * (dot.detach() > -0.999).float()
+            valid_mask = valid_mask > 0.0
+
+            al = torch.acos(dot[valid_mask])
+            loss = torch.mean(al)
+
+        elif self.loss_type == 'NLL_vMF':
+            dot = torch.cosine_similarity(pred_norm, gt_norm, dim=1)
+
+            valid_mask = gt_norm_mask[:, 0, :, :].float() \
+                         * (dot.detach() < 0.999).float() \
+                         * (dot.detach() > -0.999).float()
+            valid_mask = valid_mask > 0.0
+
+            dot = dot[valid_mask]
+            kappa = pred_kappa[:, 0, :, :][valid_mask]
+
+            loss_pixelwise = - torch.log(kappa) \
+                             - (kappa * (dot - 1)) \
+                             + torch.log(1 - torch.exp(- 2 * kappa))
+            loss = torch.mean(loss_pixelwise)
+
+        elif self.loss_type == 'NLL_ours':
+            dot = torch.cosine_similarity(pred_norm, gt_norm, dim=1)
+
+            valid_mask = gt_norm_mask[:, 0, :, :].float() \
+                         * (dot.detach() < 0.999).float() \
+                         * (dot.detach() > -0.999).float()
+            valid_mask = valid_mask > 0.0
+
+            dot = dot[valid_mask]
+            kappa = pred_kappa[:, 0, :, :][valid_mask]
+
+            loss_pixelwise = - torch.log(torch.square(kappa) + 1) \
+                             + kappa * torch.acos(dot) \
+                             + torch.log(1 + torch.exp(-kappa * np.pi))
+            loss = torch.mean(loss_pixelwise)
+
+        else:
+            raise Exception('invalid loss type')
+
+        return loss
+
+
+    def forward_UG(self, pred_list, coord_list, gt_norm, gt_norm_mask):
+        loss = 0.0
+        for (pred, coord) in zip(pred_list, coord_list):
+            if coord is None:
+                pred = F.interpolate(pred, size=[gt_norm.size(2), gt_norm.size(3)], mode='bilinear', align_corners=True)
+                pred_norm, pred_kappa = pred[:, 0:3, :, :], pred[:, 3:, :, :]
+
+                if self.loss_type == 'UG_NLL_vMF':
+                    dot = torch.cosine_similarity(pred_norm, gt_norm, dim=1)
+
+                    valid_mask = gt_norm_mask[:, 0, :, :].float() \
+                                * (dot.detach() < 0.999).float() \
+                                * (dot.detach() > -0.999).float()
+                    valid_mask = valid_mask > 0.5
+
+                    # mask
+                    dot = dot[valid_mask]
+                    kappa = pred_kappa[:, 0, :, :][valid_mask]
+
+                    loss_pixelwise = - torch.log(kappa) \
+                                     - (kappa * (dot - 1)) \
+                                     + torch.log(1 - torch.exp(- 2 * kappa))
+                    loss = loss + torch.mean(loss_pixelwise)
+
+                elif self.loss_type == 'UG_NLL_ours':
+                    dot = torch.cosine_similarity(pred_norm, gt_norm, dim=1)
+
+                    valid_mask = gt_norm_mask[:, 0, :, :].float() \
+                                * (dot.detach() < 0.999).float() \
+                                * (dot.detach() > -0.999).float()
+                    valid_mask = valid_mask > 0.5
+
+                    dot = dot[valid_mask]
+                    kappa = pred_kappa[:, 0, :, :][valid_mask]
+
+                    loss_pixelwise = - torch.log(torch.square(kappa) + 1) \
+                                     + kappa * torch.acos(dot) \
+                                     + torch.log(1 + torch.exp(-kappa * np.pi))
+                    loss = loss + torch.mean(loss_pixelwise)
+
+                else:
+                    raise Exception
+
+            else:
+                # coord: B, 1, N, 2
+                # pred: B, 4, N
+                gt_norm_ = F.grid_sample(gt_norm, coord, mode='nearest', align_corners=True)  # (B, 3, 1, N)
+                gt_norm_mask_ = F.grid_sample(gt_norm_mask.float(), coord, mode='nearest', align_corners=True)  # (B, 1, 1, N)
+                gt_norm_ = gt_norm_[:, :, 0, :]  # (B, 3, N)
+                gt_norm_mask_ = gt_norm_mask_[:, :, 0, :] > 0.5  # (B, 1, N)
+
+                pred_norm, pred_kappa = pred[:, 0:3, :], pred[:, 3:, :]
+
+                if self.loss_type == 'UG_NLL_vMF':
+                    dot = torch.cosine_similarity(pred_norm, gt_norm_, dim=1)  # (B, N)
+
+                    valid_mask = gt_norm_mask_[:, 0, :].float() \
+                                 * (dot.detach() < 0.999).float() \
+                                 * (dot.detach() > -0.999).float()
+                    valid_mask = valid_mask > 0.5
+
+                    dot = dot[valid_mask]
+                    kappa = pred_kappa[:, 0, :][valid_mask]
+
+                    loss_pixelwise = - torch.log(kappa) \
+                                     - (kappa * (dot - 1)) \
+                                     + torch.log(1 - torch.exp(- 2 * kappa))
+                    loss = loss + torch.mean(loss_pixelwise)
+
+                elif self.loss_type == 'UG_NLL_ours':
+                    dot = torch.cosine_similarity(pred_norm, gt_norm_, dim=1)  # (B, N)
+
+                    valid_mask = gt_norm_mask_[:, 0, :].float() \
+                                 * (dot.detach() < 0.999).float() \
+                                 * (dot.detach() > -0.999).float()
+                    valid_mask = valid_mask > 0.5
+
+                    dot = dot[valid_mask]
+                    kappa = pred_kappa[:, 0, :][valid_mask]
+
+                    loss_pixelwise = - torch.log(torch.square(kappa) + 1) \
+                                     + kappa * torch.acos(dot) \
+                                     + torch.log(1 + torch.exp(-kappa * np.pi))
+                    loss = loss + torch.mean(loss_pixelwise)
+
+                else:
+                    raise Exception
+        return loss

prismer/experts/normal/utils/utils.py

@@ -0,0 +1,191 @@
+import os
+import numpy as np
+from PIL import Image
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
+
+
+# convert arg line to args
+def convert_arg_line_to_args(arg_line):
+    for arg in arg_line.split():
+        if not arg.strip():
+            continue
+        yield str(arg)
+
+
+# save args
+def save_args(args, filename):
+    with open(filename, 'w') as f:
+        for arg in vars(args):
+            f.write('{}: {}\n'.format(arg, getattr(args, arg)))
+
+
+# concatenate images
+def concat_image(image_path_list, concat_image_path):
+    imgs = [Image.open(i).convert("RGB").resize((640, 480), resample=Image.BILINEAR) for i in image_path_list]
+    imgs_list = []
+    for i in range(len(imgs)):
+        img = imgs[i]
+        imgs_list.append(np.asarray(img))
+
+        H, W, _ = np.asarray(img).shape
+        imgs_list.append(255 * np.ones((H, 20, 3)).astype('uint8'))
+
+    imgs_comb = np.hstack(imgs_list[:-1])
+    imgs_comb = Image.fromarray(imgs_comb)
+    imgs_comb.save(concat_image_path)
+
+
+# load model
+def load_checkpoint(fpath, model):
+    ckpt = torch.load(fpath, map_location='cpu')['model']
+
+    load_dict = {}
+    for k, v in ckpt.items():
+        if k.startswith('module.'):
+            k_ = k.replace('module.', '')
+            load_dict[k_] = v
+        else:
+            load_dict[k] = v
+
+    model.load_state_dict(load_dict)
+    return model
+
+
+# compute normal errors
+def compute_normal_errors(total_normal_errors):
+    metrics = {
+        'mean': np.average(total_normal_errors),
+        'median': np.median(total_normal_errors),
+        'rmse': np.sqrt(np.sum(total_normal_errors * total_normal_errors) / total_normal_errors.shape),
+        'a1': 100.0 * (np.sum(total_normal_errors < 5) / total_normal_errors.shape[0]),
+        'a2': 100.0 * (np.sum(total_normal_errors < 7.5) / total_normal_errors.shape[0]),
+        'a3': 100.0 * (np.sum(total_normal_errors < 11.25) / total_normal_errors.shape[0]),
+        'a4': 100.0 * (np.sum(total_normal_errors < 22.5) / total_normal_errors.shape[0]),
+        'a5': 100.0 * (np.sum(total_normal_errors < 30) / total_normal_errors.shape[0])
+    }
+    return metrics
+
+
+# log normal errors
+def log_normal_errors(metrics, where_to_write, first_line):
+    print(first_line)
+    print("mean median rmse 5 7.5 11.25 22.5 30")
+    print("%.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f" % (
+        metrics['mean'], metrics['median'], metrics['rmse'],
+        metrics['a1'], metrics['a2'], metrics['a3'], metrics['a4'], metrics['a5']))
+
+    with open(where_to_write, 'a') as f:
+        f.write('%s\n' % first_line)
+        f.write("mean median rmse 5 7.5 11.25 22.5 30\n")
+        f.write("%.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f\n\n" % (
+            metrics['mean'], metrics['median'], metrics['rmse'],
+            metrics['a1'], metrics['a2'], metrics['a3'], metrics['a4'], metrics['a5']))
+
+
+# makedir
+def makedir(dirpath):
+    if not os.path.exists(dirpath):
+        os.makedirs(dirpath)
+
+
+# makedir from list
+def make_dir_from_list(dirpath_list):
+    for dirpath in dirpath_list:
+        makedir(dirpath)
+
+
+
+########################################################################################################################
+# Visualization
+########################################################################################################################
+
+
+# unnormalize image
+__imagenet_stats = {'mean': [0.485, 0.456, 0.406], 'std': [0.229, 0.224, 0.225]}
+def unnormalize(img_in):
+    img_out = np.zeros(img_in.shape)
+    for ich in range(3):
+        img_out[:, :, ich] = img_in[:, :, ich] * __imagenet_stats['std'][ich]
+        img_out[:, :, ich] += __imagenet_stats['mean'][ich]
+    img_out = (img_out * 255).astype(np.uint8)
+    return img_out
+
+
+# kappa to exp error (only applicable to AngMF distribution)
+def kappa_to_alpha(pred_kappa):
+    alpha = ((2 * pred_kappa) / ((pred_kappa ** 2.0) + 1)) \
+            + ((np.exp(- pred_kappa * np.pi) * np.pi) / (1 + np.exp(- pred_kappa * np.pi)))
+    alpha = np.degrees(alpha)
+    return alpha
+
+
+# normal vector to rgb values
+def norm_to_rgb(norm):
+    # norm: (B, H, W, 3)
+    norm_rgb = ((norm[0, ...] + 1) * 0.5) * 255
+    norm_rgb = np.clip(norm_rgb, a_min=0, a_max=255)
+    norm_rgb = norm_rgb.astype(np.uint8)
+    return norm_rgb
+
+
+# visualize during training
+def visualize(args, img, gt_norm, gt_norm_mask, norm_out_list, total_iter):
+    B, _, H, W = gt_norm.shape
+
+    pred_norm_list = []
+    pred_kappa_list = []
+    for norm_out in norm_out_list:
+        norm_out = F.interpolate(norm_out, size=[gt_norm.size(2), gt_norm.size(3)], mode='nearest')
+        pred_norm = norm_out[:, :3, :, :]  # (B, 3, H, W)
+        pred_norm = pred_norm.detach().cpu().permute(0, 2, 3, 1).numpy()  # (B, H, W, 3)
+        pred_norm_list.append(pred_norm)
+
+        pred_kappa = norm_out[:, 3:, :, :]  # (B, 1, H, W)
+        pred_kappa = pred_kappa.detach().cpu().permute(0, 2, 3, 1).numpy()  # (B, H, W, 1)
+        pred_kappa_list.append(pred_kappa)
+
+    # to numpy arrays
+    img = img.detach().cpu().permute(0, 2, 3, 1).numpy()                     # (B, H, W, 3)
+    gt_norm = gt_norm.detach().cpu().permute(0, 2, 3, 1).numpy()             # (B, H, W, 3)
+    gt_norm_mask = gt_norm_mask.detach().cpu().permute(0, 2, 3, 1).numpy()   # (B, H, W, 1)
+
+    # input image
+    target_path = '%s/%08d_img.jpg' % (args.exp_vis_dir, total_iter)
+    img = unnormalize(img[0, ...])
+    plt.imsave(target_path, img)
+
+    # gt norm
+    gt_norm_rgb = ((gt_norm[0, ...] + 1) * 0.5) * 255
+    gt_norm_rgb = np.clip(gt_norm_rgb, a_min=0, a_max=255)
+    gt_norm_rgb = gt_norm_rgb.astype(np.uint8)
+
+    target_path = '%s/%08d_gt_norm.jpg' % (args.exp_vis_dir, total_iter)
+    plt.imsave(target_path, gt_norm_rgb * gt_norm_mask[0, ...])
+
+    # pred_norm
+    for i in range(len(pred_norm_list)):
+        pred_norm = pred_norm_list[i]
+        pred_norm_rgb = norm_to_rgb(pred_norm)
+        target_path = '%s/%08d_pred_norm_%d.jpg' % (args.exp_vis_dir, total_iter, i)
+        plt.imsave(target_path, pred_norm_rgb)
+
+        pred_kappa = pred_kappa_list[i]
+        pred_alpha = kappa_to_alpha(pred_kappa)
+        target_path = '%s/%08d_pred_alpha_%d.jpg' % (args.exp_vis_dir, total_iter, i)
+        plt.imsave(target_path, pred_alpha[0, :, :, 0], vmin=0, vmax=60, cmap='jet')
+
+        # error in angles
+        DP = np.sum(gt_norm * pred_norm, axis=3, keepdims=True)  # (B, H, W, 1)
+        DP = np.clip(DP, -1, 1)
+        E = np.degrees(np.arccos(DP))  # (B, H, W, 1)
+        E = E * gt_norm_mask
+        target_path = '%s/%08d_pred_error_%d.jpg' % (args.exp_vis_dir, total_iter, i)
+        plt.imsave(target_path, E[0, :, :, 0], vmin=0, vmax=60, cmap='jet')

prismer/experts/obj_detection/configs/Base-CRCNN-COCO.yaml

@@ -0,0 +1,48 @@
+MODEL:
+  META_ARCHITECTURE: "GeneralizedRCNN"
+  BACKBONE:
+    NAME: "build_p67_resnet_fpn_backbone"
+  WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl"
+  RESNETS:
+    OUT_FEATURES: ["res3", "res4", "res5"]
+    DEPTH: 50
+  FPN:
+    IN_FEATURES: ["res3", "res4", "res5"]
+  ANCHOR_GENERATOR:
+    SIZES: !!python/object/apply:eval ["[[x, x * 2**(1.0/3), x * 2**(2.0/3) ] for x in [32, 64, 128, 256, 512 ]]"]
+    ASPECT_RATIOS: [[0.5, 1.0, 2.0]]
+  RPN:
+    IN_FEATURES: ["p3", "p4", "p5", "p6", "p7"]
+    PRE_NMS_TOPK_TRAIN: 2000
+    PRE_NMS_TOPK_TEST: 1000
+    POST_NMS_TOPK_TRAIN: 2000
+    POST_NMS_TOPK_TEST: 1000
+  ROI_HEADS:
+    NUM_CLASSES: 80
+    NAME: CustomCascadeROIHeads
+    IN_FEATURES: ["p3", "p4", "p5"]
+    SCORE_THRESH_TEST: 0.0001
+  ROI_BOX_HEAD:
+    NAME: "FastRCNNConvFCHead"
+    NUM_FC: 2
+    POOLER_RESOLUTION: 7
+    CLS_AGNOSTIC_BBOX_REG: True
+DATASETS:
+  TRAIN: ("coco_2017_train",)
+  TEST: ("coco_2017_val",)
+TEST:
+  DETECTIONS_PER_IMAGE: 300
+SOLVER:
+  IMS_PER_BATCH: 16
+  BASE_LR: 0.01
+  STEPS: (60000, 80000)
+  MAX_ITER: 90000
+  CHECKPOINT_PERIOD: 1000000
+  WARMUP_ITERS: 4000
+  WARMUP_FACTOR: 0.00025
+  CLIP_GRADIENTS:
+    ENABLED: True
+INPUT:
+  MIN_SIZE_TRAIN: (640, 672, 704, 736, 768, 800)
+VERSION: 2
+OUTPUT_DIR: "output/UniDet/auto"

prismer/experts/obj_detection/configs/O365_CRFR50_CAS_2x.yaml

@@ -0,0 +1,15 @@
+_BASE_: "Base-CRCNN-COCO"
+MODEL:
+  ROI_HEADS:
+    NUM_CLASSES: 365
+DATASETS:
+  TRAIN: ("objects365_train",)
+  TEST: ("objects365_val",)
+SOLVER:
+  IMS_PER_BATCH: 16
+  BASE_LR: 0.01
+  STEPS: (120000, 160000,)
+  MAX_ITER: 180000
+  CHECKPOINT_PERIOD: 120000
+DATALOADER:
+  SAMPLER_TRAIN: "ClassAwareSampler"

prismer/experts/obj_detection/configs/OID_CRFR50_CAS_2x.yaml

@@ -0,0 +1,22 @@
+_BASE_: "Base-CRCNN-COCO"
+MODEL:
+  ROI_HEADS:
+    NUM_CLASSES: 500
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    USE_EQL_LOSS: True
+    EQL_FREQ_CAT: 200
+    EQL_CAT_INFO: 'datasets/oid/annotations/openimages_challenge_2019_train_v2_cat_info.json'
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+DATASETS:
+  TRAIN: ("oid_train",)
+  TEST: ("oid_val_expanded",)
+SOLVER:
+  IMS_PER_BATCH: 16
+  BASE_LR: 0.01
+  STEPS: (120000, 160000,)
+  MAX_ITER: 180000
+  CHECKPOINT_PERIOD: 120000
+DATALOADER:
+  SAMPLER_TRAIN: "ClassAwareSampler"

prismer/experts/obj_detection/configs/Partitioned_COIM_R50_6x+2x.yaml

@@ -0,0 +1,28 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  META_ARCHITECTURE: "SplitClassifierRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: -1
+    NAME: "MultiDatasetCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+DATASETS:
+  TRAIN: ("coco_2017_train","objects365_train","oid_train","mapillary_960_train")
+  TEST: ('coco_2017_val','oid_val_expanded','objects365_val',"mapillary_val")
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['coco', 'objects365', 'oid', 'mapillary']
+  NUM_CLASSES: [80, 365, 500, 37]
+  DATA_RATIO: [1, 1, 1, 1]
+  USE_CAS: [False, True, True, False]
+SOLVER:
+  IMS_PER_BATCH: 16
+  BASE_LR: 0.001
+  STEPS: (160000,)
+  MAX_ITER: 180000

prismer/experts/obj_detection/configs/Partitioned_COI_R50_2x.yaml

@@ -0,0 +1,29 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  META_ARCHITECTURE: "SplitClassifierRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: -1
+    NAME: "MultiDatasetCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+DATASETS:
+  TRAIN: ("coco_2017_train","objects365_train","oid_train",)
+  TEST: ('coco_2017_val','oid_val_expanded','objects365_val',)
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['coco', 'objects365', 'oid']
+  NUM_CLASSES: [80, 365, 500]
+  DATA_RATIO: [1, 1, 1]
+  USE_CAS: [False, True, True]
+SOLVER:
+  IMS_PER_BATCH: 16
+  STEPS: (120000, 160000)
+  MAX_ITER: 180000
+  CHECKPOINT_PERIOD: 120000
+

prismer/experts/obj_detection/configs/Partitioned_COI_R50_6x.yaml

@@ -0,0 +1,28 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  META_ARCHITECTURE: "SplitClassifierRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: -1
+    NAME: "MultiDatasetCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+DATASETS:
+  TRAIN: ("coco_2017_train","objects365_train","oid_train",)
+  TEST: ('coco_2017_val','oid_val_expanded','objects365_val',)
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['coco', 'objects365', 'oid']
+  NUM_CLASSES: [80, 365, 500]
+  DATA_RATIO: [1, 1, 1]
+  USE_CAS: [False, True, True]
+SOLVER:
+  IMS_PER_BATCH: 16
+  STEPS: (480000, 500000)
+  MAX_ITER: 540000
+  CHECKPOINT_PERIOD: 120000

prismer/experts/obj_detection/configs/Partitioned_COI_R50_8x.yaml

@@ -0,0 +1,28 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  META_ARCHITECTURE: "SplitClassifierRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: -1
+    NAME: "MultiDatasetCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+DATASETS:
+  TRAIN: ("coco_2017_train","objects365_train","oid_train",)
+  TEST: ('coco_2017_val','oid_val_expanded','objects365_val',)
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['coco', 'objects365', 'oid']
+  NUM_CLASSES: [80, 365, 500]
+  DATA_RATIO: [1, 1, 1]
+  USE_CAS: [False, True, True]
+SOLVER:
+  IMS_PER_BATCH: 16
+  STEPS: (660000, 700000)
+  MAX_ITER: 720000
+  CHECKPOINT_PERIOD: 120000

prismer/experts/obj_detection/configs/Partitioned_COI_RS101_2x.yaml

@@ -0,0 +1,46 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  BACKBONE:
+    NAME: "build_p67_resnest_fpn_backbone"
+  WEIGHTS: "https://hangzh.s3-us-west-1.amazonaws.com/encoding/models/resnest101_detectron-486f69a8.pth"
+  PIXEL_MEAN: [123.68, 116.779, 103.939]
+  PIXEL_STD: [58.393, 57.12, 57.375]
+  RESNETS:
+    DEPTH: 101
+    STRIDE_IN_1X1: False
+    RADIX: 2
+    NORM: "SyncBN"
+  FPN:
+    NORM: "SyncBN"
+  META_ARCHITECTURE: "SplitClassifierRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: -1
+    NAME: "MultiDatasetCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+    NAME: "FastRCNNConvFCHead"
+    NUM_CONV: 4
+    NUM_FC: 1
+    NORM: "SyncBN"
+INPUT:
+  FORMAT: "RGB"
+DATASETS:
+  TRAIN: ("coco_2017_train","objects365_train","oid_train",)
+  TEST: ('coco_2017_val','oid_val_expanded','objects365_val',)
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['coco', 'objects365', 'oid']
+  NUM_CLASSES: [80, 365, 500]
+  DATA_RATIO: [1, 1, 1]
+  USE_CAS: [False, True, True]
+SOLVER:
+  IMS_PER_BATCH: 16
+  STEPS: (120000, 160000)
+  MAX_ITER: 180000
+  CHECKPOINT_PERIOD: 120000

prismer/experts/obj_detection/configs/Unified_human_OCI_R50_2x.yaml

@@ -0,0 +1,29 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  META_ARCHITECTURE: "UnifiedRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: 659
+    NAME: "UnifiedCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+DATASETS:
+  TRAIN: ("coco_2017_train","objects365_train","oid_train",)
+  TEST: ('coco_2017_val','oid_val_expanded','objects365_val',)
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['coco', 'objects365', 'oid']
+  NUM_CLASSES: [80, 365, 500]
+  DATA_RATIO: [1, 1, 1]
+  USE_CAS: [False, True, True]
+  UNIFIED_LABEL_FILE: 'datasets/label_spaces/manual.json'
+SOLVER:
+  IMS_PER_BATCH: 16
+  STEPS: (120000,160000)
+  MAX_ITER: 180000
+  CHECKPOINT_PERIOD: 120000

prismer/experts/obj_detection/configs/Unified_learned_OCIM_R50_6x+2x.yaml

@@ -0,0 +1,35 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  META_ARCHITECTURE: "UnifiedRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: 722
+    NAME: "UnifiedCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+DATASETS:
+  TRAIN: ("objects365_train","coco_2017_train","oid_train","mapillary_960_train")
+  TEST: ("coco_2017_val", "objects365_val", "oid_val_v2_expanded","mapillary_val")
+  # TEST: ('voc_cocoformat_test','viper_val', 'scannet_val','wilddash_public',
+  #   'kitti_train','crowdhuman_val', 'cityscapes_cocoformat_val',)
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['objects365', 'coco', 'oid', 'mapillary']
+  NUM_CLASSES: [365, 80, 500, 37]
+  DATA_RATIO: [1, 1, 1, 1]
+  USE_CAS: [True, False, True, False]
+  UNIFIED_LABEL_FILE: 'experts/obj_detection/datasets/label_spaces/learned_mAP+M.json'
+  # MATCH_NOVEL_CLASSES_FILE: 'datasets/label_spaces/mAP_val+M_722_4d_labelmap_test.json'
+  # UNIFIED_EVAL: True
+  # UNIFIED_NOVEL_CLASSES_EVAL: True
+  # UNIFY_LABEL_TEST: False
+SOLVER:
+  IMS_PER_BATCH: 16
+  BASE_LR: 0.001
+  STEPS: (160000,)
+  MAX_ITER: 180000

prismer/experts/obj_detection/configs/Unified_learned_OCIM_RS200_6x+2x.yaml

@@ -0,0 +1,46 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  BACKBONE:
+    NAME: "build_p67_resnest_fpn_backbone"
+  PIXEL_MEAN: [123.68, 116.779, 103.939]
+  PIXEL_STD: [58.393, 57.12, 57.375]
+  RESNETS:
+    DEPTH: 200
+    STRIDE_IN_1X1: False
+    RADIX: 2
+    NORM: "SyncBN"
+  FPN:
+    NORM: "SyncBN"
+  META_ARCHITECTURE: "UnifiedRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: 722
+    NAME: "UnifiedCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+    NAME: "FastRCNNConvFCHead"
+    NUM_CONV: 4
+    NUM_FC: 1
+    NORM: "SyncBN"
+INPUT:
+  FORMAT: "RGB"
+DATASETS:
+  TRAIN: ("objects365_train","coco_2017_train","oid_train","mapillary_960_train")
+  TEST: ("coco_2017_val", "objects365_val", "oid_val_v2_expanded","mapillary_val")
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['objects365', 'coco', 'oid', 'mapillary']
+  NUM_CLASSES: [365, 80, 500, 37]
+  DATA_RATIO: [1, 1, 1, 1]
+  USE_CAS: [True, False, True, False]
+  UNIFIED_LABEL_FILE: 'experts/obj_detection/datasets/label_spaces/learned_mAP+M.json'
+SOLVER:
+  IMS_PER_BATCH: 16
+  STEPS: (480000, 500000)
+  MAX_ITER: 540000
+  CHECKPOINT_PERIOD: 120000

prismer/experts/obj_detection/configs/Unified_learned_OCI_R50_2x.yaml

@@ -0,0 +1,29 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  META_ARCHITECTURE: "UnifiedRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: 701
+    NAME: "UnifiedCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+DATASETS:
+  TRAIN: ("coco_2017_train","objects365_train","oid_train",)
+  TEST: ('coco_2017_val','oid_val_expanded','objects365_val',)
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['coco', 'objects365', 'oid']
+  NUM_CLASSES: [80, 365, 500]
+  DATA_RATIO: [1, 1, 1]
+  USE_CAS: [False, True, True]
+  UNIFIED_LABEL_FILE: 'datasets/label_spaces/learned_mAP.json'
+SOLVER:
+  IMS_PER_BATCH: 16
+  STEPS: (120000,160000)
+  MAX_ITER: 180000
+  CHECKPOINT_PERIOD: 120000

prismer/experts/obj_detection/configs/Unified_learned_OCI_R50_6x.yaml

@@ -0,0 +1,29 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  META_ARCHITECTURE: "UnifiedRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: 701
+    NAME: "UnifiedCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+DATASETS:
+  TRAIN: ("coco_2017_train","objects365_train","oid_train",)
+  TEST: ('coco_2017_val','oid_val_expanded','objects365_val',)
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['coco', 'objects365', 'oid']
+  NUM_CLASSES: [80, 365, 500]
+  DATA_RATIO: [1, 1, 1]
+  USE_CAS: [False, True, True]
+  UNIFIED_LABEL_FILE: 'datasets/label_spaces/learned_mAP.json'
+SOLVER:
+  IMS_PER_BATCH: 16
+  STEPS: (480000, 500000)
+  MAX_ITER: 540000
+  CHECKPOINT_PERIOD: 120000

prismer/experts/obj_detection/configs/Unified_learned_OCI_R50_8x.yaml

@@ -0,0 +1,29 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  META_ARCHITECTURE: "UnifiedRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: 722
+    NAME: "UnifiedCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+DATASETS:
+  TRAIN: ("coco_2017_train","objects365_train","oid_train",)
+  TEST: ('coco_2017_val','oid_val_expanded','objects365_val',)
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['coco', 'objects365', 'oid']
+  NUM_CLASSES: [80, 365, 500]
+  DATA_RATIO: [1, 1, 1]
+  USE_CAS: [False, True, True]
+  UNIFIED_LABEL_FILE: 'datasets/label_spaces/learned_mAP+M.json'
+SOLVER:
+  IMS_PER_BATCH: 16
+  STEPS: (660000, 700000)
+  MAX_ITER: 720000
+  CHECKPOINT_PERIOD: 120000

prismer/experts/obj_detection/configs/Unified_learned_OCI_RS200_6x.yaml

@@ -0,0 +1,46 @@
+_BASE_: "Base-CRCNN-COCO.yaml"
+MODEL:
+  BACKBONE:
+    NAME: "build_p67_resnest_fpn_backbone"
+  PIXEL_MEAN: [123.68, 116.779, 103.939]
+  PIXEL_STD: [58.393, 57.12, 57.375]
+  RESNETS:
+    DEPTH: 200
+    STRIDE_IN_1X1: False
+    RADIX: 2
+    NORM: "SyncBN"
+  FPN:
+    NORM: "SyncBN"
+  META_ARCHITECTURE: "UnifiedRCNN"
+  ROI_HEADS:
+    NUM_CLASSES: 701
+    NAME: "UnifiedCascadeROIHeads"
+  ROI_BOX_HEAD:
+    USE_SIGMOID_CE: True
+    # USE_EQL_LOSS: True
+    HIERARCHY_IGNORE: True
+    HIERARCHY_POS_PARENTS: True
+    NAME: "FastRCNNConvFCHead"
+    NUM_CONV: 4
+    NUM_FC: 1
+    NORM: "SyncBN"
+INPUT:
+  FORMAT: "RGB"
+DATASETS:
+  TRAIN: ("coco_2017_train","objects365_train","oid_train",)
+  TEST: ('coco_2017_val','oid_val_expanded','objects365_val',)
+DATALOADER:                                                                                                                                                                                                                                                                                                                                                                 
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+  NUM_WORKERS: 1
+MULTI_DATASET:
+  ENABLED: True
+  DATASETS: ['coco', 'objects365', 'oid']
+  NUM_CLASSES: [80, 365, 500]
+  DATA_RATIO: [1, 1, 1]
+  USE_CAS: [False, True, True]
+  UNIFIED_LABEL_FILE: 'datasets/label_spaces/learned_mAP.json'
+SOLVER:
+  IMS_PER_BATCH: 16
+  STEPS: (480000, 500000)
+  MAX_ITER: 540000
+  CHECKPOINT_PERIOD: 120000