| 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)) |
|
|
