def test_cnn_preprocessor_augment_off(dataset_df):
"""should return same image each time"""
dataset = CnnPreprocessor(dataset_df)
dataset.augmentation_off()
sample1 = dataset[0]["X"].numpy()
sample2 = dataset[0]["X"].numpy()
assert np.array_equal(sample1, sample2)
def test_cnn_preprocessor(dataset_df):
"""should return tensor and labels"""
dataset = CnnPreprocessor(dataset_df)
dataset.augmentation_off()
sample1 = dataset[0]["X"]
assert sample1.numpy().shape == (3, 224, 224)
assert dataset[0]["y"].numpy().shape == (2,)