From dfcb5203dc025d22012f90c224ba344abc579939 Mon Sep 17 00:00:00 2001
From: Fuou Marinas <57849916+FuouM@users.noreply.github.com>
Date: Sun, 30 Jun 2024 22:26:08 +0700
Subject: [PATCH] Fix find_best_frame

used to take the number of color channels (3) instead of number of frames
---
 utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/utils.py b/utils.py
index 12894cd..acc77d9 100644
--- a/utils.py
+++ b/utils.py
@@ -120,7 +120,7 @@ def tensor_to_numpy(tensor):
     return tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()
 
 
-def find_best_frame(source_image, driving_video: list, cpu=False) -> int:
+def find_best_frame(source_image, driving_video, cpu=False) -> int:
     if not face_alignment_available:
         print("Face Alignment module were not found. Best frame set to 0.")
         return 0
@@ -139,7 +139,7 @@ def find_best_frame(source_image, driving_video: list, cpu=False) -> int:
     norm = float("inf")
     frame_num = 0
 
-    num_frames = driving_video.shape[1]
+    num_frames = driving_video.shape[2]
 
     for i in tqdm(range(num_frames)):
         # Extract single frame from driving video