深度学习--离线数据增强

最近做项目遇见数据集背景非常单一，为了增加模型的返回能里，只能自己做一些数据增强来增加背景的多样性。代码如下：

import numpy as np
import cv2


def create_mask(box, height, width):
    """
    创建一个全零的掩码图像，目标区域是255(白色)，北京是0(黑色)
    :param box: 坐标框
    :param height: 图片高
    :param width: 图片宽
    :return:
    """
    mask = np.zeros((height, width), dtype=np.uint8)

    # 在掩码图像上绘制矩形
    for b in box:
        x_min, y_min, x_max, y_max = b
        mask[y_min:y_max, x_min:x_max] = 255

    return mask


def blend_images_with_mask(image1, image2, mask, alpha=0.3, beta=0.7):
    """

    :param image1:
    :param image2:
    :param mask:
    :param alpha:
    :param beta:
    :return:
    """
    # 根据mask将目标从其中抠出来，除了目标区域其余都是0
    obj_masked = cv2.bitwise_or(image1, image1, mask=mask)
    # 根据mask将image1和image2中目标位置删除置为0(目标区域是黑色)
    image1_masked = cv2.bitwise_and(image1, image1, mask=cv2.bitwise_not(mask))
    image2_masked = cv2.bitwise_and(image2, image2, mask=cv2.bitwise_not(mask))

    # image1和image2融合
    blended = cv2.addWeighted(image1_masked, alpha, image2_masked, beta, 0)

    # 将目标放回融合后图像对应位置
    result = cv2.bitwise_or(blended, obj_masked)

    return result


if __name__ == '__main__':
    # 示例用法
    image1 = cv2.imread('image1.jpg')  # 原图
    image2 = cv2.imread('image2.jpg')  # 背景图
    # 保证两张图shape一致，这里只是粗暴的直接resize成了一样尺寸，
    # TODO:后续可以实现使用等比例缩放，多余区域填灰条的方式resize
    image2 = cv2.resize(image2, image1.shape[:2][::-1], )

    # 假设box是一个形状为 (N, 4) 的 numpy 数组，每行代表一个目标的坐标 [x_min, y_min, x_max, y_max]
    # TODO:这里的数据是直接复制过来的，可以改为从文件中读取
    box = np.array([[958, 302, 1046, 416],
                    [871, 316, 975, 464],
                    [1626, 445, 1676, 551],
                    ])  # 两个目标的坐标

    # TODO:添加代码将image1原图整图移动，保证目标不总是出现在同一个位置

    # 向外扩展一些，多保留一些原图背景，不会显得太突兀
    # TODO:后续改成随机向外扩展，实现完整代码，包括外扩后边界检查
    roi_box = box + np.array([-50, -50, 50, 50])

    # 获取图像宽高
    img_h, img_w = image1.shape[:2]
    # 创建mask掩码
    mask = create_mask(roi_box, img_h, img_w)

    # 设置融合比例
    # TODO:设置成随机
    alpha = 0.2  # image1 的比例
    beta = 0.8  # image2 的比例

    # 进行图像融合
    result = blend_images_with_mask(image1, image2, mask, alpha, beta)

    cv2.namedWindow('Blended Image', 0)
    cv2.resizeWindow('Blended Image', 1920 // 2, 1080 // 2)
    # 显示结果
    cv2.imshow('Blended Image', result)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

image1

image2

obj_mask

image1_masked

image2_masked

blended

result