Update app.py

app.py

@@ -1,18 +1,30 @@
 import numpy as np
 import gradio as gr
 import cv2
-from black import to_black
+import os
+import argparse
+from inference import Predictor
+#from black import to_black
 
-# def to_black(image, transfer_style):
-#     if transfer_style == "Hayao":
-#          output = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)# 转换为灰度图像
-#          return output
-#     elif transfer_style == "Shinkai":
-#         return image
-#     elif transfer_style == "Kon Satoshi":
-#         return image
-#     else:
-#         return image
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--weight', type=str, default='GeneratorV2_train_photo_Paprika_init.pt')
+    parser.add_argument('--device', type=str, default='cpu', help='Device, cuda or cpu')
+
+    return parser.parse_args()
+
+def transfer(image, transfer_style):
+    if transfer_style == "Hayao":
+         #output = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)# 转换为灰度图像
+        predictor = Predictor(args.weight, args.device)
+        anime_img = predictor.transform_image(args.src, save_path)
+        return anime_img
+    elif transfer_style == "Shinkai":
+        return image
+    elif transfer_style == "Kon Satoshi":
+        return image
+    else:
+        return image
 
 def clear_output(input_widget):
     input_widget = np.array([])
@@ -37,7 +49,7 @@ with gr.Blocks() as demo:
                 video_button = gr.Button("Transfer")
                 clear_video_button = gr.Button("Clear")
 
-    image_button.click(to_black, inputs=[image_input,image_dropdown], outputs=image_output)
+    image_button.click(transfer, inputs=[image_input,image_dropdown], outputs=image_output)
     video_button.click(to_black, inputs=[video_input,video_dropdown], outputs=video_output)
     clear_image_button.click(clear_output, inputs=image_input,outputs=image_output)
     clear_video_button.click(clear_output, inputs=video_input,outputs=video_output)