mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-09 00:09:45 +00:00
Aryan
512a8a6854
* update * update * update * update * update * update * update * update * update * update * update * update * update * update * update
105 lines
3.0 KiB
Python
105 lines
3.0 KiB
Python
# Run: python3 tests/test_dataset.py
|
|
|
|
import sys
|
|
|
|
|
|
def test_video_dataset():
|
|
from dataset import VideoDataset
|
|
|
|
dataset_dirs = VideoDataset(
|
|
data_root="assets/tests/",
|
|
caption_column="prompts.txt",
|
|
video_column="videos.txt",
|
|
max_num_frames=49,
|
|
id_token=None,
|
|
random_flip=None,
|
|
)
|
|
dataset_csv = VideoDataset(
|
|
data_root="assets/tests/",
|
|
dataset_file="assets/tests/metadata.csv",
|
|
caption_column="caption",
|
|
video_column="video",
|
|
max_num_frames=49,
|
|
id_token=None,
|
|
random_flip=None,
|
|
)
|
|
|
|
assert len(dataset_dirs) == 1
|
|
assert len(dataset_csv) == 1
|
|
assert dataset_dirs[0]["video"].shape == (49, 3, 480, 720)
|
|
assert (dataset_dirs[0]["video"] == dataset_csv[0]["video"]).all()
|
|
|
|
print(dataset_dirs[0]["video"].shape)
|
|
|
|
|
|
def test_video_dataset_with_resizing():
|
|
from dataset import VideoDatasetWithResizing
|
|
|
|
dataset_dirs = VideoDatasetWithResizing(
|
|
data_root="assets/tests/",
|
|
caption_column="prompts.txt",
|
|
video_column="videos.txt",
|
|
max_num_frames=49,
|
|
id_token=None,
|
|
random_flip=None,
|
|
)
|
|
dataset_csv = VideoDatasetWithResizing(
|
|
data_root="assets/tests/",
|
|
dataset_file="assets/tests/metadata.csv",
|
|
caption_column="caption",
|
|
video_column="video",
|
|
max_num_frames=49,
|
|
id_token=None,
|
|
random_flip=None,
|
|
)
|
|
|
|
assert len(dataset_dirs) == 1
|
|
assert len(dataset_csv) == 1
|
|
assert dataset_dirs[0]["video"].shape == (48, 3, 480, 720) # Changes due to T2V frame bucket sampling
|
|
assert (dataset_dirs[0]["video"] == dataset_csv[0]["video"]).all()
|
|
|
|
print(dataset_dirs[0]["video"].shape)
|
|
|
|
|
|
def test_video_dataset_with_bucket_sampler():
|
|
import torch
|
|
from dataset import BucketSampler, VideoDatasetWithResizing
|
|
from torch.utils.data import DataLoader
|
|
|
|
dataset_dirs = VideoDatasetWithResizing(
|
|
data_root="assets/tests/",
|
|
caption_column="prompts_multi.txt",
|
|
video_column="videos_multi.txt",
|
|
max_num_frames=49,
|
|
id_token=None,
|
|
random_flip=None,
|
|
)
|
|
sampler = BucketSampler(dataset_dirs, batch_size=8)
|
|
|
|
def collate_fn(data):
|
|
captions = [x["prompt"] for x in data[0]]
|
|
videos = [x["video"] for x in data[0]]
|
|
videos = torch.stack(videos)
|
|
return captions, videos
|
|
|
|
dataloader = DataLoader(dataset_dirs, batch_size=1, sampler=sampler, collate_fn=collate_fn)
|
|
first = False
|
|
|
|
for captions, videos in dataloader:
|
|
if not first:
|
|
assert len(captions) == 8 and isinstance(captions[0], str)
|
|
assert videos.shape == (8, 48, 3, 480, 720)
|
|
first = True
|
|
else:
|
|
assert len(captions) == 8 and isinstance(captions[0], str)
|
|
assert videos.shape == (8, 48, 3, 256, 360)
|
|
break
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.path.append("./training")
|
|
|
|
test_video_dataset()
|
|
test_video_dataset_with_resizing()
|
|
test_video_dataset_with_bucket_sampler()
|