limit number of frames

evaluate.py

@@ -66,7 +66,7 @@ def get_score(item, image_paths, video_path, metrics, train_steps=100, inference
         images.append([img])
     
     gt_frames = extract_frames(video_path, fps)
-    #gt_frames = gt_frames[:2]
+    gt_frames = gt_frames[:1000]
     for f in gt_frames:
         f.thumbnail((512,512))
     

main.py

@@ -85,7 +85,7 @@ debug = False
 save_model = True
 should_gen_vid = False
 max_batch_size = 8
-
+max_frame_count = 1000
 
 def save_temp_imgs(imgs):
     os.makedirs('temp', exist_ok=True)
@@ -245,6 +245,9 @@ def extract_frames(video_path, fps):
         if not success: break
         index_in += 1
 
+        if frame_count > max_frame_count:
+            break
+
         out_due = int(index_in / fps_in * fps_out)
         if out_due > index_out:
             success, frame = video_capture.retrieve()
@@ -1271,7 +1274,7 @@ def run_eval(images_orig, video_path, train_steps=100, inference_steps=10, fps=1
     in_img, in_pose, train_imgs, train_poses = prepare_inputs_train(images, bg_remove, dwpose, rembg_session)
     in_img, target_poses, in_pose, _, _ = prepare_inputs_inference(in_img, video_path, [], fps, dwpose, rembg_session, bg_remove, resize_inputs, is_app)
     
-    #target_poses = target_poses[:2]
+    target_poses = target_poses[:max_frame_count]
     #train_steps = 3
     
     finetune = False