reverted model convention

profjsb · Jul 25, 2019 · b9c49de · b9c49de
1 parent e9d0195
commit b9c49de
Show file tree

Hide file tree

Showing 7 changed files with 7 additions and 12 deletions.
diff --git a/deepCR/model.py b/deepCR/model.py
@@ -23,7 +23,7 @@
 
 class deepCR():
 
-    def __init__(self, mask='ACS-WFC-2-32', inpaint='ACS-WFC-2-32', device='CPU',
+    def __init__(self, mask='ACS-WFC-F606W-2-32', inpaint='ACS-WFC-F606W-2-32', device='CPU',
                  model_dir=default_model_path):
 
         """

diff --git a/deepCR/test/test_model.py b/deepCR/test/test_model.py
@@ -8,7 +8,7 @@
 
 def test_deepCR_serial():
 
-    mdl = model.deepCR(mask='ACS-WFC-2-32', device='CPU')
+    mdl = model.deepCR(mask='ACS-WFC-F606W-2-32', device='CPU')
     in_im = np.ones((299, 299))
     out = mdl.clean(in_im)
     assert (out[0].shape, out[1].shape) == (in_im.shape, in_im.shape)
@@ -18,7 +18,7 @@ def test_deepCR_serial():
 
 def test_deepCR_parallel():
 
-    mdl = model.deepCR(mask='ACS-WFC-2-32', device='CPU')
+    mdl = model.deepCR(mask='ACS-WFC-F606W-2-32', device='CPU')
     in_im = np.ones((299, 299))
     out = mdl.clean(in_im, parallel=True)
     assert (out[0].shape, out[1].shape) == (in_im.shape, in_im.shape)
@@ -39,7 +39,7 @@ def test_deepCR_parallel():
         assert par_runtime < ser_runtime
 
 def test_seg():
-    mdl = model.deepCR(mask='ACS-WFC-2-32', inpaint='ACS-WFC-2-32', device='CPU')
+    mdl = model.deepCR(mask='ACS-WFC-F606W-2-32', inpaint='ACS-WFC-F606W-2-32', device='CPU')
     in_im = np.ones((500, 1000))
     out = mdl.clean(in_im, segment=True)
     assert (out[0].shape, out[1].shape) == (in_im.shape, in_im.shape)

diff --git a/learned_models/__init__.py b/learned_models/__init__.py
@@ -3,16 +3,11 @@
 
 __all__ = ('mask_dict', 'inpaint_dict')
 
-mask_dict = {'ACS-WFC-2-32': [UNet2Sigmoid, (1, 1, 32), 100],
-             'ACS-WFC-2-4': [UNet2Sigmoid, (1, 1, 4), 100],
-             'ACS-WFC-F606W-2-32': [UNet2Sigmoid, (1, 1, 32), 100],
+mask_dict = {'ACS-WFC-F606W-2-32': [UNet2Sigmoid, (1, 1, 32), 100],
              'ACS-WFC-F606W-2-4': [UNet2Sigmoid, (1, 1, 4), 100],
              'example_model': [UNet2Sigmoid, (1, 1, 32), 100]}
 
-inpaint_dict = {'ACS-WFC-3-32': [UNet3, (2, 1, 32)],
-                'ACS-WFC-2-32': [UNet2, (2, 1, 32)],
-                'ACS-WFC-F606W-3-32': [UNet3, (2, 1, 32)],
-                'ACS-WFC-F606W-2-32': [UNet2, (2, 1, 32)]
-                }
+inpaint_dict = {'ACS-WFC-F606W-3-32': [UNet3, (2, 1, 32)],
+                'ACS-WFC-F606W-2-32': [UNet2, (2, 1, 32)]}
 
 default_model_path = path.join(path.dirname(__file__))
diff --git a/learned_models/inpaint/ACS-WFC-2-32.pth b/learned_models/inpaint/ACS-WFC-2-32.pth
diff --git a/learned_models/inpaint/ACS-WFC-3-32.pth b/learned_models/inpaint/ACS-WFC-3-32.pth
diff --git a/learned_models/mask/ACS-WFC-2-32.pth b/learned_models/mask/ACS-WFC-2-32.pth
diff --git a/learned_models/mask/ACS-WFC-2-4.pth b/learned_models/mask/ACS-WFC-2-4.pth