Coverage for src/transforms.py: 55%

1from torch import nn 

2import torch 

3import numpy as np 

4import kornia.augmentation as kaug 

5from kornia import image_to_tensor 

6from random import choice, random 

7 

8 

9class Preprocess(nn.Module): 

10 """Module to perform pre-process using Kornia on torch tensors.""" 

11 

12 @torch.no_grad() # disable gradients for efficiency 

13 def forward(self, x): 

14 x_tmp = np.array(x) 

15 x_out = image_to_tensor(x_tmp, keepdim=True) 

16 return x_out.float() / 255.0 

17 

18 

19class Transforms(nn.Module): 

20 def __init__(self, prob: float = 0.5, test: bool = False, image_size=(224, 224)): 

21 """ 

22 Class handles transformation of images: 

23 Params: 

24 prob (floating) - probability of applying transformation 

25 test (boolean) - specify if the current transformation works on train or test set 

26 """ 

27 super().__init__() 

28 self.prob = prob 

29 self.test = test 

30 

31 self.train_transforms = [ 

32 kaug.RandomHorizontalFlip(p=1.0), 

33 kaug.RandomVerticalFlip(p=1.0), 

34 kaug.RandomRotation(degrees=30, p=1.0), 

35 # kaug.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=1.0), # This tranformation is making some issues while images are not in RGB or grayscale 

36 kaug.RandomAffine(degrees=15, translate=(0.1, 0.1), scale=(0.8, 1.2), shear=10, p=1.0), 

37 kaug.RandomPerspective(distortion_scale=0.2, p=1.0), 

38 kaug.RandomGaussianNoise(mean=0.0, std=0.1, p=1.0), 

39 ] 

40 

41 self.resize = kaug.Resize(size=image_size) 

42 

43 @torch.no_grad() 

44 def forward(self, x): 

45 

46 if self.test or random() < self.prob: 

47 result = self.resize(x) 

48 return result.squeeze() 

49 

50 else: 

51 transform = choice( 

52 self.train_transforms) # In paper method applies only one transformation at once, might change in the future 

53 result = transform(self.resize(x)) 

54 return result.squeeze()