1from fastai.callback.core import Callback 

2 

3 

4# Can this be a transform? 

5class DihedralCallback(Callback): 

6 def before_batch(self): 

7 """ 

8 x: (batch_size, c, d, h, w) 

9 """ 

10 xb = self.xb[0] 

11 yb = self.yb[0] 

12 

13 k = random.randint(0,7) 

14 

15 if k in [1,3,4,7]: 

16 xb = xb.flip(-1) 

17 yb = yb.flip(-1) 

18 

19 if k in [2,4,5,7]: 

20 xb = xb.flip(-2) 

21 yb = yb.flip(-2) 

22 

23 if k in [3,5,6,7]: 

24 xb = xb.transpose(-1,-2) 

25 yb = yb.transpose(-1,-2) 

26 

27 self.learn.xb = (xb,) 

28 self.learn.yb = (yb,)