简介

内容

这里粘贴了自己写的一个类，主要参考第22行即可。

 15 class ElasticTransform(object):
 16     def __init__(self, mode='train', probability=0.6): 
 17         self.mode = mode
 18         self.probability = probability  
 19 
 20     def __call__(self, sample):
 21         if self.mode == 'train':        
 22             if round(np.random.uniform(0, 1), 1) <= self.probability:
 23                 image, target = sample['image'], sample['label']
 24                 sigma = 0.1 * np.random.randint(2, 5) # sigma and points are experience parameter via experiment
 25                 points = np.random.randint(3, 8)
 26                 img_aug = deform_grid(image, target, sigma=sigma, points=points)
 27 
 28                 sample['image'] = img_aug[0]    
 29                 sample['label'] = img_aug[1]    
 30                 return sample
 31             else:
 32                 return sample
 33 
 34         if self.mode == 'test' or self.mode == 'infer': 
 35             return sample