Update util.py

util.py

@@ -158,10 +158,17 @@ def scatter_pitch(ax, t, f, s=60, marker="o", color="limegreen", label="Pitch"):
     ax.legend(loc="upper left")
 
 
-# Load video frame
-def load_frame(video_path):
-    vr = decord.VideoReader(video_path, num_threads=1)
-    frame = PIL.Image.fromarray(vr[0].asnumpy())
+def load_frame(video_path, video_backend="decord"):
+    if video_backend == "decord":
+        vr = decord.VideoReader(video_path, num_threads=1)
+        frame = PIL.Image.fromarray(vr[0].asnumpy())
+    elif video_backend == "torchvision":
+        import torchvision.io as tio
+        video, _, _ = tio.read_video(video_path, pts_unit="sec")
+        frame = video[0]
+        frame = PIL.Image.fromarray(frame.numpy())
+    else:
+        raise ValueError(f"Unknown video backend: {video_backend}")
     frame = audio_loader.crop_or_pad_to_size(frame, size=(270, 480))
     return frame