dewarpNet矫正扭曲

dewarp/models/__init__.py

@@ -0,0 +1,24 @@
+from dewarp.models.densenetccnl import dnetccnl
+from dewarp.models.unetnc import UnetGenerator
+
+
+def get_model(name, n_classes=1, in_channels=3):
+    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 == 'unetnc':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    else:
+        model = model(n_classes=n_classes)
+    return model
+
+
+def _get_model_instance(name):
+    try:
+        return {
+            'dnetccnl': dnetccnl,
+            'unetnc': UnetGenerator,
+        }[name]
+    except:
+        print('Model {} not available'.format(name))