最近做项目遇见数据集背景非常单一,为了增加模型的返回能里,只能自己做一些数据增强来增加背景的多样性。代码如下:
import numpy as np
import cv2def 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 = bmask[y_min:y_max, x_min:x_max] = 255return maskdef 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将目标从其中抠出来,除了目标区域其余都是0obj_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 resultif __name__ == '__main__':# 示例用法image1 = cv2.imread('image1.jpg') # 原图image2 = cv2.imread('image2.jpg') # 背景图# 保证两张图shape一致,这里只是粗暴的直接resize成了一样尺寸,# TODO:后续可以实现使用等比例缩放,多余区域填灰条的方式resizeimage2 = 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