use original aspect ration, use ffmpeg

- pytorch 2.x compatibility
- use original aspect ratio
- use ffmpeg

FGT_codes/tool/video_inpainting.py

@@ -28,6 +28,7 @@ import copy
 import glob
 import cv2
 import argparse
+import ffmpeg
 
 
 def to_tensor(img):
@@ -206,13 +207,9 @@ def initialize_LAFC(args, device):
     pkg = import_module("LAFC.models.{}".format(model))
     model = pkg.Model(configs)
     if not torch.cuda.is_available():
-        state = torch.load(
-        checkpoint, map_location=lambda storage, loc: storage
-        )
+        state = torch.load(checkpoint, map_location="cpu")
     else:
-        state = torch.load(
-            checkpoint, map_location=lambda storage, loc: storage.cuda(device)
-        )
+        state = torch.load(checkpoint, map_location=device)
     model.load_state_dict(state["model_state_dict"])
     model = model.to(device)
     return model, configs
@@ -230,13 +227,9 @@ def initialize_FGT(args, device):
     net = import_module("FGT.models.{}".format(model))
     model = net.Model(configs).to(device)
     if not torch.cuda.is_available():
-        state = torch.load(
-        checkpoint, map_location=lambda storage, loc: storage
-        )
+        state = torch.load(checkpoint, map_location="cpu")
     else:
-        state = torch.load(
-            checkpoint, map_location=lambda storage, loc: storage.cuda(device)
-        )
+        state = torch.load(checkpoint, map_location=device)
     model.load_state_dict(state["model_state_dict"])
     return model, configs
 
@@ -427,7 +420,7 @@ def save_results(outdir, comp_frames):
         cv2.imwrite(out_path, comp_frames[i][:, :, ::-1])
 
 
-def video_inpainting(args, imgArr, imgMaskArr):
+def video_inpainting(args, original_frame_list, imgArr, imgMaskArr):
     #device = torch.device("cuda:{}".format(args.gpu))
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     print(args)
@@ -739,6 +732,9 @@ def video_inpainting(args, imgArr, imgMaskArr):
     flows = np2tensor(videoFlowF, near="t")
     flows = norm_flows(flows).to(device)
 
+    original_frame = original_frame_list[0]
+    orig_h, orig_w = original_frame.shape[:2]
+
     for f in range(0, video_length, neighbor_stride):
         neighbor_ids = [
             i
@@ -770,14 +766,34 @@ def video_inpainting(args, imgArr, imgMaskArr):
                     comp_frames[idx].astype(np.float32) * 0.5
                     + comp.astype(np.float32) * 0.5
                 )
+                
     if args.vis_frame:
         save_results(args.outroot, comp_frames)
+
     create_dir(args.outroot)
+
     for i in range(len(comp_frames)):
-        comp_frames[i] = comp_frames[i].astype(np.uint8)
-    imageio.mimwrite(
-        os.path.join(args.outroot, args.outfilename), comp_frames, fps=args.out_fps, quality=8
+        if comp_frames[i] is not None:
+            comp_frames[i] = cv2.resize(
+                comp_frames[i].astype(np.uint8),
+                (orig_w, orig_h),
+                interpolation=cv2.INTER_LANCZOS4
+            )
+    output_path = os.path.join(args.outroot, args.outfilename)
+    process = (
+        ffmpeg
+        .input('pipe:', format='rawvideo', pix_fmt='rgb24',
+            s=f'{orig_w}x{orig_h}', framerate=args.out_fps)
+        .output(output_path, vcodec='libx264', crf=18, pix_fmt='yuv420p')
+        .overwrite_output()
+        .run_async(pipe_stdin=True)
     )
+
+    for frame in comp_frames:
+        process.stdin.write(frame.tobytes())
+    process.stdin.close()
+    process.wait()
+
     print(f"Done, please check your result in {args.outroot} ")