torch 数据增强-随机擦除（随机遮挡）-两种方案

Serafina ·

更新时间:2024-09-20

· 903 次阅读

第一种
原文

class RandomErasing(object):
    def __init__(self, p=0.5, sl=0.02, sh=0.4, r1=0.3, r2=3):
        self.p = p
        self.sl = sl
        self.sh = sh
        self.r1 = r1
        self.r2 = r2
    def __call__(self, img):
        if np.random.rand() > self.p:
            return img
        img = np.array(img)
        while True:
            img_h, img_w, img_c = img.shape
            img_area = img_h * img_w
            mask_area = np.random.uniform(self.sl, self.sh) * img_area
            mask_aspect_ratio = np.random.uniform(self.r1, self.r2)
            mask_w = int(np.sqrt(mask_area / mask_aspect_ratio))
            mask_h = int(np.sqrt(mask_area * mask_aspect_ratio))
            mask = np.random.rand(mask_h, mask_w, img_c) * 255
            left = np.random.randint(0, img_w)
            top = np.random.randint(0, img_h)
            right = left + mask_w
            bottom = top + mask_h
            if right <= img_w and bottom <= img_h:
                break
        img[top:bottom, left:right, :] = mask
        return Image.fromarray(img)

效果：
在这里插入图片描述
第二种：

原文

class RandomErasing(object):
def init(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=(0.485, 0.456, 0.406)):
self.probability = probability
self.mean = mean
self.sl = sl
self.sh = sh
self.rl = r1
self.rh = 1./r1

def __call__(self, sample):                                                                    
    image, landmarks_crop = sample['image_crop'], sample['landmarks_crop']                     
    image = image / 0.0078125 + 127.5                                                          
    image = self._random_erase(image)                                                          
    image_crop = image.astype(int)                                                             
    # image_crop = (image_crop - 127.5) * 0.0078125                                            
    return {'image_crop': torch.from_numpy(image_crop),                                        
            'landmarks_crop': landmarks_crop,                                                  
            'img_name': sample['img_name'],                                                    
            'image_ori': sample['image_ori'],                                                  
            'landmarks_ori': sample['landmarks_ori'],                                          
            'landmarks_train': sample['landmarks_train'],                                      
            }                                                                                  
    # return Image.fromarray(img)                                                              
def _random_erase(self,image):                                                                 
    image = np.asarray(image).copy()                                                           
    if np.random.random() > self.probability:                                                  
        return image                                                                           
    h, w = image.shape[:2]                                                                     
    image_area = h * w                                                                         
    for _ in range(20):                                                                        
        mask_area = np.random.uniform(self.sl, self.sh) * image_area                           
        aspect_ratio = np.random.uniform(self.rl, self.rh)                                     
        mask_h = int(np.sqrt(mask_area * aspect_ratio))                                        
        mask_w = int(np.sqrt(mask_area / aspect_ratio))                                        
        if mask_w < w and mask_h < h:                                                          
            x0 = np.random.randint(0, w - mask_w)                                              
            y0 = np.random.randint(0, h - mask_h)                                              
            x1 = x0 + mask_w                                                                   
            y1 = y0 + mask_h                                                                   
            image[y0:y1, x0:x1] = np.random.uniform(0, 1)                                      
            break                                                                              
    return image

效果：
在这里插入图片描述