|
|
import pathlib |
|
|
import tempfile |
|
|
import unittest |
|
|
|
|
|
import torch |
|
|
from PIL import Image |
|
|
|
|
|
from finetrainers.data import ( |
|
|
ImageCaptionFilePairDataset, |
|
|
ImageFileCaptionFileListDataset, |
|
|
ImageFolderDataset, |
|
|
ValidationDataset, |
|
|
VideoCaptionFilePairDataset, |
|
|
VideoFileCaptionFileListDataset, |
|
|
VideoFolderDataset, |
|
|
VideoWebDataset, |
|
|
initialize_dataset, |
|
|
) |
|
|
from finetrainers.utils import find_files |
|
|
|
|
|
from .utils import create_dummy_directory_structure |
|
|
|
|
|
|
|
|
class DatasetTesterMixin: |
|
|
num_data_files = None |
|
|
directory_structure = None |
|
|
caption = "A cat ruling the world" |
|
|
metadata_extension = None |
|
|
|
|
|
def setUp(self): |
|
|
if self.num_data_files is None: |
|
|
raise ValueError("num_data_files is not defined") |
|
|
if self.directory_structure is None: |
|
|
raise ValueError("dataset_structure is not defined") |
|
|
|
|
|
self.tmpdir = tempfile.TemporaryDirectory() |
|
|
create_dummy_directory_structure( |
|
|
self.directory_structure, self.tmpdir, self.num_data_files, self.caption, self.metadata_extension |
|
|
) |
|
|
|
|
|
def tearDown(self): |
|
|
self.tmpdir.cleanup() |
|
|
|
|
|
|
|
|
class ImageDatasetTesterMixin(DatasetTesterMixin): |
|
|
metadata_extension = "jpg" |
|
|
|
|
|
|
|
|
class VideoDatasetTesterMixin(DatasetTesterMixin): |
|
|
metadata_extension = "mp4" |
|
|
|
|
|
|
|
|
class ImageCaptionFilePairDatasetFastTests(ImageDatasetTesterMixin, unittest.TestCase): |
|
|
num_data_files = 3 |
|
|
directory_structure = [ |
|
|
"0.jpg", |
|
|
"1.jpg", |
|
|
"2.jpg", |
|
|
"0.txt", |
|
|
"1.txt", |
|
|
"2.txt", |
|
|
] |
|
|
|
|
|
def setUp(self): |
|
|
super().setUp() |
|
|
self.dataset = ImageCaptionFilePairDataset(self.tmpdir.name, infinite=False) |
|
|
|
|
|
def test_getitem(self): |
|
|
iterator = iter(self.dataset) |
|
|
for _ in range(self.num_data_files): |
|
|
item = next(iterator) |
|
|
self.assertEqual(item["caption"], self.caption) |
|
|
self.assertTrue(torch.is_tensor(item["image"])) |
|
|
self.assertEqual(item["image"].shape, (3, 64, 64)) |
|
|
|
|
|
def test_initialize_dataset(self): |
|
|
dataset = initialize_dataset(self.tmpdir.name, "image", infinite=False) |
|
|
self.assertIsInstance(dataset, ImageCaptionFilePairDataset) |
|
|
|
|
|
|
|
|
class ImageFileCaptionFileListDatasetFastTests(ImageDatasetTesterMixin, unittest.TestCase): |
|
|
num_data_files = 3 |
|
|
directory_structure = [ |
|
|
"prompts.txt", |
|
|
"images.txt", |
|
|
"images/", |
|
|
"images/0.jpg", |
|
|
"images/1.jpg", |
|
|
"images/2.jpg", |
|
|
] |
|
|
|
|
|
def setUp(self): |
|
|
super().setUp() |
|
|
self.dataset = ImageFileCaptionFileListDataset(self.tmpdir.name, infinite=False) |
|
|
|
|
|
def test_getitem(self): |
|
|
iterator = iter(self.dataset) |
|
|
for i in range(3): |
|
|
item = next(iterator) |
|
|
self.assertEqual(item["caption"], self.caption) |
|
|
self.assertTrue(torch.is_tensor(item["image"])) |
|
|
self.assertEqual(item["image"].shape, (3, 64, 64)) |
|
|
|
|
|
def test_initialize_dataset(self): |
|
|
dataset = initialize_dataset(self.tmpdir.name, "image", infinite=False) |
|
|
self.assertIsInstance(dataset, ImageFileCaptionFileListDataset) |
|
|
|
|
|
|
|
|
class ImageFolderDatasetFastTests___CSV(ImageDatasetTesterMixin, unittest.TestCase): |
|
|
num_data_files = 3 |
|
|
directory_structure = [ |
|
|
"metadata.csv", |
|
|
"0.jpg", |
|
|
"1.jpg", |
|
|
"2.jpg", |
|
|
] |
|
|
|
|
|
def setUp(self): |
|
|
super().setUp() |
|
|
self.dataset = ImageFolderDataset(self.tmpdir.name, infinite=False) |
|
|
|
|
|
def test_getitem(self): |
|
|
iterator = iter(self.dataset) |
|
|
for _ in range(3): |
|
|
item = next(iterator) |
|
|
self.assertIn("caption", item) |
|
|
self.assertEqual(item["caption"], self.caption) |
|
|
self.assertTrue(torch.is_tensor(item["image"])) |
|
|
|
|
|
def test_initialize_dataset(self): |
|
|
dataset = initialize_dataset(self.tmpdir.name, "image", infinite=False) |
|
|
self.assertIsInstance(dataset, ImageFolderDataset) |
|
|
|
|
|
|
|
|
class ImageFolderDatasetFastTests___JSONL(ImageDatasetTesterMixin, unittest.TestCase): |
|
|
num_data_files = 3 |
|
|
directory_structure = [ |
|
|
"metadata.jsonl", |
|
|
"0.jpg", |
|
|
"1.jpg", |
|
|
"2.jpg", |
|
|
] |
|
|
|
|
|
def setUp(self): |
|
|
super().setUp() |
|
|
self.dataset = ImageFolderDataset(self.tmpdir.name, infinite=False) |
|
|
|
|
|
def test_getitem(self): |
|
|
iterator = iter(self.dataset) |
|
|
for _ in range(3): |
|
|
item = next(iterator) |
|
|
self.assertIn("caption", item) |
|
|
self.assertEqual(item["caption"], self.caption) |
|
|
self.assertTrue(torch.is_tensor(item["image"])) |
|
|
|
|
|
def test_initialize_dataset(self): |
|
|
dataset = initialize_dataset(self.tmpdir.name, "image", infinite=False) |
|
|
self.assertIsInstance(dataset, ImageFolderDataset) |
|
|
|
|
|
|
|
|
class VideoCaptionFilePairDatasetFastTests(VideoDatasetTesterMixin, unittest.TestCase): |
|
|
num_data_files = 3 |
|
|
directory_structure = [ |
|
|
"0.mp4", |
|
|
"1.mp4", |
|
|
"2.mp4", |
|
|
"0.txt", |
|
|
"1.txt", |
|
|
"2.txt", |
|
|
] |
|
|
|
|
|
def setUp(self): |
|
|
super().setUp() |
|
|
self.dataset = VideoCaptionFilePairDataset(self.tmpdir.name, infinite=False) |
|
|
|
|
|
def test_getitem(self): |
|
|
iterator = iter(self.dataset) |
|
|
for _ in range(self.num_data_files): |
|
|
item = next(iterator) |
|
|
self.assertEqual(item["caption"], self.caption) |
|
|
self.assertTrue(torch.is_tensor(item["video"])) |
|
|
self.assertEqual(len(item["video"]), 4) |
|
|
self.assertEqual(item["video"][0].shape, (3, 64, 64)) |
|
|
|
|
|
def test_initialize_dataset(self): |
|
|
dataset = initialize_dataset(self.tmpdir.name, "video", infinite=False) |
|
|
self.assertIsInstance(dataset, VideoCaptionFilePairDataset) |
|
|
|
|
|
|
|
|
class VideoFileCaptionFileListDatasetFastTests(VideoDatasetTesterMixin, unittest.TestCase): |
|
|
num_data_files = 3 |
|
|
directory_structure = [ |
|
|
"prompts.txt", |
|
|
"videos.txt", |
|
|
"videos/", |
|
|
"videos/0.mp4", |
|
|
"videos/1.mp4", |
|
|
"videos/2.mp4", |
|
|
] |
|
|
|
|
|
def setUp(self): |
|
|
super().setUp() |
|
|
self.dataset = VideoFileCaptionFileListDataset(self.tmpdir.name, infinite=False) |
|
|
|
|
|
def test_getitem(self): |
|
|
iterator = iter(self.dataset) |
|
|
for _ in range(3): |
|
|
item = next(iterator) |
|
|
self.assertEqual(item["caption"], self.caption) |
|
|
self.assertTrue(torch.is_tensor(item["video"])) |
|
|
self.assertEqual(len(item["video"]), 4) |
|
|
self.assertEqual(item["video"][0].shape, (3, 64, 64)) |
|
|
|
|
|
def test_initialize_dataset(self): |
|
|
dataset = initialize_dataset(self.tmpdir.name, "video", infinite=False) |
|
|
self.assertIsInstance(dataset, VideoFileCaptionFileListDataset) |
|
|
|
|
|
|
|
|
class VideoFolderDatasetFastTests___CSV(VideoDatasetTesterMixin, unittest.TestCase): |
|
|
num_data_files = 3 |
|
|
directory_structure = [ |
|
|
"metadata.csv", |
|
|
"0.mp4", |
|
|
"1.mp4", |
|
|
"2.mp4", |
|
|
] |
|
|
|
|
|
def setUp(self): |
|
|
super().setUp() |
|
|
self.dataset = VideoFolderDataset(self.tmpdir.name, infinite=False) |
|
|
|
|
|
def test_getitem(self): |
|
|
iterator = iter(self.dataset) |
|
|
for _ in range(3): |
|
|
item = next(iterator) |
|
|
self.assertIn("caption", item) |
|
|
self.assertEqual(item["caption"], self.caption) |
|
|
self.assertTrue(torch.is_tensor(item["video"])) |
|
|
self.assertEqual(len(item["video"]), 4) |
|
|
self.assertEqual(item["video"][0].shape, (3, 64, 64)) |
|
|
|
|
|
def test_initialize_dataset(self): |
|
|
dataset = initialize_dataset(self.tmpdir.name, "video", infinite=False) |
|
|
self.assertIsInstance(dataset, VideoFolderDataset) |
|
|
|
|
|
|
|
|
class VideoFolderDatasetFastTests___JSONL(VideoDatasetTesterMixin, unittest.TestCase): |
|
|
num_data_files = 3 |
|
|
directory_structure = [ |
|
|
"metadata.jsonl", |
|
|
"0.mp4", |
|
|
"1.mp4", |
|
|
"2.mp4", |
|
|
] |
|
|
|
|
|
def setUp(self): |
|
|
super().setUp() |
|
|
self.dataset = VideoFolderDataset(self.tmpdir.name, infinite=False) |
|
|
|
|
|
def test_getitem(self): |
|
|
iterator = iter(self.dataset) |
|
|
for _ in range(3): |
|
|
item = next(iterator) |
|
|
self.assertIn("caption", item) |
|
|
self.assertEqual(item["caption"], self.caption) |
|
|
self.assertTrue(torch.is_tensor(item["video"])) |
|
|
self.assertEqual(len(item["video"]), 4) |
|
|
self.assertEqual(item["video"][0].shape, (3, 64, 64)) |
|
|
|
|
|
def test_initialize_dataset(self): |
|
|
dataset = initialize_dataset(self.tmpdir.name, "video", infinite=False) |
|
|
self.assertIsInstance(dataset, VideoFolderDataset) |
|
|
|
|
|
|
|
|
class ImageWebDatasetFastTests(unittest.TestCase): |
|
|
|
|
|
pass |
|
|
|
|
|
|
|
|
class VideoWebDatasetFastTests(unittest.TestCase): |
|
|
def setUp(self): |
|
|
self.num_data_files = 15 |
|
|
self.dataset = VideoWebDataset("finetrainers/dummy-squish-wds", infinite=False) |
|
|
|
|
|
def test_getitem(self): |
|
|
for index, item in enumerate(self.dataset): |
|
|
if index > 2: |
|
|
break |
|
|
self.assertIn("caption", item) |
|
|
self.assertIn("video", item) |
|
|
self.assertTrue(torch.is_tensor(item["video"])) |
|
|
self.assertEqual(len(item["video"]), 121) |
|
|
self.assertEqual(item["video"][0].shape, (3, 720, 1280)) |
|
|
|
|
|
def test_initialize_dataset(self): |
|
|
dataset = initialize_dataset("finetrainers/dummy-squish-wds", "video", infinite=False) |
|
|
self.assertIsInstance(dataset, VideoWebDataset) |
|
|
|
|
|
|
|
|
class DatasetUtilsFastTests(unittest.TestCase): |
|
|
def test_find_files_depth_0(self): |
|
|
with tempfile.TemporaryDirectory() as tmpdir: |
|
|
file1 = tempfile.NamedTemporaryFile(dir=tmpdir, suffix=".txt", delete=False) |
|
|
file2 = tempfile.NamedTemporaryFile(dir=tmpdir, suffix=".txt", delete=False) |
|
|
file3 = tempfile.NamedTemporaryFile(dir=tmpdir, suffix=".txt", delete=False) |
|
|
|
|
|
files = find_files(tmpdir, "*.txt") |
|
|
self.assertEqual(len(files), 3) |
|
|
self.assertIn(file1.name, files) |
|
|
self.assertIn(file2.name, files) |
|
|
self.assertIn(file3.name, files) |
|
|
|
|
|
def test_find_files_depth_n(self): |
|
|
with tempfile.TemporaryDirectory() as tmpdir: |
|
|
dir1 = tempfile.TemporaryDirectory(dir=tmpdir) |
|
|
dir2 = tempfile.TemporaryDirectory(dir=dir1.name) |
|
|
file1 = tempfile.NamedTemporaryFile(dir=dir1.name, suffix=".txt", delete=False) |
|
|
file2 = tempfile.NamedTemporaryFile(dir=dir2.name, suffix=".txt", delete=False) |
|
|
|
|
|
files = find_files(tmpdir, "*.txt", depth=1) |
|
|
self.assertEqual(len(files), 1) |
|
|
self.assertIn(file1.name, files) |
|
|
self.assertNotIn(file2.name, files) |
|
|
|
|
|
files = find_files(tmpdir, "*.txt", depth=2) |
|
|
self.assertEqual(len(files), 2) |
|
|
self.assertIn(file1.name, files) |
|
|
self.assertIn(file2.name, files) |
|
|
self.assertNotIn(dir1.name, files) |
|
|
self.assertNotIn(dir2.name, files) |
|
|
|
|
|
|
|
|
class ValidationDatasetFastTests(unittest.TestCase): |
|
|
def setUp(self): |
|
|
num_data_files = 3 |
|
|
|
|
|
self.tmpdir = tempfile.TemporaryDirectory() |
|
|
metadata_filename = pathlib.Path(self.tmpdir.name) / "metadata.csv" |
|
|
|
|
|
with open(metadata_filename, "w") as f: |
|
|
f.write("caption,image_path,video_path\n") |
|
|
for i in range(num_data_files): |
|
|
Image.new("RGB", (64, 64)).save((pathlib.Path(self.tmpdir.name) / f"{i}.jpg").as_posix()) |
|
|
f.write(f"test caption,{self.tmpdir.name}/{i}.jpg,\n") |
|
|
|
|
|
self.dataset = ValidationDataset(metadata_filename.as_posix()) |
|
|
|
|
|
def tearDown(self): |
|
|
self.tmpdir.cleanup() |
|
|
|
|
|
def test_getitem(self): |
|
|
for i, data in enumerate(self.dataset): |
|
|
self.assertEqual(data["image_path"], f"{self.tmpdir.name}/{i}.jpg") |
|
|
self.assertIsInstance(data["image"], Image.Image) |
|
|
self.assertEqual(data["image"].size, (64, 64)) |
|
|
|
