Add ICNet

meetps · Apr 20, 2018 · 730b5af · 730b5af
1 parent 6d73a6e
commit 730b5af
Show file tree

Hide file tree

Showing 3 changed files with 560 additions and 33 deletions.
diff --git a/ptsemseg/models/__init__.py b/ptsemseg/models/__init__.py
@@ -4,6 +4,7 @@
 from ptsemseg.models.segnet import *
 from ptsemseg.models.unet import *
 from ptsemseg.models.pspnet import *
+from ptsemseg.models.icnet import *
 from ptsemseg.models.linknet import *
 from ptsemseg.models.frrn import *
 
@@ -34,6 +35,11 @@ def get_model(name, n_classes, version=None):
     elif name == 'pspnet':
         model = model(n_classes=n_classes, version=version)
 
+    elif name == 'icnet':
+        model = model(n_classes=n_classes, with_bn=False, version=version)
+    elif name == 'icnetBN':
+        model = model(n_classes=n_classes, with_bn=True, version=version)
+
     else:
         model = model(n_classes=n_classes)
 
@@ -48,6 +54,8 @@ def _get_model_instance(name):
             'unet': unet,
             'segnet': segnet,
             'pspnet': pspnet,
+			'icnet': icnet,
+			'icnetBN': icnet,
             'linknet': linknet,
             'frrnA': frrn,
             'frrnB': frrn,