diff --git a/docs/en/benchmark.md b/docs/en/benchmark.md
index 75e3f657b..21582493b 100644
--- a/docs/en/benchmark.md
+++ b/docs/en/benchmark.md
@@ -1291,6 +1291,18 @@ Users can directly test the performance through [how_to_evaluate_a_model.md](tut
| 70.92 |
$MMSEG_DIR/configs/fastscnn/fast_scnn_lr0.12_8x4_160k_cityscapes.py |
+
+ | UNet |
+ Cityscapes |
+ mIoU |
+ 69.10 |
+ - |
+ 69.10 |
+ 69.10 |
+ 68.95 |
+ - |
+ $MMSEG_DIR/configs/unet/fcn_unet_s5-d16_4x4_512x1024_160k_cityscapes.py |
+
diff --git a/docs/en/codebases/mmseg.md b/docs/en/codebases/mmseg.md
index c56e27cb0..53a0d312c 100644
--- a/docs/en/codebases/mmseg.md
+++ b/docs/en/codebases/mmseg.md
@@ -8,13 +8,14 @@ Please refer to [get_started.md](https://github.com/open-mmlab/mmsegmentation/bl
### List of MMSegmentation models supported by MMDeploy
-| Model | OnnxRuntime | TensorRT | NCNN | PPLNN | OpenVino | Model config |
-| :-------------------------- | :---------: | :------: | :---: | :---: | :------: | :--------------------------------------------------------------------------------------: |
-| FCN | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/fcn) |
-| PSPNet[*](#static_shape) | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/pspnet) |
-| DeepLabV3 | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3) |
-| DeepLabV3+ | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3plus) |
-| Fast-SCNN[*](#static_shape) | Y | Y | N | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/fastscnn) |
+| Model | OnnxRuntime | TensorRT | NCNN | PPLNN | OpenVino | Model config |
+|:----------------------------|:-----------:|:--------:|:----:|:-----:|:--------:|:----------------------------------------------------------------------------------------:|
+| FCN | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/fcn) |
+| PSPNet[*](#static_shape) | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/pspnet) |
+| DeepLabV3 | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3) |
+| DeepLabV3+ | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3plus) |
+| Fast-SCNN[*](#static_shape) | Y | Y | N | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/fastscnn) |
+| UNet | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/unet) |
### Reminder
diff --git a/docs/en/supported_models.md b/docs/en/supported_models.md
index 03dad9d40..da7368edb 100644
--- a/docs/en/supported_models.md
+++ b/docs/en/supported_models.md
@@ -2,42 +2,43 @@
The table below lists the models that are guaranteed to be exportable to other backends.
-| Model | Codebase | OnnxRuntime | TensorRT | NCNN | PPLNN | OpenVINO | Model config |
-| :------------------------ | :--------------- | :---------: | :------: | :---: | :---: | :------: | :--------------------------------------------------------------------------------------------: |
-| RetinaNet | MMDetection | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/retinanet) |
-| Faster R-CNN | MMDetection | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/faster_rcnn) |
-| YOLOv3 | MMDetection | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/yolo) |
-| YOLOX | MMDetection | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/yolox) |
-| FCOS | MMDetection | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/fcos) |
-| FSAF | MMDetection | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/fsaf) |
-| Mask R-CNN | MMDetection | Y | Y | N | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/mask_rcnn) |
-| SSD[*](#note) | MMDetection | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/ssd) |
-| FoveaBox | MMDetection | Y | N | N | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/foveabox) |
-| ATSS | MMDetection | Y | Y | N | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/atss) |
-| Cascade R-CNN | MMDetection | Y | Y | N | Y | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/cascade_rcnn) |
-| Cascade Mask R-CNN | MMDetection | Y | N | N | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/cascade_rcnn) |
-| VFNet | MMDetection | N | N | N | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/vfnet) |
-| ResNet | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/resnet) |
-| ResNeXt | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/resnext) |
-| SE-ResNet | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/seresnet) |
-| MobileNetV2 | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/mobilenet_v2) |
-| ShuffleNetV1 | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/shufflenet_v1) |
-| ShuffleNetV2 | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/shufflenet_v2) |
-| FCN | MMSegmentation | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/fcn) |
-| PSPNet[*static](#note) | MMSegmentation | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/pspnet) |
-| DeepLabV3 | MMSegmentation | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3) |
-| DeepLabV3+ | MMSegmentation | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3plus) |
-| Fast-SCNN[*static](#note) | MMSegmentation | Y | Y | N | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/fastscnn) |
-| SRCNN | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/srcnn) |
-| ESRGAN | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/esrgan) |
-| SRGAN | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/srresnet_srgan) |
-| SRResNet | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/srresnet_srgan) |
-| Real-ESRGAN | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/real_esrgan) |
-| EDSR | MMEditing | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/edsr) |
-| RDN | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/rdn) |
-| DBNet | MMOCR | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmocr/tree/main/configs/textdet/dbnet) |
-| CRNN | MMOCR | Y | Y | Y | Y | N | [config](https://github.com/open-mmlab/mmocr/tree/main/configs/textrecog/crnn) |
-| SAR | MMOCR | Y | N | N | N | N | [config](https://github.com/open-mmlab/mmocr/tree/main/configs/textrecog/sar) |
+| Model | Codebase | OnnxRuntime | TensorRT | NCNN | PPLNN | OpenVINO | Model config |
+|:--------------------------|:-----------------|:-----------:|:--------:|:----:|:-----:|:--------:|:----------------------------------------------------------------------------------------------:|
+| RetinaNet | MMDetection | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/retinanet) |
+| Faster R-CNN | MMDetection | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/faster_rcnn) |
+| YOLOv3 | MMDetection | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/yolo) |
+| YOLOX | MMDetection | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/yolox) |
+| FCOS | MMDetection | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/fcos) |
+| FSAF | MMDetection | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/fsaf) |
+| Mask R-CNN | MMDetection | Y | Y | N | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/mask_rcnn) |
+| SSD[*](#note) | MMDetection | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/ssd) |
+| FoveaBox | MMDetection | Y | N | N | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/foveabox) |
+| ATSS | MMDetection | Y | Y | N | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/atss) |
+| Cascade R-CNN | MMDetection | Y | Y | N | Y | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/cascade_rcnn) |
+| Cascade Mask R-CNN | MMDetection | Y | N | N | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/cascade_rcnn) |
+| VFNet | MMDetection | N | N | N | N | Y | [config](https://github.com/open-mmlab/mmdetection/tree/master/configs/vfnet) |
+| ResNet | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/resnet) |
+| ResNeXt | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/resnext) |
+| SE-ResNet | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/seresnet) |
+| MobileNetV2 | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/mobilenet_v2) |
+| ShuffleNetV1 | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/shufflenet_v1) |
+| ShuffleNetV2 | MMClassification | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmclassification/tree/master/configs/shufflenet_v2) |
+| FCN | MMSegmentation | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/fcn) |
+| PSPNet[*static](#note) | MMSegmentation | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/pspnet) |
+| DeepLabV3 | MMSegmentation | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3) |
+| DeepLabV3+ | MMSegmentation | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3plus) |
+| Fast-SCNN[*static](#note) | MMSegmentation | Y | Y | N | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/fastscnn) |
+| UNet | MMSegmentation | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/unet) |
+| SRCNN | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/srcnn) |
+| ESRGAN | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/esrgan) |
+| SRGAN | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/srresnet_srgan) |
+| SRResNet | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/srresnet_srgan) |
+| Real-ESRGAN | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/real_esrgan) |
+| EDSR | MMEditing | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/edsr) |
+| RDN | MMEditing | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmediting/tree/master/configs/restorers/rdn) |
+| DBNet | MMOCR | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmocr/tree/main/configs/textdet/dbnet) |
+| CRNN | MMOCR | Y | Y | Y | Y | N | [config](https://github.com/open-mmlab/mmocr/tree/main/configs/textrecog/crnn) |
+| SAR | MMOCR | Y | N | N | N | N | [config](https://github.com/open-mmlab/mmocr/tree/main/configs/textrecog/sar) |
### Note
diff --git a/docs/zh_cn/benchmark.md b/docs/zh_cn/benchmark.md
index 237a3daaf..c9bc38a09 100644
--- a/docs/zh_cn/benchmark.md
+++ b/docs/zh_cn/benchmark.md
@@ -1292,6 +1292,18 @@ GPU: ncnn, TensorRT, PPLNN
70.92 |
$MMSEG_DIR/configs/fastscnn/fast_scnn_lr0.12_8x4_160k_cityscapes.py |
+
+ | UNet |
+ Cityscapes |
+ mIoU |
+ 69.10 |
+ - |
+ 69.10 |
+ 69.10 |
+ 68.95 |
+ - |
+ $MMSEG_DIR/configs/unet/fcn_unet_s5-d16_4x4_512x1024_160k_cityscapes.py |
+
diff --git a/mmdeploy/codebase/mmdet/deploy/object_detection.py b/mmdeploy/codebase/mmdet/deploy/object_detection.py
index 493f84381..f6f17e8ab 100644
--- a/mmdeploy/codebase/mmdet/deploy/object_detection.py
+++ b/mmdeploy/codebase/mmdet/deploy/object_detection.py
@@ -101,9 +101,8 @@ def create_input(self,
"""Create input for detector.
Args:
- task (Task): Specifying task type.
- imgs (Any): Input image(s), accpeted data type are `str`,
- `np.ndarray`, `torch.Tensor`.
+ imgs (str|np.ndarray): Input image(s), accpeted data type are
+ `str`, `np.ndarray`.
input_shape (list[int]): A list of two integer in (width, height)
format specifying input shape. Defaults to `None`.
diff --git a/mmdeploy/codebase/mmdet/deploy/object_detection_model.py b/mmdeploy/codebase/mmdet/deploy/object_detection_model.py
index 8ddebd02b..35a5d6378 100644
--- a/mmdeploy/codebase/mmdet/deploy/object_detection_model.py
+++ b/mmdeploy/codebase/mmdet/deploy/object_detection_model.py
@@ -45,8 +45,13 @@ class End2EndModel(BaseBackendModel):
"""End to end model for inference of detection.
Args:
+ backend (Backend): The backend enum, specifying backend type.
+ backend_files (Sequence[str]): Paths to all required backend files
+ (e.g. '.onnx' for ONNX Runtime, '.param' and '.bin' for ncnn).
+ device (str): A string specifying device type.
class_names (Sequence[str]): A list of string specifying class names.
- device_id (int): An integer represents device index.
+ deploy_cfg (str|mmcv.Config): Deployment config file or loaded Config
+ object.
"""
def __init__(self, backend: Backend, backend_files: Sequence[str],
@@ -286,11 +291,15 @@ class PartitionSingleStageModel(End2EndModel):
"""Partitioned single stage detection model.
Args:
- model_file (str): The path of input model file.
+ backend (Backend): The backend enum, specifying backend type.
+ backend_files (Sequence[str]): Paths to all required backend files
+ (e.g. '.onnx' for ONNX Runtime, '.param' and '.bin' for ncnn).
+ device (str): A string specifying device type.
class_names (Sequence[str]): A list of string specifying class names.
- model_cfg: (str | mmcv.Config): Input model config.
- deploy_cfg: (str | mmcv.Config): Input deployment config.
- device_id (int): An integer represents device index.
+ model_cfg (str|mmcv.Config): Input model config file or Config
+ object.
+ deploy_cfg (str|mmcv.Config): Deployment config file or loaded Config
+ object.
"""
def __init__(self, backend: Backend, backend_files: Sequence[str],
@@ -366,10 +375,15 @@ class PartitionTwoStageModel(End2EndModel):
"""Partitioned two stage detection model.
Args:
+ backend (Backend): The backend enum, specifying backend type.
+ backend_files (Sequence[str]): Paths to all required backend files
+ (e.g. '.onnx' for ONNX Runtime, '.param' and '.bin' for ncnn).
+ device (str): A string specifying device type.
class_names (Sequence[str]): A list of string specifying class names.
- model_cfg: (str | mmcv.Config): Input model config.
- deploy_cfg: (str | mmcv.Config): Input deployment config.
- device_id (int): An integer represents device index.
+ model_cfg (str|mmcv.Config): Input model config file or Config
+ object.
+ deploy_cfg (str|mmcv.Config): Deployment config file or loaded Config
+ object.
"""
def __init__(self, backend: Backend, backend_files: Sequence[str],
@@ -543,11 +557,15 @@ class NCNNEnd2EndModel(End2EndModel):
and its output is different from original mmdet style of `dets`, `labels`.
Args:
- model_file (str): The path of input model file.
+ backend (Backend): The backend enum, specifying backend type.
+ backend_files (Sequence[str]): Paths to all required backend files
+ (e.g. '.onnx' for ONNX Runtime, '.param' and '.bin' for ncnn).
+ device (str): A string specifying device type.
class_names (Sequence[str]): A list of string specifying class names.
- model_cfg: (str | mmcv.Config): Input model config.
- deploy_cfg: (str | mmcv.Config): Input deployment config.
- device_id (int): An integer represents device index.
+ model_cfg (str|mmcv.Config): Input model config file or Config
+ object.
+ deploy_cfg (str|mmcv.Config): Deployment config file or loaded Config
+ object.
"""
def __init__(self, backend: Backend, backend_files: Sequence[str],
@@ -691,7 +709,7 @@ def build_object_detection_model(model_files: Sequence[str],
device (str): Device to input model
Returns:
- DeployBaseDetector: Detector for a configured backend.
+ End2EndModel: Detector for a configured backend.
"""
# load cfg if necessary
deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg)
diff --git a/mmdeploy/codebase/mmseg/deploy/segmentation_model.py b/mmdeploy/codebase/mmseg/deploy/segmentation_model.py
index 5b86ad8e4..46e078903 100644
--- a/mmdeploy/codebase/mmseg/deploy/segmentation_model.py
+++ b/mmdeploy/codebase/mmseg/deploy/segmentation_model.py
@@ -183,12 +183,12 @@ def get_classes_palette_from_config(model_cfg: Union[str, mmcv.Config]):
module_dict = DATASETS.module_dict
data_cfg = model_cfg.data
- if 'train' in data_cfg:
- module = module_dict[data_cfg.train.type]
- elif 'val' in data_cfg:
+ if 'val' in data_cfg:
module = module_dict[data_cfg.val.type]
elif 'test' in data_cfg:
module = module_dict[data_cfg.test.type]
+ elif 'train' in data_cfg:
+ module = module_dict[data_cfg.train.type]
else:
raise RuntimeError(f'No dataset config found in: {model_cfg}')