Update preprocess.py

preprocess.py

@@ -26,6 +26,11 @@ fa = [face_detection.FaceAlignment(face_detection.LandmarksType._2D, flip_input=
 
 template = 'ffmpeg -loglevel panic -y -i {} -strict -2 {}'
 
+def _progress(generated, to_generate):
+    progress((generated, to_generate))
+
+_progress(0, 1)
+
 def process_video_file(vfile, args, gpu_id):
     video_stream = cv2.VideoCapture(vfile)
     frames = []
@@ -41,7 +46,7 @@ def process_video_file(vfile, args, gpu_id):
     os.makedirs(fulldir, exist_ok=True)
     batches = [frames[i:i + args.batch_size] for i in range(0, len(frames), args.batch_size)]
     i = -1
-    for fb in batches:
+    for fb in tqdm(batches, desc='Processing Video Frames'):
         preds = fa[gpu_id].get_detections_for_batch(np.asarray(fb))
         for j, f in enumerate(preds):
             i += 1
@@ -49,6 +54,7 @@ def process_video_file(vfile, args, gpu_id):
                 continue
             x1, y1, x2, y2 = f
             cv2.imwrite(path.join(fulldir, '{}.jpg'.format(i)), fb[j][y1:y2, x1:x2])
+            _progress(i + 1, len(batches) * args.batch_size)
 
 def process_audio_file(vfile, args):
     vidname = os.path.basename(vfile).split('.')[0]
@@ -58,6 +64,7 @@ def process_audio_file(vfile, args):
     wavpath = path.join(fulldir, 'audio.wav')
     command = template.format(vfile, wavpath)
     subprocess.call(command, shell=True)
+    _progress(1, 1)
 
 def mp_handler(job):
     vfile, args, gpu_id = job
@@ -70,20 +77,21 @@ def mp_handler(job):
 
 def main(args):
     print('Started processing for {} with {} GPUs'.format(args.data_root, args.ngpu))
-filelist = glob(path.join(args.data_root, '*.mp4'))
+    filelist = glob(path.join(args.data_root, '*.mp4'))
 
-jobs = [(vfile, args, i % args.ngpu) for i, vfile in enumerate(filelist)]
-p = ThreadPoolExecutor(args.ngpu)
-futures = [p.submit(mp_handler, j) for j in jobs]
-_ = [r.result() for r in tqdm(as_completed(futures), total=len(futures))]
+    jobs = [(vfile, args, i % args.ngpu) for i, vfile in enumerate(filelist)]
+    p = ThreadPoolExecutor(args.ngpu)
+    futures = [p.submit(mp_handler, j) for j in jobs]
+    _ = [r.result() for r in tqdm(as_completed(futures), total=len(futures))]
+    _progress(1, 1)
 
 print('Dumping audios...')
-
-for vfile in tqdm(filelist):
+for vfile in tqdm(filelist, desc='Processing Audio Files'):
     try:
         process_audio_file(vfile, args)
     except KeyboardInterrupt:
         exit(0)
     except:
         traceback.print_exc()
-        continue
+        continue
+    _progress(1, len(filelist))