From 069d385ce04957b950f383e1c676388e68be6c65 Mon Sep 17 00:00:00 2001
From: plucky <plucky.pan@ubnt.com>
Date: Tue, 1 Jan 2019 18:20:43 +0800
Subject: [PATCH] support preview or not when train(resolve cannot connect to X
 server)

---
 README.md              | 3 +++
 main.py                | 6 ++++--
 mainscripts/Trainer.py | 9 +++++----
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 1794088..bf29df7 100644
--- a/README.md
+++ b/README.md
@@ -176,6 +176,9 @@ Video tutorial: https://www.youtube.com/watch?v=K98nTNjXkq8
 
 Windows 10 consumes % of VRAM even if card unused for video output.
 
+### For Mac Users
+Check out [DockerCPU.md](DockerCPU.md) for more detailed instructions.
+
 ### **Problem of the year**:
 
 algorithm of overlaying neural face onto video face located in ConverterMasked.py.
diff --git a/main.py b/main.py
index ca62958..0bc322c 100644
--- a/main.py
+++ b/main.py
@@ -80,6 +80,7 @@ if __name__ == "__main__":
             model_path=arguments.model_dir, 
             model_name=arguments.model_name,
             debug              = arguments.debug,
+            preview           = arguments.preview,
             #**options
             batch_size         = arguments.batch_size,
             write_preview_history = arguments.write_preview_history,
@@ -106,8 +107,9 @@ if __name__ == "__main__":
     train_parser.add_argument('--force-best-gpu-idx', type=int, dest="force_best_gpu_idx", default=-1, help="Force to choose this GPU idx as best(worst).")
     train_parser.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="MultiGPU option. It will select only same best(worst) GPU models.")
     train_parser.add_argument('--force-gpu-idxs', type=str, dest="force_gpu_idxs", default=None, help="Override final GPU idxs. Example: 0,1,2.")
-    train_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.")    
-    
+    train_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.")
+    train_parser.add_argument('--preview', action="store_true",dest="preview", default=False, help="Show preview.")
+
     train_parser.set_defaults (func=process_train)
     
     def process_convert(arguments):
diff --git a/mainscripts/Trainer.py b/mainscripts/Trainer.py
index 96b4542..2c6ae83 100644
--- a/mainscripts/Trainer.py
+++ b/mainscripts/Trainer.py
@@ -277,13 +277,14 @@ def previewThread (input_queue, output_queue):
             
     cv2.destroyAllWindows()
     
-def main (training_data_src_dir, training_data_dst_dir, model_path, model_name, **in_options):
-    print ("Running trainer.\r\n")
+def main (training_data_src_dir, training_data_dst_dir, model_path, model_name,preview, **in_options):
+    print ("Running trainer(preview=%s).\r\n" % (preview))
     
     output_queue = queue.Queue()
     input_queue = queue.Queue()
     import threading
     thread = threading.Thread(target=trainerThread, args=(output_queue, input_queue, training_data_src_dir, training_data_dst_dir, model_path, model_name), kwargs=in_options )
     thread.start()
-    
-    previewThread (input_queue, output_queue)
\ No newline at end of file
+
+    if preview:
+        previewThread (input_queue, output_queue)
\ No newline at end of file