chore: import upstream snapshot with attribution

dit/object_detection/README.md

@@ -0,0 +1,120 @@
+# DiT for Object Detection
+
+This folder contains Mask R-CNN Cascade Mask R-CNN running instructions on top of [Detectron2](https://github.com/facebookresearch/detectron2) for PubLayNet and ICDAR 2019 cTDaR.
+
+## Usage
+
+### Inference
+
+The quickest way to try out DiT for document layout analysis is the web demo: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/nielsr/dit-document-layout-analysis).
+
+One can run inference using the `inference.py` script. It can be run as follows (from the root of the unilm repository):
+
+```
+python ./dit/object_detection/inference.py \
+--image_path ./dit/object_detection/publaynet_example.jpeg \
+--output_file_name output.jpg \
+--config ./dit/object_detection/publaynet_configs/maskrcnn/maskrcnn_dit_base.yaml \
+--opts MODEL.WEIGHTS https://layoutlm.blob.core.windows.net/dit/dit-fts/publaynet_dit-b_mrcnn.pth \
+```
+
+Make sure that the configuration file (YAML) and PyTorch checkpoint match. The example above uses DiT-base with the Mask R-CNN framework fine-tuned on PubLayNet.
+
+### Data Preparation
+
+**PubLayNet**
+
+Download the data from this [link](https://dax-cdn.cdn.appdomain.cloud/dax-publaynet/1.0.0/publaynet.tar.gz?_ga=2.218138265.1825957955.1646384196-1495010506.1633610665) (~96GB). Then extract it to `PATH-to-PubLayNet`.
+
+A soft link needs to be created to make the data accessible for the program:`ln -s PATH-to-PubLayNet publaynet_data`.
+
+**ICDAR 2019 cTDaR**
+
+Download the data from this [link](https://github.com/cndplab-founder/ICDAR2019_cTDaR) (~4GB). Assume path to this repository is named as `PATH-to-ICDARrepo`.
+
+Then run `python convert_to_coco_format.py --root_dir=PATH-to-ICDARrepo --target_dir=PATH-toICDAR`. Now the path to processed data is `PATH-to-ICDAR`.
+
+Run the following command to get the adaptively binarized images for archival subset.
+
+```
+cp -r PATH-to-ICDAR/trackA_archival PATH-to-ICDAR/at_trackA_archival
+python adaptive_binarize.py --root_dir PATH-to-ICDAR/at_trackA_archival
+```
+
+The binarized archival subset will be in `PATH-to-ICDAR/at_trackA_archival`.
+
+According to the subset you want to evaluate/fine-tune, a soft link should be created:`ln -s PATH-to-ICDAR/trackA_modern data` or `ln -s PATH-to-ICDAR/at_trackA_archival data`.
+
+### Evaluation
+
+Following commands provide two examples to evaluate the fine-tuned checkpoints.
+
+The config files can be found in `icdar19_configs` and `publaynet_configs`.
+
+1) Evaluate the fine-tuned checkpoint of DiT-Base with Mask R-CNN on PublayNet:
+```bash
+python train_net.py --config-file publaynet_configs/maskrcnn/maskrcnn_dit_base.yaml --eval-only --num-gpus 8 MODEL.WEIGHTS <finetuned_checkpoint_file_path or link> OUTPUT_DIR <your_output_dir> 
+``` 
+
+2) Evaluate the fine-tuned checkpoint of DiT-Large with Cascade Mask R-CNN on ICDAR 2019 cTDaR archival subset (make sure you have created a soft link from `PATH-to-ICDAR/at_trackA_archival` to `data`):
+```bash
+python train_net.py --config-file icdar19_configs/cascade/cascade_dit_large.yaml --eval-only --num-gpus 8 MODEL.WEIGHTS <finetuned_checkpoint_file_path or link> OUTPUT_DIR <your_output_dir> 
+```
+
+**Note**: We have fixed the **bug** in the [ICDAR2019 measurement tool](https://github.com/cndplab-founder/ctdar_measurement_tool) during integrating the tool into our code. If you use the tool to get the evaluation score, please modify the [code](https://github.com/cndplab-founder/ctdar_measurement_tool/blob/738456d3164a838ffaeefe7d1b5e64f3a4368a0e/evaluate.py#L146
+) as follows:
+```bash
+    ...
+    # print(each_file)
+
+# for file in gt_file_lst:
+#     if file.split(".") != "xml":
+#         gt_file_lst.remove(file)
+#     # print(gt_file_lst)
+
+# Comment the code above and add the code below
+for i in range(len(gt_file_lst) - 1, -1, -1):
+    if gt_file_lst[i].split(".")[-1] != "xml":
+        del gt_file_lst[i]
+
+if len(gt_file_lst) > 0:
+    ...
+```
+
+### Training
+The following commands provide two examples to train the Mask R-CNN/Cascade Mask R-CNN with DiT backbone on 8 32GB Nvidia V100 GPUs.
+
+1) Fine-tune DiT-Base with Cascade Mask R-CNN on PublayNet:
+```bash
+python train_net.py --config-file publaynet_configs/cascade/cascade_dit_base.yaml --num-gpus 8 MODEL.WEIGHTS <DiT-Base_file_path or link> OUTPUT_DIR <your_output_dir> 
+``` 
+
+
+2) Fine-tune DiT-Large with Mask R-CNN on ICDAR 2019 cTDaR modern:
+```bash
+python train_net.py --config-file icdar19_configs/markrcnn/maskrcnn_dit_large.yaml --num-gpus 8 MODEL.WEIGHTS <DiT-Large_file_path or link> OUTPUT_DIR <your_output_dir> 
+``` 
+
+
+
+[Detectron2's document](https://detectron2.readthedocs.io/en/latest/tutorials/getting_started.html) may help you for more details.
+
+
+## Citation
+
+If you find this repository useful, please consider citing our work:
+```
+@misc{li2022dit,
+    title={DiT: Self-supervised Pre-training for Document Image Transformer},
+    author={Junlong Li and Yiheng Xu and Tengchao Lv and Lei Cui and Cha Zhang and Furu Wei},
+    year={2022},
+    eprint={2203.02378},
+    archivePrefix={arXiv},
+    primaryClass={cs.CV}
+}
+```
+
+
+
+## Acknowledgment
+Thanks to [Detectron2](https://github.com/facebookresearch/detectron2) for Mask R-CNN and Cascade Mask R-CNN implementation.

dit/object_detection/adaptive_binarize.py

@@ -0,0 +1,33 @@
+import argparse
+import os
+import cv2
+import tqdm
+
+
+def convert(fn):
+    # given a file name, convert it into binary and store at the same position
+    img = cv2.imread(fn)
+    gim = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    gim = cv2.adaptiveThreshold(gim, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 45, 11)
+    g3im = cv2.cvtColor(gim, cv2.COLOR_GRAY2BGR)
+    cv2.imwrite(fn, g3im)
+
+
+
+if __name__ == '__main__':
+    """
+    Now only feasible for trackA_XX
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root_dir', default="../datasets/icdar2019/at_trackA_archival")
+    args = parser.parse_args()
+
+    for fdname in os.listdir(args.root_dir):
+        if fdname.endswith(".json"):
+            continue
+        ffdname = os.path.join(args.root_dir, fdname)
+        for file in tqdm.tqdm(os.listdir(ffdname)):
+            if file.endswith(".xml"):
+                continue
+            ffile = os.path.join(ffdname, file)
+            convert(ffile)

dit/object_detection/convert_to_coco_format.py

@@ -0,0 +1,104 @@
+import os
+from PIL import Image
+import xml.etree.ElementTree as ET
+import numpy as np
+import json
+from PIL import Image
+from shutil import copyfile
+
+
+def convert(ROOT, TRACK, SPLIT):
+    coco_data = {
+        "images": [],
+        "annotations": [],
+        "categories": [{"id": 1, "name": "table"}, ],
+    }
+    DATA_DIR = f"{ROOT}/{TRACK}/{SPLIT}"
+    prefix = "cTDaR_t0" if TRACK == "trackA_archival" else "cTDaR_t1"
+    print(TRACK, SPLIT, prefix)
+    table_count = 0
+    for file in sorted(os.listdir(DATA_DIR)):
+        if file.startswith(prefix) and file.endswith(".jpg"):
+            img = Image.open(os.path.join(DATA_DIR, file))
+            coco_data["images"].append(
+                {
+                    "file_name": file,
+                    "height": img.height,
+                    "width": img.width,
+                    "id": int(file[7:-4]),
+                }
+            )
+        elif file.startswith(prefix) and file.endswith(".xml"):
+            # print(file)
+            tree = ET.parse(os.path.join(DATA_DIR, file))
+            root = tree.getroot()
+            assert len(root.findall("./table/Coords")) > 0
+            for table_id in range(len(root.findall("./table/Coords"))):
+                four_points = root.findall("./table/Coords")[table_id].attrib["points"]
+                four_points = list(map(lambda x: x.split(","), four_points.split()))
+                four_points = [[int(j) for j in i] for i in four_points]
+                segmentation = [j for i in four_points for j in i]
+                bbox = [
+                    four_points[0][0],
+                    four_points[0][1],
+                    four_points[2][0] - four_points[0][0],
+                    four_points[2][1] - four_points[0][1],
+                ]
+                coco_data["annotations"].append(
+                    {
+                        "segmentation": [segmentation],
+                        "area": bbox[2] * bbox[3],
+                        "iscrowd": 0,
+                        "image_id": int(file[7:-4]),
+                        "bbox": bbox,
+                        "category_id": 1,
+                        "id": table_count,
+                    }
+                )
+                table_count += 1
+
+    with open(f"{ROOT}/{TRACK}/{SPLIT}.json", "w") as f:
+        json.dump(coco_data, f)
+
+
+def clean_img(DATA_DIR):
+    for file in sorted(os.listdir(DATA_DIR)):
+        if file.endswith(".JPG"):
+            os.rename(os.path.join(DATA_DIR, file), os.path.join(DATA_DIR, file.replace(".JPG", ".jpg")))
+        elif file.endswith(".TIFF"):
+            img = Image.open(os.path.join(DATA_DIR, file))
+            img.save(os.path.join(DATA_DIR, file.replace(".TIFF", ".jpg")))
+            os.remove(os.path.join(DATA_DIR, file))
+        elif file.endswith(".png"):
+            img = Image.open(os.path.join(DATA_DIR, file))
+            img.save(os.path.join(DATA_DIR, file.replace(".png", ".jpg")))
+            os.remove(os.path.join(DATA_DIR, file))
+
+
+if __name__ == '__main__':
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root_dir', required=True)
+    parser.add_argument('--target_dir', required=True)
+    args = parser.parse_args()
+
+    test_data_dir = os.path.join(args.root_dir, 'test', 'TRACKA')
+    test_gt_dir = os.path.join(args.root_dir, 'test_ground_truth', 'TRACKA')
+    training_data_dir = os.path.join(args.root_dir, 'training', 'TRACKA', 'ground_truth')
+
+    raw_datas = {"train": [training_data_dir], "test": [test_data_dir, test_gt_dir]}
+
+    TRACKS = ["trackA_modern", "trackA_archival"]
+    SPLITS = ["train", "test"]
+    for track in TRACKS:
+        prefix = "cTDaR_t0" if track == "trackA_archival" else "cTDaR_t1"
+        for split in SPLITS:
+            os.makedirs(os.path.join(args.target_dir, track, split))
+            for source_dir in raw_datas[split]:
+                for fn in os.listdir(source_dir):
+                    if fn.startswith(prefix):
+                        ffn = os.path.join(source_dir, fn)
+                        copyfile(ffn, os.path.join(args.target_dir, track, split, fn))
+            clean_img(os.path.join(args.target_dir, track, split))
+            convert(args.target_dir, track, split)

dit/object_detection/ditod/__init__.py

@@ -0,0 +1,16 @@
+# --------------------------------------------------------------------------------
+# MPViT: Multi-Path Vision Transformer for Dense Prediction
+# Copyright (c) 2022 Electronics and Telecommunications Research Institute (ETRI).
+# All Rights Reserved.
+# Written by Youngwan Lee
+# This source code is licensed(Dual License(GPL3.0 & Commercial)) under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------------------------------
+
+from .config import add_vit_config
+from .backbone import build_vit_fpn_backbone
+from .dataset_mapper import DetrDatasetMapper
+from .mycheckpointer import MyDetectionCheckpointer
+from .icdar_evaluation import ICDAREvaluator
+from .mytrainer import MyTrainer
+from .table_evaluation import calc_table_score

dit/object_detection/ditod/backbone.py

@@ -0,0 +1,156 @@
+# --------------------------------------------------------------------------------
+# VIT: Multi-Path Vision Transformer for Dense Prediction
+# Copyright (c) 2022 Electronics and Telecommunications Research Institute (ETRI).
+# All Rights Reserved.
+# Written by Youngwan Lee
+# This source code is licensed(Dual License(GPL3.0 & Commercial)) under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------------------------------
+# References:
+# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# CoaT: https://github.com/mlpc-ucsd/CoaT
+# --------------------------------------------------------------------------------
+
+
+import torch
+
+from detectron2.layers import (
+    ShapeSpec,
+)
+from detectron2.modeling import Backbone, BACKBONE_REGISTRY, FPN
+from detectron2.modeling.backbone.fpn import LastLevelP6P7, LastLevelMaxPool
+
+from .beit import beit_base_patch16, dit_base_patch16, dit_large_patch16, beit_large_patch16
+from .deit import deit_base_patch16, mae_base_patch16
+
+__all__ = [
+    "build_vit_fpn_backbone",
+]
+
+
+class VIT_Backbone(Backbone):
+    """
+    Implement VIT backbone.
+    """
+
+    def __init__(self, name, out_features, drop_path, img_size, pos_type, model_kwargs):
+        super().__init__()
+        self._out_features = out_features
+        if 'base' in name:
+            self._out_feature_strides = {"layer3": 4, "layer5": 8, "layer7": 16, "layer11": 32}
+        else:
+            self._out_feature_strides = {"layer7": 4, "layer11": 8, "layer15": 16, "layer23": 32}
+
+        if name == 'beit_base_patch16':
+            model_func = beit_base_patch16
+            self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
+        elif name == 'dit_base_patch16':
+            model_func = dit_base_patch16
+            self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
+        elif name == "deit_base_patch16":
+            model_func = deit_base_patch16
+            self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
+        elif name == "mae_base_patch16":
+            model_func = mae_base_patch16
+            self._out_feature_channels = {"layer3": 768, "layer5": 768, "layer7": 768, "layer11": 768}
+        elif name == "dit_large_patch16":
+            model_func = dit_large_patch16
+            self._out_feature_channels = {"layer7": 1024, "layer11": 1024, "layer15": 1024, "layer23": 1024}
+        elif name == "beit_large_patch16":
+            model_func = beit_large_patch16
+            self._out_feature_channels = {"layer7": 1024, "layer11": 1024, "layer15": 1024, "layer23": 1024}
+        else:
+            raise ValueError("Unsupported VIT name yet.")
+
+        if 'beit' in name or 'dit' in name:
+            if pos_type == "abs":
+                self.backbone = model_func(img_size=img_size,
+                                           out_features=out_features,
+                                           drop_path_rate=drop_path,
+                                           use_abs_pos_emb=True,
+                                           **model_kwargs)
+            elif pos_type == "shared_rel":
+                self.backbone = model_func(img_size=img_size,
+                                           out_features=out_features,
+                                           drop_path_rate=drop_path,
+                                           use_shared_rel_pos_bias=True,
+                                           **model_kwargs)
+            elif pos_type == "rel":
+                self.backbone = model_func(img_size=img_size,
+                                           out_features=out_features,
+                                           drop_path_rate=drop_path,
+                                           use_rel_pos_bias=True,
+                                           **model_kwargs)
+            else:
+                raise ValueError()
+        else:
+            self.backbone = model_func(img_size=img_size,
+                                       out_features=out_features,
+                                       drop_path_rate=drop_path,
+                                       **model_kwargs)
+
+    def forward(self, x):
+        """
+        Args:
+            x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.
+
+        Returns:
+            dict[str->Tensor]: names and the corresponding features
+        """
+        assert x.dim() == 4, f"VIT takes an input of shape (N, C, H, W). Got {x.shape} instead!"
+        return self.backbone.forward_features(x)
+
+    def output_shape(self):
+        return {
+            name: ShapeSpec(
+                channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
+            )
+            for name in self._out_features
+        }
+
+
+def build_VIT_backbone(cfg):
+    """
+    Create a VIT instance from config.
+
+    Args:
+        cfg: a detectron2 CfgNode
+
+    Returns:
+        A VIT backbone instance.
+    """
+    # fmt: off
+    name = cfg.MODEL.VIT.NAME
+    out_features = cfg.MODEL.VIT.OUT_FEATURES
+    drop_path = cfg.MODEL.VIT.DROP_PATH
+    img_size = cfg.MODEL.VIT.IMG_SIZE
+    pos_type = cfg.MODEL.VIT.POS_TYPE
+
+    model_kwargs = eval(str(cfg.MODEL.VIT.MODEL_KWARGS).replace("`", ""))
+
+    return VIT_Backbone(name, out_features, drop_path, img_size, pos_type, model_kwargs)
+
+
+@BACKBONE_REGISTRY.register()
+def build_vit_fpn_backbone(cfg, input_shape: ShapeSpec):
+    """
+    Create a VIT w/ FPN backbone.
+
+    Args:
+        cfg: a detectron2 CfgNode
+
+    Returns:
+        backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
+    """
+    bottom_up = build_VIT_backbone(cfg)
+    in_features = cfg.MODEL.FPN.IN_FEATURES
+    out_channels = cfg.MODEL.FPN.OUT_CHANNELS
+    backbone = FPN(
+        bottom_up=bottom_up,
+        in_features=in_features,
+        out_channels=out_channels,
+        norm=cfg.MODEL.FPN.NORM,
+        top_block=LastLevelMaxPool(),
+        fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
+    )
+    return backbone

dit/object_detection/ditod/beit.py

@@ -0,0 +1,671 @@
+""" Vision Transformer (ViT) in PyTorch
+
+A PyTorch implement of Vision Transformers as described in
+'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale' - https://arxiv.org/abs/2010.11929
+
+The official jax code is released and available at https://github.com/google-research/vision_transformer
+
+Status/TODO:
+* Models updated to be compatible with official impl. Args added to support backward compat for old PyTorch weights.
+* Weights ported from official jax impl for 384x384 base and small models, 16x16 and 32x32 patches.
+* Trained (supervised on ImageNet-1k) my custom 'small' patch model to 77.9, 'base' to 79.4 top-1 with this code.
+* Hopefully find time and GPUs for SSL or unsupervised pretraining on OpenImages w/ ImageNet fine-tune in future.
+
+Acknowledgments:
+* The paper authors for releasing code and weights, thanks!
+* I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch ... check it out
+for some einops/einsum fun
+* Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT
+* Bert reference code checks against Huggingface Transformers and Tensorflow Bert
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import warnings
+import math
+import torch
+from functools import partial
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import drop_path, to_2tuple, trunc_normal_
+
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'mean': (0.5, 0.5, 0.5), 'std': (0.5, 0.5, 0.5),
+        **kwargs
+    }
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+    def extra_repr(self) -> str:
+        return 'p={}'.format(self.drop_prob)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        # x = self.drop(x)
+        # commit this for the orignal BERT implement
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
+            proj_drop=0., window_size=None, attn_head_dim=None):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+
+        if window_size:
+            self.window_size = window_size
+            self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+            self.relative_position_bias_table = nn.Parameter(
+                torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+            # cls to token & token 2 cls & cls to cls
+
+            # get pair-wise relative position index for each token inside the window
+            coords_h = torch.arange(window_size[0])
+            coords_w = torch.arange(window_size[1])
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+            coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+            relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
+            relative_coords[:, :, 1] += window_size[1] - 1
+            relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+            relative_position_index = \
+                torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
+            relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+            relative_position_index[0, 0:] = self.num_relative_distance - 3
+            relative_position_index[0:, 0] = self.num_relative_distance - 2
+            relative_position_index[0, 0] = self.num_relative_distance - 1
+
+            self.register_buffer("relative_position_index", relative_position_index)
+
+            # trunc_normal_(self.relative_position_bias_table, std=.0)
+        else:
+            self.window_size = None
+            self.relative_position_bias_table = None
+            self.relative_position_index = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(all_head_dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, rel_pos_bias=None, training_window_size=None):
+        B, N, C = x.shape
+        qkv_bias = None
+        if self.q_bias is not None:
+            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+        # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        if self.relative_position_bias_table is not None:
+            if training_window_size == self.window_size:
+                relative_position_bias = \
+                    self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                        self.window_size[0] * self.window_size[1] + 1,
+                        self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+                relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+                attn = attn + relative_position_bias.unsqueeze(0)
+            else:
+                training_window_size = tuple(training_window_size.tolist())
+                new_num_relative_distance = (2 * training_window_size[0] - 1) * (2 * training_window_size[1] - 1) + 3
+                # new_num_relative_dis 为 所有可能的相对位置选项，包含cls-cls，tok-cls，与cls-tok
+                new_relative_position_bias_table = F.interpolate(
+                    self.relative_position_bias_table[:-3, :].permute(1, 0).view(1, self.num_heads,
+                                                                                 2 * self.window_size[0] - 1,
+                                                                                 2 * self.window_size[1] - 1),
+                    size=(2 * training_window_size[0] - 1, 2 * training_window_size[1] - 1), mode='bicubic',
+                    align_corners=False)
+                new_relative_position_bias_table = new_relative_position_bias_table.view(self.num_heads,
+                                                                                         new_num_relative_distance - 3).permute(
+                    1, 0)
+                new_relative_position_bias_table = torch.cat(
+                    [new_relative_position_bias_table, self.relative_position_bias_table[-3::]], dim=0)
+
+                # get pair-wise relative position index for each token inside the window
+                coords_h = torch.arange(training_window_size[0])
+                coords_w = torch.arange(training_window_size[1])
+                coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+                coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+                relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+                relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+                relative_coords[:, :, 0] += training_window_size[0] - 1  # shift to start from 0
+                relative_coords[:, :, 1] += training_window_size[1] - 1
+                relative_coords[:, :, 0] *= 2 * training_window_size[1] - 1
+                relative_position_index = \
+                    torch.zeros(size=(training_window_size[0] * training_window_size[1] + 1,) * 2,
+                                dtype=relative_coords.dtype)
+                relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+                relative_position_index[0, 0:] = new_num_relative_distance - 3
+                relative_position_index[0:, 0] = new_num_relative_distance - 2
+                relative_position_index[0, 0] = new_num_relative_distance - 1
+
+                relative_position_bias = \
+                    new_relative_position_bias_table[relative_position_index.view(-1)].view(
+                        training_window_size[0] * training_window_size[1] + 1,
+                        training_window_size[0] * training_window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+                relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+                attn = attn + relative_position_bias.unsqueeze(0)
+
+        if rel_pos_bias is not None:
+            attn = attn + rel_pos_bias
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+                 window_size=None, attn_head_dim=None):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if init_values is not None:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+    def forward(self, x, rel_pos_bias=None, training_window_size=None):
+        if self.gamma_1 is None:
+            x = x + self.drop_path(
+                self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, training_window_size=training_window_size))
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias,
+                                                            training_window_size=training_window_size))
+            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=[224, 224], patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches_w = self.patch_shape[0]
+        self.num_patches_h = self.patch_shape[1]
+        # the so-called patch_shape is the patch shape during pre-training
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x, position_embedding=None, **kwargs):
+        # FIXME look at relaxing size constraints
+        # assert H == self.img_size[0] and W == self.img_size[1], \
+        #     f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x)
+        Hp, Wp = x.shape[2], x.shape[3]
+
+        if position_embedding is not None:
+            # interpolate the position embedding to the corresponding size
+            position_embedding = position_embedding.view(1, self.patch_shape[0], self.patch_shape[1], -1).permute(0, 3,
+                                                                                                                  1, 2)
+            position_embedding = F.interpolate(position_embedding, size=(Hp, Wp), mode='bicubic')
+            x = x + position_embedding
+
+        x = x.flatten(2).transpose(1, 2)
+        return x, (Hp, Wp)
+
+
+class HybridEmbed(nn.Module):
+    """ CNN Feature Map Embedding
+    Extract feature map from CNN, flatten, project to embedding dim.
+    """
+
+    def __init__(self, backbone, img_size=[224, 224], feature_size=None, in_chans=3, embed_dim=768):
+        super().__init__()
+        assert isinstance(backbone, nn.Module)
+        img_size = to_2tuple(img_size)
+        self.img_size = img_size
+        self.backbone = backbone
+        if feature_size is None:
+            with torch.no_grad():
+                # FIXME this is hacky, but most reliable way of determining the exact dim of the output feature
+                # map for all networks, the feature metadata has reliable channel and stride info, but using
+                # stride to calc feature dim requires info about padding of each stage that isn't captured.
+                training = backbone.training
+                if training:
+                    backbone.eval()
+                o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))[-1]
+                feature_size = o.shape[-2:]
+                feature_dim = o.shape[1]
+                backbone.train(training)
+        else:
+            feature_size = to_2tuple(feature_size)
+            feature_dim = self.backbone.feature_info.channels()[-1]
+        self.num_patches = feature_size[0] * feature_size[1]
+        self.proj = nn.Linear(feature_dim, embed_dim)
+
+    def forward(self, x):
+        x = self.backbone(x)[-1]
+        x = x.flatten(2).transpose(1, 2)
+        x = self.proj(x)
+        return x
+
+
+class RelativePositionBias(nn.Module):
+
+    def __init__(self, window_size, num_heads):
+        super().__init__()
+        self.window_size = window_size
+        self.num_heads = num_heads
+        self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+        # cls to token & token 2 cls & cls to cls
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(window_size[0])
+        coords_w = torch.arange(window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+        relative_position_index = \
+            torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
+        relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        relative_position_index[0, 0:] = self.num_relative_distance - 3
+        relative_position_index[0:, 0] = self.num_relative_distance - 2
+        relative_position_index[0, 0] = self.num_relative_distance - 1
+
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        # trunc_normal_(self.relative_position_bias_table, std=.02)
+
+    def forward(self, training_window_size):
+        if training_window_size == self.window_size:
+            relative_position_bias = \
+                self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                    self.window_size[0] * self.window_size[1] + 1,
+                    self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        else:
+            training_window_size = tuple(training_window_size.tolist())
+            new_num_relative_distance = (2 * training_window_size[0] - 1) * (2 * training_window_size[1] - 1) + 3
+            # new_num_relative_dis 为 所有可能的相对位置选项，包含cls-cls，tok-cls，与cls-tok
+            new_relative_position_bias_table = F.interpolate(
+                self.relative_position_bias_table[:-3, :].permute(1, 0).view(1, self.num_heads,
+                                                                             2 * self.window_size[0] - 1,
+                                                                             2 * self.window_size[1] - 1),
+                size=(2 * training_window_size[0] - 1, 2 * training_window_size[1] - 1), mode='bicubic',
+                align_corners=False)
+            new_relative_position_bias_table = new_relative_position_bias_table.view(self.num_heads,
+                                                                                     new_num_relative_distance - 3).permute(
+                1, 0)
+            new_relative_position_bias_table = torch.cat(
+                [new_relative_position_bias_table, self.relative_position_bias_table[-3::]], dim=0)
+
+            # get pair-wise relative position index for each token inside the window
+            coords_h = torch.arange(training_window_size[0])
+            coords_w = torch.arange(training_window_size[1])
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+            coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+            relative_coords[:, :, 0] += training_window_size[0] - 1  # shift to start from 0
+            relative_coords[:, :, 1] += training_window_size[1] - 1
+            relative_coords[:, :, 0] *= 2 * training_window_size[1] - 1
+            relative_position_index = \
+                torch.zeros(size=(training_window_size[0] * training_window_size[1] + 1,) * 2,
+                            dtype=relative_coords.dtype)
+            relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+            relative_position_index[0, 0:] = new_num_relative_distance - 3
+            relative_position_index[0:, 0] = new_num_relative_distance - 2
+            relative_position_index[0, 0] = new_num_relative_distance - 1
+
+            relative_position_bias = \
+                new_relative_position_bias_table[relative_position_index.view(-1)].view(
+                    training_window_size[0] * training_window_size[1] + 1,
+                    training_window_size[0] * training_window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+
+        return relative_position_bias
+
+
+class BEiT(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+
+    def __init__(self,
+                 img_size=[224, 224],
+                 patch_size=16,
+                 in_chans=3,
+                 num_classes=80,
+                 embed_dim=768,
+                 depth=12,
+                 num_heads=12,
+                 mlp_ratio=4.,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 drop_rate=0.,
+                 attn_drop_rate=0.,
+                 drop_path_rate=0.,
+                 hybrid_backbone=None,
+                 norm_layer=None,
+                 init_values=None,
+                 use_abs_pos_emb=False,
+                 use_rel_pos_bias=False,
+                 use_shared_rel_pos_bias=False,
+                 use_checkpoint=True,
+                 pretrained=None,
+                 out_features=None,
+                 ):
+
+        super(BEiT, self).__init__()
+
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.use_checkpoint = use_checkpoint
+
+        if hybrid_backbone is not None:
+            self.patch_embed = HybridEmbed(
+                hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim)
+        else:
+            self.patch_embed = PatchEmbed(
+                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+        self.out_features = out_features
+        self.out_indices = [int(name[5:]) for name in out_features]
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        # self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        if use_abs_pos_emb:
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        else:
+            self.pos_embed = None
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        self.use_shared_rel_pos_bias = use_shared_rel_pos_bias
+        if use_shared_rel_pos_bias:
+            self.rel_pos_bias = RelativePositionBias(window_size=self.patch_embed.patch_shape, num_heads=num_heads)
+        else:
+            self.rel_pos_bias = None
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.use_rel_pos_bias = use_rel_pos_bias
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None)
+            for i in range(depth)])
+
+        # trunc_normal_(self.mask_token, std=.02)
+
+        if patch_size == 16:
+            self.fpn1 = nn.Sequential(
+                nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
+                # nn.SyncBatchNorm(embed_dim),
+                nn.BatchNorm2d(embed_dim),
+                nn.GELU(),
+                nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
+            )
+
+            self.fpn2 = nn.Sequential(
+                nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
+            )
+
+            self.fpn3 = nn.Identity()
+
+            self.fpn4 = nn.MaxPool2d(kernel_size=2, stride=2)
+        elif patch_size == 8:
+            self.fpn1 = nn.Sequential(
+                nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
+            )
+
+            self.fpn2 = nn.Identity()
+
+            self.fpn3 = nn.Sequential(
+                nn.MaxPool2d(kernel_size=2, stride=2),
+            )
+
+            self.fpn4 = nn.Sequential(
+                nn.MaxPool2d(kernel_size=4, stride=4),
+            )
+
+        if self.pos_embed is not None:
+            trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+        self.fix_init_weight()
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    '''
+    def init_weights(self):
+        """Initialize the weights in backbone.
+
+        Args:
+            pretrained (str, optional): Path to pre-trained weights.
+                Defaults to None.
+        """
+        logger = get_root_logger()
+
+        if self.pos_embed is not None:
+            trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+        self.fix_init_weight()
+
+        if self.init_cfg is None:
+            logger.warn(f'No pre-trained weights for '
+                        f'{self.__class__.__name__}, '
+                        f'training start from scratch')
+        else:
+            assert 'checkpoint' in self.init_cfg, f'Only support ' \
+                                                  f'specify `Pretrained` in ' \
+                                                  f'`init_cfg` in ' \
+                                                  f'{self.__class__.__name__} '
+            logger.info(f"Will load ckpt from {self.init_cfg['checkpoint']}")
+            load_checkpoint(self,
+                            filename=self.init_cfg['checkpoint'],
+                            strict=False,
+                            logger=logger,
+                            beit_spec_expand_rel_pos = self.use_rel_pos_bias,
+                            )
+    '''
+
+    def get_num_layers(self):
+        return len(self.blocks)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def forward_features(self, x):
+        B, C, H, W = x.shape
+        x, (Hp, Wp) = self.patch_embed(x, self.pos_embed[:, 1:, :] if self.pos_embed is not None else None)
+        # Hp, Wp are HW for patches
+        batch_size, seq_len, _ = x.size()
+
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        if self.pos_embed is not None:
+            cls_tokens = cls_tokens + self.pos_embed[:, :1, :]
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = self.pos_drop(x)
+
+        features = []
+        training_window_size = torch.tensor([Hp, Wp])
+
+        rel_pos_bias = self.rel_pos_bias(training_window_size) if self.rel_pos_bias is not None else None
+
+        for i, blk in enumerate(self.blocks):
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x, rel_pos_bias, training_window_size)
+            else:
+                x = blk(x, rel_pos_bias=rel_pos_bias, training_window_size=training_window_size)
+            if i in self.out_indices:
+                xp = x[:, 1:, :].permute(0, 2, 1).reshape(B, -1, Hp, Wp)
+                features.append(xp.contiguous())
+
+        ops = [self.fpn1, self.fpn2, self.fpn3, self.fpn4]
+        for i in range(len(features)):
+            features[i] = ops[i](features[i])
+
+        feat_out = {}
+
+        for name, value in zip(self.out_features, features):
+            feat_out[name] = value
+
+        return feat_out
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        return x
+
+
+def beit_base_patch16(pretrained=False, **kwargs):
+    model = BEiT(
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        init_values=None,
+        **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def beit_large_patch16(pretrained=False, **kwargs):
+    model = BEiT(
+        patch_size=16,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        init_values=None,
+        **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def dit_base_patch16(pretrained=False, **kwargs):
+    model = BEiT(
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        init_values=0.1,
+        **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def dit_large_patch16(pretrained=False, **kwargs):
+    model = BEiT(
+        patch_size=16,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        init_values=1e-5,
+        **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+if __name__ == '__main__':
+    model = BEiT(use_checkpoint=True, use_shared_rel_pos_bias=True)
+    model = model.to("cuda:0")
+    input1 = torch.rand(2, 3, 512, 762).to("cuda:0")
+    input2 = torch.rand(2, 3, 800, 1200).to("cuda:0")
+    input3 = torch.rand(2, 3, 720, 1000).to("cuda:0")
+    output1 = model(input1)
+    output2 = model(input2)
+    output3 = model(input3)
+    print("all done")

dit/object_detection/ditod/config.py

@@ -0,0 +1,32 @@
+from detectron2.config import CfgNode as CN
+
+
+def add_vit_config(cfg):
+    """
+    Add config for VIT.
+    """
+    _C = cfg
+
+    _C.MODEL.VIT = CN()
+
+    # CoaT model name.
+    _C.MODEL.VIT.NAME = ""
+
+    # Output features from CoaT backbone.
+    _C.MODEL.VIT.OUT_FEATURES = ["layer3", "layer5", "layer7", "layer11"]
+
+    _C.MODEL.VIT.IMG_SIZE = [224, 224]
+
+    _C.MODEL.VIT.POS_TYPE = "shared_rel"
+
+    _C.MODEL.VIT.DROP_PATH = 0.
+
+    _C.MODEL.VIT.MODEL_KWARGS = "{}"
+
+    _C.SOLVER.OPTIMIZER = "ADAMW"
+
+    _C.SOLVER.BACKBONE_MULTIPLIER = 1.0
+
+    _C.AUG = CN()
+
+    _C.AUG.DETR = False

dit/object_detection/ditod/dataset_mapper.py

@@ -0,0 +1,124 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# from https://github.com/facebookresearch/detr/blob/main/d2/detr/dataset_mapper.py
+
+
+import copy
+import logging
+
+import numpy as np
+import torch
+
+from detectron2.data import detection_utils as utils
+from detectron2.data import transforms as T
+
+__all__ = ["DetrDatasetMapper"]
+
+
+def build_transform_gen(cfg, is_train):
+    """
+    Create a list of :class:`TransformGen` from config.
+    Returns:
+        list[TransformGen]
+    """
+    if is_train:
+        min_size = cfg.INPUT.MIN_SIZE_TRAIN
+        max_size = cfg.INPUT.MAX_SIZE_TRAIN
+        sample_style = cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING
+    else:
+        min_size = cfg.INPUT.MIN_SIZE_TEST
+        max_size = cfg.INPUT.MAX_SIZE_TEST
+        sample_style = "choice"
+    if sample_style == "range":
+        assert len(min_size) == 2, "more than 2 ({}) min_size(s) are provided for ranges".format(len(min_size))
+
+    logger = logging.getLogger(__name__)
+    tfm_gens = []
+    if is_train:
+        tfm_gens.append(T.RandomFlip())
+    tfm_gens.append(T.ResizeShortestEdge(min_size, max_size, sample_style))
+    if is_train:
+        logger.info("TransformGens used in training: " + str(tfm_gens))
+    return tfm_gens
+
+
+class DetrDatasetMapper:
+    """
+    A callable which takes a dataset dict in Detectron2 Dataset format,
+    and map it into a format used by DETR.
+
+    The callable currently does the following:
+
+    1. Read the image from "file_name"
+    2. Applies geometric transforms to the image and annotation
+    3. Find and applies suitable cropping to the image and annotation
+    4. Prepare image and annotation to Tensors
+    """
+
+    def __init__(self, cfg, is_train=True):
+        if cfg.INPUT.CROP.ENABLED and is_train:
+            self.crop_gen = [
+                T.ResizeShortestEdge([400, 500, 600], sample_style="choice"),
+                T.RandomCrop(cfg.INPUT.CROP.TYPE, cfg.INPUT.CROP.SIZE),
+            ]
+        else:
+            self.crop_gen = None
+
+        self.mask_on = cfg.MODEL.MASK_ON
+        self.tfm_gens = build_transform_gen(cfg, is_train)
+        logging.getLogger(__name__).info(
+            "Full TransformGens used in training: {}, crop: {}".format(str(self.tfm_gens), str(self.crop_gen))
+        )
+
+        self.img_format = cfg.INPUT.FORMAT
+        self.is_train = is_train
+
+    def __call__(self, dataset_dict):
+        """
+        Args:
+            dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.
+
+        Returns:
+            dict: a format that builtin models in detectron2 accept
+        """
+        dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below
+        image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
+        utils.check_image_size(dataset_dict, image)
+
+        if self.crop_gen is None:
+            image, transforms = T.apply_transform_gens(self.tfm_gens, image)
+        else:
+            if np.random.rand() > 0.5:
+                image, transforms = T.apply_transform_gens(self.tfm_gens, image)
+            else:
+                image, transforms = T.apply_transform_gens(
+                    self.tfm_gens[:-1] + self.crop_gen + self.tfm_gens[-1:], image
+                )
+
+        image_shape = image.shape[:2]  # h, w
+
+        # Pytorch's dataloader is efficient on torch.Tensor due to shared-memory,
+        # but not efficient on large generic data structures due to the use of pickle & mp.Queue.
+        # Therefore it's important to use torch.Tensor.
+        dataset_dict["image"] = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1)))
+
+        if not self.is_train:
+            # USER: Modify this if you want to keep them for some reason.
+            dataset_dict.pop("annotations", None)
+            return dataset_dict
+
+        if "annotations" in dataset_dict:
+            # USER: Modify this if you want to keep them for some reason.
+            for anno in dataset_dict["annotations"]:
+                if not self.mask_on:
+                    anno.pop("segmentation", None)
+                anno.pop("keypoints", None)
+
+            # USER: Implement additional transformations if you have other types of data
+            annos = [
+                utils.transform_instance_annotations(obj, transforms, image_shape)
+                for obj in dataset_dict.pop("annotations")
+                if obj.get("iscrowd", 0) == 0
+            ]
+            instances = utils.annotations_to_instances(annos, image_shape)
+            dataset_dict["instances"] = utils.filter_empty_instances(instances)
+        return dataset_dict

dit/object_detection/ditod/deit.py

@@ -0,0 +1,476 @@
+"""
+Mostly copy-paste from DINO and timm library:
+https://github.com/facebookresearch/dino
+https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+"""
+import warnings
+
+import math
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import trunc_normal_, drop_path, to_2tuple
+from functools import partial
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'mean': (0.5, 0.5, 0.5), 'std': (0.5, 0.5, 0.5),
+        **kwargs
+    }
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+    def extra_repr(self) -> str:
+        return 'p={}'.format(self.drop_prob)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        q, k, v = self.qkv(x).reshape(B, N, 3, self.num_heads,
+                                      C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(
+            drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
+                       act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+
+        self.window_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+
+        self.num_patches_w, self.num_patches_h = self.window_size
+
+        self.num_patches = self.window_size[0] * self.window_size[1]
+        self.img_size = img_size
+        self.patch_size = patch_size
+
+        self.proj = nn.Conv2d(in_chans, embed_dim,
+                              kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        x = self.proj(x)
+        return x
+
+
+class HybridEmbed(nn.Module):
+    """ CNN Feature Map Embedding
+    Extract feature map from CNN, flatten, project to embedding dim.
+    """
+
+    def __init__(self, backbone, img_size=224, feature_size=None, in_chans=3, embed_dim=768):
+        super().__init__()
+        assert isinstance(backbone, nn.Module)
+        img_size = to_2tuple(img_size)
+        self.img_size = img_size
+        self.backbone = backbone
+        if feature_size is None:
+            with torch.no_grad():
+                # FIXME this is hacky, but most reliable way of determining the exact dim of the output feature
+                # map for all networks, the feature metadata has reliable channel and stride info, but using
+                # stride to calc feature dim requires info about padding of each stage that isn't captured.
+                training = backbone.training
+                if training:
+                    backbone.eval()
+                o = self.backbone(torch.zeros(
+                    1, in_chans, img_size[0], img_size[1]))[-1]
+                feature_size = o.shape[-2:]
+                feature_dim = o.shape[1]
+                backbone.train(training)
+        else:
+            feature_size = to_2tuple(feature_size)
+            feature_dim = self.backbone.feature_info.channels()[-1]
+        self.num_patches = feature_size[0] * feature_size[1]
+        self.proj = nn.Linear(feature_dim, embed_dim)
+
+    def forward(self, x):
+        x = self.backbone(x)[-1]
+        x = x.flatten(2).transpose(1, 2)
+        x = self.proj(x)
+        return x
+
+
+class ViT(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+
+    def __init__(self,
+                 model_name='vit_base_patch16_224',
+                 img_size=384,
+                 patch_size=16,
+                 in_chans=3,
+                 embed_dim=1024,
+                 depth=24,
+                 num_heads=16,
+                 num_classes=19,
+                 mlp_ratio=4.,
+                 qkv_bias=True,
+                 qk_scale=None,
+                 drop_rate=0.1,
+                 attn_drop_rate=0.,
+                 drop_path_rate=0.,
+                 hybrid_backbone=None,
+                 norm_layer=partial(nn.LayerNorm, eps=1e-6),
+                 norm_cfg=None,
+                 pos_embed_interp=False,
+                 random_init=False,
+                 align_corners=False,
+                 use_checkpoint=False,
+                 num_extra_tokens=1,
+                 out_features=None,
+                 **kwargs,
+                 ):
+
+        super(ViT, self).__init__()
+        self.model_name = model_name
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+        self.depth = depth
+        self.num_heads = num_heads
+        self.num_classes = num_classes
+        self.mlp_ratio = mlp_ratio
+        self.qkv_bias = qkv_bias
+        self.qk_scale = qk_scale
+        self.drop_rate = drop_rate
+        self.attn_drop_rate = attn_drop_rate
+        self.drop_path_rate = drop_path_rate
+        self.hybrid_backbone = hybrid_backbone
+        self.norm_layer = norm_layer
+        self.norm_cfg = norm_cfg
+        self.pos_embed_interp = pos_embed_interp
+        self.random_init = random_init
+        self.align_corners = align_corners
+        self.use_checkpoint = use_checkpoint
+        self.num_extra_tokens = num_extra_tokens
+        self.out_features = out_features
+        self.out_indices = [int(name[5:]) for name in out_features]
+
+        # self.num_stages = self.depth
+        # self.out_indices = tuple(range(self.num_stages))
+
+        if self.hybrid_backbone is not None:
+            self.patch_embed = HybridEmbed(
+                self.hybrid_backbone, img_size=self.img_size, in_chans=self.in_chans, embed_dim=self.embed_dim)
+        else:
+            self.patch_embed = PatchEmbed(
+                img_size=self.img_size, patch_size=self.patch_size, in_chans=self.in_chans, embed_dim=self.embed_dim)
+        self.num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+
+        if self.num_extra_tokens == 2:
+            self.dist_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+
+        self.pos_embed = nn.Parameter(torch.zeros(
+            1, self.num_patches + self.num_extra_tokens, self.embed_dim))
+        self.pos_drop = nn.Dropout(p=self.drop_rate)
+
+        # self.num_extra_tokens = self.pos_embed.shape[-2] - self.num_patches
+        dpr = [x.item() for x in torch.linspace(0, self.drop_path_rate,
+                                                self.depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=self.embed_dim, num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, qkv_bias=self.qkv_bias,
+                qk_scale=self.qk_scale,
+                drop=self.drop_rate, attn_drop=self.attn_drop_rate, drop_path=dpr[i], norm_layer=self.norm_layer)
+            for i in range(self.depth)])
+
+        # NOTE as per official impl, we could have a pre-logits representation dense layer + tanh here
+        # self.repr = nn.Linear(embed_dim, representation_size)
+        # self.repr_act = nn.Tanh()
+
+        if patch_size == 16:
+            self.fpn1 = nn.Sequential(
+                nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
+                nn.SyncBatchNorm(embed_dim),
+                nn.GELU(),
+                nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
+            )
+
+            self.fpn2 = nn.Sequential(
+                nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
+            )
+
+            self.fpn3 = nn.Identity()
+
+            self.fpn4 = nn.MaxPool2d(kernel_size=2, stride=2)
+        elif patch_size == 8:
+            self.fpn1 = nn.Sequential(
+                nn.ConvTranspose2d(embed_dim, embed_dim, kernel_size=2, stride=2),
+            )
+
+            self.fpn2 = nn.Identity()
+
+            self.fpn3 = nn.Sequential(
+                nn.MaxPool2d(kernel_size=2, stride=2),
+            )
+
+            self.fpn4 = nn.Sequential(
+                nn.MaxPool2d(kernel_size=4, stride=4),
+            )
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        if self.num_extra_tokens==2:
+            trunc_normal_(self.dist_token, std=0.2)
+        self.apply(self._init_weights)
+        # self.fix_init_weight()
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    '''
+    def init_weights(self):
+        logger = get_root_logger()
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+        if self.init_cfg is None:
+            logger.warn(f'No pre-trained weights for '
+                        f'{self.__class__.__name__}, '
+                        f'training start from scratch')
+        else:
+            assert 'checkpoint' in self.init_cfg, f'Only support ' \
+                                                  f'specify `Pretrained` in ' \
+                                                  f'`init_cfg` in ' \
+                                                  f'{self.__class__.__name__} '
+            logger.info(f"Will load ckpt from {self.init_cfg['checkpoint']}")
+            load_checkpoint(self, filename=self.init_cfg['checkpoint'], strict=False, logger=logger)
+    '''
+
+    def get_num_layers(self):
+        return len(self.blocks)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def _conv_filter(self, state_dict, patch_size=16):
+        """ convert patch embedding weight from manual patchify + linear proj to conv"""
+        out_dict = {}
+        for k, v in state_dict.items():
+            if 'patch_embed.proj.weight' in k:
+                v = v.reshape((v.shape[0], 3, patch_size, patch_size))
+            out_dict[k] = v
+        return out_dict
+
+    def to_2D(self, x):
+        n, hw, c = x.shape
+        h = w = int(math.sqrt(hw))
+        x = x.transpose(1, 2).reshape(n, c, h, w)
+        return x
+
+    def to_1D(self, x):
+        n, c, h, w = x.shape
+        x = x.reshape(n, c, -1).transpose(1, 2)
+        return x
+
+    def interpolate_pos_encoding(self, x, w, h):
+        npatch = x.shape[1] - self.num_extra_tokens
+        N = self.pos_embed.shape[1] - self.num_extra_tokens
+        if npatch == N and w == h:
+            return self.pos_embed
+
+        class_ORdist_pos_embed = self.pos_embed[:, 0:self.num_extra_tokens]
+
+        patch_pos_embed = self.pos_embed[:, self.num_extra_tokens:]
+
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size[0]
+        h0 = h // self.patch_embed.patch_size[1]
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode='bicubic',
+        )
+        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+        return torch.cat((class_ORdist_pos_embed, patch_pos_embed), dim=1)
+
+    def prepare_tokens(self, x, mask=None):
+        B, nc, w, h = x.shape
+        # patch linear embedding
+        x = self.patch_embed(x)
+
+        # mask image modeling
+        if mask is not None:
+            x = self.mask_model(x, mask)
+        x = x.flatten(2).transpose(1, 2)
+
+        # add the [CLS] token to the embed patch tokens
+        all_tokens = [self.cls_token.expand(B, -1, -1)]
+
+        if self.num_extra_tokens == 2:
+            dist_tokens = self.dist_token.expand(B, -1, -1)
+            all_tokens.append(dist_tokens)
+        all_tokens.append(x)
+
+        x = torch.cat(all_tokens, dim=1)
+
+        # add positional encoding to each token
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        return self.pos_drop(x)
+
+    def forward_features(self, x):
+        # print(f"==========shape of x is {x.shape}==========")
+        B, _, H, W = x.shape
+        Hp, Wp = H // self.patch_size, W // self.patch_size
+        x = self.prepare_tokens(x)
+
+        features = []
+        for i, blk in enumerate(self.blocks):
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x)
+            if i in self.out_indices:
+                xp = x[:, self.num_extra_tokens:, :].permute(0, 2, 1).reshape(B, -1, Hp, Wp)
+                features.append(xp.contiguous())
+
+        ops = [self.fpn1, self.fpn2, self.fpn3, self.fpn4]
+        for i in range(len(features)):
+            features[i] = ops[i](features[i])
+
+        feat_out = {}
+
+        for name, value in zip(self.out_features, features):
+            feat_out[name] = value
+
+        return feat_out
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        return x
+
+
+def deit_base_patch16(pretrained=False, **kwargs):
+    model = ViT(
+        patch_size=16,
+        drop_rate=0.,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        num_classes=1000,
+        mlp_ratio=4.,
+        qkv_bias=True,
+        use_checkpoint=True,
+        num_extra_tokens=2,
+        **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def mae_base_patch16(pretrained=False, **kwargs):
+    model = ViT(
+        patch_size=16,
+        drop_rate=0.,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        num_classes=1000,
+        mlp_ratio=4.,
+        qkv_bias=True,
+        use_checkpoint=True,
+        num_extra_tokens=1,
+        **kwargs)
+    model.default_cfg = _cfg()
+    return model

dit/object_detection/ditod/icdar_evaluation.py

@@ -0,0 +1,100 @@
+import copy
+import itertools
+import os
+import os.path as osp
+import shutil
+from collections import OrderedDict
+from xml.dom.minidom import Document
+
+import detectron2.utils.comm as comm
+import torch
+from detectron2.evaluation import COCOEvaluator
+from detectron2.utils.file_io import PathManager
+
+from .table_evaluation.evaluate import calc_table_score
+
+
+class ICDAREvaluator(COCOEvaluator):
+    def evaluate(self, img_ids=None):
+        """
+        Args:
+            img_ids: a list of image IDs to evaluate on. Default to None for the whole dataset
+        """
+        if self._distributed:
+            comm.synchronize()
+            predictions = comm.gather(self._predictions, dst=0)
+            predictions = list(itertools.chain(*predictions))
+
+            if not comm.is_main_process():
+                return {}
+        else:
+            predictions = self._predictions
+
+        if len(predictions) == 0:
+            self._logger.warning("[COCOEvaluator] Did not receive valid predictions.")
+            return {}
+
+        if self._output_dir:
+            PathManager.mkdirs(self._output_dir)
+            file_path = os.path.join(self._output_dir, "instances_predictions.pth")
+            with PathManager.open(file_path, "wb") as f:
+                torch.save(predictions, f)
+
+        self._results = OrderedDict()
+        if "proposals" in predictions[0]:
+            self._eval_box_proposals(predictions)
+        if "instances" in predictions[0]:
+            self._eval_predictions(predictions, img_ids=img_ids)
+            self.evaluate_table(predictions)
+        # Copy so the caller can do whatever with results
+        return copy.deepcopy(self._results)
+
+    def evaluate_table(self, predictions):
+        xml_dir = self.convert_to_xml(predictions)
+        results = calc_table_score(xml_dir)
+        self._results["wF1"] = results['wF1']
+
+    def convert_to_xml(self, predictions):
+        output_dir = osp.join(self._output_dir, "xml_results")
+        if os.path.exists(output_dir):
+            shutil.rmtree(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        coco_results = list(itertools.chain(*[x["instances"] for x in predictions]))
+        results_dict = {}
+        for result in coco_results:
+            if result["score"] < 0.7:
+                continue
+            image_id = result["image_id"]
+            if image_id not in results_dict:
+                results_dict[image_id] = []
+
+            results_dict[image_id].append(result)
+
+        for image_id, tables in results_dict.items():
+            file_name = f"cTDaR_t{image_id:05d}.jpg"
+            doc = Document()
+            root = doc.createElement('document')
+            root.setAttribute('filename', file_name)
+            doc.appendChild(root)
+            for table_id, table in enumerate(tables, start=1):
+                nodeManager = doc.createElement('table')
+                nodeManager.setAttribute('id', str(table_id))
+                bbox = list(map(int, table['bbox']))
+                bbox_str = '{},{} {},{} {},{} {},{}'.format(bbox[0], bbox[1],
+                                                            bbox[0], bbox[1] + bbox[3],
+                                                            bbox[0] + bbox[2], bbox[1] + bbox[3],
+                                                            bbox[0] + bbox[2], bbox[1])
+                nodeCoords = doc.createElement('Coords')
+                nodeCoords.setAttribute('points', bbox_str)
+                nodeManager.appendChild(nodeCoords)
+                root.appendChild(nodeManager)
+            filename = '{}-result.xml'.format(file_name[:-4])
+            fp = open(os.path.join(output_dir, filename), 'w')
+            doc.writexml(fp, indent='', addindent='\t', newl='\n', encoding="utf-8")
+            fp.flush()
+            fp.close()
+        return output_dir
+
+
+if __name__ == '__main__':
+    pass

dit/object_detection/ditod/mycheckpointer.py

@@ -0,0 +1,257 @@
+from detectron2.checkpoint import DetectionCheckpointer
+
+from typing import Any
+import torch
+import torch.nn as nn
+from fvcore.common.checkpoint import _IncompatibleKeys, _strip_prefix_if_present, TORCH_VERSION, quantization, \
+    ObserverBase, FakeQuantizeBase
+from torch import distributed as dist
+from scipy import interpolate
+import numpy as np
+import torch.nn.functional as F
+
+
+def append_prefix(k):
+    prefix = 'backbone.bottom_up.backbone.'
+    return prefix + k if not k.startswith(prefix) else k
+
+
+def modify_ckpt_state(model, state_dict, logger=None):
+    # reshape absolute position embedding for Swin
+    if state_dict.get(append_prefix('absolute_pos_embed')) is not None:
+        absolute_pos_embed = state_dict[append_prefix('absolute_pos_embed')]
+        N1, L, C1 = absolute_pos_embed.size()
+        N2, C2, H, W = model.backbone.bottom_up.backbone.absolute_pos_embed.size()
+        if N1 != N2 or C1 != C2 or L != H * W:
+            logger.warning("Error in loading absolute_pos_embed, pass")
+        else:
+            state_dict[append_prefix('absolute_pos_embed')] = absolute_pos_embed.view(N2, H, W, C2).permute(0, 3, 1, 2)
+
+    def get_dist_info():
+        if dist.is_available() and dist.is_initialized():
+            rank = dist.get_rank()
+            world_size = dist.get_world_size()
+        else:
+            rank = 0
+            world_size = 1
+        return rank, world_size
+
+    rank, _ = get_dist_info()
+    all_keys = list(state_dict.keys())
+    for key in all_keys:
+        if "relative_position_index" in key:
+            state_dict.pop(key)
+
+        if "relative_position_bias_table" in key:
+            rel_pos_bias = state_dict[key]
+            src_num_pos, num_attn_heads = rel_pos_bias.size()
+            if key not in model.state_dict():
+                continue
+            dst_num_pos, _ = model.state_dict()[key].size()
+            dst_patch_shape = model.backbone.bottom_up.backbone.patch_embed.patch_shape
+            if dst_patch_shape[0] != dst_patch_shape[1]:
+                raise NotImplementedError()
+            num_extra_tokens = dst_num_pos - (dst_patch_shape[0] * 2 - 1) * (dst_patch_shape[1] * 2 - 1)
+            src_size = int((src_num_pos - num_extra_tokens) ** 0.5)
+            dst_size = int((dst_num_pos - num_extra_tokens) ** 0.5)
+            if src_size != dst_size:
+                if rank == 0:
+                    print("Position interpolate for %s from %dx%d to %dx%d" % (
+                        key, src_size, src_size, dst_size, dst_size))
+                extra_tokens = rel_pos_bias[-num_extra_tokens:, :]
+                rel_pos_bias = rel_pos_bias[:-num_extra_tokens, :]
+
+                def geometric_progression(a, r, n):
+                    return a * (1.0 - r ** n) / (1.0 - r)
+
+                left, right = 1.01, 1.5
+                while right - left > 1e-6:
+                    q = (left + right) / 2.0
+                    gp = geometric_progression(1, q, src_size // 2)
+                    if gp > dst_size // 2:
+                        right = q
+                    else:
+                        left = q
+
+                # if q > 1.13492:
+                #     q = 1.13492
+
+                dis = []
+                cur = 1
+                for i in range(src_size // 2):
+                    dis.append(cur)
+                    cur += q ** (i + 1)
+
+                r_ids = [-_ for _ in reversed(dis)]
+
+                x = r_ids + [0] + dis
+                y = r_ids + [0] + dis
+
+                t = dst_size // 2.0
+                dx = np.arange(-t, t + 0.1, 1.0)
+                dy = np.arange(-t, t + 0.1, 1.0)
+                if rank == 0:
+                    print("x = {}".format(x))
+                    print("dx = {}".format(dx))
+
+                all_rel_pos_bias = []
+
+                for i in range(num_attn_heads):
+                    z = rel_pos_bias[:, i].view(src_size, src_size).float().numpy()
+                    f = interpolate.interp2d(x, y, z, kind='cubic')
+                    all_rel_pos_bias.append(
+                        torch.Tensor(f(dx, dy)).contiguous().view(-1, 1).to(rel_pos_bias.device))
+
+                rel_pos_bias = torch.cat(all_rel_pos_bias, dim=-1)
+                new_rel_pos_bias = torch.cat((rel_pos_bias, extra_tokens), dim=0)
+                state_dict[key] = new_rel_pos_bias
+
+    if append_prefix('pos_embed') in state_dict:
+        pos_embed_checkpoint = state_dict[append_prefix('pos_embed')]
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.backbone.bottom_up.backbone.patch_embed.num_patches
+        num_extra_tokens = model.backbone.bottom_up.backbone.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        # new_size = int(num_patches ** 0.5)
+        new_size_w = model.backbone.bottom_up.backbone.patch_embed.num_patches_w
+        new_size_h = model.backbone.bottom_up.backbone.patch_embed.num_patches_h
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size_h or orig_size != new_size_w:
+            if rank == 0:
+                print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size_w, new_size_h))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size_w, new_size_h), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            state_dict[append_prefix('pos_embed')] = new_pos_embed
+
+    # interpolate position bias table if needed
+    relative_position_bias_table_keys = [k for k in state_dict.keys() if "relative_position_bias_table" in k]
+    for table_key in relative_position_bias_table_keys:
+        table_pretrained = state_dict[table_key]
+        if table_key not in model.state_dict():
+            continue
+        table_current = model.state_dict()[table_key]
+        L1, nH1 = table_pretrained.size()
+        L2, nH2 = table_current.size()
+        if nH1 != nH2:
+            logger.warning(f"Error in loading {table_key}, pass")
+        else:
+            if L1 != L2:
+                S1 = int(L1 ** 0.5)
+                S2 = int(L2 ** 0.5)
+                table_pretrained_resized = F.interpolate(
+                    table_pretrained.permute(1, 0).view(1, nH1, S1, S1),
+                    size=(S2, S2), mode='bicubic')
+                state_dict[table_key] = table_pretrained_resized.view(nH2, L2).permute(1, 0)
+
+    if append_prefix('rel_pos_bias.relative_position_bias_table') in state_dict and \
+            model.backbone.bottom_up.backbone.use_rel_pos_bias and \
+            not model.backbone.bottom_up.backbone.use_shared_rel_pos_bias and \
+            append_prefix('blocks.0.attn.relative_position_bias_table') not in state_dict:
+        logger.info("[BEIT] Expand the shared relative position embedding to each transformer block. ")
+        num_layers = model.backbone.bottom_up.backbone.get_num_layers()
+        rel_pos_bias = state_dict[append_prefix("rel_pos_bias.relative_position_bias_table")]
+        for i in range(num_layers):
+            state_dict["blocks.%d.attn.relative_position_bias_table" % i] = rel_pos_bias.clone()
+        state_dict.pop(append_prefix("rel_pos_bias.relative_position_bias_table"))
+
+    return state_dict
+
+
+class MyDetectionCheckpointer(DetectionCheckpointer):
+    def _load_model(self, checkpoint: Any) -> _IncompatibleKeys:
+        """
+        Load weights from a checkpoint.
+
+        Args:
+            checkpoint (Any): checkpoint contains the weights.
+
+        Returns:
+            ``NamedTuple`` with ``missing_keys``, ``unexpected_keys``,
+                and ``incorrect_shapes`` fields:
+                * **missing_keys** is a list of str containing the missing keys
+                * **unexpected_keys** is a list of str containing the unexpected keys
+                * **incorrect_shapes** is a list of (key, shape in checkpoint, shape in model)
+
+            This is just like the return value of
+            :func:`torch.nn.Module.load_state_dict`, but with extra support
+            for ``incorrect_shapes``.
+        """
+        checkpoint_state_dict = checkpoint.pop("model")
+        self._convert_ndarray_to_tensor(checkpoint_state_dict)
+
+        # if the state_dict comes from a model that was wrapped in a
+        # DataParallel or DistributedDataParallel during serialization,
+        # remove the "module" prefix before performing the matching.
+        _strip_prefix_if_present(checkpoint_state_dict, "module.")
+
+        # workaround https://github.com/pytorch/pytorch/issues/24139
+        model_state_dict = self.model.state_dict()
+        incorrect_shapes = []
+
+        # rename the para in checkpoint_state_dict
+        # some bug here, do not support re load
+
+        checkpoint_state_dict = {
+            append_prefix(k): checkpoint_state_dict[k]
+            for k in checkpoint_state_dict.keys()
+        }
+
+        checkpoint_state_dict = modify_ckpt_state(self.model, checkpoint_state_dict, logger=self.logger)
+
+        for k in list(checkpoint_state_dict.keys()):
+            if k in model_state_dict:
+                model_param = model_state_dict[k]
+                # Allow mismatch for uninitialized parameters
+                if TORCH_VERSION >= (1, 8) and isinstance(
+                        model_param, nn.parameter.UninitializedParameter
+                ):
+                    continue
+                shape_model = tuple(model_param.shape)
+                shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
+                if shape_model != shape_checkpoint:
+
+                    has_observer_base_classes = (
+                            TORCH_VERSION >= (1, 8)
+                            and hasattr(quantization, "ObserverBase")
+                            and hasattr(quantization, "FakeQuantizeBase")
+                    )
+                    if has_observer_base_classes:
+                        # Handle the special case of quantization per channel observers,
+                        # where buffer shape mismatches are expected.
+                        def _get_module_for_key(
+                                model: torch.nn.Module, key: str
+                        ) -> torch.nn.Module:
+                            # foo.bar.param_or_buffer_name -> [foo, bar]
+                            key_parts = key.split(".")[:-1]
+                            cur_module = model
+                            for key_part in key_parts:
+                                cur_module = getattr(cur_module, key_part)
+                            return cur_module
+
+                        cls_to_skip = (
+                            ObserverBase,
+                            FakeQuantizeBase,
+                        )
+                        target_module = _get_module_for_key(self.model, k)
+                        if isinstance(target_module, cls_to_skip):
+                            # Do not remove modules with expected shape mismatches
+                            # them from the state_dict loading. They have special logic
+                            # in _load_from_state_dict to handle the mismatches.
+                            continue
+
+                    incorrect_shapes.append((k, shape_checkpoint, shape_model))
+                    checkpoint_state_dict.pop(k)
+        incompatible = self.model.load_state_dict(checkpoint_state_dict, strict=False)
+        return _IncompatibleKeys(
+            missing_keys=incompatible.missing_keys,
+            unexpected_keys=incompatible.unexpected_keys,
+            incorrect_shapes=incorrect_shapes,
+        )

dit/object_detection/ditod/mytrainer.py

@@ -0,0 +1,744 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+"""
+This file contains components with some default boilerplate logic user may need
+in training / testing. They will not work for everyone, but many users may find them useful.
+
+The behavior of functions/classes in this file is subject to change,
+since they are meant to represent the "common default behavior" people need in their projects.
+"""
+
+import argparse
+import logging
+import os
+import sys
+import weakref
+from collections import OrderedDict
+from typing import Optional
+import torch
+from fvcore.nn.precise_bn import get_bn_modules
+from omegaconf import OmegaConf
+from torch.nn.parallel import DistributedDataParallel
+
+import detectron2.data.transforms as T
+from detectron2.checkpoint import DetectionCheckpointer
+from detectron2.config import CfgNode, LazyConfig
+from detectron2.data import (
+    MetadataCatalog,
+    build_detection_test_loader,
+    build_detection_train_loader,
+)
+from detectron2.evaluation import (
+    DatasetEvaluator,
+    inference_on_dataset,
+    print_csv_format,
+    verify_results,
+)
+from detectron2.modeling import build_model
+from detectron2.solver import build_lr_scheduler, build_optimizer
+from detectron2.utils import comm
+from detectron2.utils.collect_env import collect_env_info
+from detectron2.utils.env import seed_all_rng
+from detectron2.utils.events import CommonMetricPrinter, JSONWriter, TensorboardXWriter
+from detectron2.utils.file_io import PathManager
+from detectron2.utils.logger import setup_logger
+
+from detectron2.engine import hooks
+from detectron2.engine.train_loop import AMPTrainer, SimpleTrainer, TrainerBase
+
+from .mycheckpointer import MyDetectionCheckpointer
+from typing import Any, Dict, List, Set
+import itertools
+from detectron2.solver.build import maybe_add_gradient_clipping
+from .dataset_mapper import DetrDatasetMapper
+from .icdar_evaluation import ICDAREvaluator
+from detectron2.evaluation import COCOEvaluator
+
+__all__ = [
+    "create_ddp_model",
+    "default_argument_parser",
+    "default_setup",
+    "default_writers",
+    "DefaultPredictor",
+    "MyTrainer",
+]
+
+
+def create_ddp_model(model, *, fp16_compression=False, **kwargs):
+    """
+    Create a DistributedDataParallel model if there are >1 processes.
+
+    Args:
+        model: a torch.nn.Module
+        fp16_compression: add fp16 compression hooks to the ddp object.
+            See more at https://pytorch.org/docs/stable/ddp_comm_hooks.html#torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_hook
+        kwargs: other arguments of :module:`torch.nn.parallel.DistributedDataParallel`.
+    """  # noqa
+    if comm.get_world_size() == 1:
+        return model
+    if "device_ids" not in kwargs:
+        kwargs["device_ids"] = [comm.get_local_rank()]
+    ddp = DistributedDataParallel(model, **kwargs)
+    if fp16_compression:
+        from torch.distributed.algorithms.ddp_comm_hooks import default as comm_hooks
+
+        ddp.register_comm_hook(state=None, hook=comm_hooks.fp16_compress_hook)
+    return ddp
+
+
+def default_argument_parser(epilog=None):
+    """
+    Create a parser with some common arguments used by detectron2 users.
+
+    Args:
+        epilog (str): epilog passed to ArgumentParser describing the usage.
+
+    Returns:
+        argparse.ArgumentParser:
+    """
+    parser = argparse.ArgumentParser(
+        epilog=epilog
+        or f"""
+Examples:
+
+Run on single machine:
+    $ {sys.argv[0]} --num-gpus 8 --config-file cfg.yaml
+
+Change some config options:
+    $ {sys.argv[0]} --config-file cfg.yaml MODEL.WEIGHTS /path/to/weight.pth SOLVER.BASE_LR 0.001
+
+Run on multiple machines:
+    (machine0)$ {sys.argv[0]} --machine-rank 0 --num-machines 2 --dist-url <URL> [--other-flags]
+    (machine1)$ {sys.argv[0]} --machine-rank 1 --num-machines 2 --dist-url <URL> [--other-flags]
+""",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--config-file", default="", metavar="FILE", help="path to config file")
+    parser.add_argument(
+        "--resume",
+        action="store_true",
+        help="Whether to attempt to resume from the checkpoint directory. "
+        "See documentation of `MyTrainer.resume_or_load()` for what it means.",
+    )
+    parser.add_argument("--eval-only", action="store_true", help="perform evaluation only")
+    parser.add_argument("--num-gpus", type=int, default=1, help="number of gpus *per machine*")
+    parser.add_argument("--num-machines", type=int, default=1, help="total number of machines")
+    parser.add_argument(
+        "--machine-rank", type=int, default=0, help="the rank of this machine (unique per machine)"
+    )
+
+    # PyTorch still may leave orphan processes in multi-gpu training.
+    # Therefore we use a deterministic way to obtain port,
+    # so that users are aware of orphan processes by seeing the port occupied.
+    port = 2 ** 15 + 2 ** 14 + hash(os.getuid() if sys.platform != "win32" else 1) % 2 ** 14
+    parser.add_argument(
+        "--dist-url",
+        default="tcp://127.0.0.1:{}".format(port),
+        help="initialization URL for pytorch distributed backend. See "
+        "https://pytorch.org/docs/stable/distributed.html for details.",
+    )
+    parser.add_argument(
+        "opts",
+        help="""
+Modify config options at the end of the command. For Yacs configs, use
+space-separated "PATH.KEY VALUE" pairs.
+For python-based LazyConfig, use "path.key=value".
+        """.strip(),
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+    return parser
+
+
+def _try_get_key(cfg, *keys, default=None):
+    """
+    Try select keys from cfg until the first key that exists. Otherwise return default.
+    """
+    if isinstance(cfg, CfgNode):
+        cfg = OmegaConf.create(cfg.dump())
+    for k in keys:
+        none = object()
+        p = OmegaConf.select(cfg, k, default=none)
+        if p is not none:
+            return p
+    return default
+
+
+def _highlight(code, filename):
+    try:
+        import pygments
+    except ImportError:
+        return code
+
+    from pygments.lexers import Python3Lexer, YamlLexer
+    from pygments.formatters import Terminal256Formatter
+
+    lexer = Python3Lexer() if filename.endswith(".py") else YamlLexer()
+    code = pygments.highlight(code, lexer, Terminal256Formatter(style="monokai"))
+    return code
+
+
+def default_setup(cfg, args):
+    """
+    Perform some basic common setups at the beginning of a job, including:
+
+    1. Set up the detectron2 logger
+    2. Log basic information about environment, cmdline arguments, and config
+    3. Backup the config to the output directory
+
+    Args:
+        cfg (CfgNode or omegaconf.DictConfig): the full config to be used
+        args (argparse.NameSpace): the command line arguments to be logged
+    """
+    output_dir = _try_get_key(cfg, "OUTPUT_DIR", "output_dir", "train.output_dir")
+    if comm.is_main_process() and output_dir:
+        PathManager.mkdirs(output_dir)
+
+    rank = comm.get_rank()
+    setup_logger(output_dir, distributed_rank=rank, name="fvcore")
+    logger = setup_logger(output_dir, distributed_rank=rank)
+
+    logger.info("Rank of current process: {}. World size: {}".format(rank, comm.get_world_size()))
+    logger.info("Environment info:\n" + collect_env_info())
+
+    logger.info("Command line arguments: " + str(args))
+    if hasattr(args, "config_file") and args.config_file != "":
+        logger.info(
+            "Contents of args.config_file={}:\n{}".format(
+                args.config_file,
+                _highlight(PathManager.open(args.config_file, "r").read(), args.config_file),
+            )
+        )
+
+    if comm.is_main_process() and output_dir:
+        # Note: some of our scripts may expect the existence of
+        # config.yaml in output directory
+        path = os.path.join(output_dir, "config.yaml")
+        if isinstance(cfg, CfgNode):
+            logger.info("Running with full config:\n{}".format(_highlight(cfg.dump(), ".yaml")))
+            with PathManager.open(path, "w") as f:
+                f.write(cfg.dump())
+        else:
+            LazyConfig.save(cfg, path)
+        logger.info("Full config saved to {}".format(path))
+
+    # make sure each worker has a different, yet deterministic seed if specified
+    seed = _try_get_key(cfg, "SEED", "train.seed", default=-1)
+    seed_all_rng(None if seed < 0 else seed + rank)
+
+    # cudnn benchmark has large overhead. It shouldn't be used considering the small size of
+    # typical validation set.
+    if not (hasattr(args, "eval_only") and args.eval_only):
+        torch.backends.cudnn.benchmark = _try_get_key(
+            cfg, "CUDNN_BENCHMARK", "train.cudnn_benchmark", default=False
+        )
+
+
+def default_writers(output_dir: str, max_iter: Optional[int] = None):
+    """
+    Build a list of :class:`EventWriter` to be used.
+    It now consists of a :class:`CommonMetricPrinter`,
+    :class:`TensorboardXWriter` and :class:`JSONWriter`.
+
+    Args:
+        output_dir: directory to store JSON metrics and tensorboard events
+        max_iter: the total number of iterations
+
+    Returns:
+        list[EventWriter]: a list of :class:`EventWriter` objects.
+    """
+    PathManager.mkdirs(output_dir)
+    return [
+        # It may not always print what you want to see, since it prints "common" metrics only.
+        CommonMetricPrinter(max_iter),
+        JSONWriter(os.path.join(output_dir, "metrics.json")),
+        TensorboardXWriter(output_dir),
+    ]
+
+
+class DefaultPredictor:
+    """
+    Create a simple end-to-end predictor with the given config that runs on
+    single device for a single input image.
+
+    Compared to using the model directly, this class does the following additions:
+
+    1. Load checkpoint from `cfg.MODEL.WEIGHTS`.
+    2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.
+    3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.
+    4. Take one input image and produce a single output, instead of a batch.
+
+    This is meant for simple demo purposes, so it does the above steps automatically.
+    This is not meant for benchmarks or running complicated inference logic.
+    If you'd like to do anything more complicated, please refer to its source code as
+    examples to build and use the model manually.
+
+    Attributes:
+        metadata (Metadata): the metadata of the underlying dataset, obtained from
+            cfg.DATASETS.TEST.
+
+    Examples:
+    ::
+        pred = DefaultPredictor(cfg)
+        inputs = cv2.imread("input.jpg")
+        outputs = pred(inputs)
+    """
+
+    def __init__(self, cfg):
+        self.cfg = cfg.clone()  # cfg can be modified by model
+        self.model = build_model(self.cfg)
+        self.model.eval()
+        if len(cfg.DATASETS.TEST):
+            self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
+
+        checkpointer = DetectionCheckpointer(self.model)
+        checkpointer.load(cfg.MODEL.WEIGHTS)
+
+        self.aug = T.ResizeShortestEdge(
+            [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
+        )
+
+        self.input_format = cfg.INPUT.FORMAT
+        assert self.input_format in ["RGB", "BGR"], self.input_format
+
+    def __call__(self, original_image):
+        """
+        Args:
+            original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).
+
+        Returns:
+            predictions (dict):
+                the output of the model for one image only.
+                See :doc:`/tutorials/models` for details about the format.
+        """
+        with torch.no_grad():  # https://github.com/sphinx-doc/sphinx/issues/4258
+            # Apply pre-processing to image.
+            if self.input_format == "RGB":
+                # whether the model expects BGR inputs or RGB
+                original_image = original_image[:, :, ::-1]
+            height, width = original_image.shape[:2]
+            image = self.aug.get_transform(original_image).apply_image(original_image)
+            image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
+
+            inputs = {"image": image, "height": height, "width": width}
+            predictions = self.model([inputs])[0]
+            return predictions
+
+
+class MyTrainer(TrainerBase):
+    """
+    A trainer with default training logic. It does the following:
+
+    1. Create a :class:`SimpleTrainer` using model, optimizer, dataloader
+       defined by the given config. Create a LR scheduler defined by the config.
+    2. Load the last checkpoint or `cfg.MODEL.WEIGHTS`, if exists, when
+       `resume_or_load` is called.
+    3. Register a few common hooks defined by the config.
+
+    It is created to simplify the **standard model training workflow** and reduce code boilerplate
+    for users who only need the standard training workflow, with standard features.
+    It means this class makes *many assumptions* about your training logic that
+    may easily become invalid in a new research. In fact, any assumptions beyond those made in the
+    :class:`SimpleTrainer` are too much for research.
+
+    The code of this class has been annotated about restrictive assumptions it makes.
+    When they do not work for you, you're encouraged to:
+
+    1. Overwrite methods of this class, OR:
+    2. Use :class:`SimpleTrainer`, which only does minimal SGD training and
+       nothing else. You can then add your own hooks if needed. OR:
+    3. Write your own training loop similar to `tools/plain_train_net.py`.
+
+    See the :doc:`/tutorials/training` tutorials for more details.
+
+    Note that the behavior of this class, like other functions/classes in
+    this file, is not stable, since it is meant to represent the "common default behavior".
+    It is only guaranteed to work well with the standard models and training workflow in detectron2.
+    To obtain more stable behavior, write your own training logic with other public APIs.
+
+    Examples:
+    ::
+        trainer = MyTrainer(cfg)
+        trainer.resume_or_load()  # load last checkpoint or MODEL.WEIGHTS
+        trainer.train()
+
+    Attributes:
+        scheduler:
+        checkpointer (DetectionCheckpointer):
+        cfg (CfgNode):
+    """
+
+    def __init__(self, cfg):
+        """
+        Args:
+            cfg (CfgNode):
+        """
+        super().__init__()
+        logger = logging.getLogger("detectron2")
+        if not logger.isEnabledFor(logging.INFO):  # setup_logger is not called for d2
+            setup_logger()
+        cfg = MyTrainer.auto_scale_workers(cfg, comm.get_world_size())
+
+        self.cfg = cfg
+
+        # Assume these objects must be constructed in this order.
+        model = self.build_model(cfg)
+        optimizer = self.build_optimizer(cfg, model)
+        data_loader = self.build_train_loader(cfg)
+
+        model = create_ddp_model(model, broadcast_buffers=False)
+        self._trainer = (AMPTrainer if cfg.SOLVER.AMP.ENABLED else SimpleTrainer)(
+            model, data_loader, optimizer
+        )
+
+        self.scheduler = self.build_lr_scheduler(cfg, optimizer)
+        self.checkpointer = MyDetectionCheckpointer(
+            # Assume you want to save checkpoints together with logs/statistics
+            model,
+            cfg.OUTPUT_DIR,
+            trainer=weakref.proxy(self),
+        )
+        self.start_iter = 0
+        self.max_iter = cfg.SOLVER.MAX_ITER
+        self.cfg = cfg
+
+        self.register_hooks(self.build_hooks())
+
+    def resume_or_load(self, resume=True):
+        """
+        If `resume==True` and `cfg.OUTPUT_DIR` contains the last checkpoint (defined by
+        a `last_checkpoint` file), resume from the file. Resuming means loading all
+        available states (eg. optimizer and scheduler) and update iteration counter
+        from the checkpoint. ``cfg.MODEL.WEIGHTS`` will not be used.
+
+        Otherwise, this is considered as an independent training. The method will load model
+        weights from the file `cfg.MODEL.WEIGHTS` (but will not load other states) and start
+        from iteration 0.
+
+        Args:
+            resume (bool): whether to do resume or not
+        """
+        self.checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS, resume=resume)
+        if resume and self.checkpointer.has_checkpoint():
+            # The checkpoint stores the training iteration that just finished, thus we start
+            # at the next iteration
+            self.start_iter = self.iter + 1
+
+    def build_hooks(self):
+        """
+        Build a list of default hooks, including timing, evaluation,
+        checkpointing, lr scheduling, precise BN, writing events.
+
+        Returns:
+            list[HookBase]:
+        """
+        cfg = self.cfg.clone()
+        cfg.defrost()
+        cfg.DATALOADER.NUM_WORKERS = 0  # save some memory and time for PreciseBN
+
+        ret = [
+            hooks.IterationTimer(),
+            hooks.LRScheduler(),
+            hooks.PreciseBN(
+                # Run at the same freq as (but before) evaluation.
+                cfg.TEST.EVAL_PERIOD,
+                self.model,
+                # Build a new data loader to not affect training
+                self.build_train_loader(cfg),
+                cfg.TEST.PRECISE_BN.NUM_ITER,
+            )
+            if cfg.TEST.PRECISE_BN.ENABLED and get_bn_modules(self.model)
+            else None,
+        ]
+
+        # Do PreciseBN before checkpointer, because it updates the model and need to
+        # be saved by checkpointer.
+        # This is not always the best: if checkpointing has a different frequency,
+        # some checkpoints may have more precise statistics than others.
+        if comm.is_main_process():
+            ret.append(hooks.PeriodicCheckpointer(self.checkpointer, cfg.SOLVER.CHECKPOINT_PERIOD))
+
+        def test_and_save_results():
+            self._last_eval_results = self.test(self.cfg, self.model)
+            return self._last_eval_results
+
+        # Do evaluation after checkpointer, because then if it fails,
+        # we can use the saved checkpoint to debug.
+        ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
+
+        if comm.is_main_process():
+            # Here the default print/log frequency of each writer is used.
+            # run writers in the end, so that evaluation metrics are written
+            ret.append(hooks.PeriodicWriter(self.build_writers(), period=20))
+        return ret
+
+    def build_writers(self):
+        """
+        Build a list of writers to be used using :func:`default_writers()`.
+        If you'd like a different list of writers, you can overwrite it in
+        your trainer.
+
+        Returns:
+            list[EventWriter]: a list of :class:`EventWriter` objects.
+        """
+        return default_writers(self.cfg.OUTPUT_DIR, self.max_iter)
+
+    def train(self):
+        """
+        Run training.
+
+        Returns:
+            OrderedDict of results, if evaluation is enabled. Otherwise None.
+        """
+        super().train(self.start_iter, self.max_iter)
+        if len(self.cfg.TEST.EXPECTED_RESULTS) and comm.is_main_process():
+            assert hasattr(
+                self, "_last_eval_results"
+            ), "No evaluation results obtained during training!"
+            verify_results(self.cfg, self._last_eval_results)
+            return self._last_eval_results
+
+    def run_step(self):
+        self._trainer.iter = self.iter
+        self._trainer.run_step()
+
+    @classmethod
+    def build_model(cls, cfg):
+        """
+        Returns:
+            torch.nn.Module:
+
+        It now calls :func:`detectron2.modeling.build_model`.
+        Overwrite it if you'd like a different model.
+        """
+        model = build_model(cfg)
+        logger = logging.getLogger(__name__)
+        logger.info("Model:\n{}".format(model))
+        return model
+
+    @classmethod
+    def build_optimizer(cls, cfg, model):
+        params: List[Dict[str, Any]] = []
+        memo: Set[torch.nn.parameter.Parameter] = set()
+        for key, value in model.named_parameters(recurse=True):
+            if not value.requires_grad:
+                continue
+            # Avoid duplicating parameters
+            if value in memo:
+                continue
+            memo.add(value)
+            lr = cfg.SOLVER.BASE_LR
+            weight_decay = cfg.SOLVER.WEIGHT_DECAY
+            if "backbone" in key:
+                lr = lr * cfg.SOLVER.BACKBONE_MULTIPLIER
+            params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}]
+
+        def maybe_add_full_model_gradient_clipping(optim):  # optim: the optimizer class
+            # detectron2 doesn't have full model gradient clipping now
+            clip_norm_val = cfg.SOLVER.CLIP_GRADIENTS.CLIP_VALUE
+            enable = (
+                    cfg.SOLVER.CLIP_GRADIENTS.ENABLED
+                    and cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model"
+                    and clip_norm_val > 0.0
+            )
+
+            class FullModelGradientClippingOptimizer(optim):
+                def step(self, closure=None):
+                    all_params = itertools.chain(*[x["params"] for x in self.param_groups])
+                    torch.nn.utils.clip_grad_norm_(all_params, clip_norm_val)
+                    super().step(closure=closure)
+
+            return FullModelGradientClippingOptimizer if enable else optim
+
+        optimizer_type = cfg.SOLVER.OPTIMIZER
+        if optimizer_type == "SGD":
+            optimizer = maybe_add_full_model_gradient_clipping(torch.optim.SGD)(
+                params, cfg.SOLVER.BASE_LR, momentum=cfg.SOLVER.MOMENTUM
+            )
+        elif optimizer_type == "ADAMW":
+            optimizer = maybe_add_full_model_gradient_clipping(torch.optim.AdamW)(
+                params, cfg.SOLVER.BASE_LR
+            )
+        else:
+            raise NotImplementedError(f"no optimizer type {optimizer_type}")
+        if not cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model":
+            optimizer = maybe_add_gradient_clipping(cfg, optimizer)
+        return optimizer
+
+    @classmethod
+    def build_lr_scheduler(cls, cfg, optimizer):
+        """
+        It now calls :func:`detectron2.solver.build_lr_scheduler`.
+        Overwrite it if you'd like a different scheduler.
+        """
+        return build_lr_scheduler(cfg, optimizer)
+
+    @classmethod
+    def build_train_loader(cls, cfg):
+        if cfg.AUG.DETR:
+            mapper = DetrDatasetMapper(cfg, is_train=True)
+        else:
+            mapper = None
+        return build_detection_train_loader(cfg, mapper=mapper)
+
+    @classmethod
+    def build_test_loader(cls, cfg, dataset_name):
+        """
+        Returns:
+            iterable
+
+        It now calls :func:`detectron2.data.build_detection_test_loader`.
+        Overwrite it if you'd like a different data loader.
+        """
+        return build_detection_test_loader(cfg, dataset_name)
+
+    @classmethod
+    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
+        if output_folder is None:
+            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
+        if 'icdar' not in dataset_name:
+            return COCOEvaluator(dataset_name, output_dir=output_folder)
+        else:
+            return ICDAREvaluator(dataset_name, output_dir=output_folder)
+
+    @classmethod
+    def test(cls, cfg, model, evaluators=None):
+        """
+        Evaluate the given model. The given model is expected to already contain
+        weights to evaluate.
+
+        Args:
+            cfg (CfgNode):
+            model (nn.Module):
+            evaluators (list[DatasetEvaluator] or None): if None, will call
+                :meth:`build_evaluator`. Otherwise, must have the same length as
+                ``cfg.DATASETS.TEST``.
+
+        Returns:
+            dict: a dict of result metrics
+        """
+        logger = logging.getLogger(__name__)
+        if isinstance(evaluators, DatasetEvaluator):
+            evaluators = [evaluators]
+        if evaluators is not None:
+            assert len(cfg.DATASETS.TEST) == len(evaluators), "{} != {}".format(
+                len(cfg.DATASETS.TEST), len(evaluators)
+            )
+
+        results = OrderedDict()
+        for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
+            data_loader = cls.build_test_loader(cfg, dataset_name)
+            # When evaluators are passed in as arguments,
+            # implicitly assume that evaluators can be created before data_loader.
+            if evaluators is not None:
+                evaluator = evaluators[idx]
+            else:
+                try:
+                    evaluator = cls.build_evaluator(cfg, dataset_name)
+                except NotImplementedError:
+                    logger.warn(
+                        "No evaluator found. Use `MyTrainer.test(evaluators=)`, "
+                        "or implement its `build_evaluator` method."
+                    )
+                    results[dataset_name] = {}
+                    continue
+            results_i = inference_on_dataset(model, data_loader, evaluator)
+            results[dataset_name] = results_i
+            if comm.is_main_process():
+                assert isinstance(
+                    results_i, dict
+                ), "Evaluator must return a dict on the main process. Got {} instead.".format(
+                    results_i
+                )
+                logger.info("Evaluation results for {} in csv format:".format(dataset_name))
+                print_csv_format(results_i)
+
+        if len(results) == 1:
+            results = list(results.values())[0]
+        return results
+
+    @staticmethod
+    def auto_scale_workers(cfg, num_workers: int):
+        """
+        When the config is defined for certain number of workers (according to
+        ``cfg.SOLVER.REFERENCE_WORLD_SIZE``) that's different from the number of
+        workers currently in use, returns a new cfg where the total batch size
+        is scaled so that the per-GPU batch size stays the same as the
+        original ``IMS_PER_BATCH // REFERENCE_WORLD_SIZE``.
+
+        Other config options are also scaled accordingly:
+        * training steps and warmup steps are scaled inverse proportionally.
+        * learning rate are scaled proportionally, following :paper:`ImageNet in 1h`.
+
+        For example, with the original config like the following:
+
+        .. code-block:: yaml
+
+            IMS_PER_BATCH: 16
+            BASE_LR: 0.1
+            REFERENCE_WORLD_SIZE: 8
+            MAX_ITER: 5000
+            STEPS: (4000,)
+            CHECKPOINT_PERIOD: 1000
+
+        When this config is used on 16 GPUs instead of the reference number 8,
+        calling this method will return a new config with:
+
+        .. code-block:: yaml
+
+            IMS_PER_BATCH: 32
+            BASE_LR: 0.2
+            REFERENCE_WORLD_SIZE: 16
+            MAX_ITER: 2500
+            STEPS: (2000,)
+            CHECKPOINT_PERIOD: 500
+
+        Note that both the original config and this new config can be trained on 16 GPUs.
+        It's up to user whether to enable this feature (by setting ``REFERENCE_WORLD_SIZE``).
+
+        Returns:
+            CfgNode: a new config. Same as original if ``cfg.SOLVER.REFERENCE_WORLD_SIZE==0``.
+        """
+        old_world_size = cfg.SOLVER.REFERENCE_WORLD_SIZE
+        if old_world_size == 0 or old_world_size == num_workers:
+            return cfg
+        cfg = cfg.clone()
+        frozen = cfg.is_frozen()
+        cfg.defrost()
+
+        assert (
+            cfg.SOLVER.IMS_PER_BATCH % old_world_size == 0
+        ), "Invalid REFERENCE_WORLD_SIZE in config!"
+        scale = num_workers / old_world_size
+        bs = cfg.SOLVER.IMS_PER_BATCH = int(round(cfg.SOLVER.IMS_PER_BATCH * scale))
+        lr = cfg.SOLVER.BASE_LR = cfg.SOLVER.BASE_LR * scale
+        max_iter = cfg.SOLVER.MAX_ITER = int(round(cfg.SOLVER.MAX_ITER / scale))
+        warmup_iter = cfg.SOLVER.WARMUP_ITERS = int(round(cfg.SOLVER.WARMUP_ITERS / scale))
+        cfg.SOLVER.STEPS = tuple(int(round(s / scale)) for s in cfg.SOLVER.STEPS)
+        cfg.TEST.EVAL_PERIOD = int(round(cfg.TEST.EVAL_PERIOD / scale))
+        cfg.SOLVER.CHECKPOINT_PERIOD = int(round(cfg.SOLVER.CHECKPOINT_PERIOD / scale))
+        cfg.SOLVER.REFERENCE_WORLD_SIZE = num_workers  # maintain invariant
+        logger = logging.getLogger(__name__)
+        logger.info(
+            f"Auto-scaling the config to batch_size={bs}, learning_rate={lr}, "
+            f"max_iter={max_iter}, warmup={warmup_iter}."
+        )
+
+        if frozen:
+            cfg.freeze()
+        return cfg
+
+
+# Access basic attributes from the underlying trainer
+for _attr in ["model", "data_loader", "optimizer"]:
+    setattr(
+        MyTrainer,
+        _attr,
+        property(
+            # getter
+            lambda self, x=_attr: getattr(self._trainer, x),
+            # setter
+            lambda self, value, x=_attr: setattr(self._trainer, x, value),
+        ),
+    )

dit/object_detection/ditod/table_evaluation/__init__.py

@@ -0,0 +1 @@
+from .evaluate import calc_table_score

dit/object_detection/ditod/table_evaluation/data_structure.py

@@ -0,0 +1,469 @@
+"""
+Data structures used by the evaluation process.
+Yu Fang - March 2019
+"""
+
+from collections import Iterable
+
+import numpy as np
+from shapely.geometry import Polygon
+
+
+# helper functions
+def flatten(lis):
+    for item in lis:
+        if isinstance(item, Iterable) and not isinstance(item, str):
+            for x in flatten(item):
+                yield x
+        else:
+            yield item
+
+# derived from https://blog.csdn.net/u012433049/article/details/82909484
+def compute_poly_iou(list1, list2):
+    a1 = np.array(list1, dtype=int).reshape(-1, 2)
+    poly1 = Polygon(a1)
+    poly1_clean = poly1.buffer(0)
+
+    a2 = np.array(list2, dtype=int).reshape(-1, 2)
+    poly2 = Polygon(a2)
+    poly2_clean = poly2.buffer(0)
+
+    try:
+        # iou = poly1.intersection(poly2).area / poly1.union(poly2).area
+        iou = poly1_clean.intersection(poly2_clean).area / poly1_clean.union(poly2_clean).area
+    except ZeroDivisionError:
+        iou = 0
+    return iou
+
+
+class Cell(object):
+    # @:param start_row : start row index of the Cell
+    # @:param start_col : start column index of the Cell
+    # @:param end-row : end row index of the Cell
+    # @:param end-col : end column index of the Cell
+    # @:param cell_box: bounding-box of the Cell (coordinates are saved as a string)
+    # @:param content_box: bounding-box of the text content within Cell (unused variable)
+    # @:param cell_id: unique id of the Cell
+
+    def __init__(self, table_id, start_row, start_col, cell_box, end_row, end_col, content_box=""):
+        self._start_row = int(start_row)
+        self._start_col = int(start_col)
+        self._cell_box = cell_box
+        self._content_box = content_box
+        self._table_id = table_id    # the table_id this cell belongs to
+        # self._cell_name = cell_id    # specify the cell using passed-in cell_id
+        self._cell_id = id(self)
+        # self._region = region
+
+        # check for end-row and end-col special case
+        if end_row == -1:
+            self._end_row = self.start_row
+        else:
+            self._end_row = int(end_row)
+        if end_col == -1:
+            self._end_col = self._start_col
+        else:
+            self._end_col = int(end_col)
+
+    @property
+    def start_row(self):
+        return self._start_row
+
+    @property
+    def start_col(self):
+        return self._start_col
+
+    @property
+    def end_row(self):
+        return self._end_row
+
+    @property
+    def end_col(self):
+        return self._end_col
+
+    @property
+    def cell_box(self):
+        return self._cell_box
+
+    @property
+    def content_box(self):
+        return self._content_box
+
+    @property
+    def cell_id(self):
+        return self._cell_id
+
+    @property
+    def table_id(self):
+        return self._table_id
+
+    def __str__(self):
+        return "CELL row=[%d, %d] col=[%d, %d] (coords=%s)" %(self.start_row, self.end_row
+                                                              , self.start_col, self.end_col
+                                                              , self.cell_box)
+
+    # return the IoU value of two cell blocks
+    def compute_cell_iou(self, another_cell):
+        cell_box_1_temp = []
+        for el in self.cell_box.split():
+            cell_box_1_temp.append((el.split(",")))
+        cell_box_1 = list(flatten(cell_box_1_temp))
+        cell_box_1 = [int(x) for x in cell_box_1]
+
+        cell_box_2_temp = []
+        for el in another_cell.cell_box.split():
+            cell_box_2_temp.append((el.split(",")))
+        cell_box_2 = list(flatten(cell_box_2_temp))
+        cell_box_2 = [int(x) for x in cell_box_2]
+
+        return compute_poly_iou(cell_box_1, cell_box_2)
+
+    # check if the two cell object denotes same cell area in table
+    def check_same(self, another_cell):
+        return self._start_row == another_cell.start_row and self._end_row == another_cell.end_row and \
+               self._start_col == another_cell.start_col and self._end_col == another_cell.end_col
+
+
+# Note: currently save the relation with two cell object involved,
+# can be replaced by cell_id in follow-up memory clean up
+class AdjRelation:
+
+    DIR_HORIZ = 1
+    DIR_VERT = 2
+
+    def __init__(self, fromText, toText, direction):
+        # @param: fromText, toText are Cell objects （may be changed to cell-ID for further development）
+        self._fromText = fromText
+        self._toText = toText
+        self._direction = direction
+
+    @property
+    def fromText(self):
+        return self._fromText
+
+    @property
+    def toText(self):
+        return self._toText
+
+    @property
+    def direction(self):
+        return self._direction
+
+    def __str__(self):
+        if self.direction == self.DIR_VERT:
+            dir = "vertical"
+        else:
+            dir = "horizontal"
+        return 'ADJ_RELATION: ' + str(self._fromText) + '  ' + str(self._toText) + '    ' + dir
+
+    def isEqual(self, otherRelation):
+        return self.fromText.cell_id == otherRelation.fromText.cell_id and \
+               self.toText.cell_id == otherRelation.toText.cell_id and self.direction == otherRelation.direction
+
+
+class Table:
+
+    def __init__(self, tableNode):
+        self._root = tableNode
+        self._id = id(self)
+        self._table_coords = ""
+        self._maxRow = 0    # PS: indexing from 0
+        self._maxCol = 0
+        self._cells = []    # save a table as list of <Cell>s
+        self.adj_relations = []    # save the adj_relations for the table
+        self.parsed = False
+        self.found = False    # check if the find_adj_relations() has been called once
+
+        self.parse_table()
+
+    def __str__(self):
+        return "TABLE object - {} row x {} col".format(self._maxRow+1, self._maxCol+1)
+
+    @property
+    def id(self):
+        return self._id
+
+    @property
+    def table_coords(self):
+        return self._table_coords
+
+    @property
+    def table_cells(self):
+        return self._cells
+
+    # parse input xml to cell lists
+    def parse_table(self):
+        # get the table bbox
+        self._table_coords = str(self._root.getElementsByTagName("Coords")[0].getAttribute("points"))
+
+        # get info for each cell
+        cells = self._root.getElementsByTagName("cell")
+        max_row = max_col = 0
+        for cell in cells:
+            sr = cell.getAttribute("start-row")
+            sc = cell.getAttribute("start-col")
+            cell_id = cell.getAttribute("id")
+            b_points = str(cell.getElementsByTagName("Coords")[0].getAttribute("points"))
+            # try:
+            #     try:
+            #         text = cell.getElementsByTagName("content")[0].firstChild.nodeValue
+            #     except AttributeError:
+            #         text = ""
+            # except IndexError:
+            #     text = "initialized cell as no content"
+            er = cell.getAttribute("end-row") if cell.hasAttribute("end-row") else -1
+            ec = cell.getAttribute("end-col") if cell.hasAttribute("end-col") else -1
+            new_cell = Cell(table_id=str(self.id), start_row=sr, start_col=sc, cell_box=b_points,
+                            end_row=er, end_col=ec)
+            max_row = max(max_row, int(sr), int(er))
+            max_col = max(max_col, int(sc), int(ec))
+            self._cells.append(new_cell)
+        self._maxCol = max_col
+        self._maxRow = max_row
+        self.parsed = True
+
+    # generate a table-like structure for finding adj_relations
+    def convert_2d(self):
+        table = [[0 for x in range(self._maxCol+1)] for y in range(self._maxRow+1)]    # init blank cell with int 0
+        for cell in self._cells:
+            cur_row = cell.start_row
+            while cur_row <= cell.end_row:
+                cur_col = cell.start_col
+                while cur_col <= cell.end_col:
+                    temp = table[cur_row][cur_col]
+                    if temp == 0:
+                        table[cur_row][cur_col] = cell
+                    elif type(temp) == list:
+                        temp.append(cell)
+                        table[cur_row][cur_col] = temp
+                    else:
+                        table[cur_row][cur_col] = [temp, cell]
+                    cur_col += 1
+                cur_row += 1
+
+        return table
+
+    def find_adj_relations(self):
+        if self.found:
+            return self.adj_relations
+        else:
+            # if len(self._cells) == 0:
+            if self.parsed == False:
+                # fix: cases where there's no cell in table?
+                print("table is not parsed for further steps.")
+                self.parse_table()
+                self.find_adj_relations()
+            else:
+                retVal = []
+                tab = self.convert_2d()
+
+                # find horizontal relations
+                for r in range(self._maxRow+1):
+                    for c_from in range(self._maxCol):
+                        temp_pos = tab[r][c_from]
+                        if temp_pos == 0:
+                            continue
+                        elif type(temp_pos) == list:
+                            for cell in temp_pos:
+                                c_to = c_from + 1
+                                if tab[r][c_to] != 0:
+                                    # find relation between two adjacent cells
+                                    if type(tab[r][c_to]) == list:
+                                        for cell_to in tab[r][c_to]:
+                                            if cell != cell_to and (not cell.check_same(cell_to)):
+                                                adj_relation = AdjRelation(cell, cell_to, AdjRelation.DIR_HORIZ)
+                                                retVal.append(adj_relation)
+                                    else:
+                                        if cell != tab[r][c_to]:
+                                            adj_relation = AdjRelation(cell, tab[r][c_to], AdjRelation.DIR_HORIZ)
+                                            retVal.append(adj_relation)
+                                else:
+                                    # find the next non-blank cell, if exists
+                                    for temp in range(c_from + 1, self._maxCol + 1):
+                                        if tab[r][temp] != 0:
+                                            if type(tab[r][temp]) == list:
+                                                for cell_to in tab[r][temp]:
+                                                    adj_relation = AdjRelation(cell, cell_to,
+                                                                               AdjRelation.DIR_HORIZ)
+                                                    retVal.append(adj_relation)
+                                            else:
+                                                adj_relation = AdjRelation(cell, tab[r][temp],
+                                                                           AdjRelation.DIR_HORIZ)
+                                                retVal.append(adj_relation)
+                                            break
+                        else:
+                            c_to = c_from + 1
+                            if tab[r][c_to] != 0:
+                                # find relation between two adjacent cells
+                                if type(tab[r][c_to]) == list:
+                                    for cell_to in tab[r][c_to]:
+                                        if temp_pos != cell_to:
+                                            adj_relation = AdjRelation(temp_pos, cell_to, AdjRelation.DIR_HORIZ)
+                                            retVal.append(adj_relation)
+                                else:
+                                    if temp_pos != tab[r][c_to]:
+                                        adj_relation = AdjRelation(temp_pos, tab[r][c_to], AdjRelation.DIR_HORIZ)
+                                        retVal.append(adj_relation)
+                            else:
+                                # find the next non-blank cell, if exists
+                                for temp in range(c_from + 1, self._maxCol + 1):
+                                    if tab[r][temp] != 0:
+                                        if type(tab[r][temp]) == list:
+                                            for cell_to in tab[r][temp]:
+                                                adj_relation = AdjRelation(temp_pos, cell_to,
+                                                                           AdjRelation.DIR_HORIZ)
+                                                retVal.append(adj_relation)
+                                        else:
+                                            adj_relation = AdjRelation(temp_pos, tab[r][temp], AdjRelation.DIR_HORIZ)
+                                            retVal.append(adj_relation)
+                                        break
+
+                # find vertical relations
+                for c in range(self._maxCol+1):
+                    for r_from in range(self._maxRow):
+                        temp_pos = tab[r_from][c]
+                        if temp_pos == 0:
+                            continue
+                        elif type(temp_pos) == list:
+                            for cell in temp_pos:
+                                r_to = r_from + 1
+                                if tab[r_to][c] != 0:
+                                    # find relation between two adjacent cells
+                                    if type(tab[r_to][c]) == list:
+                                        for cell_to in tab[r_to][c]:
+                                            if cell != cell_to and (not cell.check_same(cell_to)):
+                                                adj_relation = AdjRelation(cell, cell_to, AdjRelation.DIR_VERT)
+                                                retVal.append(adj_relation)
+                                    else:
+                                        if cell != tab[r_to][c]:
+                                            adj_relation = AdjRelation(cell, tab[r_to][c], AdjRelation.DIR_VERT)
+                                            retVal.append(adj_relation)
+                                else:
+                                    # find the next non-blank cell, if exists
+                                    for temp in range(r_from + 1, self._maxRow + 1):
+                                        if tab[temp][c] != 0:
+                                            if type(tab[temp][c]) == list:
+                                                for cell_to in tab[temp][c]:
+                                                    adj_relation = AdjRelation(cell, cell_to,
+                                                                               AdjRelation.DIR_VERT)
+                                                    retVal.append(adj_relation)
+                                            else:
+                                                adj_relation = AdjRelation(cell, tab[temp][c],
+                                                                           AdjRelation.DIR_VERT)
+                                                retVal.append(adj_relation)
+                                            break
+                        else:
+                            r_to = r_from + 1
+                            if tab[r_to][c] != 0:
+                                # find relation between two adjacent cells
+                                if type(tab[r_to][c]) == list:
+                                    for cell_to in tab[r_to][c]:
+                                        if temp_pos != cell_to:
+                                            adj_relation = AdjRelation(temp_pos, cell_to, AdjRelation.DIR_VERT)
+                                            retVal.append(adj_relation)
+                                else:
+                                    if temp_pos != tab[r_to][c]:
+                                        adj_relation = AdjRelation(temp_pos, tab[r_to][c], AdjRelation.DIR_VERT)
+                                        retVal.append(adj_relation)
+                            else:
+                                # find the next non-blank cell, if exists
+                                for temp in range(r_from + 1, self._maxRow + 1):
+                                    if tab[temp][c] != 0:
+                                        if type(tab[temp][c]) == list:
+                                            for cell_to in tab[temp][c]:
+                                                adj_relation = AdjRelation(temp_pos, cell_to, AdjRelation.DIR_VERT)
+                                                retVal.append(adj_relation)
+                                        else:
+                                            adj_relation = AdjRelation(temp_pos, tab[temp][c], AdjRelation.DIR_VERT)
+                                            retVal.append(adj_relation)
+                                        break
+
+                # eliminate duplicates
+                repeat = True
+                while repeat:
+                    repeat = False
+                    duplicates = []
+
+                    for ar1 in retVal:
+                        for ar2 in retVal:
+                            if ar1 != ar2:
+                                if ar1.direction == ar2.direction and ar1.fromText == ar2.fromText and\
+                                        ar1.toText == ar2.toText:
+                                    duplicates.append(ar2)
+                                    break
+                        else:
+                            continue
+                        break
+
+                    if len(duplicates) > 0:
+                        repeat = True
+                        retVal.remove(duplicates[0])
+
+                self.found = True
+                self.adj_relations = retVal
+            return self.adj_relations
+
+    # compute the IOU of table, pass-in var is another Table object
+    def compute_table_iou(self, another_table):
+        table_box_1_temp = []
+        for el in self.table_coords.split():
+            table_box_1_temp.append((el.split(",")))
+        table_box_1 = list(flatten(table_box_1_temp))
+        table_box_1 = [int(x) for x in table_box_1]
+
+        table_box_2_temp = []
+        for el in another_table.table_coords.split():
+            table_box_2_temp.append((el.split(",")))
+        table_box_2 = list(flatten(table_box_2_temp))
+        table_box_2 = [int(x) for x in table_box_2]
+
+        return compute_poly_iou(table_box_1, table_box_2)
+
+    # find the cell mapping of tables as dictionary, pass-in var is another table and the desired IOU value
+    def find_cell_mapping(self, target_table, iou_value):
+        mapped_cell = []    # store the matches as tuples - (gt, result) mind the order of table when passing in
+        for cell_1 in self.table_cells:
+            for cell_2 in target_table.table_cells:
+                if cell_1.compute_cell_iou(cell_2) >= iou_value:
+                    mapped_cell.append((cell_1, cell_2))
+                    break
+        ret = dict(mapped_cell)
+        # print(ret)
+        return ret
+
+    # to print a table cell mapping
+    @classmethod
+    def printCellMapping(cls, dMappedCell):
+        print("-"*25)
+        for cell1, cell2 in dMappedCell.items():
+            print("  ", cell1, " --> ", cell2)
+
+    # to print a table set of adjacency relations
+    @classmethod
+    def printAdjacencyRelationList(cls, lAdjRel, title=""):
+        print("--- %s "%title + "-"*25)
+        for adj in lAdjRel:
+            print(adj)
+
+
+class ResultStructure:
+
+    def __init__(self, truePos, gtTotal, resTotal):
+        self._truePos = truePos
+        self._gtTotal = gtTotal
+        self._resTotal = resTotal
+
+    @property
+    def truePos(self):
+        return self._truePos
+
+    @property
+    def gtTotal(self):
+        return self._gtTotal
+
+    @property
+    def resTotal(self):
+        return self._resTotal
+
+    def __str__(self):
+        return "true: {}, gt: {}, res: {}".format(self._truePos, self._gtTotal, self._resTotal)

dit/object_detection/ditod/table_evaluation/evaluate.py

@@ -0,0 +1,404 @@
+"""
+Evaluation of -.tar.gz file.
+Yu Fang - March 2019
+"""
+
+import os
+import xml.dom.minidom
+
+# from eval import eval
+
+reg_gt_path = os.path.abspath("data/test")
+reg_gt_path_archival = os.path.abspath("data/test")
+reg_gt_path_modern = os.path.abspath("data/test")
+str_gt_path_1 = os.path.abspath("data/test")
+str_gt_path_2 = os.path.abspath("data/test")
+str_gt_path_archival = os.path.abspath("data/test")
+str_gt_path_modern = os.path.abspath("data/test")
+
+import xml.dom.minidom
+# from functools import cmp_to_key
+from os.path import join as osj
+from .data_structure import *
+
+
+class eval:
+    STR = "-str"
+    REG = "-reg"
+    DEFAULT_ENCODING = "UTF-8"
+    # reg_gt_path = "./annotations/trackA/"
+    # str_gt_path = "./annotations/trackB/"
+    # reg_gt_path = os.path.abspath("data/test")
+    # reg_gt_path_archival = os.path.abspath("data/test")
+    # reg_gt_path_modern = os.path.abspath("data/test")
+    # str_gt_path_1 = os.path.abspath("data/test")
+    # str_gt_path_2 = os.path.abspath("data/test")
+    # str_gt_path_archival = os.path.abspath("data/test")
+    # str_gt_path_modern = os.path.abspath("data/test")
+
+    # dummyDom = xml.dom.minidom.parse("./dummyXML.xml")
+
+    def __init__(self, track, res_path):
+        self.return_result = None
+        self.reg = True
+        self.str = False
+
+        self.resultFile = res_path
+        self.inPrefix = os.path.split(res_path)[-1].split(".")[0][:-7]
+
+        if track == "-trackA":
+            self.reg = True
+            self.GTFile = osj(reg_gt_path, self.inPrefix + ".xml")
+            # self.GTFile = osj(self.reg_gt_path, self.inPrefix)
+        elif track == "-trackA1":  # archival documents
+            self.reg = True
+            self.GTFile = osj(reg_gt_path_archival, self.inPrefix + ".xml")
+        elif track == "-trackA2":  # modern documents
+            self.reg = True
+            self.GTFile = osj(reg_gt_path_modern, self.inPrefix + ".xml")
+        elif track == "-trackB1":
+            self.str = True
+            self.GTFile = osj(str_gt_path_1, self.inPrefix + ".xml")
+            # self.GTFile = osj(self.str_gt_path_1, self.inPrefix)
+        elif track == "-trackB2":
+            self.str = True
+            self.GTFile = osj(str_gt_path_2, self.inPrefix + ".xml")
+            # print(self.GTFile)
+            # self.GTFile = osj(self.str_gt_path_2, self.inPrefix)
+        elif track == "-trackB2_a":
+            self.str = True
+            self.GTFile = osj(str_gt_path_archival, self.inPrefix + ".xml")
+        elif track == "-trackB2_m":
+            self.str = True
+            self.GTFile = osj(str_gt_path_modern, self.inPrefix + ".xml")
+        else:
+            print(track)
+            print("Not a valid track, please check your spelling.")
+
+        # self.resultFile = res_path
+        # self.inPrefix = os.path.split(res_path)[-1].split("-")[0]
+
+        # if self.str:
+        #     # self.GTFile = osj(self.str_gt_path, self.inPrefix + "-str.xml")
+        #     self.GTFile = osj(self.str_gt_path, self.inPrefix + ".xml")
+        # elif self.reg:
+        #     # self.GTFile = osj(self.reg_gt_path, self.inPrefix + "-reg.xml")
+        #     self.GTFile = osj(self.reg_gt_path, self.inPrefix + ".xml")
+        # else:
+        #     print("Not a valid track, please check your spelling.")
+
+        self.gene_ret_lst()
+
+    @property
+    def result(self):
+        return self.return_result
+
+    def gene_ret_lst(self):
+        ret_lst = []
+        for iou in [0.6, 0.7, 0.8, 0.9]:
+            temp = self.compute_retVal(iou)
+            ret_lst.append(temp)
+            # ret_lst.append(self.compute_retVal(iou))
+
+        ret_lst.append(self.inPrefix + ".xml")
+        # ret_lst.append(self.inPrefix)
+        # print("Done processing {}\n".format(self.resultFile))
+        self.return_result = ret_lst
+
+    def compute_retVal(self, iou):
+        gt_dom = xml.dom.minidom.parse(self.GTFile)
+        # incorrect submission format handling
+        try:
+            result_dom = xml.dom.minidom.parse(self.resultFile)
+        except Exception as e:
+            # result_dom = xml.dom.minidom.parse(dummyDom)
+            gt_tables = eval.get_table_list(gt_dom)
+            retVal = ResultStructure(truePos=0, gtTotal=len(gt_tables), resTotal=0)
+            return retVal
+
+        # result_dom = xml.dom.minidom.parse(self.resultFile)
+        if self.reg:
+            ret = self.evaluate_result_reg(gt_dom, result_dom, iou)
+            return ret
+        if self.str:
+            ret = self.evaluate_result_str(gt_dom, result_dom, iou)
+            return ret
+
+    @staticmethod
+    def get_table_list(dom):
+        """
+        return a list of Table objects corresponding to the table element of the DOM.
+        """
+        return [Table(_nd) for _nd in dom.documentElement.getElementsByTagName("table")]
+
+    @staticmethod
+    def evaluate_result_reg(gt_dom, result_dom, iou_value):
+        # parse the tables in input elements
+        gt_tables = eval.get_table_list(gt_dom)
+        result_tables = eval.get_table_list(result_dom)
+        # duplicate result table list
+        remaining_tables = result_tables.copy()
+
+        # map the tables in gt and result file
+        table_matches = []  # @param: table_matches - list of mapping of tables in gt and res file, in order (gt, res)
+        for gtt in gt_tables:
+            for rest in remaining_tables:
+                if gtt.compute_table_iou(rest) >= iou_value:
+                    remaining_tables.remove(rest)
+                    table_matches.append((gtt, rest))
+                    break
+
+        assert len(table_matches) <= len(gt_tables)
+        assert len(table_matches) <= len(result_tables)
+
+        retVal = ResultStructure(truePos=len(table_matches), gtTotal=len(gt_tables), resTotal=len(result_tables))
+        return retVal
+
+    @staticmethod
+    def evaluate_result_str(gt_dom, result_dom, iou_value, table_iou_value=0.8):
+        # parse the tables in input elements
+        gt_tables = eval.get_table_list(gt_dom)
+        result_tables = eval.get_table_list(result_dom)
+
+        # duplicate result table list
+        remaining_tables = result_tables.copy()
+        gt_remaining = gt_tables.copy()
+
+        # map the tables in gt and result file
+        table_matches = []  # @param: table_matches - list of mapping of tables in gt and res file, in order (gt, res)
+        for gtt in gt_remaining:
+            for rest in remaining_tables:
+                # note: for structural analysis, use 0.8 for table mapping
+                if gtt.compute_table_iou(rest) >= table_iou_value:
+                    table_matches.append((gtt, rest))
+                    remaining_tables.remove(rest)  # unsafe... should be ok with the break below
+                    gt_remaining.remove(gtt)
+                    break
+
+        total_gt_relation, total_res_relation, total_correct_relation = 0, 0, 0
+        for gt_table, ress_table in table_matches:
+
+            # set up the cell mapping for matching tables
+            cell_mapping = gt_table.find_cell_mapping(ress_table, iou_value)
+            # set up the adj relations, convert the one for result table to a dictionary for faster searching
+            gt_AR = gt_table.find_adj_relations()
+            total_gt_relation += len(gt_AR)
+
+            res_AR = ress_table.find_adj_relations()
+            total_res_relation += len(res_AR)
+
+            if False:  # for DEBUG
+                Table.printCellMapping(cell_mapping)
+                Table.printAdjacencyRelationList(gt_AR, "GT")
+                Table.printAdjacencyRelationList(res_AR, "run")
+
+            # Now map GT adjacency relations to result
+            lMappedAR = []
+            for ar in gt_AR:
+                try:
+                    resFromCell = cell_mapping[ar.fromText]
+                    resToCell = cell_mapping[ar.toText]
+                    # make a mapped adjacency relation
+                    lMappedAR.append(AdjRelation(resFromCell, resToCell, ar.direction))
+                except:
+                    # no mapping is possible
+                    pass
+
+            # compare two list of adjacency relation
+            correct_dect = 0
+            for ar1 in res_AR:
+                for ar2 in lMappedAR:
+                    if ar1.isEqual(ar2):
+                        correct_dect += 1
+                        break
+
+            total_correct_relation += correct_dect
+
+        # handle gt_relations in unmatched gt table
+        for gtt_remain in gt_remaining:
+            total_gt_relation += len(gtt_remain.find_adj_relations())
+
+        # handle gt_relation in unmatched res table
+        for res_remain in remaining_tables:
+            total_res_relation += len(res_remain.find_adj_relations())
+
+        retVal = ResultStructure(truePos=total_correct_relation, gtTotal=total_gt_relation, resTotal=total_res_relation)
+        return retVal
+
+# calculate the gt adj_relations of the missing file
+# @param: file_lst - list of missing ground truth file
+# @param: cur_gt_num - current total of ground truth objects (tables / cells)
+def process_missing_files(track, gt_file_lst, cur_gt_num):
+    if track in ["-trackA", "-trackA1", "-trackA2"]:
+        gt_file_lst_full = [osj(reg_gt_path, filename) for filename in gt_file_lst]
+        for file in gt_file_lst_full:
+            if os.path.split(file)[-1].split(".")[-1] == "xml":
+                gt_dom = xml.dom.minidom.parse(file)
+                gt_root = gt_dom.documentElement
+                # tables = []
+                table_elements = gt_root.getElementsByTagName("table")
+                for res_table in table_elements:
+                    # t = Table(res_table)
+                    # tables.append(t)
+                    cur_gt_num += 1
+        return cur_gt_num
+    elif track == "-trackB1":
+        gt_file_lst_full = [osj(str_gt_path_1, filename) for filename in gt_file_lst]
+        for file in gt_file_lst_full:
+            if os.path.split(file)[-1].split(".")[-1] == "xml":
+                gt_dom = xml.dom.minidom.parse(file)
+                gt_root = gt_dom.documentElement
+                tables = []
+                table_elements = gt_root.getElementsByTagName("table")
+                for res_table in table_elements:
+                    t = Table(res_table)
+                    tables.append(t)
+                for table in tables:
+                    cur_gt_num += len(table.find_adj_relations())
+        return cur_gt_num
+    elif track == "-trackB2":
+        gt_file_lst_full = [osj(str_gt_path_2, filename) for filename in gt_file_lst]
+        for file in gt_file_lst_full:
+            if os.path.split(file)[-1].split(".")[-1] == "xml":
+                gt_dom = xml.dom.minidom.parse(file)
+                gt_root = gt_dom.documentElement
+                tables = []
+                table_elements = gt_root.getElementsByTagName("table")
+                for res_table in table_elements:
+                    t = Table(res_table)
+                    tables.append(t)
+                for table in tables:
+                    cur_gt_num += len(table.find_adj_relations())
+        return cur_gt_num
+
+def calc(F1):
+    sum_a = 0.6 * F1[0] + 0.7 * F1[1] + 0.8 * F1[2] + 0.9 * F1[3]
+    sum_b = 0.6 + 0.7 + 0.8 + 0.9
+
+    return sum_a / sum_b
+
+def calc_table_score(result_path):
+    # measure = eval(*sys.argv[1:])
+
+    gt_file_lst = os.listdir(reg_gt_path_archival)
+    track = "-trackA1"
+    untar_path = result_path
+
+    res_lst = []
+    for root, files, dirs in os.walk(untar_path):
+        for name in dirs:
+            if name.split(".")[-1] == "xml":
+                cur_filepath = osj(os.path.abspath(root), name)
+                res_lst.append(eval(track, cur_filepath))
+                # printing for debug
+                # print("Processing... {}".format(name))
+    # print("DONE WITH FILE PROCESSING\n")
+    # note: results are stored as list of each when iou at [0.6, 0.7, 0.8, 0.9, gt_filename]
+    # gt number should be the same for all files
+    gt_num = 0
+    correct_six, res_six = 0, 0
+    correct_seven, res_seven = 0, 0
+    correct_eight, res_eight = 0, 0
+    correct_nine, res_nine = 0, 0
+
+
+    for each_file in res_lst:
+        # print(each_file)
+        try:
+            gt_file_lst.remove(each_file.result[-1])
+            if each_file.result[-1].replace('.xml', '.jpg') in gt_file_lst:
+                gt_file_lst.remove(each_file.result[-1].replace('.xml', '.jpg'))
+            correct_six += each_file.result[0].truePos
+            gt_num += each_file.result[0].gtTotal
+            res_six += each_file.result[0].resTotal
+            # print("{} {} {}".format(each_file.result[0].truePos, each_file.result[0].gtTotal, each_file.result[0].resTotal))
+
+            correct_seven += each_file.result[1].truePos
+            res_seven += each_file.result[1].resTotal
+
+            correct_eight += each_file.result[2].truePos
+            res_eight += each_file.result[2].resTotal
+
+            correct_nine += each_file.result[3].truePos
+            res_nine += each_file.result[3].resTotal
+        except:
+            print("Error occur in processing result list.")
+            print(each_file.result[-1])
+            break
+            # print(each_file.result[-1])
+            # print(each_file)
+
+    # for file in gt_file_lst:
+    #     if file.split(".") != "xml":
+    #         gt_file_lst.remove(file)
+    #     # print(gt_file_lst)
+
+    for i in range(len(gt_file_lst) - 1, -1, -1):
+        if gt_file_lst[i].split(".")[-1] != "xml":
+            del gt_file_lst[i]
+
+    if len(gt_file_lst) > 0:
+        print("\nWarning: missing result annotations for file: {}\n".format(gt_file_lst))
+        gt_total = process_missing_files(track, gt_file_lst, gt_num)
+    else:
+        gt_total = gt_num
+
+
+    try:
+        # print("Evaluation of {}".format(track.replace("-", "")))
+        # iou @ 0.6
+        p_six = correct_six / res_six
+        r_six = correct_six / gt_total
+        f1_six = 2 * p_six * r_six / (p_six + r_six)
+        print("IOU @ 0.6 -\nprecision: {}\nrecall: {}\nf1: {}".format(p_six, r_six, f1_six))
+        print("correct: {}, gt: {}, res: {}\n".format(correct_six, gt_total, res_six))
+
+        # iou @ 0.7
+        p_seven = correct_seven / res_seven
+        r_seven = correct_seven / gt_total
+        f1_seven = 2 * p_seven * r_seven / (p_seven + r_seven)
+        print("IOU @ 0.7 -\nprecision: {}\nrecall: {}\nf1: {}".format(p_seven, r_seven, f1_seven))
+        print("correct: {}, gt: {}, res: {}\n".format(correct_seven, gt_total, res_seven))
+
+        # iou @ 0.8
+        p_eight = correct_eight / res_eight
+        r_eight = correct_eight / gt_total
+        f1_eight = 2 * p_eight * r_eight / (p_eight + r_eight)
+        print("IOU @ 0.8 -\nprecision: {}\nrecall: {}\nf1: {}".format(p_eight, r_eight, f1_eight))
+        print("correct: {}, gt: {}, res: {}\n".format(correct_eight, gt_total, res_eight))
+
+        # iou @ 0.9
+        p_nine = correct_nine / res_nine
+        r_nine = correct_nine / gt_total
+        f1_nine = 2 * p_nine * r_nine / (p_nine + r_nine)
+        print("IOU @ 0.9 -\nprecision: {}\nrecall: {}\nf1: {}".format(p_nine, r_nine, f1_nine))
+        print("correct: {}, gt: {}, res: {}".format(correct_nine, gt_total, res_nine))
+
+        F1 = [f1_six, f1_seven, f1_eight, f1_nine]
+        wF1 = calc(F1)
+
+        print("Average weight F1: {}".format(wF1))
+
+        return {
+            'p_six':p_six * 100,
+            "r_six":r_six * 100,
+            "f1_six":f1_six * 100,
+            "p_seven":p_seven * 100,
+            "r_seven":r_seven * 100,
+            "f1_seven":f1_seven * 100,
+            "p_eight":p_eight * 100,
+            "r_eight":r_eight * 100,
+            "f1_eight":f1_eight * 100,
+            "p_nine":p_nine * 100,
+            "r_nine":r_nine * 100,
+            "f1_nine":f1_nine * 100,
+            "wF1":wF1 * 100
+        }
+    except ZeroDivisionError:
+        print(
+            "Error: zero devision error found, (possible that no adjacency relations are found), please check the file input.")
+        return {"wF1": 0}
+
+
+if __name__=="__main__":
+    pass

dit/object_detection/icdar19_configs/Base-RCNN-FPN.yaml

@@ -0,0 +1,69 @@
+MODEL:
+  MASK_ON: True
+  META_ARCHITECTURE: "GeneralizedRCNN"
+  PIXEL_MEAN: [127.5, 127.5, 127.5]
+  PIXEL_STD: [127.5, 127.5, 127.5]
+  BACKBONE:
+    NAME: "build_vit_fpn_backbone"
+  VIT:
+    OUT_FEATURES: ["layer3", "layer5", "layer7", "layer11"]
+    DROP_PATH: 0.1
+    IMG_SIZE: [224,224]
+    POS_TYPE: "abs"
+  FPN:
+    IN_FEATURES: ["layer3", "layer5", "layer7", "layer11"]
+  ANCHOR_GENERATOR:
+    SIZES: [[32], [64], [128], [256], [512]]  # One size for each in feature map
+    ASPECT_RATIOS: [[0.5, 1.0, 2.0]]  # Three aspect ratios (same for all in feature maps)
+  RPN:
+    IN_FEATURES: ["p2", "p3", "p4", "p5", "p6"]
+    PRE_NMS_TOPK_TRAIN: 2000  # Per FPN level
+    PRE_NMS_TOPK_TEST: 1000  # Per FPN level
+    # Detectron1 uses 2000 proposals per-batch,
+    # (See "modeling/rpn/rpn_outputs.py" for details of this legacy issue)
+    # which is approximately 1000 proposals per-image since the default batch size for FPN is 2.
+    POST_NMS_TOPK_TRAIN: 1000
+    POST_NMS_TOPK_TEST: 1000
+  ROI_HEADS:
+    NAME: "StandardROIHeads"
+    IN_FEATURES: ["p2", "p3", "p4", "p5"]
+    NUM_CLASSES: 1
+  ROI_BOX_HEAD:
+    NAME: "FastRCNNConvFCHead"
+    NUM_FC: 2
+    POOLER_RESOLUTION: 7
+  ROI_MASK_HEAD:
+    NAME: "MaskRCNNConvUpsampleHead"
+    NUM_CONV: 4
+    POOLER_RESOLUTION: 14
+DATASETS:
+  TRAIN: ("icdar2019_train",)
+  TEST: ("icdar2019_test",)
+SOLVER:
+  LR_SCHEDULER_NAME: "WarmupCosineLR"
+  AMP:
+    ENABLED: True
+  OPTIMIZER: "ADAMW"
+  BACKBONE_MULTIPLIER: 1.0
+  CLIP_GRADIENTS:
+    ENABLED: True
+    CLIP_TYPE: "full_model"
+    CLIP_VALUE: 1.0
+    NORM_TYPE: 2.0
+  WARMUP_FACTOR: 0.01
+  BASE_LR: 0.0002
+  WEIGHT_DECAY: 0.05
+  IMS_PER_BATCH: 32
+INPUT:
+  CROP:
+    ENABLED: True
+    TYPE: "absolute_range"
+    SIZE: (384, 600)
+  MIN_SIZE_TRAIN: (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800)
+  FORMAT: "RGB"
+DATALOADER:
+  FILTER_EMPTY_ANNOTATIONS: False
+VERSION: 2
+AUG:
+  DETR: True
+SEED: 42

dit/object_detection/icdar19_configs/cascade/cascade_dit_base.yaml

@@ -0,0 +1,21 @@
+_BASE_: "../Base-RCNN-FPN.yaml"
+MODEL:
+  PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
+  PIXEL_STD: [ 127.5, 127.5, 127.5 ]
+  WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.pth"
+  VIT:
+    NAME: "dit_base_patch16"
+  ROI_HEADS:
+    NAME: CascadeROIHeads
+  ROI_BOX_HEAD:
+    CLS_AGNOSTIC_BBOX_REG: True
+  RPN:
+    POST_NMS_TOPK_TRAIN: 2000
+SOLVER:
+  WARMUP_ITERS: 1000
+  IMS_PER_BATCH: 16
+  MAX_ITER: 60000
+  CHECKPOINT_PERIOD: 1000
+  BASE_LR: 0.00005
+TEST:
+  EVAL_PERIOD: 1000

dit/object_detection/icdar19_configs/cascade/cascade_dit_large.yaml

@@ -0,0 +1,25 @@
+_BASE_: "../Base-RCNN-FPN.yaml"
+MODEL:
+  PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
+  PIXEL_STD: [ 127.5, 127.5, 127.5 ]
+  WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-large-224-p16-500k-d7a2fb.pth"
+  VIT:
+    NAME: "dit_large_patch16"
+    OUT_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
+    DROP_PATH: 0.2
+  FPN:
+    IN_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
+  ROI_HEADS:
+    NAME: CascadeROIHeads
+  ROI_BOX_HEAD:
+    CLS_AGNOSTIC_BBOX_REG: True
+  RPN:
+    POST_NMS_TOPK_TRAIN: 2000
+SOLVER:
+  WARMUP_ITERS: 1000
+  IMS_PER_BATCH: 16
+  MAX_ITER: 60000
+  CHECKPOINT_PERIOD: 1000
+  BASE_LR: 0.00005
+TEST:
+  EVAL_PERIOD: 1000

dit/object_detection/icdar19_configs/maskrcnn/maskrcnn_dit_base.yaml

@@ -0,0 +1,15 @@
+_BASE_: "../Base-RCNN-FPN.yaml"
+MODEL:
+  PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
+  PIXEL_STD: [ 127.5, 127.5, 127.5 ]
+  WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.pth"
+  VIT:
+    NAME: "dit_base_patch16"
+SOLVER:
+  WARMUP_ITERS: 1000
+  IMS_PER_BATCH: 16
+  MAX_ITER: 60000
+  CHECKPOINT_PERIOD: 1000
+  BASE_LR: 0.00005
+TEST:
+  EVAL_PERIOD: 1000

dit/object_detection/icdar19_configs/maskrcnn/maskrcnn_dit_large.yaml

@@ -0,0 +1,19 @@
+_BASE_: "../Base-RCNN-FPN.yaml"
+MODEL:
+  PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
+  PIXEL_STD: [ 127.5, 127.5, 127.5 ]
+  WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-large-224-p16-500k-d7a2fb.pth"
+  VIT:
+    NAME: "dit_large_patch16"
+    OUT_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
+    DROP_PATH: 0.2
+  FPN:
+    IN_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
+SOLVER:
+  WARMUP_ITERS: 1000
+  IMS_PER_BATCH: 16
+  MAX_ITER: 60000
+  CHECKPOINT_PERIOD: 1000
+  BASE_LR: 0.00005
+TEST:
+  EVAL_PERIOD: 1000

dit/object_detection/inference.py

@@ -0,0 +1,80 @@
+import argparse
+
+import cv2
+
+from ditod import add_vit_config
+
+import torch
+
+from detectron2.config import get_cfg
+from detectron2.utils.visualizer import ColorMode, Visualizer
+from detectron2.data import MetadataCatalog
+from detectron2.engine import DefaultPredictor
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Detectron2 inference script")
+    parser.add_argument(
+        "--image_path",
+        help="Path to input image",
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        "--output_file_name",
+        help="Name of the output visualization file.",
+        type=str,
+    )
+    parser.add_argument(
+        "--config-file",
+        default="configs/quick_schedules/mask_rcnn_R_50_FPN_inference_acc_test.yaml",
+        metavar="FILE",
+        help="path to config file",
+    )
+    parser.add_argument(
+        "--opts",
+        help="Modify config options using the command-line 'KEY VALUE' pairs",
+        default=[],
+        nargs=argparse.REMAINDER,
+    )
+
+    args = parser.parse_args()
+
+    # Step 1: instantiate config
+    cfg = get_cfg()
+    add_vit_config(cfg)
+    cfg.merge_from_file(args.config_file)
+    
+    # Step 2: add model weights URL to config
+    cfg.merge_from_list(args.opts)
+    
+    # Step 3: set device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    cfg.MODEL.DEVICE = device
+
+    # Step 4: define model
+    predictor = DefaultPredictor(cfg)
+    
+    # Step 5: run inference
+    img = cv2.imread(args.image_path)
+
+    md = MetadataCatalog.get(cfg.DATASETS.TEST[0])
+    if cfg.DATASETS.TEST[0]=='icdar2019_test':
+        md.set(thing_classes=["table"])
+    else:
+        md.set(thing_classes=["text","title","list","table","figure"])
+
+    output = predictor(img)["instances"]
+    v = Visualizer(img[:, :, ::-1],
+                    md,
+                    scale=1.0,
+                    instance_mode=ColorMode.SEGMENTATION)
+    result = v.draw_instance_predictions(output.to("cpu"))
+    result_image = result.get_image()[:, :, ::-1]
+
+    # step 6: save
+    cv2.imwrite(args.output_file_name, result_image)
+
+if __name__ == '__main__':
+    main()
+

dit/object_detection/publaynet_configs/Base-RCNN-FPN.yaml

@@ -0,0 +1,69 @@
+MODEL:
+  MASK_ON: True
+  META_ARCHITECTURE: "GeneralizedRCNN"
+  PIXEL_MEAN: [123.675, 116.280, 103.530]
+  PIXEL_STD: [58.395, 57.120, 57.375]
+  BACKBONE:
+    NAME: "build_vit_fpn_backbone"
+  VIT:
+    OUT_FEATURES: ["layer3", "layer5", "layer7", "layer11"]
+    DROP_PATH: 0.1
+    IMG_SIZE: [224,224]
+    POS_TYPE: "abs"
+  FPN:
+    IN_FEATURES: ["layer3", "layer5", "layer7", "layer11"]
+  ANCHOR_GENERATOR:
+    SIZES: [[32], [64], [128], [256], [512]]  # One size for each in feature map
+    ASPECT_RATIOS: [[0.5, 1.0, 2.0]]  # Three aspect ratios (same for all in feature maps)
+  RPN:
+    IN_FEATURES: ["p2", "p3", "p4", "p5", "p6"]
+    PRE_NMS_TOPK_TRAIN: 2000  # Per FPN level
+    PRE_NMS_TOPK_TEST: 1000  # Per FPN level
+    # Detectron1 uses 2000 proposals per-batch,
+    # (See "modeling/rpn/rpn_outputs.py" for details of this legacy issue)
+    # which is approximately 1000 proposals per-image since the default batch size for FPN is 2.
+    POST_NMS_TOPK_TRAIN: 1000
+    POST_NMS_TOPK_TEST: 1000
+  ROI_HEADS:
+    NAME: "StandardROIHeads"
+    IN_FEATURES: ["p2", "p3", "p4", "p5"]
+    NUM_CLASSES: 5
+  ROI_BOX_HEAD:
+    NAME: "FastRCNNConvFCHead"
+    NUM_FC: 2
+    POOLER_RESOLUTION: 7
+  ROI_MASK_HEAD:
+    NAME: "MaskRCNNConvUpsampleHead"
+    NUM_CONV: 4
+    POOLER_RESOLUTION: 14
+DATASETS:
+  TRAIN: ("publaynet_train",)
+  TEST: ("publaynet_val",)
+SOLVER:
+  LR_SCHEDULER_NAME: "WarmupCosineLR"
+  AMP:
+    ENABLED: True
+  OPTIMIZER: "ADAMW"
+  BACKBONE_MULTIPLIER: 1.0
+  CLIP_GRADIENTS:
+    ENABLED: True
+    CLIP_TYPE: "full_model"
+    CLIP_VALUE: 1.0
+    NORM_TYPE: 2.0
+  WARMUP_FACTOR: 0.01
+  BASE_LR: 0.0004
+  WEIGHT_DECAY: 0.05
+  IMS_PER_BATCH: 32
+INPUT:
+  CROP:
+    ENABLED: True
+    TYPE: "absolute_range"
+    SIZE: (384, 600)
+  MIN_SIZE_TRAIN: (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800)
+  FORMAT: "RGB"
+DATALOADER:
+  FILTER_EMPTY_ANNOTATIONS: False
+VERSION: 2
+AUG:
+  DETR: True
+SEED: 42

dit/object_detection/publaynet_configs/cascade/cascade_dit_base.yaml

@@ -0,0 +1,20 @@
+_BASE_: "../Base-RCNN-FPN.yaml"
+MODEL:
+  PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
+  PIXEL_STD: [ 127.5, 127.5, 127.5 ]
+  WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.pth"
+  VIT:
+    NAME: "dit_base_patch16"
+  ROI_HEADS:
+    NAME: CascadeROIHeads
+  ROI_BOX_HEAD:
+    CLS_AGNOSTIC_BBOX_REG: True
+  RPN:
+    POST_NMS_TOPK_TRAIN: 2000
+SOLVER:
+  WARMUP_ITERS: 1000
+  IMS_PER_BATCH: 16
+  MAX_ITER: 60000
+  CHECKPOINT_PERIOD: 2000
+TEST:
+  EVAL_PERIOD: 2000

dit/object_detection/publaynet_configs/cascade/cascade_dit_large.yaml

@@ -0,0 +1,28 @@
+_BASE_: "../Base-RCNN-FPN.yaml"
+MODEL:
+  PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
+  PIXEL_STD: [ 127.5, 127.5, 127.5 ]
+  WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-large-224-p16-500k-d7a2fb.pth"
+  VIT:
+    NAME: "dit_large_patch16"
+    OUT_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
+    DROP_PATH: 0.2
+  FPN:
+    IN_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
+  ROI_HEADS:
+    NAME: CascadeROIHeads
+  ROI_BOX_HEAD:
+    CLS_AGNOSTIC_BBOX_REG: True
+  RPN:
+    POST_NMS_TOPK_TRAIN: 2000
+SOLVER:
+  WARMUP_ITERS: 1000
+  IMS_PER_BATCH: 16
+  MAX_ITER: 60000
+  CHECKPOINT_PERIOD: 2000
+  BASE_LR: 0.0001
+  STEPS: (40000, 53333)
+  AMP:
+    ENABLED: False
+TEST:
+  EVAL_PERIOD: 2000

dit/object_detection/publaynet_configs/maskrcnn/maskrcnn_dit_base.yaml

@@ -0,0 +1,15 @@
+_BASE_: "../Base-RCNN-FPN.yaml"
+MODEL:
+  PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
+  PIXEL_STD: [ 127.5, 127.5, 127.5 ]
+  WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.pth"
+  VIT:
+    NAME: "dit_base_patch16"
+SOLVER:
+  WARMUP_ITERS: 1000
+  IMS_PER_BATCH: 16
+  MAX_ITER: 60000
+  CHECKPOINT_PERIOD: 2000
+TEST:
+  EVAL_PERIOD: 2000
+OUTPUT_DIR: $AMLT_OUTPUT_DIR

dit/object_detection/publaynet_configs/maskrcnn/maskrcnn_dit_large.yaml

@@ -0,0 +1,22 @@
+_BASE_: "../Base-RCNN-FPN.yaml"
+MODEL:
+  PIXEL_MEAN: [ 127.5, 127.5, 127.5 ]
+  PIXEL_STD: [ 127.5, 127.5, 127.5 ]
+  WEIGHTS: "https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-large-224-p16-500k-d7a2fb.pth"
+  VIT:
+    NAME: "dit_large_patch16"
+    OUT_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
+    DROP_PATH: 0.2
+  FPN:
+    IN_FEATURES: [ "layer7", "layer11", "layer15", "layer23" ]
+SOLVER:
+  WARMUP_ITERS: 1000
+  IMS_PER_BATCH: 16
+  MAX_ITER: 60000
+  CHECKPOINT_PERIOD: 2000
+  BASE_LR: 0.0001
+  AMP:
+    ENABLED: False
+TEST:
+  EVAL_PERIOD: 2000
+OUTPUT_DIR: "output/publaynet/mask_rcnn/dit_base_multistep_3x_ms"

dit/object_detection/train_net.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python
+# --------------------------------------------------------------------------------
+# MPViT: Multi-Path Vision Transformer for Dense Prediction
+# Copyright (c) 2022 Electronics and Telecommunications Research Institute (ETRI).
+# All Rights Reserved.
+# Written by Youngwan Lee
+# --------------------------------------------------------------------------------
+
+"""
+Detection Training Script for MPViT.
+"""
+
+import os
+import itertools
+
+import torch
+
+from typing import Any, Dict, List, Set
+
+from detectron2.data import build_detection_train_loader
+from detectron2.checkpoint import DetectionCheckpointer
+from detectron2.config import get_cfg
+from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, launch
+from detectron2.evaluation import COCOEvaluator
+from detectron2.solver.build import maybe_add_gradient_clipping
+
+from ditod import add_vit_config
+from ditod import DetrDatasetMapper
+
+from detectron2.data.datasets import register_coco_instances
+import logging
+from detectron2.utils.logger import setup_logger
+from detectron2.utils import comm
+from detectron2.engine.defaults import create_ddp_model
+import weakref
+from detectron2.engine.train_loop import AMPTrainer, SimpleTrainer
+from ditod import MyDetectionCheckpointer, ICDAREvaluator
+from ditod import MyTrainer
+
+
+def setup(args):
+    """
+    Create configs and perform basic setups.
+    """
+    cfg = get_cfg()
+    # add_coat_config(cfg)
+    add_vit_config(cfg)
+    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+    default_setup(cfg, args)
+    return cfg
+
+
+def main(args):
+    """
+    register publaynet first
+    """
+    register_coco_instances(
+        "publaynet_train",
+        {},
+        "./publaynet_data/train.json",
+        "./publaynet_data/train"
+    )
+
+    register_coco_instances(
+        "publaynet_val",
+        {},
+        "./publaynet_data/val.json",
+        "./publaynet_data/val"
+    )
+
+    register_coco_instances(
+        "icdar2019_train",
+        {},
+        "data/train.json",
+        "data/train"
+    )
+
+    register_coco_instances(
+        "icdar2019_test",
+        {},
+        "data/test.json",
+        "data/test"
+    )
+
+    cfg = setup(args)
+
+    if args.eval_only:
+        model = MyTrainer.build_model(cfg)
+        DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
+            cfg.MODEL.WEIGHTS, resume=args.resume
+        )
+        res = MyTrainer.test(cfg, model)
+        return res
+
+    trainer = MyTrainer(cfg)
+    trainer.resume_or_load(resume=args.resume)
+    return trainer.train()
+
+
+if __name__ == "__main__":
+    parser = default_argument_parser()
+    parser.add_argument("--debug", action="store_true", help="enable debug mode")
+    args = parser.parse_args()
+    print("Command Line Args:", args)
+
+    if args.debug:
+        import debugpy
+
+        print("Enabling attach starts.")
+        debugpy.listen(address=('0.0.0.0', 9310))
+        debugpy.wait_for_client()
+        print("Enabling attach ends.")
+
+    launch(
+        main,
+        args.num_gpus,
+        num_machines=args.num_machines,
+        machine_rank=args.machine_rank,
+        dist_url=args.dist_url,
+        args=(args,),
+    )