add omg code

.gitattributes

@@ -17,10 +17,6 @@
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
@@ -33,3 +29,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,125 @@
+# 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/
+*.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/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/en/_build/
+docs/zh_cn/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# 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/
+data/
+data
+.vscode
+.idea/
+.DS_Store
+
+# custom
+*.pkl
+*.pkl.json
+*.log.json
+docs/modelzoo_statistics.md
+mmdet/.mim
+work_dirs/
+
+# Pytorch
+*.py~
+*.sh~
+
+# remove tmp folder
+tmp/

app/configs/m2_convl.py

@@ -0,0 +1,152 @@
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import MSDeformAttnPixelDecoder, CrossEntropyLoss, DiceLoss, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, ClassificationCost, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+from seg.models.detectors import Mask2formerVideo
+from seg.models.fusion_head import OMGFusionHead
+from seg.models.heads import Mask2FormerVideoHead
+from seg.models.backbones import OpenCLIPBackbone
+
+num_things_classes = 80
+num_stuff_classes = 53
+
+ov_model_name = 'convnext_large_d_320'
+ov_datasets_name = 'CocoPanopticOVDataset'
+model = dict(
+    type=Mask2formerVideo,
+    data_preprocessor=None,  # to fill
+    backbone=dict(
+        type=OpenCLIPBackbone,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='./models/omg_seg_convl.pth',
+            prefix='panoptic_head.'
+        ),
+        type=Mask2FormerVideoHead,
+        sphere_cls=True,
+        ov_classifier_name=f'{ov_model_name}_{ov_datasets_name}',
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+            positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=None  # [1.0] * num_classes + [0.1]
+        ),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean'
+        )
+    ),
+    panoptic_fusion_head=dict(
+        type=OMGFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None
+    ),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                dict(type=ClassificationCost, weight=2.0),
+                dict(
+                    type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True,
+        object_mask_thr=0.,
+    ),
+    init_cfg=None
+)

ext/cityscapes_scripts/createPanopticImgs.py

@@ -0,0 +1,194 @@
+#!/usr/bin/python
+#
+# Converts the *instanceIds.png annotations of the Cityscapes dataset
+# to COCO-style panoptic segmentation format (http://cocodataset.org/#format-data).
+# The convertion is working for 'fine' set of the annotations.
+#
+# By default with this tool uses IDs specified in labels.py. You can use flag
+# --use-train-id to get train ids for categories. 'ignoreInEval' categories are
+# removed during the conversion.
+#
+# In panoptic segmentation format image_id is used to match predictions and ground truth.
+# For cityscapes image_id has form <city>_123456_123456 and corresponds to the prefix
+# of cityscapes image files.
+#
+
+# python imports
+from __future__ import print_function, absolute_import, division, unicode_literals
+import os
+import glob
+import sys
+import argparse
+import json
+import numpy as np
+
+# Image processing
+from PIL import Image
+
+# cityscapes imports
+from ext.cityscapes_scripts.helpers.csHelpers import printError
+from ext.cityscapes_scripts.helpers.labels import id2label, labels
+
+
+import mmengine
+
+
+# The main method
+def convert2panoptic(cityscapesPath=None, outputFolder=None, useTrainId=False, setNames=["val", "train", "test"]):
+    # Where to look for Cityscapes
+    if cityscapesPath is None:
+        if 'CITYSCAPES_DATASET' in os.environ:
+            cityscapesPath = os.environ['CITYSCAPES_DATASET']
+        else:
+            cityscapesPath = 'data/cityscapes'
+        cityscapesPath = os.path.join(cityscapesPath, "gtFine")
+
+    if outputFolder is None:
+        outputFolder = cityscapesPath.replace('gtFine', "annotations")
+
+    mmengine.mkdir_or_exist(outputFolder)
+
+    categories = []
+    for label in labels:
+        if label.ignoreInEval:
+            continue
+        categories.append({'id': int(label.trainId) if useTrainId else int(label.id),
+                           'name': label.name,
+                           'color': label.color,
+                           'supercategory': label.category,
+                           'isthing': 1 if label.hasInstances else 0})
+
+    categories = sorted(categories, key=lambda x:x['id'])
+
+    for setName in setNames:
+        # how to search for all ground truth
+        searchFine   = os.path.join(cityscapesPath, setName, "*", "*_instanceIds.png")
+        # search files
+        filesFine = glob.glob(searchFine)
+        filesFine.sort()
+
+        files = filesFine
+        # quit if we did not find anything
+        if not files:
+            printError(
+                "Did not find any files for {} set using matching pattern {}. Please consult the README.".format(setName, searchFine)
+            )
+        # a bit verbose
+        print("Converting {} annotation files for {} set.".format(len(files), setName))
+
+        trainIfSuffix = "_trainId" if useTrainId else ""
+        outputBaseFile = "cityscapes_panoptic_{}{}".format(setName, trainIfSuffix)
+        outFile = os.path.join(outputFolder, "{}.json".format(outputBaseFile))
+        print("Json file with the annotations in panoptic format will be saved in {}".format(outFile))
+        panopticFolder = os.path.join(outputFolder, outputBaseFile)
+        if not os.path.isdir(panopticFolder):
+            print("Creating folder {} for panoptic segmentation PNGs".format(panopticFolder))
+            os.mkdir(panopticFolder)
+        print("Corresponding segmentations in .png format will be saved in {}".format(panopticFolder))
+
+        images = []
+        annotations = []
+        for progress, f in enumerate(files):
+
+            originalFormat = np.array(Image.open(f))
+
+            fileName = os.path.basename(f)
+            location = fileName.split('_')[0]
+            imageId = fileName.replace("_gtFine_instanceIds.png", "")
+            fileName = os.path.join(location, fileName)
+            inputFileName = fileName.replace("_gtFine_instanceIds.png", "_leftImg8bit.png")
+            outputFileName = fileName.replace("_gtFine_instanceIds.png", "_panoptic.png")
+            # image entry, id for image is its filename without extension
+            images.append({"id": imageId,
+                           "width": int(originalFormat.shape[1]),
+                           "height": int(originalFormat.shape[0]),
+                           "file_name": inputFileName})
+
+            pan_format = np.zeros(
+                (originalFormat.shape[0], originalFormat.shape[1], 3), dtype=np.uint8
+            )
+
+            segmentIds = np.unique(originalFormat)
+            segmInfo = []
+            for segmentId in segmentIds:
+                if segmentId < 1000:
+                    semanticId = segmentId
+                    isCrowd = 1
+                else:
+                    semanticId = segmentId // 1000
+                    isCrowd = 0
+                labelInfo = id2label[semanticId]
+                categoryId = labelInfo.trainId if useTrainId else labelInfo.id
+                if labelInfo.ignoreInEval:
+                    continue
+                if not labelInfo.hasInstances:
+                    isCrowd = 0
+
+                mask = originalFormat == segmentId
+                color = [segmentId % 256, segmentId // 256, segmentId // 256 // 256]
+                pan_format[mask] = color
+
+                area = np.sum(mask) # segment area computation
+
+                # bbox computation for a segment
+                hor = np.sum(mask, axis=0)
+                hor_idx = np.nonzero(hor)[0]
+                x = hor_idx[0]
+                width = hor_idx[-1] - x + 1
+                vert = np.sum(mask, axis=1)
+                vert_idx = np.nonzero(vert)[0]
+                y = vert_idx[0]
+                height = vert_idx[-1] - y + 1
+                bbox = [int(x), int(y), int(width), int(height)]
+
+                segmInfo.append({"id": int(segmentId),
+                                 "category_id": int(categoryId),
+                                 "area": int(area),
+                                 "bbox": bbox,
+                                 "iscrowd": isCrowd})
+
+            annotations.append({'image_id': imageId,
+                                'file_name': outputFileName,
+                                "segments_info": segmInfo})
+
+            mmengine.mkdir_or_exist(os.path.dirname(os.path.join(panopticFolder, outputFileName)))
+            Image.fromarray(pan_format).save(os.path.join(panopticFolder, outputFileName))
+
+            print("\rProgress: {:>3.2f} %".format((progress + 1) * 100 / len(files)), end=' ')
+            sys.stdout.flush()
+
+        print("\nSaving the json file {}".format(outFile))
+        d = {'images': images,
+             'annotations': annotations,
+             'categories': categories}
+        with open(outFile, 'w') as f:
+            json.dump(d, f, sort_keys=True, indent=4)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset-folder",
+                        dest="cityscapesPath",
+                        help="path to the Cityscapes dataset 'gtFine' folder",
+                        default=None,
+                        type=str)
+    parser.add_argument("--output-folder",
+                        dest="outputFolder",
+                        help="path to the output folder.",
+                        default=None,
+                        type=str)
+    parser.add_argument("--use-train-id", default=True,action="store_true", dest="useTrainId")
+    parser.add_argument("--set-names",
+                        dest="setNames",
+                        help="set names to which apply the function to",
+                        nargs='+',
+                        default=["val", "train"],
+                        type=str)
+    args = parser.parse_args()
+
+    convert2panoptic(args.cityscapesPath, args.outputFolder, args.useTrainId, args.setNames)
+
+
+# call the main
+if __name__ == "__main__":
+    main()

ext/cityscapes_scripts/helpers/__init__.py

@@ -0,0 +1 @@
+# empty

ext/cityscapes_scripts/helpers/annotation.py

@@ -0,0 +1,441 @@
+#!/usr/bin/python
+#
+# Classes to store, read, and write annotations
+#
+
+from __future__ import print_function, absolute_import, division
+import os
+import json
+import numpy as np
+from collections import namedtuple
+
+# get current date and time
+import datetime
+import locale
+
+from abc import ABCMeta, abstractmethod
+from .box3dImageTransform import Camera
+
+# A point in a polygon
+Point = namedtuple('Point', ['x', 'y'])
+
+
+class CsObjectType():
+    """Type of an object"""
+    POLY = 1  # polygon
+    BBOX2D = 2  # bounding box
+    BBOX3D = 3  # 3d bounding box
+    IGNORE2D = 4  # 2d ignore region
+
+
+class CsObject:
+    """Abstract base class for annotation objects"""
+    __metaclass__ = ABCMeta
+
+    def __init__(self, objType):
+        self.objectType = objType
+        # the label
+        self.label = ""
+
+        # If deleted or not
+        self.deleted = 0
+        # If verified or not
+        self.verified = 0
+        # The date string
+        self.date = ""
+        # The username
+        self.user = ""
+        # Draw the object
+        # Not read from or written to JSON
+        # Set to False if deleted object
+        # Might be set to False by the application for other reasons
+        self.draw = True
+
+    @abstractmethod
+    def __str__(self): pass
+
+    @abstractmethod
+    def fromJsonText(self, jsonText, objId=-1): pass
+
+    @abstractmethod
+    def toJsonText(self): pass
+
+    def updateDate(self):
+        try:
+            locale.setlocale(locale.LC_ALL, 'en_US.utf8')
+        except locale.Error:
+            locale.setlocale(locale.LC_ALL, 'en_US')
+        except locale.Error:
+            locale.setlocale(locale.LC_ALL, 'us_us.utf8')
+        except locale.Error:
+            locale.setlocale(locale.LC_ALL, 'us_us')
+        except Exception:
+            pass
+        self.date = datetime.datetime.now().strftime("%d-%b-%Y %H:%M:%S")
+
+    # Mark the object as deleted
+    def delete(self):
+        self.deleted = 1
+        self.draw = False
+
+
+class CsPoly(CsObject):
+    """Class that contains the information of a single annotated object as polygon"""
+
+    # Constructor
+    def __init__(self):
+        CsObject.__init__(self, CsObjectType.POLY)
+        # the polygon as list of points
+        self.polygon = []
+        # the object ID
+        self.id = -1
+
+    def __str__(self):
+        polyText = ""
+        if self.polygon:
+            if len(self.polygon) <= 4:
+                for p in self.polygon:
+                    polyText += '({},{}) '.format(p.x, p.y)
+            else:
+                polyText += '({},{}) ({},{}) ... ({},{}) ({},{})'.format(
+                    self.polygon[0].x, self.polygon[0].y,
+                    self.polygon[1].x, self.polygon[1].y,
+                    self.polygon[-2].x, self.polygon[-2].y,
+                    self.polygon[-1].x, self.polygon[-1].y)
+        else:
+            polyText = "none"
+        text = "Object: {} - {}".format(self.label, polyText)
+        return text
+
+    def fromJsonText(self, jsonText, objId=-1):
+        self.id = objId
+        self.label = str(jsonText['label'])
+        self.polygon = [Point(p[0], p[1]) for p in jsonText['polygon']]
+        if 'deleted' in jsonText.keys():
+            self.deleted = jsonText['deleted']
+        else:
+            self.deleted = 0
+        if 'verified' in jsonText.keys():
+            self.verified = jsonText['verified']
+        else:
+            self.verified = 1
+        if 'user' in jsonText.keys():
+            self.user = jsonText['user']
+        else:
+            self.user = ''
+        if 'date' in jsonText.keys():
+            self.date = jsonText['date']
+        else:
+            self.date = ''
+        if self.deleted == 1:
+            self.draw = False
+        else:
+            self.draw = True
+
+    def toJsonText(self):
+        objDict = {}
+        objDict['label'] = self.label
+        objDict['id'] = self.id
+        objDict['deleted'] = self.deleted
+        objDict['verified'] = self.verified
+        objDict['user'] = self.user
+        objDict['date'] = self.date
+        objDict['polygon'] = []
+        for pt in self.polygon:
+            objDict['polygon'].append([pt.x, pt.y])
+
+        return objDict
+
+
+class CsBbox2d(CsObject):
+    """Class that contains the information of a single annotated object as bounding box"""
+
+    # Constructor
+    def __init__(self):
+        CsObject.__init__(self, CsObjectType.BBOX2D)
+        # the polygon as list of points
+        self.bbox_amodal_xywh = []
+        self.bbox_modal_xywh = []
+
+        # the ID of the corresponding object
+        self.instanceId = -1
+        # the label of the corresponding object
+        self.label = ""
+
+    def __str__(self):
+        bboxAmodalText = ""
+        bboxAmodalText += '[(x1: {}, y1: {}), (w: {}, h: {})]'.format(
+            self.bbox_amodal_xywh[0], self.bbox_amodal_xywh[1],  self.bbox_amodal_xywh[2],  self.bbox_amodal_xywh[3])
+
+        bboxModalText = ""
+        bboxModalText += '[(x1: {}, y1: {}), (w: {}, h: {})]'.format(
+            self.bbox_modal_xywh[0], self.bbox_modal_xywh[1], self.bbox_modal_xywh[2], self.bbox_modal_xywh[3])
+
+        text = "Object: {}\n - Amodal {}\n - Modal {}".format(
+            self.label, bboxAmodalText, bboxModalText)
+        return text
+
+    def setAmodalBox(self, bbox_amodal):
+        # sets the amodal box if required
+        self.bbox_amodal_xywh = [
+            bbox_amodal[0],
+            bbox_amodal[1],
+            bbox_amodal[2] - bbox_amodal[0],
+            bbox_amodal[3] - bbox_amodal[1]
+        ]
+
+    # access 2d boxes in [xmin, ymin, xmax, ymax] format
+    @property
+    def bbox_amodal(self):
+        """Returns the 2d box as [xmin, ymin, xmax, ymax]"""
+        return [
+            self.bbox_amodal_xywh[0],
+            self.bbox_amodal_xywh[1],
+            self.bbox_amodal_xywh[0] + self.bbox_amodal_xywh[2],
+            self.bbox_amodal_xywh[1] + self.bbox_amodal_xywh[3]
+        ]
+
+    @property
+    def bbox_modal(self):
+        """Returns the 2d box as [xmin, ymin, xmax, ymax]"""
+        return [
+            self.bbox_modal_xywh[0],
+            self.bbox_modal_xywh[1],
+            self.bbox_modal_xywh[0] + self.bbox_modal_xywh[2],
+            self.bbox_modal_xywh[1] + self.bbox_modal_xywh[3]
+        ]
+
+    def fromJsonText(self, jsonText, objId=-1):
+        # try to load from cityperson format
+        if 'bbox' in jsonText.keys() and 'bboxVis' in jsonText.keys():
+            self.bbox_amodal_xywh = jsonText['bbox']
+            self.bbox_modal_xywh = jsonText['bboxVis']
+        # both modal and amodal boxes are provided
+        elif "modal" in jsonText.keys() and "amodal" in jsonText.keys():
+            self.bbox_amodal_xywh = jsonText['amodal']
+            self.bbox_modal_xywh = jsonText['modal']
+        # only amodal boxes are provided
+        else:
+            self.bbox_modal_xywh = jsonText['amodal']
+            self.bbox_amodal_xywh = jsonText['amodal']
+
+        # load label and instanceId if available
+        if 'label' in jsonText.keys() and 'instanceId' in jsonText.keys():
+            self.label = str(jsonText['label'])
+            self.instanceId = jsonText['instanceId']
+
+    def toJsonText(self):
+        objDict = {}
+        objDict['label'] = self.label
+        objDict['instanceId'] = self.instanceId
+        objDict['modal'] = self.bbox_modal_xywh
+        objDict['amodal'] = self.bbox_amodal_xywh
+
+        return objDict
+
+
+class CsBbox3d(CsObject):
+    """Class that contains the information of a single annotated object as 3D bounding box"""
+
+    # Constructor
+    def __init__(self):
+        CsObject.__init__(self, CsObjectType.BBOX3D)
+
+        self.bbox_2d = None
+
+        self.center = []
+        self.dims = []
+        self.rotation = []
+        self.instanceId = -1
+        self.label = ""
+        self.score = -1.
+
+    def __str__(self):
+        bbox2dText = str(self.bbox_2d)
+
+        bbox3dText = ""
+        bbox3dText += '\n - Center (x/y/z) [m]: {}/{}/{}'.format(
+            self.center[0], self.center[1],  self.center[2])
+        bbox3dText += '\n - Dimensions (l/w/h) [m]: {}/{}/{}'.format(
+            self.dims[0], self.dims[1],  self.dims[2])
+        bbox3dText += '\n - Rotation: {}/{}/{}/{}'.format(
+            self.rotation[0], self.rotation[1], self.rotation[2], self.rotation[3])
+
+        text = "Object: {}\n2D {}\n - 3D {}".format(
+            self.label, bbox2dText, bbox3dText)
+        return text
+
+    def fromJsonText(self, jsonText, objId=-1):
+        # load 2D box
+        self.bbox_2d = CsBbox2d()
+        self.bbox_2d.fromJsonText(jsonText['2d'])
+
+        self.center = jsonText['3d']['center']
+        self.dims = jsonText['3d']['dimensions']
+        self.rotation = jsonText['3d']['rotation']
+        self.label = jsonText['label']
+        self.score = jsonText['score']
+
+        if 'instanceId' in jsonText.keys():
+            self.instanceId = jsonText['instanceId']
+
+    def toJsonText(self):
+        objDict = {}
+        objDict['label'] = self.label
+        objDict['instanceId'] = self.instanceId
+        objDict['2d']['amodal'] = self.bbox_2d.bbox_amodal_xywh
+        objDict['2d']['modal'] = self.bbox_2d.bbox_modal_xywh
+        objDict['3d']['center'] = self.center
+        objDict['3d']['dimensions'] = self.dims
+        objDict['3d']['rotation'] = self.rotation
+
+        return objDict
+
+    @property
+    def depth(self):
+        # returns the BEV depth
+        return np.sqrt(self.center[0]**2 + self.center[1]**2).astype(int)
+
+
+class CsIgnore2d(CsObject):
+    """Class that contains the information of a single annotated 2d ignore region"""
+
+    # Constructor
+    def __init__(self):
+        CsObject.__init__(self, CsObjectType.IGNORE2D)
+
+        self.bbox_xywh = []
+        self.label = ""
+        self.instanceId = -1
+
+    def __str__(self):
+        bbox2dText = ""
+        bbox2dText += 'Ignore Region:  (x1: {}, y1: {}), (w: {}, h: {})'.format(
+            self.bbox_xywh[0], self.bbox_xywh[1], self.bbox_xywh[2], self.bbox_xywh[3])
+
+        return bbox2dText
+
+    def fromJsonText(self, jsonText, objId=-1):
+        self.bbox_xywh = jsonText['2d']
+
+        if 'label' in jsonText.keys():
+            self.label = jsonText['label']
+
+        if 'instanceId' in jsonText.keys():
+            self.instanceId = jsonText['instanceId']
+
+    def toJsonText(self):
+        objDict = {}
+        objDict['label'] = self.label
+        objDict['instanceId'] = self.instanceId
+        objDict['2d'] = self.bbox_xywh
+
+        return objDict
+
+    @property
+    def bbox(self):
+        """Returns the 2d box as [xmin, ymin, xmax, ymax]"""
+        return [
+            self.bbox_xywh[0],
+            self.bbox_xywh[1],
+            self.bbox_xywh[0] + self.bbox_xywh[2],
+            self.bbox_xywh[1] + self.bbox_xywh[3]
+        ]
+
+    # Extend api to be compatible to bbox2d
+    @property
+    def bbox_amodal_xywh(self):
+        return self.bbox_xywh
+
+    @property
+    def bbox_modal_xywh(self):
+        return self.bbox_xywh
+
+
+class Annotation:
+    """The annotation of a whole image (doesn't support mixed annotations, i.e. combining CsPoly and CsBbox2d)"""
+
+    # Constructor
+    def __init__(self, objType=CsObjectType.POLY):
+        # the width of that image and thus of the label image
+        self.imgWidth = 0
+        # the height of that image and thus of the label image
+        self.imgHeight = 0
+        # the list of objects
+        self.objects = []
+        # the camera calibration
+        self.camera = None
+        assert objType in CsObjectType.__dict__.values()
+        self.objectType = objType
+
+    def toJson(self):
+        return json.dumps(self, default=lambda o: o.__dict__, sort_keys=True, indent=4)
+
+    def fromJsonText(self, jsonText):
+        jsonDict = json.loads(jsonText)
+        self.imgWidth = int(jsonDict['imgWidth'])
+        self.imgHeight = int(jsonDict['imgHeight'])
+        self.objects = []
+        # load objects
+        if self.objectType != CsObjectType.IGNORE2D:
+            for objId, objIn in enumerate(jsonDict['objects']):
+                if self.objectType == CsObjectType.POLY:
+                    obj = CsPoly()
+                elif self.objectType == CsObjectType.BBOX2D:
+                    obj = CsBbox2d()
+                elif self.objectType == CsObjectType.BBOX3D:
+                    obj = CsBbox3d()
+                obj.fromJsonText(objIn, objId)
+                self.objects.append(obj)
+
+        # load ignores
+        if 'ignore' in jsonDict.keys():
+            for ignoreId, ignoreIn in enumerate(jsonDict['ignore']):
+                obj = CsIgnore2d()
+                obj.fromJsonText(ignoreIn, ignoreId)
+                self.objects.append(obj)
+
+        # load camera calibration
+        if 'sensor' in jsonDict.keys():
+            self.camera = Camera(fx=jsonDict['sensor']['fx'],
+                                 fy=jsonDict['sensor']['fy'],
+                                 u0=jsonDict['sensor']['u0'],
+                                 v0=jsonDict['sensor']['v0'],
+                                 sensor_T_ISO_8855=jsonDict['sensor']['sensor_T_ISO_8855'])
+
+    def toJsonText(self):
+        jsonDict = {}
+        jsonDict['imgWidth'] = self.imgWidth
+        jsonDict['imgHeight'] = self.imgHeight
+        jsonDict['objects'] = []
+        for obj in self.objects:
+            objDict = obj.toJsonText()
+            jsonDict['objects'].append(objDict)
+
+        return jsonDict
+
+    # Read a json formatted polygon file and return the annotation
+    def fromJsonFile(self, jsonFile):
+        if not os.path.isfile(jsonFile):
+            print('Given json file not found: {}'.format(jsonFile))
+            return
+        with open(jsonFile, 'r') as f:
+            jsonText = f.read()
+            self.fromJsonText(jsonText)
+
+    def toJsonFile(self, jsonFile):
+        with open(jsonFile, 'w') as f:
+            f.write(self.toJson())
+
+
+# a dummy example
+if __name__ == "__main__":
+    obj = CsPoly()
+    obj.label = 'car'
+    obj.polygon.append(Point(0, 0))
+    obj.polygon.append(Point(1, 0))
+    obj.polygon.append(Point(1, 1))
+    obj.polygon.append(Point(0, 1))
+
+    print(type(obj).__name__)
+    print(obj)

ext/cityscapes_scripts/helpers/csHelpers.py

@@ -0,0 +1,129 @@
+#!/usr/bin/python
+#
+# Various helper methods and includes for Cityscapes
+#
+
+# Python imports
+from __future__ import print_function, absolute_import, division
+import os
+import sys
+import getopt
+import glob
+import math
+import json
+from collections import namedtuple
+import logging
+import traceback
+
+# Image processing
+from PIL import Image
+from PIL import ImageDraw
+
+# Numpy for datastructures
+import numpy as np
+
+# Cityscapes modules
+# from .annotation import Annotation
+from .labels import labels, name2label, id2label, trainId2label, category2labels
+
+
+def printError(message):
+    """Print an error message and quit"""
+    print('ERROR: ' + str(message))
+    sys.exit(-1)
+
+
+class colors:
+    """Class for colors"""
+    RED = '\033[31;1m'
+    GREEN = '\033[32;1m'
+    YELLOW = '\033[33;1m'
+    BLUE = '\033[34;1m'
+    MAGENTA = '\033[35;1m'
+    CYAN = '\033[36;1m'
+    BOLD = '\033[1m'
+    UNDERLINE = '\033[4m'
+    ENDC = '\033[0m'
+
+
+def getColorEntry(val, args):
+    """Colored value output if colorized flag is activated."""
+
+    if not args.colorized:
+        return ""
+    if not isinstance(val, float) or math.isnan(val):
+        return colors.ENDC
+    if (val < .20):
+        return colors.RED
+    elif (val < .40):
+        return colors.YELLOW
+    elif (val < .60):
+        return colors.BLUE
+    elif (val < .80):
+        return colors.CYAN
+    else:
+        return colors.GREEN
+
+
+# Cityscapes files have a typical filename structure
+# <city>_<sequenceNb>_<frameNb>_<type>[_<type2>].<ext>
+# This class contains the individual elements as members
+# For the sequence and frame number, the strings are returned, including leading zeros
+CsFile = namedtuple('csFile', ['city', 'sequenceNb', 'frameNb', 'type', 'type2', 'ext'])
+
+
+def getCsFileInfo(fileName):
+    """Returns a CsFile object filled from the info in the given filename"""
+    baseName = os.path.basename(fileName)
+    parts = baseName.split('_')
+    parts = parts[:-1] + parts[-1].split('.')
+    if not parts:
+        printError('Cannot parse given filename ({}). Does not seem to be a valid Cityscapes file.'.format(fileName))
+    if len(parts) == 5:
+        csFile = CsFile(*parts[:-1], type2="", ext=parts[-1])
+    elif len(parts) == 6:
+        csFile = CsFile(*parts)
+    else:
+        printError('Found {} part(s) in given filename ({}). Expected 5 or 6.'.format(len(parts), fileName))
+
+    return csFile
+
+
+def getCoreImageFileName(filename):
+    """Returns the part of Cityscapes filenames that is common to all data types
+
+    e.g. for city_123456_123456_gtFine_polygons.json returns city_123456_123456
+    """
+    csFile = getCsFileInfo(filename)
+    return "{}_{}_{}".format(csFile.city, csFile.sequenceNb, csFile.frameNb)
+
+
+def getDirectory(fileName):
+    """Returns the directory name for the given filename
+
+    e.g.
+    fileName = "/foo/bar/foobar.txt"
+    return value is "bar"
+    Not much error checking though
+    """
+    dirName = os.path.dirname(fileName)
+    return os.path.basename(dirName)
+
+
+def ensurePath(path):
+    """Make sure that the given path exists"""
+    if not path:
+        return
+    if not os.path.isdir(path):
+        os.makedirs(path)
+
+
+def writeDict2JSON(dictName, fileName):
+    """Write a dictionary as json file"""
+    with open(fileName, 'w') as f:
+        f.write(json.dumps(dictName, default=lambda o: o.__dict__, sort_keys=True, indent=4))
+
+
+# dummy main
+if __name__ == "__main__":
+    printError("Only for include, not executable on its own.")

ext/cityscapes_scripts/helpers/labels.py

@@ -0,0 +1,182 @@
+#!/usr/bin/python
+#
+# Cityscapes labels
+#
+
+from __future__ import print_function, absolute_import, division
+from collections import namedtuple
+
+
+#--------------------------------------------------------------------------------
+# Definitions
+#--------------------------------------------------------------------------------
+
+# a label and all meta information
+Label = namedtuple( 'Label' , [
+
+    'name'        , # The identifier of this label, e.g. 'car', 'person', ... .
+                    # We use them to uniquely name a class
+
+    'id'          , # An integer ID that is associated with this label.
+                    # The IDs are used to represent the label in ground truth images
+                    # An ID of -1 means that this label does not have an ID and thus
+                    # is ignored when creating ground truth images (e.g. license plate).
+                    # Do not modify these IDs, since exactly these IDs are expected by the
+                    # evaluation server.
+
+    'trainId'     , # Feel free to modify these IDs as suitable for your method. Then create
+                    # ground truth images with train IDs, using the tools provided in the
+                    # 'preparation' folder. However, make sure to validate or submit results
+                    # to our evaluation server using the regular IDs above!
+                    # For trainIds, multiple labels might have the same ID. Then, these labels
+                    # are mapped to the same class in the ground truth images. For the inverse
+                    # mapping, we use the label that is defined first in the list below.
+                    # For example, mapping all void-type classes to the same ID in training,
+                    # might make sense for some approaches.
+                    # Max value is 255!
+
+    'category'    , # The name of the category that this label belongs to
+
+    'categoryId'  , # The ID of this category. Used to create ground truth images
+                    # on category level.
+
+    'hasInstances', # Whether this label distinguishes between single instances or not
+
+    'ignoreInEval', # Whether pixels having this class as ground truth label are ignored
+                    # during evaluations or not
+
+    'color'       , # The color of this label
+    ] )
+
+
+#--------------------------------------------------------------------------------
+# A list of all labels
+#--------------------------------------------------------------------------------
+
+# Please adapt the train IDs as appropriate for your approach.
+# Note that you might want to ignore labels with ID 255 during training.
+# Further note that the current train IDs are only a suggestion. You can use whatever you like.
+# Make sure to provide your results using the original IDs and not the training IDs.
+# Note that many IDs are ignored in evaluation and thus you never need to predict these!
+
+labels = [
+    #       name                     id    trainId   category            catId     hasInstances   ignoreInEval   color
+    Label(  'unlabeled'            ,  0 ,      255 , 'void'            , 0       , False        , True         , (  0,  0,  0) ),
+    Label(  'ego vehicle'          ,  1 ,      255 , 'void'            , 0       , False        , True         , (  0,  0,  0) ),
+    Label(  'rectification border' ,  2 ,      255 , 'void'            , 0       , False        , True         , (  0,  0,  0) ),
+    Label(  'out of roi'           ,  3 ,      255 , 'void'            , 0       , False        , True         , (  0,  0,  0) ),
+    Label(  'static'               ,  4 ,      255 , 'void'            , 0       , False        , True         , (  0,  0,  0) ),
+    Label(  'dynamic'              ,  5 ,      255 , 'void'            , 0       , False        , True         , (111, 74,  0) ),
+    Label(  'ground'               ,  6 ,      255 , 'void'            , 0       , False        , True         , ( 81,  0, 81) ),
+    Label(  'road'                 ,  7 ,        0 + 8, 'flat'            , 1       , False        , False        , (128, 64,128) ),
+    Label(  'sidewalk'             ,  8 ,        1 + 8, 'flat'            , 1       , False        , False        , (244, 35,232) ),
+    Label(  'parking'              ,  9 ,      255 , 'flat'            , 1       , False        , True         , (250,170,160) ),
+    Label(  'rail track'           , 10 ,      255 , 'flat'            , 1       , False        , True         , (230,150,140) ),
+    Label(  'building'             , 11 ,        2 + 8, 'construction'    , 2       , False        , False        , ( 70, 70, 70) ),
+    Label(  'wall'                 , 12 ,        3 + 8, 'construction'    , 2       , False        , False        , (102,102,156) ),
+    Label(  'fence'                , 13 ,        4 + 8, 'construction'    , 2       , False        , False        , (190,153,153) ),
+    Label(  'guard rail'           , 14 ,      255 , 'construction'    , 2       , False        , True         , (180,165,180) ),
+    Label(  'bridge'               , 15 ,      255 , 'construction'    , 2       , False        , True         , (150,100,100) ),
+    Label(  'tunnel'               , 16 ,      255 , 'construction'    , 2       , False        , True         , (150,120, 90) ),
+    Label(  'pole'                 , 17 ,        5 + 8, 'object'          , 3       , False        , False        , (153,153,153) ),
+    Label(  'polegroup'            , 18 ,      255 , 'object'          , 3       , False        , True         , (153,153,153) ),
+    Label(  'traffic light'        , 19 ,        6 + 8, 'object'          , 3       , False        , False        , (250,170, 30) ),
+    Label(  'traffic sign'         , 20 ,        7 + 8, 'object'          , 3       , False        , False        , (220,220,  0) ),
+    Label(  'vegetation'           , 21 ,        8 + 8, 'nature'          , 4       , False        , False        , (107,142, 35) ),
+    Label(  'terrain'              , 22 ,        9 + 8, 'nature'          , 4       , False        , False        , (152,251,152) ),
+    Label(  'sky'                  , 23 ,       10 + 8, 'sky'             , 5       , False        , False        , ( 70,130,180) ),
+    Label(  'person'               , 24 ,       11 - 11 , 'human'           , 6       , True         , False        , (220, 20, 60) ),
+    Label(  'rider'                , 25 ,       12 - 11 , 'human'           , 6       , True         , False        , (255,  0,  0) ),
+    Label(  'car'                  , 26 ,       13 - 11, 'vehicle'         , 7       , True         , False        , (  0,  0,142) ),
+    Label(  'truck'                , 27 ,       14 - 11, 'vehicle'         , 7       , True         , False        , (  0,  0, 70) ),
+    Label(  'bus'                  , 28 ,       15 - 11, 'vehicle'         , 7       , True         , False        , (  0, 60,100) ),
+    Label(  'caravan'              , 29 ,      255 , 'vehicle'         , 7       , True         , True         , (  0,  0, 90) ),
+    Label(  'trailer'              , 30 ,      255 , 'vehicle'         , 7       , True         , True         , (  0,  0,110) ),
+    Label(  'train'                , 31 ,       16 - 11, 'vehicle'         , 7       , True         , False        , (  0, 80,100) ),
+    Label(  'motorcycle'           , 32 ,       17 - 11, 'vehicle'         , 7       , True         , False        , (  0,  0,230) ),
+    Label(  'bicycle'              , 33 ,       18 - 11, 'vehicle'         , 7       , True         , False        , (119, 11, 32) ),
+    Label(  'license plate'        , -1 ,       -1 , 'vehicle'         , 7       , False        , True         , (  0,  0,142) ),
+]
+
+
+#--------------------------------------------------------------------------------
+# Create dictionaries for a fast lookup
+#--------------------------------------------------------------------------------
+
+# Please refer to the main method below for example usages!
+
+# name to label object
+name2label      = { label.name    : label for label in labels           }
+# id to label object
+id2label        = { label.id      : label for label in labels           }
+# trainId to label object
+trainId2label   = { label.trainId : label for label in reversed(labels) }
+# category to list of label objects
+category2labels = {}
+for label in labels:
+    category = label.category
+    if category in category2labels:
+        category2labels[category].append(label)
+    else:
+        category2labels[category] = [label]
+
+#--------------------------------------------------------------------------------
+# Assure single instance name
+#--------------------------------------------------------------------------------
+
+# returns the label name that describes a single instance (if possible)
+# e.g.     input     |   output
+#        ----------------------
+#          car       |   car
+#          cargroup  |   car
+#          foo       |   None
+#          foogroup  |   None
+#          skygroup  |   None
+def assureSingleInstanceName( name ):
+    # if the name is known, it is not a group
+    if name in name2label:
+        return name
+    # test if the name actually denotes a group
+    if not name.endswith("group"):
+        return None
+    # remove group
+    name = name[:-len("group")]
+    # test if the new name exists
+    if not name in name2label:
+        return None
+    # test if the new name denotes a label that actually has instances
+    if not name2label[name].hasInstances:
+        return None
+    # all good then
+    return name
+
+#--------------------------------------------------------------------------------
+# Main for testing
+#--------------------------------------------------------------------------------
+
+# just a dummy main
+if __name__ == "__main__":
+    # Print all the labels
+    print("List of cityscapes labels:")
+    print("")
+    print("    {:>21} | {:>3} | {:>7} | {:>14} | {:>10} | {:>12} | {:>12}".format( 'name', 'id', 'trainId', 'category', 'categoryId', 'hasInstances', 'ignoreInEval' ))
+    print("    " + ('-' * 98))
+    for label in labels:
+        print("    {:>21} | {:>3} | {:>7} | {:>14} | {:>10} | {:>12} | {:>12}".format( label.name, label.id, label.trainId, label.category, label.categoryId, label.hasInstances, label.ignoreInEval ))
+    print("")
+
+    print("Example usages:")
+
+    # Map from name to label
+    name = 'car'
+    id   = name2label[name].id
+    print("ID of label '{name}': {id}".format( name=name, id=id ))
+
+    # Map from ID to label
+    category = id2label[id].category
+    print("Category of label with ID '{id}': {category}".format( id=id, category=category ))
+
+    # Map from trainID to label
+    trainId = 0
+    name = trainId2label[trainId].name
+    print("Name of label with trainID '{id}': {name}".format( id=trainId, name=name ))

ext/cityscapes_scripts/helpers/labels_cityPersons.py

@@ -0,0 +1,61 @@
+#!/usr/bin/python
+#
+# CityPersons (cp) labels
+#
+
+from __future__ import print_function, absolute_import, division
+from collections import namedtuple
+
+
+#--------------------------------------------------------------------------------
+# Definitions
+#--------------------------------------------------------------------------------
+
+# a label and all meta information
+LabelCp = namedtuple( 'LabelCp' , [
+
+    'name'        , # The identifier of this label, e.g. 'pedestrian', 'rider', ... .
+                    # We use them to uniquely name a class
+
+    'id'          , # An integer ID that is associated with this label.
+                    # The IDs are used to represent the label in ground truth
+
+    'hasInstances', # Whether this label distinguishes between single instances or not
+
+    'ignoreInEval', # Whether pixels having this class as ground truth label are ignored
+                    # during evaluations or not
+
+    'color'       , # The color of this label
+    ] )
+
+
+#--------------------------------------------------------------------------------
+# A list of all labels
+#--------------------------------------------------------------------------------
+
+# The 'ignore' label covers representations of humans, e.g. people on posters, reflections etc.
+# Each annotation includes both the full bounding box (bbox) as well as a bounding box covering the visible area (bboxVis).
+# The latter is obtained automatically from the segmentation masks.  
+
+labelsCp = [
+    #         name                     id   hasInstances   ignoreInEval   color
+    LabelCp(  'ignore'               ,  0 , False        , True         , (250,170, 30) ),
+    LabelCp(  'pedestrian'           ,  1 , True         , False        , (220, 20, 60) ),
+    LabelCp(  'rider'                ,  2 , True         , False        , (  0,  0,142) ),
+    LabelCp(  'sitting person'       ,  3 , True         , False        , (107,142, 35) ),
+    LabelCp(  'person (other)'       ,  4 , True         , False        , (190,153,153) ),
+    LabelCp(  'person group'         ,  5 , False        , True         , (255,  0,  0) ),
+]
+
+
+#--------------------------------------------------------------------------------
+# Create dictionaries for a fast lookup
+#--------------------------------------------------------------------------------
+
+# Please refer to the main method below for example usages!
+
+# name to label object
+name2labelCp      = { label.name    : label for label in labelsCp }
+# id to label object
+id2labelCp        = { label.id      : label for label in labelsCp }
+

ext/cityscapes_scripts/helpers/version.py

@@ -0,0 +1,9 @@
+#!/usr/bin/env python
+
+import os
+
+with open(os.path.join(os.path.dirname(__file__), '..', 'VERSION')) as f:
+    version = f.read().strip()
+
+if __name__ == "__main__":
+    print(version)

ext/class_names/VIPSeg.py

@@ -0,0 +1,261 @@
+CLASSES = [
+    {"id": 0, "name": "wall", "isthing": 0, "color": [120, 120, 120]},
+    {"id": 1, "name": "ceiling", "isthing": 0, "color": [180, 120, 120]},
+    {"id": 2, "name": "door", "isthing": 1, "color": [6, 230, 230]},
+    {"id": 3, "name": "stair", "isthing": 0, "color": [80, 50, 50]},
+    {"id": 4, "name": "ladder", "isthing": 1, "color": [4, 200, 3]},
+    {"id": 5, "name": "escalator", "isthing": 0, "color": [120, 120, 80]},
+    {"id": 6, "name": "Playground_slide", "isthing": 0, "color": [140, 140, 140]},
+    {"id": 7, "name": "handrail_or_fence", "isthing": 0, "color": [204, 5, 255]},
+    {"id": 8, "name": "window", "isthing": 1, "color": [230, 230, 230]},
+    {"id": 9, "name": "rail", "isthing": 0, "color": [4, 250, 7]},
+    {"id": 10, "name": "goal", "isthing": 1, "color": [224, 5, 255]},
+    {"id": 11, "name": "pillar", "isthing": 0, "color": [235, 255, 7]},
+    {"id": 12, "name": "pole", "isthing": 0, "color": [150, 5, 61]},
+    {"id": 13, "name": "floor", "isthing": 0, "color": [120, 120, 70]},
+    {"id": 14, "name": "ground", "isthing": 0, "color": [8, 255, 51]},
+    {"id": 15, "name": "grass", "isthing": 0, "color": [255, 6, 82]},
+    {"id": 16, "name": "sand", "isthing": 0, "color": [143, 255, 140]},
+    {"id": 17, "name": "athletic_field", "isthing": 0, "color": [204, 255, 4]},
+    {"id": 18, "name": "road", "isthing": 0, "color": [255, 51, 7]},
+    {"id": 19, "name": "path", "isthing": 0, "color": [204, 70, 3]},
+    {"id": 20, "name": "crosswalk", "isthing": 0, "color": [0, 102, 200]},
+    {"id": 21, "name": "building", "isthing": 0, "color": [61, 230, 250]},
+    {"id": 22, "name": "house", "isthing": 0, "color": [255, 6, 51]},
+    {"id": 23, "name": "bridge", "isthing": 0, "color": [11, 102, 255]},
+    {"id": 24, "name": "tower", "isthing": 0, "color": [255, 7, 71]},
+    {"id": 25, "name": "windmill", "isthing": 0, "color": [255, 9, 224]},
+    {"id": 26, "name": "well_or_well_lid", "isthing": 0, "color": [9, 7, 230]},
+    {"id": 27, "name": "other_construction", "isthing": 0, "color": [220, 220, 220]},
+    {"id": 28, "name": "sky", "isthing": 0, "color": [255, 9, 92]},
+    {"id": 29, "name": "mountain", "isthing": 0, "color": [112, 9, 255]},
+    {"id": 30, "name": "stone", "isthing": 0, "color": [8, 255, 214]},
+    {"id": 31, "name": "wood", "isthing": 0, "color": [7, 255, 224]},
+    {"id": 32, "name": "ice", "isthing": 0, "color": [255, 184, 6]},
+    {"id": 33, "name": "snowfield", "isthing": 0, "color": [10, 255, 71]},
+    {"id": 34, "name": "grandstand", "isthing": 0, "color": [255, 41, 10]},
+    {"id": 35, "name": "sea", "isthing": 0, "color": [7, 255, 255]},
+    {"id": 36, "name": "river", "isthing": 0, "color": [224, 255, 8]},
+    {"id": 37, "name": "lake", "isthing": 0, "color": [102, 8, 255]},
+    {"id": 38, "name": "waterfall", "isthing": 0, "color": [255, 61, 6]},
+    {"id": 39, "name": "water", "isthing": 0, "color": [255, 194, 7]},
+    {"id": 40, "name": "billboard_or_Bulletin_Board", "isthing": 0, "color": [255, 122, 8]},
+    {"id": 41, "name": "sculpture", "isthing": 1, "color": [0, 255, 20]},
+    {"id": 42, "name": "pipeline", "isthing": 0, "color": [255, 8, 41]},
+    {"id": 43, "name": "flag", "isthing": 1, "color": [255, 5, 153]},
+    {"id": 44, "name": "parasol_or_umbrella", "isthing": 1, "color": [6, 51, 255]},
+    {"id": 45, "name": "cushion_or_carpet", "isthing": 0, "color": [235, 12, 255]},
+    {"id": 46, "name": "tent", "isthing": 1, "color": [160, 150, 20]},
+    {"id": 47, "name": "roadblock", "isthing": 1, "color": [0, 163, 255]},
+    {"id": 48, "name": "car", "isthing": 1, "color": [140, 140, 140]},
+    {"id": 49, "name": "bus", "isthing": 1, "color": [250, 10, 15]},
+    {"id": 50, "name": "truck", "isthing": 1, "color": [20, 255, 0]},
+    {"id": 51, "name": "bicycle", "isthing": 1, "color": [31, 255, 0]},
+    {"id": 52, "name": "motorcycle", "isthing": 1, "color": [255, 31, 0]},
+    {"id": 53, "name": "wheeled_machine", "isthing": 0, "color": [255, 224, 0]},
+    {"id": 54, "name": "ship_or_boat", "isthing": 1, "color": [153, 255, 0]},
+    {"id": 55, "name": "raft", "isthing": 1, "color": [0, 0, 255]},
+    {"id": 56, "name": "airplane", "isthing": 1, "color": [255, 71, 0]},
+    {"id": 57, "name": "tyre", "isthing": 0, "color": [0, 235, 255]},
+    {"id": 58, "name": "traffic_light", "isthing": 0, "color": [0, 173, 255]},
+    {"id": 59, "name": "lamp", "isthing": 0, "color": [31, 0, 255]},
+    {"id": 60, "name": "person", "isthing": 1, "color": [11, 200, 200]},
+    {"id": 61, "name": "cat", "isthing": 1, "color": [255, 82, 0]},
+    {"id": 62, "name": "dog", "isthing": 1, "color": [0, 255, 245]},
+    {"id": 63, "name": "horse", "isthing": 1, "color": [0, 61, 255]},
+    {"id": 64, "name": "cattle", "isthing": 1, "color": [0, 255, 112]},
+    {"id": 65, "name": "other_animal", "isthing": 1, "color": [0, 255, 133]},
+    {"id": 66, "name": "tree", "isthing": 0, "color": [255, 0, 0]},
+    {"id": 67, "name": "flower", "isthing": 0, "color": [255, 163, 0]},
+    {"id": 68, "name": "other_plant", "isthing": 0, "color": [255, 102, 0]},
+    {"id": 69, "name": "toy", "isthing": 0, "color": [194, 255, 0]},
+    {"id": 70, "name": "ball_net", "isthing": 0, "color": [0, 143, 255]},
+    {"id": 71, "name": "backboard", "isthing": 0, "color": [51, 255, 0]},
+    {"id": 72, "name": "skateboard", "isthing": 1, "color": [0, 82, 255]},
+    {"id": 73, "name": "bat", "isthing": 0, "color": [0, 255, 41]},
+    {"id": 74, "name": "ball", "isthing": 1, "color": [0, 255, 173]},
+    {"id": 75, "name": "cupboard_or_showcase_or_storage_rack", "isthing": 0, "color": [10, 0, 255]},
+    {"id": 76, "name": "box", "isthing": 1, "color": [173, 255, 0]},
+    {"id": 77, "name": "traveling_case_or_trolley_case", "isthing": 1, "color": [0, 255, 153]},
+    {"id": 78, "name": "basket", "isthing": 1, "color": [255, 92, 0]},
+    {"id": 79, "name": "bag_or_package", "isthing": 1, "color": [255, 0, 255]},
+    {"id": 80, "name": "trash_can", "isthing": 0, "color": [255, 0, 245]},
+    {"id": 81, "name": "cage", "isthing": 0, "color": [255, 0, 102]},
+    {"id": 82, "name": "plate", "isthing": 1, "color": [255, 173, 0]},
+    {"id": 83, "name": "tub_or_bowl_or_pot", "isthing": 1, "color": [255, 0, 20]},
+    {"id": 84, "name": "bottle_or_cup", "isthing": 1, "color": [255, 184, 184]},
+    {"id": 85, "name": "barrel", "isthing": 1, "color": [0, 31, 255]},
+    {"id": 86, "name": "fishbowl", "isthing": 1, "color": [0, 255, 61]},
+    {"id": 87, "name": "bed", "isthing": 1, "color": [0, 71, 255]},
+    {"id": 88, "name": "pillow", "isthing": 1, "color": [255, 0, 204]},
+    {"id": 89, "name": "table_or_desk", "isthing": 1, "color": [0, 255, 194]},
+    {"id": 90, "name": "chair_or_seat", "isthing": 1, "color": [0, 255, 82]},
+    {"id": 91, "name": "bench", "isthing": 1, "color": [0, 10, 255]},
+    {"id": 92, "name": "sofa", "isthing": 1, "color": [0, 112, 255]},
+    {"id": 93, "name": "shelf", "isthing": 0, "color": [51, 0, 255]},
+    {"id": 94, "name": "bathtub", "isthing": 0, "color": [0, 194, 255]},
+    {"id": 95, "name": "gun", "isthing": 1, "color": [0, 122, 255]},
+    {"id": 96, "name": "commode", "isthing": 1, "color": [0, 255, 163]},
+    {"id": 97, "name": "roaster", "isthing": 1, "color": [255, 153, 0]},
+    {"id": 98, "name": "other_machine", "isthing": 0, "color": [0, 255, 10]},
+    {"id": 99, "name": "refrigerator", "isthing": 1, "color": [255, 112, 0]},
+    {"id": 100, "name": "washing_machine", "isthing": 1, "color": [143, 255, 0]},
+    {"id": 101, "name": "Microwave_oven", "isthing": 1, "color": [82, 0, 255]},
+    {"id": 102, "name": "fan", "isthing": 1, "color": [163, 255, 0]},
+    {"id": 103, "name": "curtain", "isthing": 0, "color": [255, 235, 0]},
+    {"id": 104, "name": "textiles", "isthing": 0, "color": [8, 184, 170]},
+    {"id": 105, "name": "clothes", "isthing": 0, "color": [133, 0, 255]},
+    {"id": 106, "name": "painting_or_poster", "isthing": 1, "color": [0, 255, 92]},
+    {"id": 107, "name": "mirror", "isthing": 1, "color": [184, 0, 255]},
+    {"id": 108, "name": "flower_pot_or_vase", "isthing": 1, "color": [255, 0, 31]},
+    {"id": 109, "name": "clock", "isthing": 1, "color": [0, 184, 255]},
+    {"id": 110, "name": "book", "isthing": 0, "color": [0, 214, 255]},
+    {"id": 111, "name": "tool", "isthing": 0, "color": [255, 0, 112]},
+    {"id": 112, "name": "blackboard", "isthing": 0, "color": [92, 255, 0]},
+    {"id": 113, "name": "tissue", "isthing": 0, "color": [0, 224, 255]},
+    {"id": 114, "name": "screen_or_television", "isthing": 1, "color": [112, 224, 255]},
+    {"id": 115, "name": "computer", "isthing": 1, "color": [70, 184, 160]},
+    {"id": 116, "name": "printer", "isthing": 1, "color": [163, 0, 255]},
+    {"id": 117, "name": "Mobile_phone", "isthing": 1, "color": [153, 0, 255]},
+    {"id": 118, "name": "keyboard", "isthing": 1, "color": [71, 255, 0]},
+    {"id": 119, "name": "other_electronic_product", "isthing": 0, "color": [255, 0, 163]},
+    {"id": 120, "name": "fruit", "isthing": 0, "color": [255, 204, 0]},
+    {"id": 121, "name": "food", "isthing": 0, "color": [255, 0, 143]},
+    {"id": 122, "name": "instrument", "isthing": 1, "color": [0, 255, 235]},
+    {"id": 123, "name": "train", "isthing": 1, "color": [133, 255, 0]}
+]
+
+CLASSES_THING = [
+    {'id': 2, 'name': 'door', 'isthing': 1, 'color': [6, 230, 230]},
+    {'id': 4, 'name': 'ladder', 'isthing': 1, 'color': [4, 200, 3]},
+    {'id': 8, 'name': 'window', 'isthing': 1, 'color': [230, 230, 230]},
+    {'id': 10, 'name': 'goal', 'isthing': 1, 'color': [224, 5, 255]},
+    {'id': 41, 'name': 'sculpture', 'isthing': 1, 'color': [0, 255, 20]},
+    {'id': 43, 'name': 'flag', 'isthing': 1, 'color': [255, 5, 153]},
+    {'id': 44, 'name': 'parasol_or_umbrella', 'isthing': 1, 'color': [6, 51, 255]},
+    {'id': 46, 'name': 'tent', 'isthing': 1, 'color': [160, 150, 20]},
+    {'id': 47, 'name': 'roadblock', 'isthing': 1, 'color': [0, 163, 255]},
+    {'id': 48, 'name': 'car', 'isthing': 1, 'color': [140, 140, 140]},
+    {'id': 49, 'name': 'bus', 'isthing': 1, 'color': [250, 10, 15]},
+    {'id': 50, 'name': 'truck', 'isthing': 1, 'color': [20, 255, 0]},
+    {'id': 51, 'name': 'bicycle', 'isthing': 1, 'color': [31, 255, 0]},
+    {'id': 52, 'name': 'motorcycle', 'isthing': 1, 'color': [255, 31, 0]},
+    {'id': 54, 'name': 'ship_or_boat', 'isthing': 1, 'color': [153, 255, 0]},
+    {'id': 55, 'name': 'raft', 'isthing': 1, 'color': [0, 0, 255]},
+    {'id': 56, 'name': 'airplane', 'isthing': 1, 'color': [255, 71, 0]},
+    {'id': 60, 'name': 'person', 'isthing': 1, 'color': [11, 200, 200]},
+    {'id': 61, 'name': 'cat', 'isthing': 1, 'color': [255, 82, 0]},
+    {'id': 62, 'name': 'dog', 'isthing': 1, 'color': [0, 255, 245]},
+    {'id': 63, 'name': 'horse', 'isthing': 1, 'color': [0, 61, 255]},
+    {'id': 64, 'name': 'cattle', 'isthing': 1, 'color': [0, 255, 112]},
+    {'id': 65, 'name': 'other_animal', 'isthing': 1, 'color': [0, 255, 133]},
+    {'id': 72, 'name': 'skateboard', 'isthing': 1, 'color': [0, 82, 255]},
+    {'id': 74, 'name': 'ball', 'isthing': 1, 'color': [0, 255, 173]},
+    {'id': 76, 'name': 'box', 'isthing': 1, 'color': [173, 255, 0]},
+    {'id': 77, 'name': 'traveling_case_or_trolley_case', 'isthing': 1, 'color': [0, 255, 153]},
+    {'id': 78, 'name': 'basket', 'isthing': 1, 'color': [255, 92, 0]},
+    {'id': 79, 'name': 'bag_or_package', 'isthing': 1, 'color': [255, 0, 255]},
+    {'id': 82, 'name': 'plate', 'isthing': 1, 'color': [255, 173, 0]},
+    {'id': 83, 'name': 'tub_or_bowl_or_pot', 'isthing': 1, 'color': [255, 0, 20]},
+    {'id': 84, 'name': 'bottle_or_cup', 'isthing': 1, 'color': [255, 184, 184]},
+    {'id': 85, 'name': 'barrel', 'isthing': 1, 'color': [0, 31, 255]},
+    {'id': 86, 'name': 'fishbowl', 'isthing': 1, 'color': [0, 255, 61]},
+    {'id': 87, 'name': 'bed', 'isthing': 1, 'color': [0, 71, 255]},
+    {'id': 88, 'name': 'pillow', 'isthing': 1, 'color': [255, 0, 204]},
+    {'id': 89, 'name': 'table_or_desk', 'isthing': 1, 'color': [0, 255, 194]},
+    {'id': 90, 'name': 'chair_or_seat', 'isthing': 1, 'color': [0, 255, 82]},
+    {'id': 91, 'name': 'bench', 'isthing': 1, 'color': [0, 10, 255]},
+    {'id': 92, 'name': 'sofa', 'isthing': 1, 'color': [0, 112, 255]},
+    {'id': 95, 'name': 'gun', 'isthing': 1, 'color': [0, 122, 255]},
+    {'id': 96, 'name': 'commode', 'isthing': 1, 'color': [0, 255, 163]},
+    {'id': 97, 'name': 'roaster', 'isthing': 1, 'color': [255, 153, 0]},
+    {'id': 99, 'name': 'refrigerator', 'isthing': 1, 'color': [255, 112, 0]},
+    {'id': 100, 'name': 'washing_machine', 'isthing': 1, 'color': [143, 255, 0]},
+    {'id': 101, 'name': 'Microwave_oven', 'isthing': 1, 'color': [82, 0, 255]},
+    {'id': 102, 'name': 'fan', 'isthing': 1, 'color': [163, 255, 0]},
+    {'id': 106, 'name': 'painting_or_poster', 'isthing': 1, 'color': [0, 255, 92]},
+    {'id': 107, 'name': 'mirror', 'isthing': 1, 'color': [184, 0, 255]},
+    {'id': 108, 'name': 'flower_pot_or_vase', 'isthing': 1, 'color': [255, 0, 31]},
+    {'id': 109, 'name': 'clock', 'isthing': 1, 'color': [0, 184, 255]},
+    {'id': 114, 'name': 'screen_or_television', 'isthing': 1, 'color': [112, 224, 255]},
+    {'id': 115, 'name': 'computer', 'isthing': 1, 'color': [70, 184, 160]},
+    {'id': 116, 'name': 'printer', 'isthing': 1, 'color': [163, 0, 255]},
+    {'id': 117, 'name': 'Mobile_phone', 'isthing': 1, 'color': [153, 0, 255]},
+    {'id': 118, 'name': 'keyboard', 'isthing': 1, 'color': [71, 255, 0]},
+    {'id': 122, 'name': 'instrument', 'isthing': 1, 'color': [0, 255, 235]},
+    {'id': 123, 'name': 'train', 'isthing': 1, 'color': [133, 255, 0]}
+]
+
+CLASSES_STUFF = [
+    {'id': 0, 'name': 'wall', 'isthing': 0, 'color': [120, 120, 120]},
+    {'id': 1, 'name': 'ceiling', 'isthing': 0, 'color': [180, 120, 120]},
+    {'id': 3, 'name': 'stair', 'isthing': 0, 'color': [80, 50, 50]},
+    {'id': 5, 'name': 'escalator', 'isthing': 0, 'color': [120, 120, 80]},
+    {'id': 6, 'name': 'Playground_slide', 'isthing': 0, 'color': [140, 140, 140]},
+    {'id': 7, 'name': 'handrail_or_fence', 'isthing': 0, 'color': [204, 5, 255]},
+    {'id': 9, 'name': 'rail', 'isthing': 0, 'color': [4, 250, 7]},
+    {'id': 11, 'name': 'pillar', 'isthing': 0, 'color': [235, 255, 7]},
+    {'id': 12, 'name': 'pole', 'isthing': 0, 'color': [150, 5, 61]},
+    {'id': 13, 'name': 'floor', 'isthing': 0, 'color': [120, 120, 70]},
+    {'id': 14, 'name': 'ground', 'isthing': 0, 'color': [8, 255, 51]},
+    {'id': 15, 'name': 'grass', 'isthing': 0, 'color': [255, 6, 82]},
+    {'id': 16, 'name': 'sand', 'isthing': 0, 'color': [143, 255, 140]},
+    {'id': 17, 'name': 'athletic_field', 'isthing': 0, 'color': [204, 255, 4]},
+    {'id': 18, 'name': 'road', 'isthing': 0, 'color': [255, 51, 7]},
+    {'id': 19, 'name': 'path', 'isthing': 0, 'color': [204, 70, 3]},
+    {'id': 20, 'name': 'crosswalk', 'isthing': 0, 'color': [0, 102, 200]},
+    {'id': 21, 'name': 'building', 'isthing': 0, 'color': [61, 230, 250]},
+    {'id': 22, 'name': 'house', 'isthing': 0, 'color': [255, 6, 51]},
+    {'id': 23, 'name': 'bridge', 'isthing': 0, 'color': [11, 102, 255]},
+    {'id': 24, 'name': 'tower', 'isthing': 0, 'color': [255, 7, 71]},
+    {'id': 25, 'name': 'windmill', 'isthing': 0, 'color': [255, 9, 224]},
+    {'id': 26, 'name': 'well_or_well_lid', 'isthing': 0, 'color': [9, 7, 230]},
+    {'id': 27, 'name': 'other_construction', 'isthing': 0, 'color': [220, 220, 220]},
+    {'id': 28, 'name': 'sky', 'isthing': 0, 'color': [255, 9, 92]},
+    {'id': 29, 'name': 'mountain', 'isthing': 0, 'color': [112, 9, 255]},
+    {'id': 30, 'name': 'stone', 'isthing': 0, 'color': [8, 255, 214]},
+    {'id': 31, 'name': 'wood', 'isthing': 0, 'color': [7, 255, 224]},
+    {'id': 32, 'name': 'ice', 'isthing': 0, 'color': [255, 184, 6]},
+    {'id': 33, 'name': 'snowfield', 'isthing': 0, 'color': [10, 255, 71]},
+    {'id': 34, 'name': 'grandstand', 'isthing': 0, 'color': [255, 41, 10]},
+    {'id': 35, 'name': 'sea', 'isthing': 0, 'color': [7, 255, 255]},
+    {'id': 36, 'name': 'river', 'isthing': 0, 'color': [224, 255, 8]},
+    {'id': 37, 'name': 'lake', 'isthing': 0, 'color': [102, 8, 255]},
+    {'id': 38, 'name': 'waterfall', 'isthing': 0, 'color': [255, 61, 6]},
+    {'id': 39, 'name': 'water', 'isthing': 0, 'color': [255, 194, 7]},
+    {'id': 40, 'name': 'billboard_or_Bulletin_Board', 'isthing': 0, 'color': [255, 122, 8]},
+    {'id': 42, 'name': 'pipeline', 'isthing': 0, 'color': [255, 8, 41]},
+    {'id': 45, 'name': 'cushion_or_carpet', 'isthing': 0, 'color': [235, 12, 255]},
+    {'id': 53, 'name': 'wheeled_machine', 'isthing': 0, 'color': [255, 224, 0]},
+    {'id': 57, 'name': 'tyre', 'isthing': 0, 'color': [0, 235, 255]},
+    {'id': 58, 'name': 'traffic_light', 'isthing': 0, 'color': [0, 173, 255]},
+    {'id': 59, 'name': 'lamp', 'isthing': 0, 'color': [31, 0, 255]},
+    {'id': 66, 'name': 'tree', 'isthing': 0, 'color': [255, 0, 0]},
+    {'id': 67, 'name': 'flower', 'isthing': 0, 'color': [255, 163, 0]},
+    {'id': 68, 'name': 'other_plant', 'isthing': 0, 'color': [255, 102, 0]},
+    {'id': 69, 'name': 'toy', 'isthing': 0, 'color': [194, 255, 0]},
+    {'id': 70, 'name': 'ball_net', 'isthing': 0, 'color': [0, 143, 255]},
+    {'id': 71, 'name': 'backboard', 'isthing': 0, 'color': [51, 255, 0]},
+    {'id': 73, 'name': 'bat', 'isthing': 0, 'color': [0, 255, 41]},
+    {'id': 75, 'name': 'cupboard_or_showcase_or_storage_rack', 'isthing': 0, 'color': [10, 0, 255]},
+    {'id': 80, 'name': 'trash_can', 'isthing': 0, 'color': [255, 0, 245]},
+    {'id': 81, 'name': 'cage', 'isthing': 0, 'color': [255, 0, 102]},
+    {'id': 93, 'name': 'shelf', 'isthing': 0, 'color': [51, 0, 255]},
+    {'id': 94, 'name': 'bathtub', 'isthing': 0, 'color': [0, 194, 255]},
+    {'id': 98, 'name': 'other_machine', 'isthing': 0, 'color': [0, 255, 10]},
+    {'id': 103, 'name': 'curtain', 'isthing': 0, 'color': [255, 235, 0]},
+    {'id': 104, 'name': 'textiles', 'isthing': 0, 'color': [8, 184, 170]},
+    {'id': 105, 'name': 'clothes', 'isthing': 0, 'color': [133, 0, 255]},
+    {'id': 110, 'name': 'book', 'isthing': 0, 'color': [0, 214, 255]},
+    {'id': 111, 'name': 'tool', 'isthing': 0, 'color': [255, 0, 112]},
+    {'id': 112, 'name': 'blackboard', 'isthing': 0, 'color': [92, 255, 0]},
+    {'id': 113, 'name': 'tissue', 'isthing': 0, 'color': [0, 224, 255]},
+    {'id': 119, 'name': 'other_electronic_product', 'isthing': 0, 'color': [255, 0, 163]},
+    {'id': 120, 'name': 'fruit', 'isthing': 0, 'color': [255, 204, 0]},
+    {'id': 121, 'name': 'food', 'isthing': 0, 'color': [255, 0, 143]}
+]
+
+COCO_THINGS = [itm['name'] for itm in CLASSES_THING]
+COCO_STUFF = [itm['name'] for itm in CLASSES_STUFF]
+COCO_CLASSES = [*COCO_THINGS, *COCO_STUFF]
+PALETTE = [*[itm['color'] for itm in CLASSES_THING], *[itm['color'] for itm in CLASSES_STUFF]]

ext/davis2017/__init__.py

@@ -0,0 +1,3 @@
+from __future__ import absolute_import
+
+__version__ = '0.1.0'

ext/davis2017/davis.py

@@ -0,0 +1,122 @@
+import os
+from glob import glob
+from collections import defaultdict
+import numpy as np
+from PIL import Image
+
+
+class DAVIS(object):
+    SUBSET_OPTIONS = ['train', 'val', 'test-dev', 'test-challenge']
+    TASKS = ['semi-supervised', 'unsupervised']
+    DATASET_WEB = 'https://davischallenge.org/davis2017/code.html'
+    VOID_LABEL = 255
+
+    def __init__(self, root, task='unsupervised', subset='val', sequences='all', resolution='480p', codalab=False):
+        """
+        Class to read the DAVIS dataset
+        :param root: Path to the DAVIS folder that contains JPEGImages, Annotations, etc. folders.
+        :param task: Task to load the annotations, choose between semi-supervised or unsupervised.
+        :param subset: Set to load the annotations
+        :param sequences: Sequences to consider, 'all' to use all the sequences in a set.
+        :param resolution: Specify the resolution to use the dataset, choose between '480' and 'Full-Resolution'
+        """
+        if subset not in self.SUBSET_OPTIONS:
+            raise ValueError(f'Subset should be in {self.SUBSET_OPTIONS}')
+        if task not in self.TASKS:
+            raise ValueError(f'The only tasks that are supported are {self.TASKS}')
+
+        self.task = task
+        self.subset = subset
+        self.root = root
+        self.img_path = os.path.join(self.root, 'JPEGImages', resolution)
+        annotations_folder = 'Annotations' if task == 'semi-supervised' else 'Annotations_unsupervised'
+        self.mask_path = os.path.join(self.root, annotations_folder, resolution)
+        year = '2019' if task == 'unsupervised' and (subset == 'test-dev' or subset == 'test-challenge') else '2017'
+        self.imagesets_path = os.path.join(self.root, 'ImageSets', year)
+
+        self._check_directories()
+
+        if sequences == 'all':
+            with open(os.path.join(self.imagesets_path, f'{self.subset}.txt'), 'r') as f:
+                tmp = f.readlines()
+            sequences_names = [x.strip() for x in tmp]
+        else:
+            sequences_names = sequences if isinstance(sequences, list) else [sequences]
+        self.sequences = defaultdict(dict)
+
+        for seq in sequences_names:
+            images = np.sort(glob(os.path.join(self.img_path, seq, '*.jpg'))).tolist()
+            if len(images) == 0 and not codalab:
+                raise FileNotFoundError(f'Images for sequence {seq} not found.')
+            self.sequences[seq]['images'] = images
+            masks = np.sort(glob(os.path.join(self.mask_path, seq, '*.png'))).tolist()
+            masks.extend([-1] * (len(images) - len(masks)))
+            self.sequences[seq]['masks'] = masks
+
+    def _check_directories(self):
+        if not os.path.exists(self.root):
+            raise FileNotFoundError(f'DAVIS not found in the specified directory, download it from {self.DATASET_WEB}')
+        if not os.path.exists(os.path.join(self.imagesets_path, f'{self.subset}.txt')):
+            raise FileNotFoundError(f'Subset sequences list for {self.subset} not found, download the missing subset '
+                                    f'for the {self.task} task from {self.DATASET_WEB}')
+        if self.subset in ['train', 'val'] and not os.path.exists(self.mask_path):
+            raise FileNotFoundError(f'Annotations folder for the {self.task} task not found, download it from {self.DATASET_WEB}')
+
+    def get_frames(self, sequence):
+        for img, msk in zip(self.sequences[sequence]['images'], self.sequences[sequence]['masks']):
+            image = np.array(Image.open(img))
+            mask = None if msk is None else np.array(Image.open(msk))
+            yield image, mask
+
+    def _get_all_elements(self, sequence, obj_type):
+        obj = np.array(Image.open(self.sequences[sequence][obj_type][0]))
+        all_objs = np.zeros((len(self.sequences[sequence][obj_type]), *obj.shape))
+        obj_id = []
+        for i, obj in enumerate(self.sequences[sequence][obj_type]):
+            all_objs[i, ...] = np.array(Image.open(obj))
+            obj_id.append(''.join(obj.split('/')[-1].split('.')[:-1]))
+        return all_objs, obj_id
+
+    def get_all_images(self, sequence):
+        return self._get_all_elements(sequence, 'images')
+
+    def get_all_masks(self, sequence, separate_objects_masks=False):
+        masks, masks_id = self._get_all_elements(sequence, 'masks')
+        masks_void = np.zeros_like(masks)
+
+        # Separate void and object masks
+        for i in range(masks.shape[0]):
+            masks_void[i, ...] = masks[i, ...] == 255
+            masks[i, masks[i, ...] == 255] = 0
+
+        if separate_objects_masks:
+            num_objects = int(np.max(masks[0, ...]))
+            tmp = np.ones((num_objects, *masks.shape))
+            tmp = tmp * np.arange(1, num_objects + 1)[:, None, None, None]
+            masks = (tmp == masks[None, ...])
+            masks = masks > 0
+        return masks, masks_void, masks_id
+
+    def get_sequences(self):
+        for seq in self.sequences:
+            yield seq
+
+
+if __name__ == '__main__':
+    from matplotlib import pyplot as plt
+
+    only_first_frame = True
+    subsets = ['train', 'val']
+
+    for s in subsets:
+        dataset = DAVIS(root='/home/csergi/scratch2/Databases/DAVIS2017_private', subset=s)
+        for seq in dataset.get_sequences():
+            g = dataset.get_frames(seq)
+            img, mask = next(g)
+            plt.subplot(2, 1, 1)
+            plt.title(seq)
+            plt.imshow(img)
+            plt.subplot(2, 1, 2)
+            plt.imshow(mask)
+            plt.show(block=True)
+

ext/davis2017/evaluation.py

@@ -0,0 +1,110 @@
+import sys
+from tqdm import tqdm
+import warnings
+warnings.filterwarnings("ignore", category=RuntimeWarning)
+
+import numpy as np
+from ext.davis2017.davis import DAVIS
+from ext.davis2017.metrics import db_eval_boundary, db_eval_iou
+from ext.davis2017 import utils
+from ext.davis2017.results import Results
+from scipy.optimize import linear_sum_assignment
+
+
+class DAVISEvaluation(object):
+    def __init__(self, davis_root, task, gt_set, sequences='all', codalab=False):
+        """
+        Class to evaluate DAVIS sequences from a certain set and for a certain task
+        :param davis_root: Path to the DAVIS folder that contains JPEGImages, Annotations, etc. folders.
+        :param task: Task to compute the evaluation, chose between semi-supervised or unsupervised.
+        :param gt_set: Set to compute the evaluation
+        :param sequences: Sequences to consider for the evaluation, 'all' to use all the sequences in a set.
+        """
+        self.davis_root = davis_root
+        self.task = task
+        self.dataset = DAVIS(root=davis_root, task=task, subset=gt_set, sequences=sequences, codalab=codalab)
+
+    @staticmethod
+    def _evaluate_semisupervised(all_gt_masks, all_res_masks, all_void_masks, metric):
+        if all_res_masks.shape[0] > all_gt_masks.shape[0]:
+            sys.stdout.write("\nIn your PNG files there is an index higher than the number of objects in the sequence!")
+            sys.exit()
+        elif all_res_masks.shape[0] < all_gt_masks.shape[0]:
+            zero_padding = np.zeros((all_gt_masks.shape[0] - all_res_masks.shape[0], *all_res_masks.shape[1:]))
+            all_res_masks = np.concatenate([all_res_masks, zero_padding], axis=0)
+        j_metrics_res, f_metrics_res = np.zeros(all_gt_masks.shape[:2]), np.zeros(all_gt_masks.shape[:2])
+        for ii in range(all_gt_masks.shape[0]):
+            if 'J' in metric:
+                j_metrics_res[ii, :] = db_eval_iou(all_gt_masks[ii, ...], all_res_masks[ii, ...], all_void_masks)
+            if 'F' in metric:
+                f_metrics_res[ii, :] = db_eval_boundary(all_gt_masks[ii, ...], all_res_masks[ii, ...], all_void_masks)
+        return j_metrics_res, f_metrics_res
+
+    @staticmethod
+    def _evaluate_unsupervised(all_gt_masks, all_res_masks, all_void_masks, metric, max_n_proposals=20):
+        if all_res_masks.shape[0] > max_n_proposals:
+            sys.stdout.write(f"\nIn your PNG files there is an index higher than the maximum number ({max_n_proposals}) of proposals allowed!")
+            sys.exit()
+        elif all_res_masks.shape[0] < all_gt_masks.shape[0]:
+            zero_padding = np.zeros((all_gt_masks.shape[0] - all_res_masks.shape[0], *all_res_masks.shape[1:]))
+            all_res_masks = np.concatenate([all_res_masks, zero_padding], axis=0)
+        j_metrics_res = np.zeros((all_res_masks.shape[0], all_gt_masks.shape[0], all_gt_masks.shape[1]))
+        f_metrics_res = np.zeros((all_res_masks.shape[0], all_gt_masks.shape[0], all_gt_masks.shape[1]))
+        for ii in range(all_gt_masks.shape[0]):
+            for jj in range(all_res_masks.shape[0]):
+                if 'J' in metric:
+                    j_metrics_res[jj, ii, :] = db_eval_iou(all_gt_masks[ii, ...], all_res_masks[jj, ...], all_void_masks)
+                if 'F' in metric:
+                    f_metrics_res[jj, ii, :] = db_eval_boundary(all_gt_masks[ii, ...], all_res_masks[jj, ...], all_void_masks)
+        if 'J' in metric and 'F' in metric:
+            all_metrics = (np.mean(j_metrics_res, axis=2) + np.mean(f_metrics_res, axis=2)) / 2
+        else:
+            all_metrics = np.mean(j_metrics_res, axis=2) if 'J' in metric else np.mean(f_metrics_res, axis=2)
+        row_ind, col_ind = linear_sum_assignment(-all_metrics)
+        return j_metrics_res[row_ind, col_ind, :], f_metrics_res[row_ind, col_ind, :]
+
+    def evaluate(self, res_path, metric=('J', 'F'), debug=False):
+        metric = metric if isinstance(metric, tuple) or isinstance(metric, list) else [metric]
+        if 'T' in metric:
+            raise ValueError('Temporal metric not supported!')
+        if 'J' not in metric and 'F' not in metric:
+            raise ValueError('Metric possible values are J for IoU or F for Boundary')
+
+        # Containers
+        metrics_res = {}
+        if 'J' in metric:
+            metrics_res['J'] = {"M": [], "R": [], "D": [], "M_per_object": {}}
+        if 'F' in metric:
+            metrics_res['F'] = {"M": [], "R": [], "D": [], "M_per_object": {}}
+
+        # Sweep all sequences
+        results = Results(root_dir=res_path)
+        for seq in tqdm(list(self.dataset.get_sequences())):
+            all_gt_masks, all_void_masks, all_masks_id = self.dataset.get_all_masks(seq, True)
+            if self.task == 'semi-supervised':
+                all_gt_masks, all_masks_id = all_gt_masks[:, 1:-1, :, :], all_masks_id[1:-1]
+            all_res_masks = results.read_masks(seq, all_masks_id)
+            if self.task == 'unsupervised':
+                j_metrics_res, f_metrics_res = self._evaluate_unsupervised(all_gt_masks, all_res_masks, all_void_masks, metric)
+            elif self.task == 'semi-supervised':
+                j_metrics_res, f_metrics_res = self._evaluate_semisupervised(all_gt_masks, all_res_masks, None, metric)
+            for ii in range(all_gt_masks.shape[0]):
+                seq_name = f'{seq}_{ii+1}'
+                if 'J' in metric:
+                    [JM, JR, JD] = utils.db_statistics(j_metrics_res[ii])
+                    metrics_res['J']["M"].append(JM)
+                    metrics_res['J']["R"].append(JR)
+                    metrics_res['J']["D"].append(JD)
+                    metrics_res['J']["M_per_object"][seq_name] = JM
+                if 'F' in metric:
+                    [FM, FR, FD] = utils.db_statistics(f_metrics_res[ii])
+                    metrics_res['F']["M"].append(FM)
+                    metrics_res['F']["R"].append(FR)
+                    metrics_res['F']["D"].append(FD)
+                    metrics_res['F']["M_per_object"][seq_name] = FM
+
+            # Show progress
+            if debug:
+                sys.stdout.write(seq + '\n')
+                sys.stdout.flush()
+        return metrics_res

ext/davis2017/metrics.py

@@ -0,0 +1,197 @@
+import math
+import numpy as np
+import cv2
+
+
+def db_eval_iou(annotation, segmentation, void_pixels=None):
+    """ Compute region similarity as the Jaccard Index.
+    Arguments:
+        annotation   (ndarray): binary annotation   map.
+        segmentation (ndarray): binary segmentation map.
+        void_pixels  (ndarray): optional mask with void pixels
+
+    Return:
+        jaccard (float): region similarity
+    """
+    assert annotation.shape == segmentation.shape, \
+        f'Annotation({annotation.shape}) and segmentation:{segmentation.shape} dimensions do not match.'
+    annotation = annotation.astype(bool)
+    segmentation = segmentation.astype(bool)
+
+    if void_pixels is not None:
+        assert annotation.shape == void_pixels.shape, \
+            f'Annotation({annotation.shape}) and void pixels:{void_pixels.shape} dimensions do not match.'
+        void_pixels = void_pixels.astype(bool)
+    else:
+        void_pixels = np.zeros_like(segmentation)
+
+    # Intersection between all sets
+    inters = np.sum((segmentation & annotation) & np.logical_not(void_pixels), axis=(-2, -1))
+    union = np.sum((segmentation | annotation) & np.logical_not(void_pixels), axis=(-2, -1))
+
+    j = inters / union
+    if j.ndim == 0:
+        j = 1 if np.isclose(union, 0) else j
+    else:
+        j[np.isclose(union, 0)] = 1
+    return j
+
+
+def db_eval_boundary(annotation, segmentation, void_pixels=None, bound_th=0.008):
+    assert annotation.shape == segmentation.shape
+    if void_pixels is not None:
+        assert annotation.shape == void_pixels.shape
+    if annotation.ndim == 3:
+        n_frames = annotation.shape[0]
+        f_res = np.zeros(n_frames)
+        for frame_id in range(n_frames):
+            void_pixels_frame = None if void_pixels is None else void_pixels[frame_id, :, :, ]
+            f_res[frame_id] = f_measure(segmentation[frame_id, :, :, ], annotation[frame_id, :, :], void_pixels_frame, bound_th=bound_th)
+    elif annotation.ndim == 2:
+        f_res = f_measure(segmentation, annotation, void_pixels, bound_th=bound_th)
+    else:
+        raise ValueError(f'db_eval_boundary does not support tensors with {annotation.ndim} dimensions')
+    return f_res
+
+
+def f_measure(foreground_mask, gt_mask, void_pixels=None, bound_th=0.008):
+    """
+    Compute mean,recall and decay from per-frame evaluation.
+    Calculates precision/recall for boundaries between foreground_mask and
+    gt_mask using morphological operators to speed it up.
+
+    Arguments:
+        foreground_mask (ndarray): binary segmentation image.
+        gt_mask         (ndarray): binary annotated image.
+        void_pixels     (ndarray): optional mask with void pixels
+
+    Returns:
+        F (float): boundaries F-measure
+    """
+    assert np.atleast_3d(foreground_mask).shape[2] == 1
+    if void_pixels is not None:
+        void_pixels = void_pixels.astype(bool)
+    else:
+        void_pixels = np.zeros_like(foreground_mask).astype(bool)
+
+    bound_pix = bound_th if bound_th >= 1 else \
+        np.ceil(bound_th * np.linalg.norm(foreground_mask.shape))
+
+    # Get the pixel boundaries of both masks
+    fg_boundary = _seg2bmap(foreground_mask * np.logical_not(void_pixels))
+    gt_boundary = _seg2bmap(gt_mask * np.logical_not(void_pixels))
+
+    from skimage.morphology import disk
+
+    # fg_dil = binary_dilation(fg_boundary, disk(bound_pix))
+    fg_dil = cv2.dilate(fg_boundary.astype(np.uint8), disk(bound_pix).astype(np.uint8))
+    # gt_dil = binary_dilation(gt_boundary, disk(bound_pix))
+    gt_dil = cv2.dilate(gt_boundary.astype(np.uint8), disk(bound_pix).astype(np.uint8))
+
+    # Get the intersection
+    gt_match = gt_boundary * fg_dil
+    fg_match = fg_boundary * gt_dil
+
+    # Area of the intersection
+    n_fg = np.sum(fg_boundary)
+    n_gt = np.sum(gt_boundary)
+
+    # % Compute precision and recall
+    if n_fg == 0 and n_gt > 0:
+        precision = 1
+        recall = 0
+    elif n_fg > 0 and n_gt == 0:
+        precision = 0
+        recall = 1
+    elif n_fg == 0 and n_gt == 0:
+        precision = 1
+        recall = 1
+    else:
+        precision = np.sum(fg_match) / float(n_fg)
+        recall = np.sum(gt_match) / float(n_gt)
+
+    # Compute F measure
+    if precision + recall == 0:
+        F = 0
+    else:
+        F = 2 * precision * recall / (precision + recall)
+
+    return F
+
+
+def _seg2bmap(seg, width=None, height=None):
+    """
+    From a segmentation, compute a binary boundary map with 1 pixel wide
+    boundaries.  The boundary pixels are offset by 1/2 pixel towards the
+    origin from the actual segment boundary.
+    Arguments:
+        seg     : Segments labeled from 1..k.
+        width	  :	Width of desired bmap  <= seg.shape[1]
+        height  :	Height of desired bmap <= seg.shape[0]
+    Returns:
+        bmap (ndarray):	Binary boundary map.
+     David Martin <[email protected]>
+     January 2003
+    """
+
+    seg = seg.astype(bool)
+    seg[seg > 0] = 1
+
+    assert np.atleast_3d(seg).shape[2] == 1
+
+    width = seg.shape[1] if width is None else width
+    height = seg.shape[0] if height is None else height
+
+    h, w = seg.shape[:2]
+
+    ar1 = float(width) / float(height)
+    ar2 = float(w) / float(h)
+
+    assert not (
+        width > w | height > h | abs(ar1 - ar2) > 0.01
+    ), "Can" "t convert %dx%d seg to %dx%d bmap." % (w, h, width, height)
+
+    e = np.zeros_like(seg)
+    s = np.zeros_like(seg)
+    se = np.zeros_like(seg)
+
+    e[:, :-1] = seg[:, 1:]
+    s[:-1, :] = seg[1:, :]
+    se[:-1, :-1] = seg[1:, 1:]
+
+    b = seg ^ e | seg ^ s | seg ^ se
+    b[-1, :] = seg[-1, :] ^ e[-1, :]
+    b[:, -1] = seg[:, -1] ^ s[:, -1]
+    b[-1, -1] = 0
+
+    if w == width and h == height:
+        bmap = b
+    else:
+        bmap = np.zeros((height, width))
+        for x in range(w):
+            for y in range(h):
+                if b[y, x]:
+                    j = 1 + math.floor((y - 1) + height / h)
+                    i = 1 + math.floor((x - 1) + width / h)
+                    bmap[j, i] = 1
+
+    return bmap
+
+
+if __name__ == '__main__':
+    from davis2017.davis import DAVIS
+    from davis2017.results import Results
+
+    dataset = DAVIS(root='input_dir/ref', subset='val', sequences='aerobatics')
+    results = Results(root_dir='examples/osvos')
+    # Test timing F measure
+    for seq in dataset.get_sequences():
+        all_gt_masks, _, all_masks_id = dataset.get_all_masks(seq, True)
+        all_gt_masks, all_masks_id = all_gt_masks[:, 1:-1, :, :], all_masks_id[1:-1]
+        all_res_masks = results.read_masks(seq, all_masks_id)
+        f_metrics_res = np.zeros(all_gt_masks.shape[:2])
+        for ii in range(all_gt_masks.shape[0]):
+            f_metrics_res[ii, :] = db_eval_boundary(all_gt_masks[ii, ...], all_res_masks[ii, ...])
+
+    # Run using to profile code: python -m cProfile -o f_measure.prof metrics.py
+    #                            snakeviz f_measure.prof

ext/davis2017/results.py

@@ -0,0 +1,52 @@
+import os
+import numpy as np
+from PIL import Image, ImagePalette
+import sys
+
+
+davis_palette = b'\x00\x00\x00\x80\x00\x00\x00\x80\x00\x80\x80\x00\x00\x00\x80\x80\x00\x80\x00\x80\x80\x80\x80\x80@\x00\x00\xc0\x00\x00@\x80\x00\xc0\x80\x00@\x00\x80\xc0\x00\x80@\x80\x80\xc0\x80\x80\x00@\x00\x80@\x00\x00\xc0\x00\x80\xc0\x00\x00@\x80\x80@\x80\x00\xc0\x80\x80\xc0\x80@@\x00\xc0@\x00@\xc0\x00\xc0\xc0\x00@@\x80\xc0@\x80@\xc0\x80\xc0\xc0\x80\x00\x00@\x80\x00@\x00\x80@\x80\x80@\x00\x00\xc0\x80\x00\xc0\x00\x80\xc0\x80\x80\xc0@\x00@\xc0\x00@@\x80@\xc0\x80@@\x00\xc0\xc0\x00\xc0@\x80\xc0\xc0\x80\xc0\x00@@\x80@@\x00\xc0@\x80\xc0@\x00@\xc0\x80@\xc0\x00\xc0\xc0\x80\xc0\xc0@@@\xc0@@@\xc0@\xc0\xc0@@@\xc0\xc0@\xc0@\xc0\xc0\xc0\xc0\xc0 \x00\x00\xa0\x00\x00 \x80\x00\xa0\x80\x00 \x00\x80\xa0\x00\x80 \x80\x80\xa0\x80\x80`\x00\x00\xe0\x00\x00`\x80\x00\xe0\x80\x00`\x00\x80\xe0\x00\x80`\x80\x80\xe0\x80\x80 @\x00\xa0@\x00 \xc0\x00\xa0\xc0\x00 @\x80\xa0@\x80 \xc0\x80\xa0\xc0\x80`@\x00\xe0@\x00`\xc0\x00\xe0\xc0\x00`@\x80\xe0@\x80`\xc0\x80\xe0\xc0\x80 \x00@\xa0\x00@ \x80@\xa0\x80@ \x00\xc0\xa0\x00\xc0 \x80\xc0\xa0\x80\xc0`\x00@\xe0\x00@`\x80@\xe0\x80@`\x00\xc0\xe0\x00\xc0`\x80\xc0\xe0\x80\xc0 @@\xa0@@ \xc0@\xa0\xc0@ @\xc0\xa0@\xc0 \xc0\xc0\xa0\xc0\xc0`@@\xe0@@`\xc0@\xe0\xc0@`@\xc0\xe0@\xc0`\xc0\xc0\xe0\xc0\xc0\x00 \x00\x80 \x00\x00\xa0\x00\x80\xa0\x00\x00 \x80\x80 \x80\x00\xa0\x80\x80\xa0\x80@ \x00\xc0 \x00@\xa0\x00\xc0\xa0\x00@ \x80\xc0 \x80@\xa0\x80\xc0\xa0\x80\x00`\x00\x80`\x00\x00\xe0\x00\x80\xe0\x00\x00`\x80\x80`\x80\x00\xe0\x80\x80\xe0\x80@`\x00\xc0`\x00@\xe0\x00\xc0\xe0\x00@`\x80\xc0`\x80@\xe0\x80\xc0\xe0\x80\x00 @\x80 @\x00\xa0@\x80\xa0@\x00 \xc0\x80 \xc0\x00\xa0\xc0\x80\xa0\xc0@ @\xc0 @@\xa0@\xc0\xa0@@ \xc0\xc0 \xc0@\xa0\xc0\xc0\xa0\xc0\x00`@\x80`@\x00\xe0@\x80\xe0@\x00`\xc0\x80`\xc0\x00\xe0\xc0\x80\xe0\xc0@`@\xc0`@@\xe0@\xc0\xe0@@`\xc0\xc0`\xc0@\xe0\xc0\xc0\xe0\xc0  \x00\xa0 \x00 \xa0\x00\xa0\xa0\x00  \x80\xa0 \x80 \xa0\x80\xa0\xa0\x80` \x00\xe0 \x00`\xa0\x00\xe0\xa0\x00` \x80\xe0 \x80`\xa0\x80\xe0\xa0\x80 `\x00\xa0`\x00 \xe0\x00\xa0\xe0\x00 `\x80\xa0`\x80 \xe0\x80\xa0\xe0\x80``\x00\xe0`\x00`\xe0\x00\xe0\xe0\x00``\x80\xe0`\x80`\xe0\x80\xe0\xe0\x80  @\xa0 @ \xa0@\xa0\xa0@  \xc0\xa0 \xc0 \xa0\xc0\xa0\xa0\xc0` @\xe0 @`\xa0@\xe0\xa0@` \xc0\xe0 \xc0`\xa0\xc0\xe0\xa0\xc0 `@\xa0`@ \xe0@\xa0\xe0@ `\xc0\xa0`\xc0 \xe0\xc0\xa0\xe0\xc0``@\xe0`@`\xe0@\xe0\xe0@``\xc0\xe0`\xc0`\xe0\xc0\xe0\xe0\xc0'
+mose_palette = b'\x00\x00\x00\xe4\x1a\x1c7~\xb8M\xafJ\x98N\xa3\xff\x7f\x00\xff\xff3\xa6V(\xf7\x81\xbf\x99\x99\x99f\xc2\xa5\xfc\x8db\x8d\xa0\xcb\xe7\x8a\xc3\xa6\xd8T\xff\xd9/\xe5\xc4\x94\xb3\xb3\xb3\x8d\xd3\xc7\xff\xff\xb3\xbe\xba\xda\xfb\x80r\x80\xb1\xd3\xfd\xb4b\xb3\xdei\xfc\xcd\xe5\xd9\xd9\xd9\xbc\x80\xbd\xcc\xeb\xc5\xff\xedo'
+
+class Results(object):
+    def __init__(self, root_dir):
+        self.root_dir = root_dir
+
+    def _read_mask(self, sequence, frame_id):
+        try:
+            mask_path = os.path.join(self.root_dir, sequence, f'{frame_id}.png')
+            # BUGFIX
+            # There is a bug in the codebase
+            # Here is a compensation.
+            img = Image.open(mask_path)
+            if img.mode != 'P':
+                img_color = np.array(img)
+                h, w, three = img_color.shape
+                assert three == 3
+
+                img_new = np.ones((h, w), dtype=np.uint8) * 255
+                color_map_np = np.frombuffer(davis_palette, dtype=np.uint8).reshape(-1, 3).copy()
+                for i in range(10):
+                    cur_color = color_map_np[i]
+                    mask = np.all(img_color == cur_color, axis=-1)
+                    img_new[mask] = i
+                assert not np.all(img_new == 255).any()
+                img = img_new
+            # BUGFIX
+            return np.array(img)
+        except IOError as err:
+            sys.stdout.write(sequence + " frame %s not found!\n" % frame_id)
+            sys.stdout.write("The frames have to be indexed PNG files placed inside the corespondent sequence "
+                             "folder.\nThe indexes have to match with the initial frame.\n")
+            sys.stderr.write("IOError: " + err.strerror + "\n")
+            sys.exit()
+
+    def read_masks(self, sequence, masks_id):
+        mask_0 = self._read_mask(sequence, masks_id[0])
+        masks = np.zeros((len(masks_id), *mask_0.shape))
+        for ii, m in enumerate(masks_id):
+            masks[ii, ...] = self._read_mask(sequence, m)
+        num_objects = int(np.max(masks))
+        tmp = np.ones((num_objects, *masks.shape))
+        tmp = tmp * np.arange(1, num_objects + 1)[:, None, None, None]
+        masks = (tmp == masks[None, ...]) > 0
+        return masks

ext/davis2017/utils.py

@@ -0,0 +1,174 @@
+import os
+import errno
+import numpy as np
+from PIL import Image
+import warnings
+from ext.davis2017.davis import DAVIS
+
+
+def _pascal_color_map(N=256, normalized=False):
+    """
+    Python implementation of the color map function for the PASCAL VOC data set.
+    Official Matlab version can be found in the PASCAL VOC devkit
+    http://host.robots.ox.ac.uk/pascal/VOC/voc2012/index.html#devkit
+    """
+
+    def bitget(byteval, idx):
+        return (byteval & (1 << idx)) != 0
+
+    dtype = 'float32' if normalized else 'uint8'
+    cmap = np.zeros((N, 3), dtype=dtype)
+    for i in range(N):
+        r = g = b = 0
+        c = i
+        for j in range(8):
+            r = r | (bitget(c, 0) << 7 - j)
+            g = g | (bitget(c, 1) << 7 - j)
+            b = b | (bitget(c, 2) << 7 - j)
+            c = c >> 3
+
+        cmap[i] = np.array([r, g, b])
+
+    cmap = cmap / 255 if normalized else cmap
+    return cmap
+
+
+def overlay_semantic_mask(im, ann, alpha=0.5, colors=None, contour_thickness=None):
+    im, ann = np.asarray(im, dtype=np.uint8), np.asarray(ann, dtype=np.int)
+    if im.shape[:-1] != ann.shape:
+        raise ValueError('First two dimensions of `im` and `ann` must match')
+    if im.shape[-1] != 3:
+        raise ValueError('im must have three channels at the 3 dimension')
+
+    colors = colors or _pascal_color_map()
+    colors = np.asarray(colors, dtype=np.uint8)
+
+    mask = colors[ann]
+    fg = im * alpha + (1 - alpha) * mask
+
+    img = im.copy()
+    img[ann > 0] = fg[ann > 0]
+
+    if contour_thickness:  # pragma: no cover
+        import cv2
+        for obj_id in np.unique(ann[ann > 0]):
+            contours = cv2.findContours((ann == obj_id).astype(
+                np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[-2:]
+            cv2.drawContours(img, contours[0], -1, colors[obj_id].tolist(),
+                             contour_thickness)
+    return img
+
+
+def generate_obj_proposals(davis_root, subset, num_proposals, save_path):
+    dataset = DAVIS(davis_root, subset=subset, codalab=True)
+    for seq in dataset.get_sequences():
+        save_dir = os.path.join(save_path, seq)
+        if os.path.exists(save_dir):
+            continue
+        all_gt_masks, all_masks_id = dataset.get_all_masks(seq, True)
+        img_size = all_gt_masks.shape[2:]
+        num_rows = int(np.ceil(np.sqrt(num_proposals)))
+        proposals = np.zeros((num_proposals, len(all_masks_id), *img_size))
+        height_slices = np.floor(np.arange(0, img_size[0] + 1, img_size[0]/num_rows)).astype(np.uint).tolist()
+        width_slices = np.floor(np.arange(0, img_size[1] + 1, img_size[1]/num_rows)).astype(np.uint).tolist()
+        ii = 0
+        prev_h, prev_w = 0, 0
+        for h in height_slices[1:]:
+            for w in width_slices[1:]:
+                proposals[ii, :, prev_h:h, prev_w:w] = 1
+                prev_w = w
+                ii += 1
+                if ii == num_proposals:
+                    break
+            prev_h, prev_w = h, 0
+            if ii == num_proposals:
+                break
+
+        os.makedirs(save_dir, exist_ok=True)
+        for i, mask_id in enumerate(all_masks_id):
+            mask = np.sum(proposals[:, i, ...] * np.arange(1, proposals.shape[0] + 1)[:, None, None], axis=0)
+            save_mask(mask, os.path.join(save_dir, f'{mask_id}.png'))
+
+
+def generate_random_permutation_gt_obj_proposals(davis_root, subset, save_path):
+    dataset = DAVIS(davis_root, subset=subset, codalab=True)
+    for seq in dataset.get_sequences():
+        gt_masks, all_masks_id = dataset.get_all_masks(seq, True)
+        obj_swap = np.random.permutation(np.arange(gt_masks.shape[0]))
+        gt_masks = gt_masks[obj_swap, ...]
+        save_dir = os.path.join(save_path, seq)
+        os.makedirs(save_dir, exist_ok=True)
+        for i, mask_id in enumerate(all_masks_id):
+            mask = np.sum(gt_masks[:, i, ...] * np.arange(1, gt_masks.shape[0] + 1)[:, None, None], axis=0)
+            save_mask(mask, os.path.join(save_dir, f'{mask_id}.png'))
+
+
+def color_map(N=256, normalized=False):
+    def bitget(byteval, idx):
+        return ((byteval & (1 << idx)) != 0)
+
+    dtype = 'float32' if normalized else 'uint8'
+    cmap = np.zeros((N, 3), dtype=dtype)
+    for i in range(N):
+        r = g = b = 0
+        c = i
+        for j in range(8):
+            r = r | (bitget(c, 0) << 7-j)
+            g = g | (bitget(c, 1) << 7-j)
+            b = b | (bitget(c, 2) << 7-j)
+            c = c >> 3
+
+        cmap[i] = np.array([r, g, b])
+
+    cmap = cmap/255 if normalized else cmap
+    return cmap
+
+
+def save_mask(mask, img_path):
+    if np.max(mask) > 255:
+        raise ValueError('Maximum id pixel value is 255')
+    mask_img = Image.fromarray(mask.astype(np.uint8))
+    mask_img.putpalette(color_map().flatten().tolist())
+    mask_img.save(img_path)
+
+
+def db_statistics(per_frame_values):
+    """ Compute mean,recall and decay from per-frame evaluation.
+    Arguments:
+        per_frame_values (ndarray): per-frame evaluation
+
+    Returns:
+        M,O,D (float,float,float):
+            return evaluation statistics: mean,recall,decay.
+    """
+
+    # strip off nan values
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=RuntimeWarning)
+        M = np.nanmean(per_frame_values)
+        O = np.nanmean(per_frame_values > 0.5)
+
+    N_bins = 4
+    ids = np.round(np.linspace(1, len(per_frame_values), N_bins + 1) + 1e-10) - 1
+    ids = ids.astype(np.uint8)
+
+    D_bins = [per_frame_values[ids[i]:ids[i + 1] + 1] for i in range(0, 4)]
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=RuntimeWarning)
+        D = np.nanmean(D_bins[0]) - np.nanmean(D_bins[3])
+
+    return M, O, D
+
+
+def list_files(dir, extension=".png"):
+    return [os.path.splitext(file_)[0] for file_ in os.listdir(dir) if file_.endswith(extension)]
+
+
+def force_symlink(file1, file2):
+    try:
+        os.symlink(file1, file2)
+    except OSError as e:
+        if e.errno == errno.EEXIST:
+            os.remove(file2)
+        os.symlink(file1, file2)

ext/meta/sam_meta.py

@@ -0,0 +1,41 @@
+meta_dict = {
+    'vit_h': dict(
+        encoder_embed_dim=1280,
+        encoder_depth=32,
+        encoder_num_heads=16,
+        encoder_global_attn_indexes=[7, 15, 23, 31],
+        # common
+        prompt_embed_dim=256,
+        image_size=1024,
+        vit_patch_size=16,
+        image_embedding_size=64
+    ),
+    'vit_l': dict(
+        encoder_embed_dim=1024,
+        encoder_depth=24,
+        encoder_num_heads=16,
+        encoder_global_attn_indexes=[5, 11, 17, 23],
+        # common
+        prompt_embed_dim=256,
+        image_size=1024,
+        vit_patch_size=16,
+        image_embedding_size=64
+    ),
+    'vit_b': dict(
+        encoder_embed_dim=768,
+        encoder_depth=12,
+        encoder_num_heads=12,
+        encoder_global_attn_indexes=[2, 5, 8, 11],
+        # common
+        prompt_embed_dim=256,
+        image_size=1024,
+        vit_patch_size=16,
+        image_embedding_size=64
+    )
+}
+
+checkpoint_dict = {
+    'vit_h': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth',
+    'vit_l': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth',
+    'vit_b': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth',
+}

ext/open_clip/__init__.py

@@ -0,0 +1,15 @@
+from .coca_model import CoCa
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from .factory import create_model, create_model_and_transforms, create_model_from_pretrained, get_tokenizer, create_loss
+from .factory import list_models, add_model_config, get_model_config, load_checkpoint
+from .loss import ClipLoss, DistillClipLoss, CoCaLoss
+from .model import CLIP, CustomTextCLIP, CLIPTextCfg, CLIPVisionCfg, \
+    convert_weights_to_lp, convert_weights_to_fp16, trace_model, get_cast_dtype, get_input_dtype
+from .openai import load_openai_model, list_openai_models
+from .pretrained import list_pretrained, list_pretrained_models_by_tag, list_pretrained_tags_by_model, \
+    get_pretrained_url, download_pretrained_from_url, is_pretrained_cfg, get_pretrained_cfg, download_pretrained
+from .push_to_hf_hub import push_pretrained_to_hf_hub, push_to_hf_hub
+from .tokenizer import SimpleTokenizer, tokenize, decode
+from .transform import image_transform, AugmentationCfg
+from .zero_shot_classifier import build_zero_shot_classifier, build_zero_shot_classifier_legacy
+from .zero_shot_metadata import OPENAI_IMAGENET_TEMPLATES, SIMPLE_IMAGENET_TEMPLATES, IMAGENET_CLASSNAMES

ext/open_clip/bpe_simple_vocab_16e6.txt.gz

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

ext/open_clip/coca_model.py

@@ -0,0 +1,458 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+import numpy as np
+from dataclasses import dataclass
+
+from .transformer import (
+    LayerNormFp32,
+    LayerNorm,
+    QuickGELU,
+    MultimodalTransformer,
+)
+from .model import CLIPTextCfg, CLIPVisionCfg, _build_vision_tower, _build_text_tower
+
+try:
+    from transformers import (
+        BeamSearchScorer,
+        LogitsProcessorList,
+        TopPLogitsWarper,
+        TopKLogitsWarper,
+        RepetitionPenaltyLogitsProcessor,
+        MinLengthLogitsProcessor,
+        MaxLengthCriteria,
+        StoppingCriteriaList
+    )
+
+    GENERATION_TYPES = {
+        "top_k": TopKLogitsWarper,
+        "top_p": TopPLogitsWarper,
+        "beam_search": "beam_search"
+    }
+    _has_transformers = True
+except ImportError as e:
+    GENERATION_TYPES = {
+        "top_k": None,
+        "top_p": None,
+        "beam_search": "beam_search"
+    }
+    _has_transformers = False
+
+
+@dataclass
+class MultimodalCfg(CLIPTextCfg):
+    mlp_ratio: int = 4
+    dim_head: int = 64
+    heads: int = 8
+    n_queries: int = 256
+    attn_pooler_heads: int = 8
+
+
+def _build_text_decoder_tower(
+        embed_dim,
+        multimodal_cfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    multimodal_cfg = MultimodalCfg(**multimodal_cfg) if isinstance(multimodal_cfg, dict) else multimodal_cfg
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+    norm_layer = (
+        LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+    )
+
+    decoder = MultimodalTransformer(
+        context_length=multimodal_cfg.context_length,
+        width=multimodal_cfg.width,
+        heads=multimodal_cfg.heads,
+        layers=multimodal_cfg.layers,
+        ls_init_value=multimodal_cfg.ls_init_value,
+        output_dim=embed_dim,
+        act_layer=act_layer,
+        norm_layer=norm_layer,
+    )
+
+    return decoder
+
+
+class CoCa(nn.Module):
+    def __init__(
+            self,
+            embed_dim,
+            multimodal_cfg: MultimodalCfg,
+            text_cfg: CLIPTextCfg,
+            vision_cfg: CLIPVisionCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            pad_id: int = 0,
+    ):
+        super().__init__()
+        multimodal_cfg = MultimodalCfg(**multimodal_cfg) if isinstance(multimodal_cfg, dict) else multimodal_cfg
+        text_cfg = CLIPTextCfg(**text_cfg) if isinstance(text_cfg, dict) else text_cfg
+        vision_cfg = CLIPVisionCfg(**vision_cfg) if isinstance(vision_cfg, dict) else vision_cfg
+
+        self.text = _build_text_tower(
+            embed_dim=embed_dim,
+            text_cfg=text_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        vocab_size = (
+            text_cfg.vocab_size  # for hf models
+            if hasattr(text_cfg, "hf_model_name") and text_cfg.hf_model_name is not None
+            else text_cfg.vocab_size
+        )
+
+        self.visual = _build_vision_tower(
+            embed_dim=embed_dim,
+            vision_cfg=vision_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        self.text_decoder = _build_text_decoder_tower(
+            vocab_size,
+            multimodal_cfg=multimodal_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        self.pad_id = pad_id
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable)
+        self.text_decoder.set_grad_checkpointing(enable)
+
+    def _encode_image(self, images, normalize=True):
+        image_latent, tokens_embs = self.visual(images)
+        image_latent = F.normalize(image_latent, dim=-1) if normalize else image_latent
+        return image_latent, tokens_embs
+
+    def _encode_text(self, text, normalize=True, embed_cls=True):
+        text = text[:, :-1] if embed_cls else text # make space for CLS token
+        text_latent, token_emb = self.text(text)
+        text_latent = F.normalize(text_latent, dim=-1) if normalize else text_latent
+        return text_latent, token_emb
+
+    def encode_image(self, images, normalize=True):
+        image_latent, _ = self._encode_image(images, normalize=normalize)
+        return image_latent
+
+    def encode_text(self, text, normalize=True, embed_cls=True):
+        text_latent, _ = self._encode_text(text, normalize=normalize, embed_cls=embed_cls)
+        return text_latent
+
+    def forward(self, image, text, embed_cls=True, image_latent=None, image_embs=None):
+        text_latent, token_embs = self._encode_text(text, embed_cls=embed_cls)
+        if image_latent is None or image_embs is None:
+            image_latent, image_embs = self._encode_image(image)
+
+        # TODO: add assertion to avoid bugs?
+        labels = text[:, -token_embs.shape[1]:]
+
+        logits = self.text_decoder(image_embs, token_embs)
+        return {
+            "image_features": image_latent,
+            "text_features": text_latent,
+            "logits": logits,
+            "labels": labels,
+            "logit_scale": self.logit_scale.exp()
+        }
+
+    def generate(
+        self,
+        image,
+        text=None,
+        seq_len=30,
+        max_seq_len=77,
+        temperature=1.,
+        generation_type="beam_search",
+        top_p=0.1,  # keep tokens in the 1 - top_p quantile
+        top_k=1,  # keeps the top_k most probable tokens
+        pad_token_id=None,
+        eos_token_id=None,
+        sot_token_id=None,
+        num_beams=6,
+        num_beam_groups=3,
+        min_seq_len=5,
+        stopping_criteria=None,
+        repetition_penalty=1.0,
+        fixed_output_length=False # if True output.shape == (batch_size, seq_len)
+    ):
+        # taking many ideas and components from HuggingFace GenerationMixin
+        # https://huggingface.co/docs/transformers/main/en/main_classes/text_generation
+        assert _has_transformers, "Please install transformers for generate functionality. `pip install transformers`."
+        assert seq_len > min_seq_len, "seq_len must be larger than min_seq_len"
+
+        with torch.no_grad():
+            sot_token_id = 49406 if sot_token_id is None else sot_token_id
+            eos_token_id = 49407 if eos_token_id is None else eos_token_id
+            pad_token_id = self.pad_id if pad_token_id is None else pad_token_id
+            logit_processor = LogitsProcessorList(
+                [
+                    MinLengthLogitsProcessor(min_seq_len, eos_token_id),
+                    RepetitionPenaltyLogitsProcessor(repetition_penalty),
+                ]
+            )
+
+            if stopping_criteria is None:
+                stopping_criteria = [MaxLengthCriteria(max_length=seq_len)]
+
+            stopping_criteria = StoppingCriteriaList(
+                stopping_criteria
+            )
+
+            device = image.device
+
+            if generation_type == "beam_search":
+                output = self._generate_beamsearch(
+                    image_inputs = image,
+                    pad_token_id=pad_token_id,
+                    eos_token_id=eos_token_id,
+                    sot_token_id=sot_token_id,
+                    num_beams=num_beams,
+                    num_beam_groups=num_beam_groups,
+                    min_seq_len=min_seq_len,
+                    stopping_criteria=stopping_criteria,
+                    logit_processor=logit_processor,
+                )
+                if fixed_output_length and output.shape[1] < seq_len:
+                    return torch.cat(
+                        (output, torch.ones(output.shape[0], seq_len-output.shape[1], device=device, dtype=output.dtype) * self.pad_id),
+                        dim=1
+                    )
+                return output
+
+            elif generation_type == "top_p":
+                logit_warper = GENERATION_TYPES[generation_type](top_p)
+            elif generation_type == "top_k":
+                logit_warper = GENERATION_TYPES[generation_type](top_k)
+            else:
+                raise ValueError(
+                    f"generation_type has to be one of "
+                    f"{'| ' + ' | '.join(list(GENERATION_TYPES.keys())) + ' |'}."
+                )
+
+            image_latent, image_embs = self._encode_image(image)
+
+            if text is None:
+                text = torch.ones((image.shape[0], 1), device=device, dtype=torch.long) * sot_token_id
+
+            was_training = self.training
+            num_dims = len(text.shape)
+
+            if num_dims == 1:
+                text = text[None, :]
+
+            cur_len = text.shape[1]
+            self.eval()
+            out = text
+
+            while True:
+                x = out[:, -max_seq_len:]
+                cur_len = x.shape[1]
+                logits = self(image, x, image_latent=image_latent, image_embs=image_embs, embed_cls=False)["logits"][:, -1]
+                mask = (out[:, -1] == eos_token_id) | (out[:, -1] == pad_token_id)
+                sample = torch.ones((out.shape[0], 1), device=device, dtype=torch.long) * pad_token_id
+
+                if mask.all():
+                    if not fixed_output_length:
+                        break
+                else:
+                    logits = logits[~mask, :]
+                    filtered_logits = logit_processor(x[~mask, :], logits)
+                    filtered_logits = logit_warper(x[~mask, :], filtered_logits)
+                    probs = F.softmax(filtered_logits / temperature, dim=-1)
+
+                    if (cur_len + 1 == seq_len):
+                        sample[~mask, :] = torch.ones((sum(~mask), 1), device=device, dtype=torch.long) * eos_token_id
+                    else:
+                        sample[~mask, :] = torch.multinomial(probs, 1)
+
+                out = torch.cat((out, sample), dim=-1)
+
+                cur_len += 1
+
+                if stopping_criteria(out, None):
+                    break
+
+            if num_dims == 1:
+                out = out.squeeze(0)
+
+            self.train(was_training)
+            return out
+
+    def _generate_beamsearch(
+            self,
+            image_inputs,
+            pad_token_id=None,
+            eos_token_id=None,
+            sot_token_id=None,
+            num_beams=6,
+            num_beam_groups=3,
+            min_seq_len=5,
+            stopping_criteria=None,
+            logit_processor=None,
+            logit_warper=None,
+    ):
+        device = image_inputs.device
+        batch_size = image_inputs.shape[0]
+        image_inputs = torch.repeat_interleave(image_inputs, num_beams, dim=0)
+        image_latent, image_embs = self._encode_image(image_inputs)
+
+        input_ids = torch.ones((batch_size * num_beams, 1), device=device, dtype=torch.long)
+        input_ids = input_ids * sot_token_id
+        beam_scorer = BeamSearchScorer(
+            batch_size=batch_size,
+            num_beams=num_beams,
+            device=device,
+            num_beam_groups=num_beam_groups,
+        )
+        # instantiate logits processors
+        logits_processor = (
+            LogitsProcessorList([MinLengthLogitsProcessor(min_seq_len, eos_token_id=eos_token_id)])
+            if logit_processor is None
+            else logit_processor
+        )
+
+        batch_size = len(beam_scorer._beam_hyps)
+        num_beams = beam_scorer.num_beams
+        num_beam_groups = beam_scorer.num_beam_groups
+        num_sub_beams = num_beams // num_beam_groups
+        batch_beam_size, cur_len = input_ids.shape
+        beam_indices = None
+
+        if num_beams * batch_size != batch_beam_size:
+            raise ValueError(
+                f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
+            )
+
+        beam_scores = torch.full((batch_size, num_beams), -1e9, dtype=torch.float, device=device)
+        # initialise score of first beam of each group with 0 and the rest with 1e-9. This ensures that the beams in
+        # the same group don't produce same tokens everytime.
+        beam_scores[:, ::num_sub_beams] = 0
+        beam_scores = beam_scores.view((batch_size * num_beams,))
+
+        while True:
+
+            # predicted tokens in cur_len step
+            current_tokens = torch.zeros(batch_size * num_beams, dtype=input_ids.dtype, device=device)
+
+            # indices which will form the beams in the next time step
+            reordering_indices = torch.zeros(batch_size * num_beams, dtype=torch.long, device=device)
+
+            # do one decoder step on all beams of all sentences in batch
+            model_inputs = prepare_inputs_for_generation(input_ids=input_ids, image_inputs=image_inputs)
+            outputs = self(
+                model_inputs['images'],
+                model_inputs['text'],
+                embed_cls=False,
+                image_latent=image_latent,
+                image_embs=image_embs
+            )
+
+            for beam_group_idx in range(num_beam_groups):
+                group_start_idx = beam_group_idx * num_sub_beams
+                group_end_idx = min(group_start_idx + num_sub_beams, num_beams)
+                group_size = group_end_idx - group_start_idx
+
+                # indices of beams of current group among all sentences in batch
+                batch_group_indices = []
+
+                for batch_idx in range(batch_size):
+                    batch_group_indices.extend(
+                        [batch_idx * num_beams + idx for idx in range(group_start_idx, group_end_idx)]
+                    )
+                group_input_ids = input_ids[batch_group_indices]
+
+                # select outputs of beams of currentg group only
+                next_token_logits = outputs['logits'][batch_group_indices, -1, :]
+                vocab_size = next_token_logits.shape[-1]
+
+                next_token_scores_processed = logits_processor(
+                    group_input_ids, next_token_logits, current_tokens=current_tokens, beam_group_idx=beam_group_idx
+                )
+                next_token_scores = next_token_scores_processed + beam_scores[batch_group_indices].unsqueeze(-1)
+                next_token_scores = next_token_scores.expand_as(next_token_scores_processed)
+
+                # reshape for beam search
+                next_token_scores = next_token_scores.view(batch_size, group_size * vocab_size)
+
+                next_token_scores, next_tokens = torch.topk(
+                    next_token_scores, 2 * group_size, dim=1, largest=True, sorted=True
+                )
+
+                next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
+                next_tokens = next_tokens % vocab_size
+
+                # stateless
+                process_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+                beam_outputs = beam_scorer.process(
+                    group_input_ids,
+                    next_token_scores,
+                    next_tokens,
+                    next_indices,
+                    pad_token_id=pad_token_id,
+                    eos_token_id=eos_token_id,
+                    beam_indices=process_beam_indices,
+                )
+                beam_scores[batch_group_indices] = beam_outputs["next_beam_scores"]
+                beam_next_tokens = beam_outputs["next_beam_tokens"]
+                beam_idx = beam_outputs["next_beam_indices"]
+
+                input_ids[batch_group_indices] = group_input_ids[beam_idx]
+                group_input_ids = torch.cat([group_input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
+                current_tokens[batch_group_indices] = group_input_ids[:, -1]
+
+                # (beam_idx // group_size) -> batch_idx
+                # (beam_idx % group_size) -> offset of idx inside the group
+                reordering_indices[batch_group_indices] = (
+                    num_beams * torch.div(beam_idx, group_size, rounding_mode="floor") + group_start_idx + (beam_idx % group_size)
+                )
+
+            input_ids = torch.cat([input_ids, current_tokens.unsqueeze(-1)], dim=-1)
+
+            # increase cur_len
+            cur_len = cur_len + 1
+            if beam_scorer.is_done or stopping_criteria(input_ids, None):
+                break
+
+        final_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+        sequence_outputs = beam_scorer.finalize(
+            input_ids,
+            beam_scores,
+            next_tokens,
+            next_indices,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            max_length=stopping_criteria.max_length,
+            beam_indices=final_beam_indices,
+        )
+        return sequence_outputs['sequences']
+
+
+def prepare_inputs_for_generation(input_ids, image_inputs, past=None, **kwargs):
+    if past:
+        input_ids = input_ids[:, -1].unsqueeze(-1)
+
+    attention_mask = kwargs.get("attention_mask", None)
+    position_ids = kwargs.get("position_ids", None)
+
+    if attention_mask is not None and position_ids is None:
+        # create position_ids on the fly for batch generation
+        position_ids = attention_mask.long().cumsum(-1) - 1
+        position_ids.masked_fill_(attention_mask == 0, 1)
+    else:
+        position_ids = None
+    return {
+        "text": input_ids,
+        "images": image_inputs,
+        "past_key_values": past,
+        "position_ids": position_ids,
+        "attention_mask": attention_mask,
+    }

ext/open_clip/constants.py

@@ -0,0 +1,2 @@
+OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)

ext/open_clip/factory.py

@@ -0,0 +1,387 @@
+import json
+import logging
+import os
+import pathlib
+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from .model import CLIP, CustomTextCLIP, convert_weights_to_lp, convert_to_custom_text_state_dict,\
+    resize_pos_embed, get_cast_dtype
+from .coca_model import CoCa
+from .loss import ClipLoss, DistillClipLoss, CoCaLoss
+from .openai import load_openai_model
+from .pretrained import is_pretrained_cfg, get_pretrained_cfg, download_pretrained,\
+    list_pretrained_tags_by_model, download_pretrained_from_hf
+from .transform import image_transform, AugmentationCfg
+from .tokenizer import HFTokenizer, tokenize
+
+
+HF_HUB_PREFIX = 'hf-hub:'
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
+_MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
+
+
+def _natural_key(string_):
+    return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
+
+
+def _rescan_model_configs():
+    global _MODEL_CONFIGS
+
+    config_ext = ('.json',)
+    config_files = []
+    for config_path in _MODEL_CONFIG_PATHS:
+        if config_path.is_file() and config_path.suffix in config_ext:
+            config_files.append(config_path)
+        elif config_path.is_dir():
+            for ext in config_ext:
+                config_files.extend(config_path.glob(f'*{ext}'))
+
+    for cf in config_files:
+        with open(cf, 'r') as f:
+            model_cfg = json.load(f)
+            if all(a in model_cfg for a in ('embed_dim', 'vision_cfg', 'text_cfg')):
+                _MODEL_CONFIGS[cf.stem] = model_cfg
+
+    _MODEL_CONFIGS = {k: v for k, v in sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0]))}
+
+
+_rescan_model_configs()  # initial populate of model config registry
+
+
+def list_models():
+    """ enumerate available model architectures based on config files """
+    return list(_MODEL_CONFIGS.keys())
+
+
+def add_model_config(path):
+    """ add model config path or file and update registry """
+    if not isinstance(path, Path):
+        path = Path(path)
+    _MODEL_CONFIG_PATHS.append(path)
+    _rescan_model_configs()
+
+
+def get_model_config(model_name):
+    if model_name in _MODEL_CONFIGS:
+        return deepcopy(_MODEL_CONFIGS[model_name])
+    else:
+        return None
+
+
+def get_tokenizer(model_name):
+    if model_name.startswith(HF_HUB_PREFIX):
+        tokenizer = HFTokenizer(model_name[len(HF_HUB_PREFIX):])
+    else:
+        config = get_model_config(model_name)
+        tokenizer = HFTokenizer(
+            config['text_cfg']['hf_tokenizer_name']) if 'hf_tokenizer_name' in config['text_cfg'] else tokenize
+    return tokenizer
+
+
+def load_state_dict(checkpoint_path: str, map_location='cpu'):
+    checkpoint = torch.load(checkpoint_path, map_location=map_location)
+    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+        state_dict = checkpoint['state_dict']
+    else:
+        state_dict = checkpoint
+    if next(iter(state_dict.items()))[0].startswith('module'):
+        state_dict = {k[7:]: v for k, v in state_dict.items()}
+    return state_dict
+
+
+def load_checkpoint(model, checkpoint_path, strict=True):
+    state_dict = load_state_dict(checkpoint_path)
+    # detect old format and make compatible with new format
+    if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
+        state_dict = convert_to_custom_text_state_dict(state_dict)
+    resize_pos_embed(state_dict, model)
+    incompatible_keys = model.load_state_dict(state_dict, strict=strict)
+    return incompatible_keys
+
+
+def create_model(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        pretrained_image: bool = False,
+        pretrained_hf: bool = True,
+        cache_dir: Optional[str] = None,
+        output_dict: Optional[bool] = None,
+        require_pretrained: bool = False,
+        logger: logging.Logger = logging,
+):
+    has_hf_hub_prefix = model_name.startswith(HF_HUB_PREFIX)
+    if has_hf_hub_prefix:
+        model_id = model_name[len(HF_HUB_PREFIX):]
+        checkpoint_path = download_pretrained_from_hf(model_id, cache_dir=cache_dir)
+        config_path = download_pretrained_from_hf(model_id, filename='open_clip_config.json', cache_dir=cache_dir)
+
+        with open(config_path, 'r', encoding='utf-8') as f:
+            config = json.load(f)
+        pretrained_cfg = config['preprocess_cfg']
+        model_cfg = config['model_cfg']
+    else:
+        model_name = model_name.replace('/', '-')  # for callers using old naming with / in ViT names
+        checkpoint_path = None
+        pretrained_cfg = {}
+        model_cfg = None
+
+    if isinstance(device, str):
+        device = torch.device(device)
+
+    if pretrained and pretrained.lower() == 'openai':
+        logger.info(f'Loading pretrained {model_name} from OpenAI.')
+        model = load_openai_model(
+            model_name,
+            precision=precision,
+            device=device,
+            cache_dir=cache_dir,
+        )
+    else:
+        model_cfg = model_cfg or get_model_config(model_name)
+        if model_cfg is not None:
+            logger.info(f'Loaded {model_name} model config.')
+        else:
+            logger.error(f'Model config for {model_name} not found; available models {list_models()}.')
+            raise RuntimeError(f'Model config for {model_name} not found.')
+
+        if force_quick_gelu:
+            # override for use of QuickGELU on non-OpenAI transformer models
+            model_cfg["quick_gelu"] = True
+
+        if force_patch_dropout is not None:
+            # override the default patch dropout value
+            model_cfg["vision_cfg"]["patch_dropout"] = force_patch_dropout
+
+        if force_image_size is not None:
+            # override model config's image size
+            model_cfg["vision_cfg"]["image_size"] = force_image_size
+
+        is_timm_model = 'timm_model_name' in model_cfg.get('vision_cfg', {})
+        if pretrained_image:
+            if is_timm_model:
+                # pretrained weight loading for timm models set via vision_cfg
+                model_cfg['vision_cfg']['timm_model_pretrained'] = True
+            else:
+                assert False, 'pretrained image towers currently only supported for timm models'
+
+        # cast_dtype set for fp16 and bf16 (manual mixed-precision), not set for 'amp' or 'pure' modes
+        cast_dtype = get_cast_dtype(precision)
+        is_hf_model = 'hf_model_name' in model_cfg.get('text_cfg', {})
+        custom_text = model_cfg.pop('custom_text', False) or force_custom_text or is_hf_model
+
+        if custom_text:
+            if is_hf_model:
+                model_cfg['text_cfg']['hf_model_pretrained'] = pretrained_hf
+            if "coca" in model_name:
+                model = CoCa(**model_cfg, cast_dtype=cast_dtype)
+            else:
+                model = CustomTextCLIP(**model_cfg, cast_dtype=cast_dtype)
+        else:
+            model = CLIP(**model_cfg, cast_dtype=cast_dtype)
+
+        if precision in ("fp16", "bf16"):
+            dtype = torch.float16 if 'fp16' in precision else torch.bfloat16
+            # manual mixed precision that matches original OpenAI behaviour
+            if is_timm_model:
+                # FIXME this is a bit janky, create timm based model in low-precision and
+                # then cast only LayerNormFp32 instances back to float32 so they don't break.
+                # Why? The convert_weights_to_lp fn only works with native models.
+                model.to(device=device, dtype=dtype)
+                from .transformer import LayerNormFp32
+                def _convert_ln(m):
+                    if isinstance(m, LayerNormFp32):
+                        m.weight.data = m.weight.data.to(torch.float32)
+                        m.bias.data = m.bias.data.to(torch.float32)
+                model.apply(_convert_ln)
+            else:
+                model.to(device=device)
+                convert_weights_to_lp(model, dtype=dtype)
+        elif precision in ("pure_fp16", "pure_bf16"):
+            dtype = torch.float16 if 'fp16' in precision else torch.bfloat16
+            model.to(device=device, dtype=dtype)
+        else:
+            model.to(device=device)
+
+        pretrained_loaded = False
+        if pretrained:
+            checkpoint_path = ''
+            pretrained_cfg = get_pretrained_cfg(model_name, pretrained)
+            if pretrained_cfg:
+                checkpoint_path = download_pretrained(pretrained_cfg, cache_dir=cache_dir)
+            elif os.path.exists(pretrained):
+                checkpoint_path = pretrained
+
+            if checkpoint_path:
+                logger.info(f'Loading pretrained {model_name} weights ({pretrained}).')
+                load_checkpoint(model, checkpoint_path)
+            else:
+                error_str = (
+                    f'Pretrained weights ({pretrained}) not found for model {model_name}.'
+                    f'Available pretrained tags ({list_pretrained_tags_by_model(model_name)}.')
+                logger.warning(error_str)
+                raise RuntimeError(error_str)
+            pretrained_loaded = True
+        elif has_hf_hub_prefix:
+            logger.info(f'Loading pretrained {model_name} weights ({pretrained}).')
+            load_checkpoint(model, checkpoint_path)
+            pretrained_loaded = True
+
+        if require_pretrained and not pretrained_loaded:
+            # callers of create_model_from_pretrained always expect pretrained weights
+            raise RuntimeError(
+                f'Pretrained weights were required for (model: {model_name}, pretrained: {pretrained}) but not loaded.')
+
+        # set image / mean metadata from pretrained_cfg if available, or use default
+        model.visual.image_mean = pretrained_cfg.get('mean', None) or OPENAI_DATASET_MEAN
+        model.visual.image_std = pretrained_cfg.get('std', None) or OPENAI_DATASET_STD
+
+    if output_dict and hasattr(model, "output_dict"):
+        model.output_dict = True
+
+    if jit:
+        model = torch.jit.script(model)
+
+    return model
+
+
+def create_loss(args):
+    if args.distill:
+        return DistillClipLoss(
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+            use_horovod=args.horovod,
+        )
+    elif "coca" in args.model.lower():
+        return CoCaLoss(
+            caption_loss_weight=args.coca_caption_loss_weight,
+            clip_loss_weight=args.coca_contrastive_loss_weight,
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+            use_horovod=args.horovod,
+        )
+    return ClipLoss(
+        local_loss=args.local_loss,
+        gather_with_grad=args.gather_with_grad,
+        cache_labels=True,
+        rank=args.rank,
+        world_size=args.world_size,
+        use_horovod=args.horovod,
+    )
+
+
+def create_model_and_transforms(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        pretrained_image: bool = False,
+        pretrained_hf: bool = True,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        aug_cfg: Optional[Union[Dict[str, Any], AugmentationCfg]] = None,
+        cache_dir: Optional[str] = None,
+        output_dict: Optional[bool] = None,
+        logger: logging.Logger = logging,
+):
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_text=force_custom_text,
+        force_patch_dropout=force_patch_dropout,
+        force_image_size=force_image_size,
+        pretrained_image=pretrained_image,
+        pretrained_hf=pretrained_hf,
+        cache_dir=cache_dir,
+        output_dict=output_dict,
+        logger=logger,
+    )
+
+    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+    image_std = image_std or getattr(model.visual, 'image_std', None)
+    preprocess_train = image_transform(
+        model.visual.image_size,
+        is_train=True,
+        mean=image_mean,
+        std=image_std,
+        aug_cfg=aug_cfg,
+    )
+    preprocess_val = image_transform(
+        model.visual.image_size,
+        is_train=False,
+        mean=image_mean,
+        std=image_std,
+    )
+
+    return model, preprocess_train, preprocess_val
+
+
+def create_model_from_pretrained(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        return_transform: bool = True,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        cache_dir: Optional[str] = None,
+        logger: logging.Logger = logging,
+):
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_text=force_custom_text,
+        force_image_size=force_image_size,
+        cache_dir=cache_dir,
+        require_pretrained=True,
+        logger=logger,
+    )
+
+    if not return_transform:
+        return model
+
+    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+    image_std = image_std or getattr(model.visual, 'image_std', None)
+    preprocess = image_transform(
+        model.visual.image_size,
+        is_train=False,
+        mean=image_mean,
+        std=image_std,
+    )
+
+    return model, preprocess

ext/open_clip/hf_configs.py

@@ -0,0 +1,56 @@
+# HF architecture dict:
+arch_dict = {
+    # https://huggingface.co/docs/transformers/model_doc/roberta#roberta
+    "roberta": {
+        "config_names": {
+            "context_length": "max_position_embeddings",
+            "vocab_size": "vocab_size",
+            "width": "hidden_size",
+            "heads": "num_attention_heads",
+            "layers": "num_hidden_layers",
+            "layer_attr": "layer",
+            "token_embeddings_attr": "embeddings"
+        },
+        "pooler": "mean_pooler",
+    },
+    # https://huggingface.co/docs/transformers/model_doc/xlm-roberta#transformers.XLMRobertaConfig
+    "xlm-roberta": {
+        "config_names": {
+            "context_length": "max_position_embeddings",
+            "vocab_size": "vocab_size",
+            "width": "hidden_size",
+            "heads": "num_attention_heads",
+            "layers": "num_hidden_layers",
+            "layer_attr": "layer",
+            "token_embeddings_attr": "embeddings"
+        },
+        "pooler": "mean_pooler",
+    },
+    # https://huggingface.co/docs/transformers/model_doc/mt5#mt5
+    "mt5": {
+        "config_names": {
+            # unlimited seqlen
+            # https://github.com/google-research/text-to-text-transfer-transformer/issues/273
+            # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/t5/modeling_t5.py#L374
+            "context_length": "",
+            "vocab_size": "vocab_size",
+            "width": "d_model",
+            "heads": "num_heads",
+            "layers": "num_layers",
+            "layer_attr": "block",
+            "token_embeddings_attr": "embed_tokens"
+        },
+        "pooler": "mean_pooler",
+    },
+    # https://huggingface.co/docs/transformers/model_doc/bert
+    "bert": {
+        "config_names": {
+            "context_length": "max_position_embeddings",
+            "vocab_size": "vocab_size",
+            "width": "hidden_size",
+            "heads": "num_attention_heads",
+            "layers": "num_hidden_layers",
+        },
+        "pooler": "cls_pooler",
+    },
+}

ext/open_clip/hf_model.py

@@ -0,0 +1,193 @@
+""" huggingface model adapter
+
+Wraps HuggingFace transformers (https://github.com/huggingface/transformers) models for use as a text tower in CLIP model.
+"""
+import re
+
+import torch
+import torch.nn as nn
+from torch import TensorType
+
+try:
+    import transformers
+    from transformers import AutoModel, AutoTokenizer, AutoConfig, PretrainedConfig
+    from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, \
+        BaseModelOutputWithPoolingAndCrossAttentions
+except ImportError as e:
+    transformers = None
+
+
+    class BaseModelOutput:
+        pass
+
+
+    class PretrainedConfig:
+        pass
+
+from .hf_configs import arch_dict
+
+
+# utils
+def _camel2snake(s):
+    return re.sub(r'(?<!^)(?=[A-Z])', '_', s).lower()
+
+
+# TODO: ?last - for gpt-like models
+_POOLERS = {}
+
+
+def register_pooler(cls):
+    """Decorator registering pooler class"""
+    _POOLERS[_camel2snake(cls.__name__)] = cls
+    return cls
+
+
+@register_pooler
+class MeanPooler(nn.Module):
+    """Mean pooling"""
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        masked_output = x.last_hidden_state * attention_mask.unsqueeze(-1)
+        return masked_output.sum(dim=1) / attention_mask.sum(-1, keepdim=True)
+
+
+@register_pooler
+class MaxPooler(nn.Module):
+    """Max pooling"""
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        masked_output = x.last_hidden_state.masked_fill(attention_mask.unsqueeze(-1), -torch.inf)
+        return masked_output.max(1).values
+
+
+@register_pooler
+class ClsPooler(nn.Module):
+    """CLS token pooling"""
+
+    def __init__(self, use_pooler_output=True):
+        super().__init__()
+        self.cls_token_position = 0
+        self.use_pooler_output = use_pooler_output
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        if (self.use_pooler_output and
+            isinstance(x, (BaseModelOutputWithPooling, BaseModelOutputWithPoolingAndCrossAttentions)) and
+            (x.pooler_output is not None)
+        ):
+            return x.pooler_output
+
+        return x.last_hidden_state[:, self.cls_token_position, :]
+
+
+@register_pooler
+class ClsLastHiddenStatePooler(nn.Module):
+    """CLS token pooling
+    NOTE: this is equivalent to ClsPooler above with use_pooler_output=False
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.cls_token_position = 0
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        return x.last_hidden_state[:, self.cls_token_position, :]
+
+
+class HFTextEncoder(nn.Module):
+    """HuggingFace model adapter"""
+    output_tokens: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            model_name_or_path: str,
+            output_dim: int,
+            config: PretrainedConfig = None,
+            pooler_type: str = None,
+            proj: str = None,
+            pretrained: bool = True,
+            output_tokens: bool = False,
+    ):
+        super().__init__()
+        self.output_tokens = output_tokens
+        self.output_dim = output_dim
+
+        # TODO: find better way to get this information
+        uses_transformer_pooler = (pooler_type == "cls_pooler")
+
+        if transformers is None:
+            raise RuntimeError("Please `pip install transformers` to use pre-trained HuggingFace models")
+        if config is None:
+            self.config = AutoConfig.from_pretrained(model_name_or_path)
+            create_func, model_args = (AutoModel.from_pretrained, model_name_or_path) if pretrained else (
+                AutoModel.from_config, self.config)
+            # TODO: do all model configs have this attribute? PretrainedConfig does so yes??
+            if hasattr(self.config, "is_encoder_decoder") and self.config.is_encoder_decoder:
+                self.transformer = create_func(model_args)
+                self.transformer = self.transformer.encoder
+            else:
+                self.transformer = create_func(model_args, add_pooling_layer=uses_transformer_pooler)
+        else:
+            self.config = config
+            self.transformer = AutoModel.from_config(config)
+        if pooler_type is None:  # get default arch pooler
+            pooler_type = (arch_dict[self.config.model_type]["pooler"])
+
+        # FIXME downstream users of OpenCLIP models use these attr, need to verify valid across all models
+        self.vocab_size = getattr(self.config, 'vocab_size', 0)
+        self.context_length = getattr(self.config, 'max_position_embeddings', 0)
+
+        self.pooler = _POOLERS[pooler_type]()
+
+        d_model = getattr(self.config, arch_dict[self.config.model_type]["config_names"]["width"])
+        if (d_model == output_dim) and (proj is None):  # do we always need a proj?
+            self.proj = nn.Identity()
+        elif proj == 'linear':
+            self.proj = nn.Linear(d_model, output_dim, bias=False)
+        elif proj == 'mlp':
+            hidden_size = (d_model + output_dim) // 2
+            self.proj = nn.Sequential(
+                nn.Linear(d_model, hidden_size, bias=False),
+                nn.GELU(),
+                nn.Linear(hidden_size, output_dim, bias=False),
+            )
+
+    def forward(self, x: TensorType):
+        attn_mask = (x != self.config.pad_token_id).long()
+        out = self.transformer(input_ids=x, attention_mask=attn_mask)
+        pooled_out = self.pooler(out, attn_mask)
+        projected = self.proj(pooled_out)
+
+        seq_len = out.last_hidden_state.shape[1]
+        tokens = (
+            out.last_hidden_state[:, torch.arange(seq_len) != self.pooler.cls_token_position, :] 
+            if type(self.pooler) == ClsPooler 
+            else out.last_hidden_state
+        )
+        
+        if self.output_tokens:
+            return projected, tokens
+        return projected
+
+    def lock(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        if not unlocked_layers:  # full freezing
+            for n, p in self.transformer.named_parameters():
+                p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+            return
+
+        encoder = self.transformer.encoder if hasattr(self.transformer, 'encoder') else self.transformer
+        layer_list = getattr(encoder, arch_dict[self.config.model_type]["config_names"]["layer_attr"])
+        print(f"Unlocking {unlocked_layers}/{len(layer_list) + 1} layers of hf model")
+        embeddings = getattr(
+            self.transformer, arch_dict[self.config.model_type]["config_names"]["token_embeddings_attr"])
+        modules = [embeddings, *layer_list][:-unlocked_layers]
+        # freeze layers
+        for module in modules:
+            for n, p in module.named_parameters():
+                p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.gradient_checkpointing_enable()
+
+    def init_parameters(self):
+        pass

ext/open_clip/loss.py

@@ -0,0 +1,216 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+try:
+    import torch.distributed.nn
+    from torch import distributed as dist
+
+    has_distributed = True
+except ImportError:
+    has_distributed = False
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+
+def gather_features(
+        image_features,
+        text_features,
+        local_loss=False,
+        gather_with_grad=False,
+        rank=0,
+        world_size=1,
+        use_horovod=False
+):
+    assert has_distributed, 'torch.distributed did not import correctly, please use a PyTorch version with support.'
+    if use_horovod:
+        assert hvd is not None, 'Please install horovod'
+        if gather_with_grad:
+            all_image_features = hvd.allgather(image_features)
+            all_text_features = hvd.allgather(text_features)
+        else:
+            with torch.no_grad():
+                all_image_features = hvd.allgather(image_features)
+                all_text_features = hvd.allgather(text_features)
+            if not local_loss:
+                # ensure grads for local rank when all_* features don't have a gradient
+                gathered_image_features = list(all_image_features.chunk(world_size, dim=0))
+                gathered_text_features = list(all_text_features.chunk(world_size, dim=0))
+                gathered_image_features[rank] = image_features
+                gathered_text_features[rank] = text_features
+                all_image_features = torch.cat(gathered_image_features, dim=0)
+                all_text_features = torch.cat(gathered_text_features, dim=0)
+    else:
+        # We gather tensors from all gpus
+        if gather_with_grad:
+            all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features), dim=0)
+            all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
+        else:
+            gathered_image_features = [torch.zeros_like(image_features) for _ in range(world_size)]
+            gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
+            dist.all_gather(gathered_image_features, image_features)
+            dist.all_gather(gathered_text_features, text_features)
+            if not local_loss:
+                # ensure grads for local rank when all_* features don't have a gradient
+                gathered_image_features[rank] = image_features
+                gathered_text_features[rank] = text_features
+            all_image_features = torch.cat(gathered_image_features, dim=0)
+            all_text_features = torch.cat(gathered_text_features, dim=0)
+
+    return all_image_features, all_text_features
+
+
+class ClipLoss(nn.Module):
+
+    def __init__(
+            self,
+            local_loss=False,
+            gather_with_grad=False,
+            cache_labels=False,
+            rank=0,
+            world_size=1,
+            use_horovod=False,
+    ):
+        super().__init__()
+        self.local_loss = local_loss
+        self.gather_with_grad = gather_with_grad
+        self.cache_labels = cache_labels
+        self.rank = rank
+        self.world_size = world_size
+        self.use_horovod = use_horovod
+
+        # cache state
+        self.prev_num_logits = 0
+        self.labels = {}
+
+    def get_ground_truth(self, device, num_logits) -> torch.Tensor:
+        # calculated ground-truth and cache if enabled
+        if self.prev_num_logits != num_logits or device not in self.labels:
+            labels = torch.arange(num_logits, device=device, dtype=torch.long)
+            if self.world_size > 1 and self.local_loss:
+                labels = labels + num_logits * self.rank
+            if self.cache_labels:
+                self.labels[device] = labels
+                self.prev_num_logits = num_logits
+        else:
+            labels = self.labels[device]
+        return labels
+
+    def get_logits(self, image_features, text_features, logit_scale):
+        if self.world_size > 1:
+            all_image_features, all_text_features = gather_features(
+                image_features, text_features,
+                self.local_loss, self.gather_with_grad, self.rank, self.world_size, self.use_horovod)
+
+            if self.local_loss:
+                logits_per_image = logit_scale * image_features @ all_text_features.T
+                logits_per_text = logit_scale * text_features @ all_image_features.T
+            else:
+                logits_per_image = logit_scale * all_image_features @ all_text_features.T
+                logits_per_text = logits_per_image.T
+        else:
+            logits_per_image = logit_scale * image_features @ text_features.T
+            logits_per_text = logit_scale * text_features @ image_features.T
+        
+        return logits_per_image, logits_per_text
+
+    def forward(self, image_features, text_features, logit_scale, output_dict=False):
+        device = image_features.device
+        logits_per_image, logits_per_text = self.get_logits(image_features, text_features, logit_scale)
+
+        labels = self.get_ground_truth(device, logits_per_image.shape[0])
+
+        total_loss = (
+            F.cross_entropy(logits_per_image, labels) +
+            F.cross_entropy(logits_per_text, labels)
+        ) / 2
+
+        return {"contrastive_loss": total_loss} if output_dict else total_loss
+
+
+class CoCaLoss(ClipLoss):
+    def __init__(
+            self,
+            caption_loss_weight,
+            clip_loss_weight,
+            pad_id=0,  # pad_token for open_clip custom tokenizer
+            local_loss=False,
+            gather_with_grad=False,
+            cache_labels=False,
+            rank=0,
+            world_size=1,
+            use_horovod=False,
+    ):
+        super().__init__(
+            local_loss=local_loss,
+            gather_with_grad=gather_with_grad,
+            cache_labels=cache_labels,
+            rank=rank,
+            world_size=world_size,
+            use_horovod=use_horovod
+        )
+
+        self.clip_loss_weight = clip_loss_weight
+        self.caption_loss_weight = caption_loss_weight
+        self.caption_loss = nn.CrossEntropyLoss(ignore_index=pad_id)
+
+    def forward(self, image_features, text_features, logits, labels, logit_scale, output_dict=False):
+        
+        clip_loss = torch.tensor(0)
+        
+        if self.clip_loss_weight:
+            clip_loss = super().forward(image_features, text_features, logit_scale)
+            clip_loss = self.clip_loss_weight * clip_loss
+
+        caption_loss = self.caption_loss(
+            logits.permute(0, 2, 1),
+            labels,
+        )
+        caption_loss = caption_loss * self.caption_loss_weight
+
+        if output_dict:
+            return {"contrastive_loss": clip_loss, "caption_loss": caption_loss}
+
+        return clip_loss, caption_loss
+
+
+class DistillClipLoss(ClipLoss):
+
+    def dist_loss(self, teacher_logits, student_logits):
+        return -(teacher_logits.softmax(dim=1) * student_logits.log_softmax(dim=1)).sum(dim=1).mean(dim=0)
+
+    def forward(
+            self,
+            image_features,
+            text_features,
+            logit_scale,
+            dist_image_features,
+            dist_text_features,
+            dist_logit_scale,
+            output_dict=False,
+    ):
+        logits_per_image, logits_per_text = \
+            self.get_logits(image_features, text_features, logit_scale)
+
+        dist_logits_per_image, dist_logits_per_text = \
+            self.get_logits(dist_image_features, dist_text_features, dist_logit_scale)
+
+        labels = self.get_ground_truth(image_features.device, logits_per_image.shape[0])
+
+        contrastive_loss = (
+            F.cross_entropy(logits_per_image, labels) +
+            F.cross_entropy(logits_per_text, labels)
+        ) / 2
+
+        distill_loss = (
+            self.dist_loss(dist_logits_per_image, logits_per_image) +
+            self.dist_loss(dist_logits_per_text, logits_per_text)
+        ) / 2
+
+        if output_dict:
+            return {"contrastive_loss": contrastive_loss, "distill_loss": distill_loss}
+
+        return contrastive_loss, distill_loss

ext/open_clip/model.py

@@ -0,0 +1,473 @@
+""" CLIP Model
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+from dataclasses import dataclass
+import logging
+import math
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from .hf_model import HFTextEncoder
+from .modified_resnet import ModifiedResNet
+from .timm_model import TimmModel
+from .transformer import LayerNormFp32, LayerNorm, QuickGELU, Attention, VisionTransformer, TextTransformer
+from .utils import to_2tuple
+
+
+@dataclass
+class CLIPVisionCfg:
+    layers: Union[Tuple[int, int, int, int], int] = 12
+    width: int = 768
+    head_width: int = 64
+    mlp_ratio: float = 4.0
+    patch_size: int = 16
+    image_size: Union[Tuple[int, int], int] = 224
+
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    patch_dropout: float = 0.  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    input_patchnorm: bool = False  # whether to use dual patchnorm - would only apply the input layernorm on each patch, as post-layernorm already exist in original clip vit design
+    global_average_pool: bool = False  # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
+    attentional_pool: bool = False  # whether to use attentional pooler in the last embedding layer
+    n_queries: int = 256  # n_queries for attentional pooler
+    attn_pooler_heads: int = 8  # n heads for attentional_pooling
+    output_tokens: bool = False
+
+    timm_model_name: str = None  # a valid model name overrides layers, width, patch_size
+    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
+    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_proj_bias: bool = False  # enable bias final projection
+    timm_drop: float = 0.  # head dropout
+    timm_drop_path: Optional[float] = None  # backbone stochastic depth
+
+
+@dataclass
+class CLIPTextCfg:
+    context_length: int = 77
+    vocab_size: int = 49408
+    width: int = 512
+    heads: int = 8
+    layers: int = 12
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    hf_model_name: str = None
+    hf_tokenizer_name: str = None
+    hf_model_pretrained: bool = True
+    proj: str = 'mlp'
+    pooler_type: str = 'mean_pooler'
+    embed_cls: bool = False
+    pad_id: int = 0
+    output_tokens: bool = False
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == 'bf16':
+        cast_dtype = torch.bfloat16
+    elif precision == 'fp16':
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def get_input_dtype(precision: str):
+    input_dtype = None
+    if precision in ('bf16', 'pure_bf16'):
+        input_dtype = torch.bfloat16
+    elif precision in ('fp16', 'pure_fp16'):
+        input_dtype = torch.float16
+    return input_dtype
+
+
+def _build_vision_tower(
+        embed_dim: int,
+        vision_cfg: CLIPVisionCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None
+):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    # OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
+    # memory efficient in recent PyTorch releases (>= 1.10).
+    # NOTE: timm models always use native GELU regardless of quick_gelu flag.
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+
+    if vision_cfg.timm_model_name:
+        visual = TimmModel(
+            vision_cfg.timm_model_name,
+            pretrained=vision_cfg.timm_model_pretrained,
+            pool=vision_cfg.timm_pool,
+            proj=vision_cfg.timm_proj,
+            proj_bias=vision_cfg.timm_proj_bias,
+            drop=vision_cfg.timm_drop,
+            drop_path=vision_cfg.timm_drop_path,
+            patch_drop=vision_cfg.patch_dropout if vision_cfg.patch_dropout > 0 else None,
+            embed_dim=embed_dim,
+            image_size=vision_cfg.image_size,
+        )
+    elif isinstance(vision_cfg.layers, (tuple, list)):
+        vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
+        visual = ModifiedResNet(
+            layers=vision_cfg.layers,
+            output_dim=embed_dim,
+            heads=vision_heads,
+            image_size=vision_cfg.image_size,
+            width=vision_cfg.width,
+        )
+    else:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        visual = VisionTransformer(
+            image_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            width=vision_cfg.width,
+            layers=vision_cfg.layers,
+            heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            ls_init_value=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            input_patchnorm=vision_cfg.input_patchnorm,
+            global_average_pool=vision_cfg.global_average_pool,
+            attentional_pool=vision_cfg.attentional_pool,
+            n_queries=vision_cfg.n_queries,
+            attn_pooler_heads=vision_cfg.attn_pooler_heads,
+            output_tokens=vision_cfg.output_tokens,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+
+    return visual
+
+
+def _build_text_tower(
+        embed_dim: int,
+        text_cfg: CLIPTextCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    if isinstance(text_cfg, dict):
+        text_cfg = CLIPTextCfg(**text_cfg)
+
+    if text_cfg.hf_model_name:
+        text = HFTextEncoder(
+            text_cfg.hf_model_name,
+            output_dim=embed_dim,
+            proj=text_cfg.proj,
+            pooler_type=text_cfg.pooler_type,
+            pretrained=text_cfg.hf_model_pretrained,
+            output_tokens=text_cfg.output_tokens,
+        )
+    else:
+        act_layer = QuickGELU if quick_gelu else nn.GELU
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+
+        text = TextTransformer(
+            context_length=text_cfg.context_length,
+            vocab_size=text_cfg.vocab_size,
+            width=text_cfg.width,
+            heads=text_cfg.heads,
+            layers=text_cfg.layers,
+            ls_init_value=text_cfg.ls_init_value,
+            output_dim=embed_dim,
+            embed_cls=text_cfg.embed_cls,
+            output_tokens=text_cfg.output_tokens,
+            pad_id=text_cfg.pad_id,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+    return text
+
+
+class CLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+
+        text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.transformer = text.transformer
+        self.context_length = text.context_length
+        self.vocab_size = text.vocab_size
+        self.token_embedding = text.token_embedding
+        self.positional_embedding = text.positional_embedding
+        self.ln_final = text.ln_final
+        self.text_projection = text.text_projection
+        self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.transformer.grad_checkpointing = enable
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return F.normalize(x, dim=-1) if normalize else x
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+
+class CustomTextCLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+        self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.context_length = self.text.context_length
+        self.vocab_size = self.text.vocab_size
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        self.text.lock(unlocked_layers, freeze_layer_norm)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable)
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        features = self.text(text)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+
+def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
+    """Convert applicable model parameters to low-precision (bf16 or fp16)"""
+
+    def _convert_weights(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.to(dtype)
+            if l.bias is not None:
+                l.bias.data = l.bias.data.to(dtype)
+
+        if isinstance(l, (nn.MultiheadAttention, Attention)):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.to(dtype)
+
+        if isinstance(l, (CLIP, TextTransformer)):
+            # convert text nn.Parameter projections
+            attr = getattr(l, "text_projection", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+        if isinstance(l, VisionTransformer):
+            # convert vision nn.Parameter projections
+            attr = getattr(l, "proj", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+    model.apply(_convert_weights)
+
+
+convert_weights_to_fp16 = convert_weights_to_lp  # backwards compat
+
+
+# used to maintain checkpoint compatibility
+def convert_to_custom_text_state_dict(state_dict: dict):
+    if 'text_projection' in state_dict:
+        # old format state_dict, move text tower -> .text
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if any(k.startswith(p) for p in (
+                'text_projection',
+                'positional_embedding',
+                'token_embedding',
+                'transformer',
+                'ln_final',
+            )):
+                k = 'text.' + k
+            new_state_dict[k] = v
+        return new_state_dict
+    return state_dict
+
+
+def build_model_from_openai_state_dict(
+        state_dict: dict,
+        quick_gelu=True,
+        cast_dtype=torch.float16,
+):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len(
+            [k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_size = vision_patch_size * grid_size
+    else:
+        counts: list = [
+            len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_size = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    vision_cfg = CLIPVisionCfg(
+        layers=vision_layers,
+        width=vision_width,
+        patch_size=vision_patch_size,
+        image_size=image_size,
+    )
+    text_cfg = CLIPTextCfg(
+        context_length=context_length,
+        vocab_size=vocab_size,
+        width=transformer_width,
+        heads=transformer_heads,
+        layers=transformer_layers,
+    )
+    model = CLIP(
+        embed_dim,
+        vision_cfg=vision_cfg,
+        text_cfg=text_cfg,
+        quick_gelu=quick_gelu,  # OpenAI models were trained with QuickGELU
+        cast_dtype=cast_dtype,
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        state_dict.pop(key, None)
+
+    convert_weights_to_fp16(model)  # OpenAI state dicts are partially converted to float16
+    model.load_state_dict(state_dict)
+    return model.eval()
+
+
+def trace_model(model, batch_size=256, device=torch.device('cpu')):
+    model.eval()
+    image_size = model.visual.image_size
+    example_images = torch.ones((batch_size, 3, image_size, image_size), device=device)
+    example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
+    model = torch.jit.trace_module(
+        model,
+        inputs=dict(
+            forward=(example_images, example_text),
+            encode_text=(example_text,),
+            encode_image=(example_images,)
+        ))
+    model.visual.image_size = image_size
+    return model
+
+
+def resize_pos_embed(state_dict, model, interpolation: str = 'bicubic', antialias: bool = True):
+    # Rescale the grid of position embeddings when loading from state_dict
+    old_pos_embed = state_dict.get('visual.positional_embedding', None)
+    if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
+        return
+    grid_size = to_2tuple(model.visual.grid_size)
+    extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
+    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    if extra_tokens:
+        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
+    else:
+        pos_emb_tok, pos_emb_img = None, old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
+
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
+    if pos_emb_tok is not None:
+        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
+    else:
+        new_pos_embed = pos_emb_img
+    state_dict['visual.positional_embedding'] = new_pos_embed

ext/open_clip/model_configs/EVA01-g-14-plus.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva_giant_patch14_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA01-g-14.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva_giant_patch14_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA02-B-16.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva02_base_patch16_clip_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA02-E-14-plus.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva02_enormous_patch14_clip_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1280,
+        "heads": 20,
+        "layers": 32
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA02-E-14.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva02_enormous_patch14_clip_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA02-L-14-336.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 336,
+        "timm_model_name": "eva02_large_patch14_clip_336",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/EVA02-L-14.json

@@ -0,0 +1,18 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 224,
+        "timm_model_name": "eva02_large_patch14_clip_224",
+        "timm_model_pretrained": false,
+        "timm_pool": "token",
+        "timm_proj": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    },
+    "custom_text": true
+}

ext/open_clip/model_configs/RN101-quickgelu.json

@@ -0,0 +1,22 @@
+{
+    "embed_dim": 512,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            23,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/RN101.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            23,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/RN50-quickgelu.json

@@ -0,0 +1,22 @@
+{
+    "embed_dim": 1024,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            6,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/RN50.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            6,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/RN50x16.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 384,
+        "layers": [
+            6,
+            8,
+            18,
+            8
+        ],
+        "width": 96,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/RN50x4.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 288,
+        "layers": [
+            4,
+            6,
+            10,
+            6
+        ],
+        "width": 80,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/RN50x64.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 448,
+        "layers": [
+            3,
+            15,
+            36,
+            10
+        ],
+        "width": 128,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/ViT-B-16-plus-240.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 240,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/ViT-B-16-plus.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/ViT-B-16.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

ext/open_clip/model_configs/ViT-B-32-plus-256.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 256,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}