From 316a705f6a8081af2de388c64fc9b4833daf3a82 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 19 Nov 2024 10:29:16 +0530 Subject: [PATCH] revert changes to training/dataset.py --- training/dataset.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/training/dataset.py b/training/dataset.py index bb8f8b9..ec47b0b 100644 --- a/training/dataset.py +++ b/training/dataset.py @@ -49,7 +49,7 @@ class VideoDataset(Dataset): self.caption_column = caption_column self.video_column = video_column self.max_num_frames = max_num_frames - self.id_token = id_token or "" + self.id_token = f"{id_token.strip()} " if id_token else "" self.height_buckets = height_buckets or HEIGHT_BUCKETS self.width_buckets = width_buckets or WIDTH_BUCKETS self.frame_buckets = frame_buckets or FRAME_BUCKETS @@ -277,9 +277,7 @@ class VideoDatasetWithResizing(VideoDataset): video_reader = decord.VideoReader(uri=path.as_posix()) video_num_frames = len(video_reader) nearest_frame_bucket = min( - [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, + self.frame_buckets, key=lambda x: abs(x - min(video_num_frames, self.max_num_frames)) ) frame_indices = list(range(0, video_num_frames, video_num_frames // nearest_frame_bucket)) @@ -344,9 +342,7 @@ class VideoDatasetWithResizeAndRectangleCrop(VideoDataset): video_reader = decord.VideoReader(uri=path.as_posix()) video_num_frames = len(video_reader) nearest_frame_bucket = min( - [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, + self.frame_buckets, key=lambda x: abs(x - min(video_num_frames, self.max_num_frames)) ) frame_indices = list(range(0, video_num_frames, video_num_frames // nearest_frame_bucket)) @@ -429,4 +425,4 @@ class BucketSampler(Sampler): random.shuffle(bucket) yield bucket del self.buckets[fhw] - self.buckets[fhw] = [] + self.buckets[fhw] = [] \ No newline at end of file