Python基于域相关实现图像增强的方法教程

from albumentations.core.transforms_interface import ImageOnlyTransform
    
class InsectAugmentation(ImageOnlyTransform):
    """
    将昆虫的图像强加到目标图像上
    -----------------------------------------------
    参数:
        insects (int): 昆虫的最大数量
        insects_folder (str): 昆虫图片文件夹的路径
    """
 
    def __init__(self, insects=2, dark_insect=False, always_apply=False, p=0.5):
        super().__init__(always_apply, p)
        self.insects = insects
        self.dark_insect = dark_insect
        self.insects_folder = "/kaggle/input/bee-augmentation/"
 
    def apply(self, image, **kwargs):
        """
        参数:
            image (PIL Image): 画昆虫的图像。
        Returns:
            PIL Image: 带昆虫的图像。
        """
        n_insects = random.randint(1, self.insects) # 在这个例子中，我用1而不是0来说明增强效果
        
        if not n_insects:
            return image
        
        height, width, _ = image.shape  # 目标图像的宽度和高度
        insects_images = [im for im in os.listdir(self.insects_folder) if 'png' in im]
        
        for _ in range(n_insects):
            insect = cv2.cvtColor(cv2.imread(os.path.join(self.insects_folder, random.choice(insects_images))), cv2.COLOR_BGR2RGB)
            insect = cv2.flip(insect, random.choice([-1, 0, 1]))
            insect = cv2.rotate(insect, random.choice([0, 1, 2]))
 
            h_height, h_width, _ = insect.shape  # 昆虫图像的宽度和高度
            roi_ho = random.randint(0, image.shape[0] - insect.shape[0])
            roi_wo = random.randint(0, image.shape[1] - insect.shape[1])
            roi = image[roi_ho:roi_ho + h_height, roi_wo:roi_wo + h_width]
 
            # 创建掩码和反掩码
            img2gray = cv2.cvtColor(insect, cv2.COLOR_BGR2GRAY)
            ret, mask = cv2.threshold(img2gray, 10, 255, cv2.THRESH_BINARY)
            mask_inv = cv2.bitwise_not(mask)
 
            # 现在黑掉的区域是昆虫
            img_bg = cv2.bitwise_and(roi, roi, mask=mask_inv)
 
            # 从昆虫图像中只选取昆虫区域。
            if self.dark_insect:
                img_bg = cv2.bitwise_and(roi, roi, mask=mask_inv)
                insect_fg = cv2.bitwise_and(img_bg, img_bg, mask=mask)
            else:
                insect_fg = cv2.bitwise_and(insect, insect, mask=mask)
 
            # 添加
            dst = cv2.add(img_bg, insect_fg, dtype=cv2.CV_64F)
 
            image[roi_ho:roi_ho + h_height, roi_wo:roi_wo + h_width] = dst
                
        return image

def NeedleAugmentation(image, n_needles=2, dark_needles=False, p=0.5, needle_folder='../input/xray-needle-augmentation'):
    aug_prob = random.random()
    if aug_prob < p:
        height, width, _ = image.shape  # 目标图像的宽度和高度
        needle_images = [im for im in os.listdir(needle_folder) if 'png' in im]
 
        for _ in range(1, n_needles):
            needle = cv2.cvtColor(cv2.imread(os.path.join(needle_folder, random.choice(needle_images))), cv2.COLOR_BGR2RGB)
            needle = cv2.flip(needle, random.choice([-1, 0, 1]))
            needle = cv2.rotate(needle, random.choice([0, 1, 2]))
 
            h_height, h_width, _ = needle.shape  # 针图像的宽度和高度
            roi_ho = random.randint(0, abs(image.shape[0] - needle.shape[0]))
            roi_wo = random.randint(0, abs(image.shape[1] - needle.shape[1]))
            roi = image[roi_ho:roi_ho + h_height, roi_wo:roi_wo + h_width]
 
            # 创建掩码和反掩码
            img2gray = cv2.cvtColor(needle, cv2.COLOR_BGR2GRAY)
            ret, mask = cv2.threshold(img2gray, 10, 255, cv2.THRESH_BINARY)
            mask_inv = cv2.bitwise_not(mask)
 
            # 现在黑掉的区域是针
            img_bg = cv2.bitwise_and(roi, roi, mask=mask_inv)
 
            # 只选取针区域。
            if dark_needles:
                img_bg = cv2.bitwise_and(roi, roi, mask=mask_inv)
                needle_fg = cv2.bitwise_and(img_bg, img_bg, mask=mask)
            else:
                needle_fg = cv2.bitwise_and(needle, needle, mask=mask)
 
            # 添加
            dst = cv2.add(img_bg, needle_fg, dtype=cv2.CV_64F)
 
            image[roi_ho:roi_ho + h_height, roi_wo:roi_wo + h_width] = dst
 
    return image