adapt PY2 PY3 (PaddlePaddle#326)

lizexu123 · Jun 3, 2020 · 0024cfc · 0024cfc
1 parent 5d2529a
commit 0024cfc
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/demo/slimfacenet/dataloader/casia.py b/demo/slimfacenet/dataloader/casia.py
@@ -13,7 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
-import imageio
+import six
+if six.PY2:
+    import scipy.misc as imgreader
+else:
+    import imageio as imgreader
 import os
 import paddle
 from paddle import fluid
@@ -54,7 +58,7 @@ def reader(self):
             target = self.label_list[index]
 
             try:
-                img = imageio.imread(img_path)
+                img = imgreader.imread(img_path)
             except:
                 continue
 

diff --git a/demo/slimfacenet/dataloader/lfw.py b/demo/slimfacenet/dataloader/lfw.py
@@ -13,7 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
-import imageio
+import six
+if six.PY2:
+    import scipy.misc as imgreader
+else:
+    import imageio as imgreader
 import paddle
 from paddle import fluid
 
@@ -32,10 +36,10 @@ def reader(self):
                 return
             index = self.shuffle_idx.pop(0)
 
-            imgl = imageio.imread(self.imgl_list[index])
+            imgl = imgreader.imread(self.imgl_list[index])
             if len(imgl.shape) == 2:
                 imgl = np.stack([imgl] * 3, 2)
-            imgr = imageio.imread(self.imgr_list[index])
+            imgr = imgreader.imread(self.imgr_list[index])
             if len(imgr.shape) == 2:
                 imgr = np.stack([imgr] * 3, 2)