From aa7acdfb1b2c5ae360472db3e2a732563514757d Mon Sep 17 00:00:00 2001 From: tianlanshidai <48075497+tianlanshidai@users.noreply.github.com> Date: Wed, 15 Jul 2020 18:16:07 +0800 Subject: [PATCH] Update model_builder.py --- pdseg/models/model_builder.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pdseg/models/model_builder.py b/pdseg/models/model_builder.py index 67b78c1a..f15c82e6 100644 --- a/pdseg/models/model_builder.py +++ b/pdseg/models/model_builder.py @@ -26,7 +26,7 @@ from loss import multi_dice_loss from loss import multi_bce_loss from lovasz_losses import lovasz_hinge from lovasz_losses import lovasz_softmax -from models.modeling import deeplab, unet, icnet, pspnet, hrnet, fast_scnn +from models.modeling import deeplab, unet, icnet, pspnet, hrnet, fast_scnn,ocnet class ModelPhase(object): @@ -85,6 +85,8 @@ def seg_model(image, class_num): logits = hrnet.hrnet(image, class_num) elif model_name == 'fast_scnn': logits = fast_scnn.fast_scnn(image, class_num) + elif model_name == 'ocnet': + logits = ocnet.ocnet(image, class_num) else: raise Exception( "unknow model name, only support unet, deeplabv3p, icnet, pspnet, hrnet, fast_scnn" -- GitLab