Update app.py

app.py

@@ -60,21 +60,6 @@ def transfer_video(video_input, transfer_style):
          #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()
-        # # 加载视频文件
-        # #video_binary = io.BytesIO(video)
-            
-        # cap = cv2.VideoCapture(video)
-        
-        # # 读取视频帧并保存到一个列表中
-        # video_frames = []
-        # while True:
-        #     ret, frame = cap.read()
-        #     if not ret:
-        #         break
-        #     video_frames.append(frame)
-        
-        # # 关闭视频文件
-        # cap.release()
         args.src = video_input
         args.out = "video.mp4"
         Predictor(args.weight).transform_video(args.src, args.out, args.batch_size, start=args.start, end=args.end)
@@ -82,19 +67,21 @@ def transfer_video(video_input, transfer_style):
         #anime_video = Predictor(args.weight).transform_video(video, args.batch_size, args.start, args.end)
         #return anime_video
     elif transfer_style == "Shinkai":
-        args = parse_args()
+        args = parse_args_video()
         args.weight = 'GeneratorV2_train_photo_Shinkai_init.pt'
-        predictor = Predictor(args.weight, args.device)
-        anime_img = predictor.transform_image(image)
-        return anime_img
+        args.src = video_input
+        args.out = "video.mp4"
+        Predictor(args.weight).transform_video(args.src, args.out, args.batch_size, start=args.start, end=args.end)
+        return args.out
     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
+        args.src = video_input
+        args.out = "video.mp4"
+        Predictor(args.weight).transform_video(args.src, args.out, args.batch_size, start=args.start, end=args.end)
+        return args.out
     else:
-        return image
+        return 0
 
 def clear_output(input_widget):
     input_widget = np.array([])