Reset

.gitignore

@@ -1 +1,162 @@
-.idea
+cache/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "prismer"]
+	path = prismer
+	url = https://github.com/nvlabs/prismer

.pre-commit-config.yaml

@@ -0,0 +1,37 @@
+exclude: patch
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.2.0
+  hooks:
+  - id: check-executables-have-shebangs
+  - id: check-json
+  - id: check-merge-conflict
+  - id: check-shebang-scripts-are-executable
+  - id: check-toml
+  - id: check-yaml
+  - id: double-quote-string-fixer
+  - id: end-of-file-fixer
+  - id: mixed-line-ending
+    args: ['--fix=lf']
+  - id: requirements-txt-fixer
+  - id: trailing-whitespace
+- repo: https://github.com/myint/docformatter
+  rev: v1.4
+  hooks:
+  - id: docformatter
+    args: ['--in-place']
+- repo: https://github.com/pycqa/isort
+  rev: 5.12.0
+  hooks:
+    - id: isort
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v0.991
+  hooks:
+    - id: mypy
+      args: ['--ignore-missing-imports']
+      additional_dependencies: ['types-python-slugify']
+- repo: https://github.com/google/yapf
+  rev: v0.32.0
+  hooks:
+  - id: yapf
+    args: ['--parallel', '--in-place']

.style.yapf

@@ -0,0 +1,5 @@
+[style]
+based_on_style = pep8
+blank_line_before_nested_class_or_def = false
+spaces_before_comment = 2
+split_before_logical_operator = true

README.md

@@ -1,11 +1,11 @@
 ---
 title: Prismer
-emoji: 🔻
-colorFrom: yellow
-colorTo: purple
+emoji: 🦀
+colorFrom: blue
+colorTo: green
 sdk: docker
-sdk_version: 3.20.1
 pinned: false
+license: other
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,34 +1,35 @@
+#!/usr/bin/env python
+
 from __future__ import annotations
 
 import os
+import shutil
+import subprocess
+
 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
+if os.getenv('SYSTEM') == 'spaces':
+    with open('patch') as f:
+        subprocess.run('patch -p1'.split(), cwd='prismer', stdin=f)
+    shutil.copytree('prismer/helpers/images',
+                    'prismer/images',
+                    dirs_exist_ok=True)
 
+from app_caption import create_demo as create_demo_caption
+from 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**.
-Please refer to our [project page](https://shikun.io/projects/prismer) or [github](https://github.com/NVlabs/prismer) for more details.
-"""
+DESCRIPTION = '# [Prismer](https://github.com/nvlabs/prismer)'
 
 if (SPACE_ID := os.getenv('SPACE_ID')) is not None:
-    description += f'For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. <a href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a>'
-
+    DESCRIPTION += f'<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. <a href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>'
 
-with gr.Blocks() as demo:
-    gr.Markdown(description)
-    with gr.Tab("Zero-shot Image Captioning"):
-        create_caption()
-    with gr.Tab("Visual Question Answering"):
-        create_vqa()
+with gr.Blocks(css='style.css') as demo:
+    gr.Markdown(DESCRIPTION)
+    with gr.Tabs():
+        with gr.TabItem('Caption'):
+            create_demo_caption()
 
 demo.queue(api_open=False).launch()
-

gradio_caption.py → app_caption.py

@@ -1,17 +1,24 @@
+#!/usr/bin/env python
+
 from __future__ import annotations
 
 import os
 import pathlib
+
 import gradio as gr
 
-from prismer_model import run_experts
+from prismer_model import Model
 
 
 def create_demo():
+    model = Model()
+
     with gr.Row():
         with gr.Column():
             image = gr.Image(label='Input', type='filepath')
-            model_name = gr.Dropdown(label='Model', choices=['prismer_base'], value='prismer_base')
+            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')
@@ -24,21 +31,28 @@ def create_demo():
                 object_detection = gr.Image(label='Object Detection')
                 ocr = gr.Image(label='OCR Detection')
 
-    inputs = [image]
-    outputs = [depth, edge, normals]
+    inputs = [image, model_name]
+    outputs = [
+        caption,
+        depth,
+        edge,
+        normals,
+        segmentation,
+        object_detection,
+        ocr,
+    ]
 
     paths = sorted(pathlib.Path('prismer/images').glob('*'))
     examples = [[path.as_posix(), 'prismer_base'] for path in paths]
     gr.Examples(examples=examples,
                 inputs=inputs,
                 outputs=outputs,
-                fn=run_experts,
+                fn=model.run_caption,
                 cache_examples=os.getenv('SYSTEM') == 'spaces')
 
-    run_button.click(fn=run_experts, inputs=inputs, outputs=outputs)
+    run_button.click(fn=model.run_caption, inputs=inputs, outputs=outputs)
 
 
 if __name__ == '__main__':
     demo = create_demo()
     demo.queue().launch()
-

gradio_vqa.py

@@ -1,33 +0,0 @@
-import gradio as gr
-from PIL import Image
-import tempfile
-
-
-def predict_depth(model, image):
-    depth = model.infer_pil(image)
-    return depth
-
-
-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")
-            ques = gr.Textbox(label="Question", placeholder="What's the number of this player?")
-            rgb = gr.Image(label="Input Image", type='pil', elem_id='img-display-input').style(height="auto")
-            submit = gr.Button("Submit")
-        with gr.Column(scale=2):
-            pred = gr.Textbox(label="Model Prediction")
-            with gr.Row():
-                depth = gr.Image(label="Depth", elem_id='img-display-output')
-                edge = gr.Image(label="Edge", elem_id='img-display-output')
-                normals = gr.Image(label="Normals", elem_id='img-display-output')
-            with gr.Row():
-                seg = gr.Image(label="Segmentation", elem_id='img-display-output')
-                obj_det = gr.Image(label="Object Detection", elem_id='img-display-output')
-                ocr_det = gr.Image(label="OCR Detection", elem_id='img-display-output')
-
-    def on_submit(im, q, model_type):
-        return pred, depth, edge, normals, seg, obj_det, ocr_det
-
-    submit.click(on_submit, inputs=[rgb, ques, model_type], outputs=[pred, depth, edge, normals, seg, obj_det, ocr_det])
-    examples = gr.Examples(examples=["examples/1.png"], inputs=[rgb])

patch

@@ -0,0 +1,82 @@
+diff --git a/dataset/caption_dataset.py b/dataset/caption_dataset.py
+index 266fdda..0cc5d3f 100644
+--- a/dataset/caption_dataset.py
++++ b/dataset/caption_dataset.py
+@@ -50,7 +50,7 @@ class Caption(Dataset):
+         elif self.dataset == 'demo':
+             img_path_split = self.data_list[index]['image'].split('/')
+             img_name = img_path_split[-2] + '/' + img_path_split[-1]
+-            image, labels, labels_info = get_expert_labels('', self.label_path, img_name, 'helpers', self.experts)
++            image, labels, labels_info = get_expert_labels('prismer', self.label_path, img_name, 'helpers', self.experts)
+ 
+         experts = self.transform(image, labels)
+         experts = post_label_process(experts, labels_info)
+diff --git a/dataset/utils.py b/dataset/utils.py
+index b368aac..418358c 100644
+--- a/dataset/utils.py
++++ b/dataset/utils.py
+@@ -5,6 +5,7 @@
+ # https://github.com/NVlabs/prismer/blob/main/LICENSE
+ 
+ import os
++import pathlib
+ import re
+ import json
+ import torch
+@@ -14,10 +15,12 @@ import torchvision.transforms as transforms
+ import torchvision.transforms.functional as transforms_f
+ from dataset.randaugment import RandAugment
+ 
+-COCO_FEATURES = torch.load('dataset/coco_features.pt')['features']
+-ADE_FEATURES = torch.load('dataset/ade_features.pt')['features']
+-DETECTION_FEATURES = torch.load('dataset/detection_features.pt')['features']
+-BACKGROUND_FEATURES = torch.load('dataset/background_features.pt')
++cur_dir = pathlib.Path(__file__).parent
++
++COCO_FEATURES = torch.load(cur_dir / 'coco_features.pt')['features']
++ADE_FEATURES = torch.load(cur_dir / 'ade_features.pt')['features']
++DETECTION_FEATURES = torch.load(cur_dir / 'detection_features.pt')['features']
++BACKGROUND_FEATURES = torch.load(cur_dir / 'background_features.pt')
+ 
+ 
+ class Transform:
+diff --git a/model/prismer.py b/model/prismer.py
+index 080253a..02362a4 100644
+--- a/model/prismer.py
++++ b/model/prismer.py
+@@ -5,6 +5,7 @@
+ # https://github.com/NVlabs/prismer/blob/main/LICENSE
+ 
+ import json
++import pathlib
+ import torch.nn as nn
+ 
+ from model.modules.vit import load_encoder
+@@ -12,6 +13,9 @@ from model.modules.roberta import load_decoder
+ from transformers import RobertaTokenizer, RobertaConfig
+ 
+ 
++cur_dir = pathlib.Path(__file__).parent
++
++
+ class Prismer(nn.Module):
+     def __init__(self, config):
+         super().__init__()
+@@ -26,7 +30,7 @@ class Prismer(nn.Module):
+             elif exp in ['obj_detection', 'ocr_detection']:
+                 self.experts[exp] = 64
+ 
+-        prismer_config = json.load(open('configs/prismer.json', 'r'))[config['prismer_model']]
++        prismer_config = json.load(open(f'{cur_dir.parent}/configs/prismer.json', 'r'))[config['prismer_model']]
+         roberta_config = RobertaConfig.from_dict(prismer_config['roberta_model'])
+ 
+         self.tokenizer = RobertaTokenizer.from_pretrained(prismer_config['roberta_model']['model_name'])
+@@ -35,7 +39,7 @@ class Prismer(nn.Module):
+ 
+         self.prepare_to_train(config['freeze'])
+         self.ignored_modules = self.get_ignored_modules(config['freeze'])
+-    
++
+     def prepare_to_train(self, mode='none'):
+         for name, params in self.named_parameters():
+             if mode == 'freeze_lang':

prismer/dataset/__init__.py

@@ -1,32 +0,0 @@
-# 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.vqa_dataset import VQA
-from dataset.caption_dataset import Caption
-
-
-def create_dataset(dataset, config):
-    if 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
-
-
-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/ade_features.pt

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

prismer/dataset/background_features.pt

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

prismer/dataset/caption_dataset.py

@@ -1,63 +0,0 @@
-# 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('prismer', 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/clip_pca.pkl

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

prismer/dataset/coco_features.pt

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

prismer/dataset/detection_features.pt

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

prismer/dataset/randaugment.py

@@ -1,265 +0,0 @@
-# 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 random
-import numpy as np
-import torch
-
-from PIL import Image, ImageOps, ImageEnhance, ImageDraw
-
-
-fillmask = {'depth': 0, 'normal': 0, 'edge': 0, 'seg_coco': 255, 'seg_ade': 255,
-            'obj_detection': 255, 'ocr_detection': 255}
-fillcolor = (0, 0, 0)
-
-
-def affine_transform(pair, affine_params):
-    img, label = pair
-    img = img.transform(img.size, Image.AFFINE, affine_params,
-                        resample=Image.BILINEAR, fillcolor=fillcolor)
-    if label is not None:
-        for exp in label:
-            label[exp] = label[exp].transform(label[exp].size, Image.AFFINE, affine_params,
-                                              resample=Image.NEAREST, fillcolor=fillmask[exp])
-    return img, label
-
-
-def ShearX(pair, v):  # [-0.3, 0.3]
-    assert -0.3 <= v <= 0.3
-    if random.random() > 0.5:
-        v = -v
-    return affine_transform(pair, (1, v, 0, 0, 1, 0))
-
-
-def ShearY(pair, v):  # [-0.3, 0.3]
-    assert -0.3 <= v <= 0.3
-    if random.random() > 0.5:
-        v = -v
-    return affine_transform(pair, (1, 0, 0, v, 1, 0))
-
-
-def TranslateX(pair, v):  # [-150, 150] => percentage: [-0.45, 0.45]
-    assert -0.45 <= v <= 0.45
-    if random.random() > 0.5:
-        v = -v
-    img, _ = pair
-    v = v * img.size[0]
-    return affine_transform(pair, (1, 0, v, 0, 1, 0))
-
-
-def TranslateY(pair, v):  # [-150, 150] => percentage: [-0.45, 0.45]
-    assert -0.45 <= v <= 0.45
-    if random.random() > 0.5:
-        v = -v
-    img, _ = pair
-    v = v * img.size[1]
-    return affine_transform(pair, (1, 0, 0, 0, 1, v))
-
-
-def TranslateXAbs(pair, v):  # [-150, 150] => percentage: [-0.45, 0.45]
-    assert 0 <= v <= 10
-    if random.random() > 0.5:
-        v = -v
-    return affine_transform(pair, (1, 0, v, 0, 1, 0))
-
-
-def TranslateYAbs(pair, v):  # [-150, 150] => percentage: [-0.45, 0.45]
-    assert 0 <= v <= 10
-    if random.random() > 0.5:
-        v = -v
-    return affine_transform(pair, (1, 0, 0, 0, 1, v))
-
-
-def Rotate(pair, v):  # [-30, 30]
-    assert -30 <= v <= 30
-    if random.random() > 0.5:
-        v = -v
-    img, label = pair
-    img = img.rotate(v, fillcolor=fillcolor)
-    if label is not None:
-        for exp in label:
-            label[exp] = label[exp].rotate(v, resample=Image.NEAREST, fillcolor=fillmask[exp])
-    return img, label
-
-
-def AutoContrast(pair, _):
-    img, label = pair
-    return ImageOps.autocontrast(img), label
-
-
-def Invert(pair, _):
-    img, label = pair
-    return ImageOps.invert(img), label
-
-
-def Equalize(pair, _):
-    img, label = pair
-    return ImageOps.equalize(img), label
-
-
-def Flip(pair, _):  # not from the paper
-    img, label = pair
-    return ImageOps.mirror(img), ImageOps.mirror(label)
-
-
-def Solarize(pair, v):  # [0, 256]
-    img, label = pair
-    assert 0 <= v <= 256
-    return ImageOps.solarize(img, v), label
-
-
-def Posterize(pair, v):  # [4, 8]
-    img, label = pair
-    assert 4 <= v <= 8
-    v = int(v)
-    return ImageOps.posterize(img, v), label
-
-
-def Posterize2(pair, v):  # [0, 4]
-    img, label = pair
-    assert 0 <= v <= 4
-    v = int(v)
-    return ImageOps.posterize(img, v), label
-
-
-def Contrast(pair, v):  # [0.1,1.9]
-    img, label = pair
-    assert 0.1 <= v <= 1.9
-    return ImageEnhance.Contrast(img).enhance(v), label
-
-
-def Color(pair, v):  # [0.1,1.9]
-    img, label = pair
-    assert 0.1 <= v <= 1.9
-    return ImageEnhance.Color(img).enhance(v), label
-
-
-def Brightness(pair, v):  # [0.1,1.9]
-    img, label = pair
-    assert 0.1 <= v <= 1.9
-    return ImageEnhance.Brightness(img).enhance(v), label
-
-
-def Sharpness(pair, v):  # [0.1,1.9]
-    img, label = pair
-    assert 0.1 <= v <= 1.9
-    return ImageEnhance.Sharpness(img).enhance(v), label
-
-
-def Cutout(pair, v):  # [0, 60] => percentage: [0, 0.2]
-    assert 0.0 <= v <= 0.2
-    if v <= 0.:
-        return pair
-    img, label = pair
-    v = v * img.size[0]
-    return CutoutAbs(img, v), label
-
-
-def CutoutAbs(img, v):  # [0, 60] => percentage: [0, 0.2]
-    # assert 0 <= v <= 20
-    if v < 0:
-        return img
-    w, h = img.size
-    x0 = np.random.uniform(w)
-    y0 = np.random.uniform(h)
-
-    x0 = int(max(0, x0 - v / 2.))
-    y0 = int(max(0, y0 - v / 2.))
-    x1 = min(w, x0 + v)
-    y1 = min(h, y0 + v)
-
-    xy = (x0, y0, x1, y1)
-    color = (125, 123, 114)
-    # color = (0, 0, 0)
-    img = img.copy()
-    ImageDraw.Draw(img).rectangle(xy, color)
-    return img
-
-
-def Identity(pair, v):
-    return pair
-
-
-def augment_list():  # 16 oeprations and their ranges
-    # https://github.com/google-research/uda/blob/master/image/randaugment/policies.py#L57
-    l = [
-        (Identity, 0., 1.0),
-        (ShearX, 0., 0.3),  # 0
-        (ShearY, 0., 0.3),  # 1
-        (TranslateX, 0., 0.33),  # 2
-        (TranslateY, 0., 0.33),  # 3
-        (Rotate, 0, 30),  # 4
-        (AutoContrast, 0, 1),  # 5
-        # (Invert, 0, 1),  # 6
-        (Equalize, 0, 1),  # 7
-        # (Solarize, 0, 110),  # 8
-        # (Posterize, 4, 8),  # 9
-        # (Color, 0.1, 1.9),  # 11
-        (Brightness, 0.1, 1.9),  # 12
-        (Sharpness, 0.1, 1.9),  # 13
-    ]
-    return l
-
-
-class Lighting(object):
-    """Lighting noise(AlexNet - style PCA - based noise)"""
-
-    def __init__(self, alphastd, eigval, eigvec):
-        self.alphastd = alphastd
-        self.eigval = torch.Tensor(eigval)
-        self.eigvec = torch.Tensor(eigvec)
-
-    def __call__(self, img):
-        if self.alphastd == 0:
-            return img
-
-        alpha = img.new().resize_(3).normal_(0, self.alphastd)
-        rgb = self.eigvec.type_as(img).clone() \
-            .mul(alpha.view(1, 3).expand(3, 3)) \
-            .mul(self.eigval.view(1, 3).expand(3, 3)) \
-            .sum(1).squeeze()
-
-        return img.add(rgb.view(3, 1, 1).expand_as(img))
-
-
-class CutoutDefault(object):
-    """
-    Reference : https://github.com/quark0/darts/blob/master/cnn/utils.py
-    """
-    def __init__(self, length):
-        self.length = length
-
-    def __call__(self, img):
-        h, w = img.size(1), img.size(2)
-        mask = np.ones((h, w), np.float32)
-        y = np.random.randint(h)
-        x = np.random.randint(w)
-
-        y1 = np.clip(y - self.length // 2, 0, h)
-        y2 = np.clip(y + self.length // 2, 0, h)
-        x1 = np.clip(x - self.length // 2, 0, w)
-        x2 = np.clip(x + self.length // 2, 0, w)
-
-        mask[y1: y2, x1: x2] = 0.
-        mask = torch.from_numpy(mask)
-        mask = mask.expand_as(img)
-        img *= mask
-        return img
-
-
-class RandAugment:
-    def __init__(self, n, m):
-        self.n = n
-        self.m = m      # [0, 10]
-        self.augment_list = augment_list()
-
-    def __call__(self, img, label):
-        pair = img, label
-        ops = random.choices(self.augment_list, k=self.n)
-        for op, minval, maxval in ops:
-            val = (float(self.m) / 10) * float(maxval - minval) + minval
-            pair = op(pair, val)
-        return pair

prismer/dataset/utils.py

@@ -1,192 +0,0 @@
-# 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
-import pathlib
-from dataset.randaugment import RandAugment
-
-
-cur_dir = pathlib.Path(__file__).parent
-
-COCO_FEATURES = torch.load(cur_dir / 'coco_features.pt')['features']
-ADE_FEATURES = torch.load(cur_dir / 'ade_features.pt')['features']
-DETECTION_FEATURES = torch.load(cur_dir / 'detection_features.pt')['features']
-BACKGROUND_FEATURES = torch.load(cur_dir / 'background_features.pt')
-
-
-class Transform:
-    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

@@ -1,53 +0,0 @@
-# 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 *
-from PIL import ImageFile
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-
-
-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

@@ -1,77 +0,0 @@
-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

@@ -1,161 +0,0 @@
-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

@@ -1,124 +0,0 @@
-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

@@ -1,16 +0,0 @@
-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

@@ -1,383 +0,0 @@
-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

@@ -1,32 +0,0 @@
-# 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

@@ -1,124 +0,0 @@
-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

@@ -1,576 +0,0 @@
-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

@@ -1,32 +0,0 @@
-# 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

@@ -1,50 +0,0 @@
-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

@@ -1,286 +0,0 @@
-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

@@ -1,56 +0,0 @@
-# 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

@@ -1,57 +0,0 @@
-# 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

@@ -1,58 +0,0 @@
-# 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

@@ -1,115 +0,0 @@
-# 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

@@ -1,86 +0,0 @@
-# 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

@@ -1,56 +0,0 @@
-# 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

@@ -1,139 +0,0 @@
-# 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

@@ -1,34 +0,0 @@
-# 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

@@ -1,22 +0,0 @@
-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

@@ -1,85 +0,0 @@
-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

@@ -1,202 +0,0 @@
-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

@@ -1,32 +0,0 @@
-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

@@ -1,140 +0,0 @@
-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

@@ -1,178 +0,0 @@
-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

@@ -1,191 +0,0 @@
-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

@@ -1,48 +0,0 @@
-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

@@ -1,15 +0,0 @@
-_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

@@ -1,22 +0,0 @@
-_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

@@ -1,28 +0,0 @@
-_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

@@ -1,29 +0,0 @@
-_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
-