init

app.py

@@ -18,6 +18,7 @@ from io import BytesIO
 sys.path.insert(0, 'animeganv2')
 
 import test1 as test
+from test1 import ImportGraph
 
 ORIGINAL_REPO_URL = 'https://github.com/TachibanaYoshino/AnimeGANv2'
 TITLE = 'TachibanaYoshino/AnimeGANv2'
@@ -48,17 +49,16 @@ def parse_args() -> argparse.Namespace:
 
 def run(
     image,
-    checkpoint_dir: str,
-) -> tuple[PIL.Image.Image]:
+    shinkai: ImportGraph,
+    hayao: ImportGraph,
+    paprika: ImportGraph,
+) -> tuple[PIL.Image.Image, PIL.Image.Image, PIL.Image.Image]:
 
-    curPath = os.path.abspath(os.path.dirname(__file__))
-    checkpoint_dir = os.path.join(curPath, 'animeganv2/checkpoint/generator_Shinkai_weight')
-    print(checkpoint_dir)
-
-    outs = test.test(checkpoint_dir, 'save_dir', image.name, True)
+    im1 = shinkai.test(image.name, True)
+    im2 = hayao.test(image.name, True)
+    im3 = paprika.test(image.name, True)
 
-
-    return PIL.Image.open(outs[0])
+    return PIL.Image.open(im1),PIL.Image.open(im2),PIL.Image.open(im3)
 
 
 def main():
@@ -66,8 +66,13 @@ def main():
 
     args = parse_args()
 
+    curPath = os.path.abspath(os.path.dirname(__file__))
+    #init
+    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, checkpoint_dir='')
+    func = functools.partial(run, shinkai=shinkai,hayao=hayao,paprika=paprika  )
     func = functools.update_wrapper(func, run)
 
     
@@ -79,7 +84,13 @@ def main():
         [
             gr.outputs.Image(
                 type='pil',
-                label='Result'),
+                label='Shinkai Result'),
+            gr.outputs.Image(
+                type='pil',
+                label='Hayao Result'),
+            gr.outputs.Image(
+                type='pil',
+                label='Paprika Result'),
         ],
         #examples=examples,
         theme=args.theme,

test1.py

@@ -8,22 +8,35 @@ import numpy as np
 from net import generator
 os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
 
-def parse_args():
-    desc = "AnimeGANv2"
-    parser = argparse.ArgumentParser(description=desc)
+class ImportGraph:
+    def __init__(self, checkpoint_dir):
+        self.graph = tf.Graph()
+        self.sess = tf.Session(graph=self.graph, config=tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options))
+        with self.graph.as_default():
+            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()
 
-    parser.add_argument('--checkpoint_dir', type=str, default='checkpoint/'+'generator_Shinkai_weight',
-                        help='Directory name to save the checkpoints')
-    parser.add_argument('--test_dir', type=str, default='dataset/test/t',
-                        help='Directory name of test photos')
-    parser.add_argument('--save_dir', type=str, default='Shinkai/t',
-                        help='what style you want to get')
-    parser.add_argument('--if_adjust_brightness', type=bool, default=True,
-                        help='adjust brightness by the real photo')
+            ckpt = tf.train.get_checkpoint_state(checkpoint_dir)  # checkpoint file information
+            if ckpt and ckpt.model_checkpoint_path:
+                ckpt_name = os.path.basename(ckpt.model_checkpoint_path)  # first line
+                saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
+                print(" [*] Success to read {}".format(os.path.join(checkpoint_dir, ckpt_name)))
+            else:
+                print(" [*] Failed to find a checkpoint")
+
+    def test(self, sample_file, if_adjust_brightness, img_size=[256,256]):
+        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)
 
-    """checking arguments"""
+        return image_path
 
-    return parser.parse_args()
 
 def stats_graph(graph):
     flops = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.float_operation())