init + inference.py патч

data/MBD/model/__init__.py

@@ -0,0 +1,50 @@
+import torchvision.models as models
+from model.densenetccnl import *
+from model.unetnc import *
+from model.gienet import *
+
+
+def get_model(name, n_classes=1, filters=64,version=None,in_channels=3, is_batchnorm=True, norm='batch', model_path=None, use_sigmoid=True, layers=3,img_size=512):
+    model = _get_model_instance(name)
+
+
+    if name == 'dnetccnl':
+        model = model(img_size=128, in_channels=in_channels, out_channels=n_classes, filters=32)
+    elif name == 'dnetccnl512':
+        model = model(img_size=img_size, in_channels=in_channels, out_channels=n_classes, filters=32)
+    elif name == 'unetnc':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'gie':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giecbam':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'gie2head':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giemask':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giemask2':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giedilated':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'bmp':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'displacement':
+        model = model(n_classes=2, num_filter=32, BatchNorm='GN', in_channels=5)
+    return model
+
+def _get_model_instance(name):
+    try:
+        return {
+            'dnetccnl': dnetccnl,
+            'dnetccnl512': dnetccnl512,
+            'unetnc': UnetGenerator,
+            'gie':GieGenerator,
+            'giecbam':GiecbamGenerator,
+            'giedilated':DilatedSingleUnet,
+            'gie2head':Gie2headGenerator,
+            'giemask':GiemaskGenerator,
+            'giemask2':Giemask2Generator,
+            'bmp':BmpGenerator,
+        }[name]
+    except:
+        print('Model {} not available'.format(name))