基于opencv的数据增强

1,旋转图像并用指定颜色填充边界

    def rotate_image(self, image, fill_color=(0, 0, 255)):
        rows, cols, _ = image.shape
        angle = random.randint(-30, 30)
        M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)

        # 添加填充颜色的边界
        rotated_image = cv2.warpAffine(image, M, (cols, rows), borderMode=cv2.BORDER_CONSTANT, borderValue=fill_color)

        return rotated_image

此处fill_color即未旋转之后背景填充

2,调整图像亮度和对比度

    def adjust_brightness_contrast(self, image):
        alpha = random.uniform(0.5, 1.5)
        beta = random.randint(10, 50)
        return cv2.convertScaleAbs(image, alpha=alpha, beta=beta)

亮度cv2.convertScaleAbs计算公式为 {亮度}=原亮度*alpha+beta

3,颜色畸变

    def color_distortion(self, image, color_matrix):
        return cv2.transform(image, color_matrix)

4,添加椒盐噪声

    def salt_and_pepper_noise(self, image, salt_prob=0.01, pepper_prob=0.01):
        noisy_image = image.copy()
        total_pixels = image.size
        num_salt = int(total_pixels * salt_prob)
        salt_coords = [np.random.randint(0, i-1, num_salt) for i in image.shape]
        noisy_image[salt_coords[0], salt_coords[1]] = 255
        num_pepper = int(total_pixels * pepper_prob)
        pepper_coords = [np.random.randint(0, i-1, num_pepper) for i in image.shape]
        noisy_image[pepper_coords[0], pepper_coords[1]] = 0
        return noisy_image

5,直方图均衡化

    def equalizeHist_image(self, image):
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        return cv2.equalizeHist(image)

6,对比度有限的自适应直方图均衡化

    def clahe_image(self, image):
        b, g, r = cv2.split(image)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(5, 5))
        clahe_b = clahe.apply(b)
        clahe_g = clahe.apply(g)
        clahe_r = clahe.apply(r)
        return cv2.merge((clahe_b, clahe_g, clahe_r))

7,图像细节增强

    def illumination_change(self, image):
        img_zero = np.zeros(image.shape, dtype=np.uint8)
        return cv2.illuminationChange(image, mask=img_zero, alpha=0.2, beta=0.4)

8,图像增强或降噪

    def enhance_reduce(self, image, strength=100):
        x, y, _ = image.shape
        radius = np.random.randint(10, int(min(x, y)), 1)
        pos_x = np.random.randint(0, (min(x, y)-radius), 1)
        pos_y = np.random.randint(0, (min(x, y)-radius), 1)
        pos_x = int(pos_x[0])
        pos_y = int(pos_y[0])
        radius = int(radius[0])
        for j in range(pos_y-radius, pos_y+radius):
            for i in range(pos_x-radius, pos_x+radius):
                distance = math.pow((pos_x-i), 2) + math.pow((pos_y-j), 2)
                distance = np.sqrt(distance)
                if distance < radius:
                    result = 1 - distance/radius
                    result = result*strength
                    if strength > 0:
                        image[i, j, 0] = min((image[i, j, 0]+result), 255)
                        image[i, j, 1] = min((image[i, j, 1]+result), 255)
                        image[i, j, 2] = min((image[i, j, 2]+result), 255)
                    else:
                        image[i, j, 0] = max((image[i, j, 0]+result), 0)
                        image[i, j, 1] = max((image[i, j, 1]+result), 0)
                        image[i, j, 2] = max((image[i, j, 2]+result), 0)
        image = image.astype(np.uint8)
        return image

9,图像遮罩

    def mask(self, image, low=10, high=50):
        x, y, _ = image.shape
        mask_size = np.random.randint(low, high, 1)
        pos_x = np.random.randint(low, (min(x, y)-high), 1)
        pos_y = np.random.randint(low, (min(x, y)-high), 1)
        pos_x = int(pos_x[0])
        pos_y = int(pos_y[0])
        mask_size = int(mask_size[0])
        image[pos_x:pos_x+mask_size, pos_y:pos_y+mask_size] = 0
        return image

依次使用

# 输入和输出文件夹路径
input_folder = r"E:\cow\shuangcheng\ce1"
output_folder = r"E:\cow\shuangcheng\ce1"

# 创建 ImageAugment 类的实例
augmenter = ImageAugment()

# 遍历ce文件夹中的所有编号文件夹
for subdir in os.listdir(input_folder):
    subdir_path = os.path.join(input_folder, subdir)
    if os.path.isdir(subdir_path):
        output_subfolder = os.path.join(output_folder, subdir)
        os.makedirs(output_subfolder, exist_ok=True)

        # 遍历每个编号文件夹中的图像
        for file in os.listdir(subdir_path):
            # 检查文件扩展名是否为PNG
            if file.lower().endswith(('.png', '.PNG')):
                # 构建图像文件的完整路径
                image_path = os.path.join(subdir_path, file)

                # 读取原图
                original_img = cv2.imread(image_path)

                # 保存原图到新的文件夹
                output_original_path = os.path.join(output_subfolder, f"original_{file}")
                cv2.imwrite(output_original_path, original_img)

                # 应用增强方法
                augmented_images = []

                rotated_image = augmenter.rotate_image(original_img)
                augmented_images.append(("rotate", rotated_image))

                brightness_contrast_image = augmenter.adjust_brightness_contrast(original_img)
                augmented_images.append(("brightness_contrast", brightness_contrast_image))

                color_distortion_matrix = np.random.rand(3, 3)
                color_distortion_image = augmenter.color_distortion(original_img, color_distortion_matrix)
                augmented_images.append(("color_distortion", color_distortion_image))

                salt_and_pepper_image = augmenter.salt_and_pepper_noise(original_img)
                augmented_images.append(("salt_and_pepper", salt_and_pepper_image))

                equalized_image = augmenter.equalizeHist_image(original_img)
                augmented_images.append(("equalizeHist", equalized_image))

                clahe_image = augmenter.clahe_image(original_img)
                augmented_images.append(("clahe", clahe_image))

                detail_enhance_image = augmenter.detailEnhance_image(original_img)
                augmented_images.append(("detail_enhance", detail_enhance_image))

                illumination_changed_image = augmenter.illumination_change(original_img)
                augmented_images.append(("illumination_change", illumination_changed_image))

                enhanced_reduced_image = augmenter.enhance_reduce(original_img)
                augmented_images.append(("enhance_reduce", enhanced_reduced_image))

                masked_image = augmenter.mask(original_img)
                augmented_images.append(("mask", masked_image))

                # 保存增强后的图像
                for method, augmented_img in augmented_images:
                    output_path = os.path.join(output_subfolder, f"{file.split('.')[0]}_{method}.png")
                    cv2.imwrite(output_path, augmented_img)