Update app.py

app.py

@@ -19,6 +19,17 @@ def parse_args():
 
     return parser.parse_args()
 
+def parse_args_video():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--weight', type=str, default='GeneratorV2_train_photo_Hayao_init.pt')
+    #parser.add_argument('--src', type=str, default='dataset/video/花.mp4', help='Path to input video')
+    #parser.add_argument('--out', type=str, default='dataset/video_Hayao/hua_hayao.mp4', help='Path to save new video')
+    parser.add_argument('--batch-size', type=int, default=4)
+    parser.add_argument('--start', type=int, default=0, help='Start time of video (second)')
+    parser.add_argument('--end', type=int, default=10, help='End time of video (second), 0 if not set')
+
+    return parser.parse_args()
+
 def transfer(image, transfer_style):
     if transfer_style == "Hayao":
          #output = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)# 转换为灰度图像
@@ -42,6 +53,29 @@ def transfer(image, transfer_style):
     else:
         return image
 
+        
+def transfer_video(video, transfer_style):
+    if transfer_style == "Hayao":
+         #output = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)# 转换为灰度图像
+        #os.system("wget https://huggingface.co/YANGYYYY/cartoonize/resolve/main/GeneratorV2_train_photo_Hayao_init.pt")
+        args = parse_args_video()
+        anime_video = Predictor(args.weight).transform_video(video, args.batch_size, start=args.start, end=args.end)
+        return anime_video
+    elif transfer_style == "Shinkai":
+        args = parse_args()
+        args.weight = 'GeneratorV2_train_photo_Shinkai_init.pt'
+        predictor = Predictor(args.weight, args.device)
+        anime_img = predictor.transform_image(image)
+        return anime_img
+    elif transfer_style == "Kon Satoshi":
+        args = parse_args()
+        args.weight = 'GeneratorV2_train_photo_Paprika_init.pt'
+        predictor = Predictor(args.weight, args.device)
+        anime_img = predictor.transform_image(image)
+        return anime_img
+    else:
+        return image
+
 def clear_output(input_widget):
     input_widget = np.array([])
 
@@ -66,7 +100,7 @@ with gr.Blocks() as demo:
                 clear_video_button = gr.Button("Clear")
 
     image_button.click(transfer, inputs=[image_input,image_dropdown], outputs=image_output)
-    video_button.click(transfer, inputs=[video_input,video_dropdown], outputs=video_output)
+    video_button.click(transfer_video, 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)
 demo.launch()