From c80bd177fdced754b2cb9e9336c6be1ce58604a5 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 19 Nov 2024 10:27:24 +0530 Subject: [PATCH] nearest_frame_bucket. --- training/dataset.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/training/dataset.py b/training/dataset.py index 7ace5c0..bb8f8b9 100644 --- a/training/dataset.py +++ b/training/dataset.py @@ -277,7 +277,9 @@ class VideoDatasetWithResizing(VideoDataset): video_reader = decord.VideoReader(uri=path.as_posix()) video_num_frames = len(video_reader) nearest_frame_bucket = min( - self.frame_buckets, key=lambda x: abs(x - min(video_num_frames, self.max_num_frames)) + [bucket for bucket in self.frame_buckets if bucket <= video_num_frames], + key=lambda x: abs(x - min(video_num_frames, self.max_num_frames)), + default=1, ) frame_indices = list(range(0, video_num_frames, video_num_frames // nearest_frame_bucket)) @@ -342,7 +344,9 @@ class VideoDatasetWithResizeAndRectangleCrop(VideoDataset): video_reader = decord.VideoReader(uri=path.as_posix()) video_num_frames = len(video_reader) nearest_frame_bucket = min( - self.frame_buckets, key=lambda x: abs(x - min(video_num_frames, self.max_num_frames)) + [bucket for bucket in self.frame_buckets if bucket <= video_num_frames], + key=lambda x: abs(x - min(video_num_frames, self.max_num_frames)), + default=1, ) frame_indices = list(range(0, video_num_frames, video_num_frames // nearest_frame_bucket))