init

app.py

@@ -45,20 +45,14 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-
-
 def run(
     image,
-    shinkai: ImportGraph,
-    hayao: ImportGraph,
-    paprika: ImportGraph,
-) -> tuple[PIL.Image.Image, PIL.Image.Image, PIL.Image.Image]:
+) -> tuple[PIL.Image.Image]:
 
-    im1 = shinkai.test('shinkai', image.name, True)
-    #im2 = hayao.test('hayao', image.name, True)
-    #im3 = paprika.test('paprika', image.name, True)
+    out = test.test(checkpoint_dir=os.path.join(curPath,'animeganv2/checkpoint/generator_Shinkai_weight'),
+              style_name='Shinkai', test_file=image.name, if_adjust_brightness=True)
 
-    return PIL.Image.open(im1),PIL.Image.open(im1),PIL.Image.open(im1)
+    return PIL.Image.open(out)
 
 
 def main():
@@ -66,13 +60,13 @@ def main():
 
     args = parse_args()
 
-    curPath = os.path.abspath(os.path.dirname(__file__))
+    #curPath = os.path.abspath(os.path.dirname(__file__))
     #init
-    shinkai = ImportGraph(checkpoint_dir=os.path.join(curPath,'animeganv2/checkpoint/generator_Shinkai_weight'))
+    #shinkai = ImportGraph(checkpoint_dir=os.path.join(curPath,'animeganv2/checkpoint/generator_Shinkai_weight'))
     #hayao = ImportGraph(checkpoint_dir=os.path.join(curPath,'animeganv2/checkpoint/generator_Hayao_weight'))
     #paprika = ImportGraph(checkpoint_dir=os.path.join(curPath,'animeganv2/checkpoint/generator_Paprika_weight'))
 
-    func = functools.partial(run, shinkai=shinkai ,hayao=None,paprika=None )
+    func = functools.partial(run)
     func = functools.update_wrapper(func, run)
 
     
@@ -84,13 +78,8 @@ def main():
         [
             gr.outputs.Image(
                 type='pil',
-                label='Shinkai Result'),
-            gr.outputs.Image(
-                type='pil',
-                label='Hayao Result'),
-            gr.outputs.Image(
-                type='pil',
-                label='Paprika Result'),
+                label='Result'),
+
         ],
         #examples=examples,
         theme=args.theme,

test1.py

@@ -53,23 +53,27 @@ def stats_graph(graph):
     # params = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.trainable_variables_parameter())
     print('FLOPs: {}'.format(flops.total_float_ops))
 
-def test(checkpoint_dir, style_name, test_dir, if_adjust_brightness, img_size=[256,256]):
+g_sess = None
+test_generated = None
+
+def test(checkpoint_dir, style_name, test_file, if_adjust_brightness, img_size=[256,256]):
+    global g_sess
+    global test_generated
+
     # tf.reset_default_graph()
     result_dir = 'results/'+style_name
     check_folder(result_dir)
-    test_files = [test_dir]
 
-    test_real = tf.placeholder(tf.float32, [1, None, None, 3], name='test')
+    if g_sess is None:
+        test_real = tf.placeholder(tf.float32, [1, None, None, 3], name='test')
 
-    with tf.variable_scope("generator", reuse=False):
-        test_generated = generator.G_net(test_real).fake
-    saver = tf.train.Saver()
+        with tf.variable_scope("generator", reuse=False):
+            test_generated = generator.G_net(test_real).fake
+        saver = tf.train.Saver()
 
-    out_paths = []
+        gpu_options = tf.GPUOptions(allow_growth=True)
+        g_sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options))
 
-    gpu_options = tf.GPUOptions(allow_growth=True)
-    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)) as sess:
-        # tf.global_variables_initializer().run()
         # load model
         ckpt = tf.train.get_checkpoint_state(checkpoint_dir)  # checkpoint file information
         if ckpt and ckpt.model_checkpoint_path:
@@ -81,22 +85,20 @@ def test(checkpoint_dir, style_name, test_dir, if_adjust_brightness, img_size=[2
             return
         # stats_graph(tf.get_default_graph())
 
-        begin = time.time()
-        for sample_file  in tqdm(test_files) :
-            # print('Processing image: ' + sample_file)
-            sample_image = np.asarray(load_test_data(sample_file, img_size))
-            image_path = os.path.join(result_dir,'{0}'.format(os.path.basename(sample_file)))
-            fake_img = sess.run(test_generated, feed_dict = {test_real : sample_image})
-            if if_adjust_brightness:
-                save_images(fake_img, image_path, sample_file)
-            else:
-                save_images(fake_img, image_path, None)
-				
-            out_paths.append(image_path)
-        end = time.time()
-        print(f'test-time: {end-begin} s')
-		
-    return out_paths
+    begin = time.time()
+    # print('Processing image: ' + sample_file)
+    sample_image = np.asarray(load_test_data(test_file, img_size))
+    image_path = os.path.join(result_dir,'{0}'.format(os.path.basename(test_file)))
+    fake_img = g_sess.run(test_generated, feed_dict = {test_real : sample_image})
+    if if_adjust_brightness:
+        save_images(fake_img, image_path, test_file)
+    else:
+        save_images(fake_img, image_path, None)
+
+    end = time.time()
+    print(f'test-time: {end-begin} s')
+
+    return image_path
     
 if __name__ == '__main__':
     arg = parse_args()