Add files using upload-large-folder tool

ktda/models/segmentors/__init__.py

@@ -0,0 +1,3 @@
+from .distill_encoder_decoder import DistillEncoderDecoder
+
+__all__ = ['DistillEncoderDecoder']

tools/analysis_tools/browse_dataset.py

@@ -0,0 +1,77 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import argparse
+import os.path as osp
+
+from mmengine.config import Config, DictAction
+from mmengine.utils import ProgressBar
+
+from mmseg.registry import DATASETS, VISUALIZERS
+from mmseg.utils import register_all_modules
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Browse a dataset')
+    parser.add_argument('config', help='train config file path')
+    parser.add_argument(
+        '--output-dir',
+        default=None,
+        type=str,
+        help='If there is no display interface, you can save it')
+    parser.add_argument('--not-show', default=False, action='store_true')
+    parser.add_argument(
+        '--show-interval',
+        type=float,
+        default=2,
+        help='the interval of show (s)')
+    parser.add_argument(
+        '--cfg-options',
+        nargs='+',
+        action=DictAction,
+        help='override some settings in the used config, the key-value pair '
+        'in xxx=yyy format will be merged into config file. If the value to '
+        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
+        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
+        'Note that the quotation marks are necessary and that no white space '
+        'is allowed.')
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    cfg = Config.fromfile(args.config)
+    if args.cfg_options is not None:
+        cfg.merge_from_dict(args.cfg_options)
+
+    # register all modules in mmdet into the registries
+    register_all_modules()
+
+    dataset = DATASETS.build(cfg.train_dataloader.dataset)
+    visualizer = VISUALIZERS.build(cfg.visualizer)
+    visualizer.dataset_meta = dataset.metainfo
+
+    progress_bar = ProgressBar(len(dataset))
+    for item in dataset:
+        img = item['inputs'].permute(1, 2, 0).numpy()
+        img = img[..., [2, 1, 0]]  # bgr to rgb
+        data_sample = item['data_samples'].numpy()
+        img_path = osp.basename(item['data_samples'].img_path)
+
+        out_file = osp.join(
+            args.output_dir,
+            osp.basename(img_path)) if args.output_dir is not None else None
+
+        visualizer.add_datasample(
+            name=osp.basename(img_path),
+            image=img,
+            data_sample=data_sample,
+            draw_gt=True,
+            draw_pred=False,
+            wait_time=args.show_interval,
+            out_file=out_file,
+            show=not args.not_show)
+        progress_bar.update()
+
+
+if __name__ == '__main__':
+    main()

tools/dataset_converters/drive.py

@@ -0,0 +1,114 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import argparse
+import os
+import os.path as osp
+import tempfile
+import zipfile
+
+import cv2
+import mmcv
+from mmengine.utils import mkdir_or_exist
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='Convert DRIVE dataset to mmsegmentation format')
+    parser.add_argument(
+        'training_path', help='the training part of DRIVE dataset')
+    parser.add_argument(
+        'testing_path', help='the testing part of DRIVE dataset')
+    parser.add_argument('--tmp_dir', help='path of the temporary directory')
+    parser.add_argument('-o', '--out_dir', help='output path')
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    training_path = args.training_path
+    testing_path = args.testing_path
+    if args.out_dir is None:
+        out_dir = osp.join('data', 'DRIVE')
+    else:
+        out_dir = args.out_dir
+
+    print('Making directories...')
+    mkdir_or_exist(out_dir)
+    mkdir_or_exist(osp.join(out_dir, 'images'))
+    mkdir_or_exist(osp.join(out_dir, 'images', 'training'))
+    mkdir_or_exist(osp.join(out_dir, 'images', 'validation'))
+    mkdir_or_exist(osp.join(out_dir, 'annotations'))
+    mkdir_or_exist(osp.join(out_dir, 'annotations', 'training'))
+    mkdir_or_exist(osp.join(out_dir, 'annotations', 'validation'))
+
+    with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
+        print('Extracting training.zip...')
+        zip_file = zipfile.ZipFile(training_path)
+        zip_file.extractall(tmp_dir)
+
+        print('Generating training dataset...')
+        now_dir = osp.join(tmp_dir, 'training', 'images')
+        for img_name in os.listdir(now_dir):
+            img = mmcv.imread(osp.join(now_dir, img_name))
+            mmcv.imwrite(
+                img,
+                osp.join(
+                    out_dir, 'images', 'training',
+                    osp.splitext(img_name)[0].replace('_training', '') +
+                    '.png'))
+
+        now_dir = osp.join(tmp_dir, 'training', '1st_manual')
+        for img_name in os.listdir(now_dir):
+            cap = cv2.VideoCapture(osp.join(now_dir, img_name))
+            ret, img = cap.read()
+            mmcv.imwrite(
+                img[:, :, 0] // 128,
+                osp.join(out_dir, 'annotations', 'training',
+                         osp.splitext(img_name)[0] + '.png'))
+
+        print('Extracting test.zip...')
+        zip_file = zipfile.ZipFile(testing_path)
+        zip_file.extractall(tmp_dir)
+
+        print('Generating validation dataset...')
+        now_dir = osp.join(tmp_dir, 'test', 'images')
+        for img_name in os.listdir(now_dir):
+            img = mmcv.imread(osp.join(now_dir, img_name))
+            mmcv.imwrite(
+                img,
+                osp.join(
+                    out_dir, 'images', 'validation',
+                    osp.splitext(img_name)[0].replace('_test', '') + '.png'))
+
+        now_dir = osp.join(tmp_dir, 'test', '1st_manual')
+        if osp.exists(now_dir):
+            for img_name in os.listdir(now_dir):
+                cap = cv2.VideoCapture(osp.join(now_dir, img_name))
+                ret, img = cap.read()
+                # The annotation img should be divided by 128, because some of
+                # the annotation imgs are not standard. We should set a
+                # threshold to convert the nonstandard annotation imgs. The
+                # value divided by 128 is equivalent to '1 if value >= 128
+                # else 0'
+                mmcv.imwrite(
+                    img[:, :, 0] // 128,
+                    osp.join(out_dir, 'annotations', 'validation',
+                             osp.splitext(img_name)[0] + '.png'))
+
+        now_dir = osp.join(tmp_dir, 'test', '2nd_manual')
+        if osp.exists(now_dir):
+            for img_name in os.listdir(now_dir):
+                cap = cv2.VideoCapture(osp.join(now_dir, img_name))
+                ret, img = cap.read()
+                mmcv.imwrite(
+                    img[:, :, 0] // 128,
+                    osp.join(out_dir, 'annotations', 'validation',
+                             osp.splitext(img_name)[0] + '.png'))
+
+        print('Removing the temporary files...')
+
+    print('Done!')
+
+
+if __name__ == '__main__':
+    main()

tools/dataset_converters/vaihingen.py

@@ -0,0 +1,156 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import argparse
+import glob
+import math
+import os
+import os.path as osp
+import tempfile
+import zipfile
+
+import mmcv
+import numpy as np
+from mmengine.utils import ProgressBar, mkdir_or_exist
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='Convert vaihingen dataset to mmsegmentation format')
+    parser.add_argument('dataset_path', help='vaihingen folder path')
+    parser.add_argument('--tmp_dir', help='path of the temporary directory')
+    parser.add_argument('-o', '--out_dir', help='output path')
+    parser.add_argument(
+        '--clip_size',
+        type=int,
+        help='clipped size of image after preparation',
+        default=512)
+    parser.add_argument(
+        '--stride_size',
+        type=int,
+        help='stride of clipping original images',
+        default=256)
+    args = parser.parse_args()
+    return args
+
+
+def clip_big_image(image_path, clip_save_dir, to_label=False):
+    # Original image of Vaihingen dataset is very large, thus pre-processing
+    # of them is adopted. Given fixed clip size and stride size to generate
+    # clipped image, the intersection　of width and height is determined.
+    # For example, given one 5120 x 5120 original image, the clip size is
+    # 512 and stride size is 256, thus it would generate 20x20 = 400 images
+    # whose size are all 512x512.
+    image = mmcv.imread(image_path)
+
+    h, w, c = image.shape
+    cs = args.clip_size
+    ss = args.stride_size
+
+    num_rows = math.ceil((h - cs) / ss) if math.ceil(
+        (h - cs) / ss) * ss + cs >= h else math.ceil((h - cs) / ss) + 1
+    num_cols = math.ceil((w - cs) / ss) if math.ceil(
+        (w - cs) / ss) * ss + cs >= w else math.ceil((w - cs) / ss) + 1
+
+    x, y = np.meshgrid(np.arange(num_cols + 1), np.arange(num_rows + 1))
+    xmin = x * cs
+    ymin = y * cs
+
+    xmin = xmin.ravel()
+    ymin = ymin.ravel()
+    xmin_offset = np.where(xmin + cs > w, w - xmin - cs, np.zeros_like(xmin))
+    ymin_offset = np.where(ymin + cs > h, h - ymin - cs, np.zeros_like(ymin))
+    boxes = np.stack([
+        xmin + xmin_offset, ymin + ymin_offset,
+        np.minimum(xmin + cs, w),
+        np.minimum(ymin + cs, h)
+    ],
+                     axis=1)
+
+    if to_label:
+        color_map = np.array([[0, 0, 0], [255, 255, 255], [255, 0, 0],
+                              [255, 255, 0], [0, 255, 0], [0, 255, 255],
+                              [0, 0, 255]])
+        flatten_v = np.matmul(
+            image.reshape(-1, c),
+            np.array([2, 3, 4]).reshape(3, 1))
+        out = np.zeros_like(flatten_v)
+        for idx, class_color in enumerate(color_map):
+            value_idx = np.matmul(class_color,
+                                  np.array([2, 3, 4]).reshape(3, 1))
+            out[flatten_v == value_idx] = idx
+        image = out.reshape(h, w)
+
+    for box in boxes:
+        start_x, start_y, end_x, end_y = box
+        clipped_image = image[start_y:end_y,
+                              start_x:end_x] if to_label else image[
+                                  start_y:end_y, start_x:end_x, :]
+        area_idx = osp.basename(image_path).split('_')[3].strip('.tif')
+        mmcv.imwrite(
+            clipped_image.astype(np.uint8),
+            osp.join(clip_save_dir,
+                     f'{area_idx}_{start_x}_{start_y}_{end_x}_{end_y}.png'))
+
+
+def main():
+    splits = {
+        'train': [
+            'area1', 'area11', 'area13', 'area15', 'area17', 'area21',
+            'area23', 'area26', 'area28', 'area3', 'area30', 'area32',
+            'area34', 'area37', 'area5', 'area7'
+        ],
+        'val': [
+            'area6', 'area24', 'area35', 'area16', 'area14', 'area22',
+            'area10', 'area4', 'area2', 'area20', 'area8', 'area31', 'area33',
+            'area27', 'area38', 'area12', 'area29'
+        ],
+    }
+
+    dataset_path = args.dataset_path
+    if args.out_dir is None:
+        out_dir = osp.join('data', 'vaihingen')
+    else:
+        out_dir = args.out_dir
+
+    print('Making directories...')
+    mkdir_or_exist(osp.join(out_dir, 'img_dir', 'train'))
+    mkdir_or_exist(osp.join(out_dir, 'img_dir', 'val'))
+    mkdir_or_exist(osp.join(out_dir, 'ann_dir', 'train'))
+    mkdir_or_exist(osp.join(out_dir, 'ann_dir', 'val'))
+
+    zipp_list = glob.glob(os.path.join(dataset_path, '*.zip'))
+    print('Find the data', zipp_list)
+
+    with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
+        for zipp in zipp_list:
+            zip_file = zipfile.ZipFile(zipp)
+            zip_file.extractall(tmp_dir)
+            src_path_list = glob.glob(os.path.join(tmp_dir, '*.tif'))
+            if 'ISPRS_semantic_labeling_Vaihingen' in zipp:
+                src_path_list = glob.glob(
+                    os.path.join(os.path.join(tmp_dir, 'top'), '*.tif'))
+            if 'ISPRS_semantic_labeling_Vaihingen_ground_truth_eroded_COMPLETE' in zipp:  # noqa
+                src_path_list = glob.glob(os.path.join(tmp_dir, '*.tif'))
+                # delete unused area9 ground truth
+                for area_ann in src_path_list:
+                    if 'area9' in area_ann:
+                        src_path_list.remove(area_ann)
+            prog_bar = ProgressBar(len(src_path_list))
+            for i, src_path in enumerate(src_path_list):
+                area_idx = osp.basename(src_path).split('_')[3].strip('.tif')
+                data_type = 'train' if area_idx in splits['train'] else 'val'
+                if 'noBoundary' in src_path:
+                    dst_dir = osp.join(out_dir, 'ann_dir', data_type)
+                    clip_big_image(src_path, dst_dir, to_label=True)
+                else:
+                    dst_dir = osp.join(out_dir, 'img_dir', data_type)
+                    clip_big_image(src_path, dst_dir, to_label=False)
+                prog_bar.update()
+
+        print('Removing the temporary files...')
+
+    print('Done!')
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main()

tools/dataset_tools/analysis_dataset.py

@@ -0,0 +1,58 @@
+from glob import glob
+from typing import Tuple,List
+import os
+import argparse
+import json
+from matplotlib import pyplot as plt
+import numpy as np
+from PIL import Image
+
+def get_args()->Tuple[str, str]:
+    """
+    Return:
+        --dataset_dir: dataset dir.
+        --save_dir: save dir.
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dataset_dir', type=str, default='data/grass')
+    parser.add_argument('--save_dir', type=str, default='dataset_num_analysis.png')
+    args = parser.parse_args()
+    return args.dataset_dir, args.save_dir
+
+def get_mask_files(dataset_dir: str)->List[str]:
+    """
+    get mask files from dataset dir.
+    Args:
+        dataset_dir: dataset dir.
+    Return:
+        mask_filenames: list of mask filenames.
+    """
+    mask_filenames = glob(os.path.join(dataset_dir, "ann_dir", "*", "*.png"))
+    return mask_filenames
+
+def main():
+    dataset_dir, save_dir = get_args()
+    mask_filenames = get_mask_files(dataset_dir)
+    statistic = {}
+    for mask_filename in mask_filenames:
+        mask = np.array(Image.open(mask_filename))
+        classes = np.unique(mask)
+        for class_ in classes:
+            class_ = int(class_)
+            if class_ not in statistic:
+                statistic[class_] = 0
+            statistic[(class_)] += int(np.sum(mask == class_))
+    
+    classes = list(statistic.keys())
+    clasees_num = list(statistic.values())
+
+    plt.title("Dataset Analysis")
+    bars = plt.bar(classes, clasees_num)
+    for bar in bars:
+        height = bar.get_height()
+        plt.text(bar.get_x() + bar.get_width() / 2, height + 5, str(height), ha='center', va='bottom')
+    plt.savefig(save_dir,dpi=300)
+    
+
+if __name__ == "__main__":
+    main()

tools/dataset_tools/process_water.py

@@ -0,0 +1,152 @@
+import shutil
+from glob import glob
+import os
+import argparse
+import numpy as np
+from rich.progress import track
+from PIL import Image
+from typing import List
+from vegseg.datasets import WaterDataset
+from sklearn.model_selection import train_test_split
+
+
+def get_args():
+    parse = argparse.ArgumentParser()
+    parse.add_argument("--raw_path", type=str)
+    parse.add_argument("--tmp_dir", type=str)
+    parse.add_argument("--save_path", type=str)
+    args = parse.parse_args()
+    return args.raw_path, args.tmp_dir, args.save_path
+
+
+def get_palette() -> List[int]:
+    """
+    get palette of dataset.
+    return:
+        palette: list of palette.
+    """
+    palette = []
+    palette_list = WaterDataset.METAINFO["palette"]
+    for palette_item in palette_list:
+        palette.extend(palette_item)
+    return palette
+
+
+def create_dataset(image_list, ann_list, image_dir, ann_dir, description="Working..."):
+    os.makedirs(image_dir, exist_ok=True)
+    os.makedirs(ann_dir, exist_ok=True)
+    for image_path, ann_path in track(
+        zip(image_list, ann_list), total=len(image_list), description=description
+    ):
+        base_name = os.path.basename(image_path)
+
+        new_image_path = os.path.join(image_dir, base_name)
+        new_ann_path = os.path.join(ann_dir, base_name)
+
+        shutil.move(image_path, new_image_path)
+        shutil.move(ann_path, new_ann_path)
+
+        mask = Image.open(new_ann_path).convert("P")
+        palette = get_palette()
+        mask.putpalette(palette)
+        mask.save(new_ann_path)
+
+
+def main():
+    classes_mapping = {
+        "CDUWD-1": 1,
+        "CDUWD-2": 2,
+        "CDUWD-3": 3,
+        "CDUWD-4": 4,
+        "CDUWD-5": 5,
+        "CDUWD-6": 0,
+    }
+
+    raw_path, tmp_dir, save_path = get_args()
+
+    all_images = glob(os.path.join(raw_path, "*", "images", "*.png"))
+
+    all_labels = [image_path.replace("images", "labels") for image_path in all_images]
+
+    target_image_dir = os.path.join(tmp_dir, "images")
+    target_label_dir = os.path.join(tmp_dir, "labels")
+
+    os.makedirs(target_image_dir, exist_ok=True)
+    os.makedirs(target_label_dir, exist_ok=True)
+
+    for image_path, label_path in track(
+        zip(all_images, all_labels), total=len(all_images), description="fuse dataset"
+    ):
+        exists_images = glob(os.path.join(target_image_dir, "*.png"))
+
+        base_name = os.path.basename(image_path)
+        if image_path not in exists_images:
+            mask = np.array(Image.open(label_path))
+
+            assert list(np.unique(mask)) in [
+                [0],
+                [1],
+                [0, 1],
+                [1, 0],
+            ], f"The mask image is not binary (it should only contain 0s and 1s),actually is {set(np.unique(mask))}"
+
+            classes_str = image_path.split(os.path.sep)[-3]
+            classes = classes_mapping[classes_str]
+            mask = np.where(mask == 1, classes, mask)
+
+            # print(classes_str)
+
+            mask = Image.fromarray(mask)
+            mask.save(os.path.join(target_label_dir, base_name))
+            shutil.copy(image_path, os.path.join(target_image_dir, base_name))
+        else:
+
+            exists_label_path = os.path.join(target_label_dir, base_name)
+            exists_mask = np.array(Image.open(exists_label_path))
+
+            mask = np.array(Image.open(label_path))
+            assert list(np.unique(mask)) in [
+                [0],
+                [1],
+                [0, 1],
+                [1, 0],
+            ], f"The mask image is not binary (it should only contain 0s and 1s),actually is {set(np.unique(mask))}"
+            classes_str = image_path.split(os.path.sep)[-3]
+            classes = classes_mapping[classes_str]
+
+            exists_mask = np.where(mask == 1, classes, exists_mask)
+
+            exists_mask = Image.fromarray(exists_mask)
+            exists_mask.save(exists_label_path)
+
+    exists_images = glob(os.path.join(target_image_dir, "*.png"))
+
+    exists_labels = [
+        image_path.replace("images", "labels") for image_path in exists_images
+    ]
+    X_train, X_test, y_train, y_test = train_test_split(
+        exists_images, exists_labels, test_size=0.2, random_state=42, shuffle=True
+    )
+
+    create_dataset(
+        X_train,
+        y_train,
+        os.path.join(save_path, "img_dir", "train"),
+        os.path.join(save_path, "ann_dir", "train"),
+        description="train dataset",
+    )
+    create_dataset(
+        X_test,
+        y_test,
+        os.path.join(save_path, "img_dir", "val"),
+        os.path.join(save_path, "ann_dir", "val"),
+        description="val dataset",
+    )
+
+    os.rmdir(target_image_dir)
+    os.rmdir(target_label_dir)
+
+
+if __name__ == "__main__":
+    # example python tools/dataset_tools/process_water.py --raw_path data/raw_water_dataset/1024 --tmp_dir data/raw_water_dataset/1024/all_dataset --save_path data/water_1024_1024
+    main()

tools/deployment/pytorch2torchscript.py

@@ -0,0 +1,185 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import argparse
+
+import numpy as np
+import torch
+import torch._C
+import torch.serialization
+from mmengine import Config
+from mmengine.runner import load_checkpoint
+from torch import nn
+
+from mmseg.models import build_segmentor
+
+torch.manual_seed(3)
+
+
+def digit_version(version_str):
+    digit_version = []
+    for x in version_str.split('.'):
+        if x.isdigit():
+            digit_version.append(int(x))
+        elif x.find('rc') != -1:
+            patch_version = x.split('rc')
+            digit_version.append(int(patch_version[0]) - 1)
+            digit_version.append(int(patch_version[1]))
+    return digit_version
+
+
+def check_torch_version():
+    torch_minimum_version = '1.8.0'
+    torch_version = digit_version(torch.__version__)
+
+    assert (torch_version >= digit_version(torch_minimum_version)), \
+        f'Torch=={torch.__version__} is not support for converting to ' \
+        f'torchscript. Please install pytorch>={torch_minimum_version}.'
+
+
+def _convert_batchnorm(module):
+    module_output = module
+    if isinstance(module, torch.nn.SyncBatchNorm):
+        module_output = torch.nn.BatchNorm2d(module.num_features, module.eps,
+                                             module.momentum, module.affine,
+                                             module.track_running_stats)
+        if module.affine:
+            module_output.weight.data = module.weight.data.clone().detach()
+            module_output.bias.data = module.bias.data.clone().detach()
+            # keep requires_grad unchanged
+            module_output.weight.requires_grad = module.weight.requires_grad
+            module_output.bias.requires_grad = module.bias.requires_grad
+        module_output.running_mean = module.running_mean
+        module_output.running_var = module.running_var
+        module_output.num_batches_tracked = module.num_batches_tracked
+    for name, child in module.named_children():
+        module_output.add_module(name, _convert_batchnorm(child))
+    del module
+    return module_output
+
+
+def _demo_mm_inputs(input_shape, num_classes):
+    """Create a superset of inputs needed to run test or train batches.
+
+    Args:
+        input_shape (tuple):
+            input batch dimensions
+        num_classes (int):
+            number of semantic classes
+    """
+    (N, C, H, W) = input_shape
+    rng = np.random.RandomState(0)
+    imgs = rng.rand(*input_shape)
+    segs = rng.randint(
+        low=0, high=num_classes - 1, size=(N, 1, H, W)).astype(np.uint8)
+    img_metas = [{
+        'img_shape': (H, W, C),
+        'ori_shape': (H, W, C),
+        'pad_shape': (H, W, C),
+        'filename': '<demo>.png',
+        'scale_factor': 1.0,
+        'flip': False,
+    } for _ in range(N)]
+    mm_inputs = {
+        'imgs': torch.FloatTensor(imgs).requires_grad_(True),
+        'img_metas': img_metas,
+        'gt_semantic_seg': torch.LongTensor(segs)
+    }
+    return mm_inputs
+
+
+def pytorch2libtorch(model,
+                     input_shape,
+                     show=False,
+                     output_file='tmp.pt',
+                     verify=False):
+    """Export Pytorch model to TorchScript model and verify the outputs are
+    same between Pytorch and TorchScript.
+
+    Args:
+        model (nn.Module): Pytorch model we want to export.
+        input_shape (tuple): Use this input shape to construct
+            the corresponding dummy input and execute the model.
+        show (bool): Whether print the computation graph. Default: False.
+        output_file (string): The path to where we store the
+            output TorchScript model. Default: `tmp.pt`.
+        verify (bool): Whether compare the outputs between
+            Pytorch and TorchScript. Default: False.
+    """
+    if isinstance(model.decode_head, nn.ModuleList):
+        num_classes = model.decode_head[-1].num_classes
+    else:
+        num_classes = model.decode_head.num_classes
+
+    mm_inputs = _demo_mm_inputs(input_shape, num_classes)
+
+    imgs = mm_inputs.pop('imgs')
+
+    # replace the original forword with forward_dummy
+    model.forward = model.forward_dummy
+    model.eval()
+    traced_model = torch.jit.trace(
+        model,
+        example_inputs=imgs,
+        check_trace=verify,
+    )
+
+    if show:
+        print(traced_model.graph)
+
+    traced_model.save(output_file)
+    print(f'Successfully exported TorchScript model: {output_file}')
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='Convert MMSeg to TorchScript')
+    parser.add_argument('config', help='test config file path')
+    parser.add_argument('--checkpoint', help='checkpoint file', default=None)
+    parser.add_argument(
+        '--show', action='store_true', help='show TorchScript graph')
+    parser.add_argument(
+        '--verify', action='store_true', help='verify the TorchScript model')
+    parser.add_argument('--output-file', type=str, default='tmp.pt')
+    parser.add_argument(
+        '--shape',
+        type=int,
+        nargs='+',
+        default=[512, 512],
+        help='input image size (height, width)')
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    check_torch_version()
+
+    if len(args.shape) == 1:
+        input_shape = (1, 3, args.shape[0], args.shape[0])
+    elif len(args.shape) == 2:
+        input_shape = (
+            1,
+            3,
+        ) + tuple(args.shape)
+    else:
+        raise ValueError('invalid input shape')
+
+    cfg = Config.fromfile(args.config)
+    cfg.model.pretrained = None
+
+    # build the model and load checkpoint
+    cfg.model.train_cfg = None
+    segmentor = build_segmentor(
+        cfg.model, train_cfg=None, test_cfg=cfg.get('test_cfg'))
+    # convert SyncBN to BN
+    segmentor = _convert_batchnorm(segmentor)
+
+    if args.checkpoint:
+        load_checkpoint(segmentor, args.checkpoint, map_location='cpu')
+
+    # convert the PyTorch model to LibTorch model
+    pytorch2libtorch(
+        segmentor,
+        input_shape,
+        show=args.show,
+        output_file=args.output_file,
+        verify=args.verify)

tools/model_converters/beit2mmseg.py

@@ -0,0 +1,56 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import argparse
+import os.path as osp
+from collections import OrderedDict
+
+import mmengine
+import torch
+from mmengine.runner import CheckpointLoader
+
+
+def convert_beit(ckpt):
+    new_ckpt = OrderedDict()
+
+    for k, v in ckpt.items():
+        if k.startswith('patch_embed'):
+            new_key = k.replace('patch_embed.proj', 'patch_embed.projection')
+            new_ckpt[new_key] = v
+        if k.startswith('blocks'):
+            new_key = k.replace('blocks', 'layers')
+            if 'norm' in new_key:
+                new_key = new_key.replace('norm', 'ln')
+            elif 'mlp.fc1' in new_key:
+                new_key = new_key.replace('mlp.fc1', 'ffn.layers.0.0')
+            elif 'mlp.fc2' in new_key:
+                new_key = new_key.replace('mlp.fc2', 'ffn.layers.1')
+            new_ckpt[new_key] = v
+        else:
+            new_key = k
+            new_ckpt[new_key] = v
+
+    return new_ckpt
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Convert keys in official pretrained beit models to'
+        'MMSegmentation style.')
+    parser.add_argument('src', help='src model path or url')
+    # The dst path must be a full path of the new checkpoint.
+    parser.add_argument('dst', help='save path')
+    args = parser.parse_args()
+
+    checkpoint = CheckpointLoader.load_checkpoint(args.src, map_location='cpu')
+    if 'state_dict' in checkpoint:
+        state_dict = checkpoint['state_dict']
+    elif 'model' in checkpoint:
+        state_dict = checkpoint['model']
+    else:
+        state_dict = checkpoint
+    weight = convert_beit(state_dict)
+    mmengine.mkdir_or_exist(osp.dirname(args.dst))
+    torch.save(weight, args.dst)
+
+
+if __name__ == '__main__':
+    main()