Support Tensorflow 2

datasets/prepare_dataset.py

@@ -4,7 +4,10 @@
 import cv2
 import numpy as np
 import json
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 from scipy.misc import imresize
 
 

datasets/prepare_div2k_dataset.py

@@ -2,7 +2,10 @@
 import argparse
 from tqdm import tqdm
 import cv2
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 
 
 def bytes_feature(value):

evaluate.py

@@ -1,4 +1,7 @@
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 import numpy as np
 import argparse
 import os

generate_header_and_model.py

@@ -1,4 +1,7 @@
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 import numpy as np
 import argparse
 import os

models/dataset.py

@@ -1,4 +1,7 @@
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 from collections import OrderedDict
 
 

models/image_warp.py

@@ -1,4 +1,7 @@
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 
 
 def image_warp(images, flow, name='image_warp'):
@@ -39,4 +42,4 @@ def gather(y_coords, x_coords):
 
         interpolated = tf.reshape(interpolated, shape)
 
-        return interpolated
+        return interpolated

models/model.py

@@ -1,5 +1,8 @@
 from abc import ABC, abstractmethod
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 import json
 from .dataset import Dataset
 

models/model_espcn.py

@@ -1,4 +1,7 @@
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 from .model import Model
 
 

models/model_srcnn.py

@@ -1,4 +1,7 @@
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 from .model import Model
 
 

models/model_vespcn.py

@@ -1,4 +1,7 @@
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 from .model import Model
 from .image_warp import image_warp
 

models/model_vsrnet.py

@@ -1,4 +1,7 @@
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 from .model import Model
 
 

train.py

@@ -1,4 +1,7 @@
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
 import argparse
 from tqdm import tqdm
 import os