Coverage for robustAI/advertrain/transforms.py: 100%

1from torchvision import transforms 

2from torchvision.transforms import Compose 

3 

4 

5class DataTransformations: 

6 """ 

7 Class to create and return training and test data transformations. 

8 

9 This class encapsulates the creation of data transformations used in training and testing. 

10 It provides methods to get composed series of transformations for both scenarios. 

11 """ 

12 

13 def __init__(self, train_prob: float = 0.5): 

14 """ 

15 Initializes the DataTransformations class. 

16 

17 Args: 

18 train_prob (float): The probability of applying transformations in training. Default is 0.5. 

19 """ 

20 self.train_prob = train_prob 

21 

22 def get_train_transforms(self) -> Compose: 

23 """ 

24 Creates and returns a series of training data transformations. 

25 

26 Returns: 

27 Compose: A composed series of transformations for training data. 

28 """ 

29 return transforms.Compose( 

30 [ 

31 transforms.Pad((0, 120)), 

32 transforms.RandomHorizontalFlip(), 

33 transforms.RandomVerticalFlip(), 

34 transforms.RandomApply( 

35 [ 

36 transforms.RandomAffine(5), 

37 transforms.RandomAffine(0, translate=(0.02, 0.02)), 

38 transforms.ColorJitter( 

39 brightness=0.5, contrast=0.2, saturation=0.2, hue=0 

40 ), 

41 ], 

42 p=self.train_prob, 

43 ), 

44 transforms.Resize(size=(64, 128)), 

45 transforms.ToTensor(), 

46 ] 

47 ) 

48 

49 def get_test_transforms(self) -> Compose: 

50 """ 

51 Creates and returns a series of test data transformations. 

52 

53 Returns: 

54 Compose: A composed series of transformations for test data. 

55 """ 

56 return transforms.Compose( 

57 [ 

58 transforms.Pad((0, 120)), 

59 transforms.Resize((64, 128)), 

60 transforms.ToTensor(), 

61 ] 

62 )