From 1b9c2a33fd0ec3118f14f9ac859d0695a4e1429c Mon Sep 17 00:00:00 2001
From: Troceleng <60510182+Troceleng@users.noreply.github.com>
Date: Mon, 30 May 2022 15:25:26 +0300
Subject: [PATCH] Important updates

Updating this to add lossless quality outputs (Using the FFV1 video codec) and offer the option to skip video preview. Also changed the default FPS to be a much more common 24000/1001 and fixed a typo (datset --> dataset).
---
 test.py | 68 +++++++++++++++++++++++++++++++++------------------------
 1 file changed, 39 insertions(+), 29 deletions(-)

diff --git a/test.py b/test.py
index 448f10c..ca16ec5 100644
--- a/test.py
+++ b/test.py
@@ -20,7 +20,9 @@
 parser.add_argument("--step", type=int, default=10)
 parser.add_argument("--num_ref", type=int, default=-1)
 parser.add_argument("--neighbor_stride", type=int, default=5)
-parser.add_argument("--savefps", type=int, default=24)
+parser.add_argument("--savefps", type=float, default=24000/1001)
+parser.add_argument("--quality", type=str, default="lossless")
+parser.add_argument("--showcase", type=str, default=True)
 
 # args for e2fgvi_hq (which can handle videos with arbitrary resolution)
 parser.add_argument("--set_size", action='store_true', default=False)
@@ -181,43 +183,51 @@ def main_worker():
     # saving videos
     print('Saving videos...')
     save_dir_name = 'results'
-    ext_name = '_results.mp4'
+    ext_name = '_results.mkv'
     save_base_name = args.video.split('/')[-1]
     save_name = save_base_name.replace(
-        '.mp4', ext_name) if args.use_mp4 else save_base_name + ext_name
+        '.mkv', ext_name) if args.use_mp4 else save_base_name + ext_name
     if not os.path.exists(save_dir_name):
         os.makedirs(save_dir_name)
     save_path = os.path.join(save_dir_name, save_name)
-    writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"),
-                             default_fps, size)
-    for f in range(video_length):
-        comp = comp_frames[f].astype(np.uint8)
-        writer.write(cv2.cvtColor(comp, cv2.COLOR_BGR2RGB))
+    if args.quality == "lossless":  # this option uses a lossless video codec for higher quality output
+        writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"FFV1"),
+                                 default_fps, size)
+        for f in range(video_length):
+            comp = comp_frames[f].astype(np.uint8)
+            writer.write(cv2.cvtColor(comp, cv2.COLOR_BGR2RGB))
+    else:
+        writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"),
+                                 default_fps, size)
+        for f in range(video_length):
+            comp = comp_frames[f].astype(np.uint8)
+            writer.write(cv2.cvtColor(comp, cv2.COLOR_BGR2RGB))
     writer.release()
     print(f'Finish test! The result video is saved in: {save_path}.')
 
     # show results
-    print('Let us enjoy the result!')
-    fig = plt.figure('Let us enjoy the result')
-    ax1 = fig.add_subplot(1, 2, 1)
-    ax1.axis('off')
-    ax1.set_title('Original Video')
-    ax2 = fig.add_subplot(1, 2, 2)
-    ax2.axis('off')
-    ax2.set_title('Our Result')
-    imdata1 = ax1.imshow(frames[0])
-    imdata2 = ax2.imshow(comp_frames[0].astype(np.uint8))
-
-    def update(idx):
-        imdata1.set_data(frames[idx])
-        imdata2.set_data(comp_frames[idx].astype(np.uint8))
-
-    fig.tight_layout()
-    anim = animation.FuncAnimation(fig,
-                                   update,
-                                   frames=len(frames),
-                                   interval=50)
-    plt.show()
+    if args.showcase == True:
+        print('Let us enjoy the result!')
+        fig = plt.figure('Let us enjoy the result')
+        ax1 = fig.add_subplot(1, 2, 1)
+        ax1.axis('off')
+        ax1.set_title('Original Video')
+        ax2 = fig.add_subplot(1, 2, 2)
+        ax2.axis('off')
+        ax2.set_title('Our Result')
+        imdata1 = ax1.imshow(frames[0])
+        imdata2 = ax2.imshow(comp_frames[0].astype(np.uint8))
+
+        def update(idx):
+            imdata1.set_data(frames[idx])
+            imdata2.set_data(comp_frames[idx].astype(np.uint8))
+
+        fig.tight_layout()
+        anim = animation.FuncAnimation(fig,
+                                       update,
+                                       frames=len(frames),
+                                       interval=50)
+        plt.show()
 
 
 if __name__ == '__main__':