一段小代码,供参考,在CV2中做一些边缘渐变的功能。
主要流程:
- 生成一个随机的mask,并归一化,使其作为另一张图片的权重。
- 对mask做高斯模糊,实现边缘渐变。
- 将加权mask与原始图片相乘。
import os
import cv2
import glob
import random
import numpy as np
class FakeMask:
def __init__(self, prop=0.99):
self.prop = prop # 做数据增强的概率
# 获取一个随机的mask,归一化后作为权重。
def _get_random_mask(self, angle=np.random.randint(45, 80), center=None, scale=1.3):
height = 100
width = 100
image = np.ones((height, width, 3), dtype=np.float32) * 255
(h, w) = image.shape[:2]
if center is None:
center = (w // 2, h // 2)
M = cv2.getRotationMatrix2D(center, angle, scale)
h = int(random.randint(12, 18) * h * 0.1)
random_mask = cv2.warpAffine(image, M, (h, int(w * 1)))
# random horizontal flip
if random.random() > 0.5:
random_mask = cv2.flip(random_mask, 1)
random_mask = random_mask / 255
return random_mask
# 对mask进行高斯模糊,使其边缘渐变
def _encode_mask(self, mask, imgw, imgh):
kh = imgh // 8 # 高斯核越小,mask边缘越明显
kw = imgw // 8
kh = kh + 1 if kh % 2 == 0 else kh
kw = kw + 1 if kw % 2 == 0 else kw
mask = cv2.GaussianBlur(mask, (kw, kh), 0) # 对mask进行高斯模糊,
mask = cv2.GaussianBlur(mask, (kw, kh), 0)
mask = cv2.resize(mask, (imgw, imgh))
return mask
# 数据增强
def transform(self, img):
if random.random() < self.prop:
imgh, imgw, _ = img.shape
random_mask = self._get_random_mask()
random_mask = self._encode_mask(random_mask, imgw, imgh)
for _ in range(2):
img = img * random_mask
return img
if __name__ == "__main__":
imgs_dir = "/data1/gy77/code/demo/aug/imgs"
save_dir = "/data1/gy77/code/demo/aug/imgs_aug"
fake_mask = FakeMask()
imgs_path = glob.glob(imgs_dir + "/*.jpg")
for img_path in imgs_path:
img_basename = os.path.basename(img_path)
img_savepath = os.path.join(save_dir, img_basename)
imgcv = cv2.imread(img_path)
imgcv = fake_mask.transform(imgcv)
cv2.imwrite(img_savepath, imgcv)
当然了,你可以微调上面参数,实现更丰富的效果。例如: