Merge pull request #74 from KUR-creative/tf2

Resolve #73 Supprot tf2
KUR-creative · Nov 16, 2019 · 321c676 · 321c676
2 parents cb43da7 + 6eb7dda
commit 321c676
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/core.py b/src/core.py
@@ -90,7 +90,7 @@ def assert_img_range(img):
     def decategorize(mask): 
         return iu.decategorize(mask, iu.rgb2wk_map)
 
-    with tf.Session() as sess:
+    with tf.compat.v1.Session() as sess:
         snet_in  = consts.snet_in('0.1.0', sess)
         snet_out = consts.snet_out('0.1.0', sess)
         def snet(img): 
@@ -169,7 +169,7 @@ def inpainted(image, segmap):
     '''
     assert (255 >= image).all(), image.max()
     assert   (image >= 0).all(), image.min()
-    with tf.Session() as sess:
+    with tf.compat.v1.Session() as sess:
         cnet_in  = consts.cnet_in('0.1.0',sess)
         cnet_out = consts.cnet_out('0.1.0',sess)
         return inpaint(