Added the encoder codec argument to crop_video_sequences

YuvalNirkin · Apr 24, 2020 · 301f2e3 · 301f2e3
1 parent 15efc29
commit 301f2e3
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/preprocess/crop_video_sequences.py b/preprocess/crop_video_sequences.py
@@ -8,7 +8,7 @@
 
 
 def main(input_path, output_dir=None, cache_path=None, seq_postfix='_dsfd_seq.pkl', resolution=256, crop_scale=2.0,
-         select='all', disable_tqdm=False):
+         select='all', disable_tqdm=False, encoder_codec='avc1'):
     cache_path = os.path.splitext(input_path)[0] + seq_postfix if cache_path is None else cache_path
     if output_dir is None:
         output_dir = os.path.splitext(input_path)[0]
@@ -47,7 +47,7 @@ def main(input_path, output_dir=None, cache_path=None, seq_postfix='_dsfd_seq.pk
 
     # For each sequence initialize output video file
     out_vids = []
-    fourcc = cv2.VideoWriter_fourcc(*'avc1')
+    fourcc = cv2.VideoWriter_fourcc(*encoder_codec)
     for seq in seq_list:
         curr_vid_name = os.path.splitext(os.path.basename(input_path))[0] + '_seq%02d.mp4' % seq.id
         curr_vid_path = os.path.join(output_dir, curr_vid_name)
@@ -139,6 +139,8 @@ def main(input_path, output_dir=None, cache_path=None, seq_postfix='_dsfd_seq.pk
                         help='selection method [all|longest]')
     parser.add_argument('-dt', '--disable_tqdm', dest='disable_tqdm', action='store_true',
                           help='if specified disables tqdm progress bar')
+    parser.add_argument('-ec', '--encoder_codec', default='avc1', metavar='STR',
+                        help='encoder codec code')
     args = parser.parse_args()
     main(args.input, args.output, args.cache, args.seq_postfix, args.resolution, args.crop_scale, args.select,
-         args.disable_tqdm)
+         args.disable_tqdm, args.encoder_codec)
diff --git a/preprocess/preprocess_video.py b/preprocess/preprocess_video.py
@@ -215,6 +215,7 @@ def __init__(self, resolution=d('resolution'), crop_scale=d('crop_scale'), gpus=
             self.smooth_seg = None
 
         # Initialize output videos format
+        self.encoder_codec = encoder_codec
         self.fourcc = cv2.VideoWriter_fourcc(*encoder_codec)
 
     def process_pose(self, input_path, output_dir, seq_file_path):
@@ -453,7 +454,8 @@ def cache(self, input_path, output_dir=None):
         if not os.path.isfile(first_cropped_path):
             if is_vid:
                 crop_video_sequences_main(input_path, output_dir, seq_file_path, self.seq_postfix, self.resolution,
-                                          self.crop_scale, select='all', disable_tqdm=False)
+                                          self.crop_scale, select='all', disable_tqdm=False,
+                                          encoder_codec=self.encoder_codec)
             else:
                 crop_image_sequences_main(input_path, output_dir, seq_file_path, self.seq_postfix, '.jpg',
                                           self.resolution, self.crop_scale)