본문 바로가기

Opencv

[Python/OpenCV] make polygone mask for image

반응형

def generate_poly_mask(img_h, img_w, height_min, height_max, width_min, width_max, num_pts):

# check number of points
if (num_pts < 3):
if (np.random.rand(1) < 0.5):
num_pts = 3
else:
num_pts = 4

# generate points
pts = []
for i in range(num_pts):
pts.append([random.randint(width_min, width_max), random.randint(height_min, height_max)])
pts_array = np.array(pts).astype('int')

pt_mean = np.mean(pts_array, axis=0)
pts_left, pts_right = [], []
for i in range(num_pts):
if (pts[i][0] < pt_mean[0]):
pts_left.append(pts[i])
else:
pts_right.append(pts[i])

pts_left_arr = np.array(pts_left)
pts_right_arr = np.array(pts_right)

dec_sort_left = np.argsort(pts_left_arr[:, 1])
dec_sort_right = np.argsort(pts_right_arr[:, 1])

final_pts = []
for i in range(len(dec_sort_left)):
final_pts.append(pts_left[dec_sort_left[i]])

for i in range(len(dec_sort_right)):
final_pts.append(pts_right[dec_sort_right[len(dec_sort_right)-i-1]])

final_pts = np.array(final_pts)

mask = np.zeros(shape=(img_h, img_w, 3))
mask = cv2.fillConvexPoly(mask, final_pts, (255, 255, 255))

mask_copy = np.copy(mask)
mask[mask_copy > 128] = 0
mask[mask_copy <= 128] = 255

kernel = np.ones((7, 7), np.float32) / 49.0
mask = cv2.filter2D(mask, -1, kernel)

return mask[:, :, 0]

def add_poly_shadow(img, height_min, height_max, width_min, width_max, num_pts):

transparency = random.uniform(0.1, 0.5)

img_h, img_w, ch = img.shape
mask = generate_poly_mask(img_h, img_w, height_min, height_max, width_min, width_max, num_pts)

yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)

yuv[:, :, 0] = yuv[:, :, 0] * transparency + (mask.astype('float') * yuv[:, :, 0] / 255.0) * (1 - transparency)
frame = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
frame[frame > 255] = 255

return np.asarray(frame, dtype=np.uint8)