From 7355f67fdcbf550973140bf52471a925084794ea Mon Sep 17 00:00:00 2001
From: deeppomf <deeppomf@gmail.com>
Date: Fri, 9 Feb 2018 23:01:07 -0500
Subject: [PATCH] update

---
 README.md       |  1 +
 src/decensor.py |  4 ++--
 src/test.py     |  2 +-
 src/train.py    | 11 ++++++-----
 4 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 6068872..5c11e46 100644
--- a/README.md
+++ b/README.md
@@ -24,6 +24,7 @@ In particular, if a vagina or penis is completely censored out, THERE IS NO HOPE
 - Python 2/3
 - TensorFlow 1.5
 - Pillow
+- tqdm
 
 # Model
 The pretrained model can be downloaded from https://drive.google.com/open?id=1mWHYSj0LDSbJQQxjR4hUMykQkVve2U3Q.
diff --git a/src/decensor.py b/src/decensor.py
index f7a9b64..862b87e 100644
--- a/src/decensor.py
+++ b/src/decensor.py
@@ -1,6 +1,6 @@
 import numpy as np
 import tensorflow as tf
-import cv2
+from PIL import Image
 import tqdm
 import os
 import matplotlib.pyplot as plt
@@ -30,7 +30,7 @@ def test():
     sess.run(init_op)
 
     saver = tf.train.Saver()
-    saver.restore(sess, '/saved_models/latest')
+    saver.restore(sess, './saved_models/latest')
 
     x_test = np.load(test_npy)
     np.random.shuffle(x_test)
diff --git a/src/test.py b/src/test.py
index f6fb519..e177d4c 100644
--- a/src/test.py
+++ b/src/test.py
@@ -1,6 +1,6 @@
 import numpy as np
 import tensorflow as tf
-import cv2
+from PIL import Image
 import tqdm
 import os
 import matplotlib.pyplot as plt
diff --git a/src/train.py b/src/train.py
index 06145e1..ab19f72 100644
--- a/src/train.py
+++ b/src/train.py
@@ -1,6 +1,6 @@
 import numpy as np
 import tensorflow as tf
-import cv2
+from PIL import Image
 import tqdm
 from model import Model
 import load
@@ -65,8 +65,8 @@ def train():
             x_batch = x_test[:BATCH_SIZE]
             completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False})
             sample = np.array((completion[0] + 1) * 127.5, dtype=np.uint8)
-            cv2.imwrite('./output/{}.jpg'.format("{0:06d}".format(sess.run(epoch))), cv2.cvtColor(sample, cv2.COLOR_RGB2BGR))
-
+            result = Image.fromarray(sample)
+            result.save('./output/{}.jpg'.format("{0:06d}".format(sess.run(epoch))))
 
             saver = tf.train.Saver()
             saver.save(sess, '/saved_model/latest', write_meta_graph=False)
@@ -106,8 +106,9 @@ def train():
             x_batch = x_test[:BATCH_SIZE]
             completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False})
             sample = np.array((completion[0] + 1) * 127.5, dtype=np.uint8)
-            cv2.imwrite('/output/{}.jpg'.format("{0:06d}".format(sess.run(epoch))), cv2.cvtColor(sample, cv2.COLOR_RGB2BGR))
-
+            result = Image.fromarray(sample)
+            result.save('./output/{}.jpg'.format("{0:06d}".format(sess.run(epoch))))
+            
             saver = tf.train.Saver()
             saver.save(sess, '/saved_model/latest', write_meta_graph=False)