def test_video_dataset(self):
if not VIDEO_DATASET_AVAILABLE:
tmp_dir = tempfile.mkdtemp()
# simulate a video
# the video dataset will check to see whether there exists a file
# with a video extension, it's enough to fake a video file here
path = os.path.join(tmp_dir, 'my_file.png')
dataset = torchvision.datasets.FakeData(size=1,
image_size=(3, 32, 32))
image, _ = dataset[0]
image.save(path)
os.rename(path, os.path.join(tmp_dir, 'my_file.avi'))
with self.assertRaises(ImportError):
dataset = LightlyDataset(from_folder=tmp_dir)
warnings.warn(
'Did not test video dataset because of missing requirements')
shutil.rmtree(tmp_dir)
return
self.create_video_dataset()
dataset = LightlyDataset(from_folder=self.input_dir)
out_dir = tempfile.mkdtemp()
dataset.dump(out_dir)
self.assertEqual(len(os.listdir(out_dir)), len(dataset))
def test_create_lightly_dataset_from_folder(self):
n_subfolders = 5
n_samples_per_subfolder = 10
n_tot_files = n_subfolders * n_samples_per_subfolder
dataset_dir, folder_names, sample_names = self.create_dataset(
n_subfolders, n_samples_per_subfolder)
dataset = LightlyDataset(from_folder=dataset_dir)
filenames = dataset.get_filenames()
fnames = []
for dir_name in folder_names:
for fname in sample_names:
fnames.append(os.path.join(dir_name, fname))
self.assertEqual(len(filenames), n_tot_files)
self.assertEqual(len(dataset), n_tot_files)
self.assertListEqual(sorted(fnames), sorted(filenames))
out_dir = tempfile.mkdtemp()
dataset.dump(out_dir)
self.assertEqual(
sum(
len(os.listdir(os.path.join(out_dir, subdir)))
for subdir in os.listdir(out_dir)),
len(dataset),
)
shutil.rmtree(dataset_dir)
shutil.rmtree(out_dir)
def test_video_dataset_available(self):
self.create_video_dataset()
dataset = LightlyDataset(input_dir=self.input_dir)
out_dir = tempfile.mkdtemp()
dataset.dump(out_dir, dataset.get_filenames()[(len(dataset) // 2):])
self.assertEqual(len(os.listdir(out_dir)), len(dataset) // 2)
for filename in os.listdir(out_dir):
self.assertIn(filename,
dataset.get_filenames()[(len(dataset) // 2):])