前言

对图像处理的一些操作还是需要掌握的嘿嘿

随机mask生成

这里代码参考MAT的一份代码: https://github.com/fenglinglwb/MAT/blob/main/datasets/mask_generator_256.py

import numpy as np
from PIL import Image, ImageDraw
import math
import random
import os


def RandomBrush(
    max_tries,
    s,
    min_num_vertex = 4,
    max_num_vertex = 18,
    mean_angle = 2*math.pi / 5,
    angle_range = 2*math.pi / 15,
    min_width = 12,
    max_width = 48):
    H, W = s, s
    average_radius = math.sqrt(H*H+W*W) / 8
    mask = Image.new('L', (W, H), 0)
    for _ in range(np.random.randint(max_tries)):
        num_vertex = np.random.randint(min_num_vertex, max_num_vertex)
        angle_min = mean_angle - np.random.uniform(0, angle_range)
        angle_max = mean_angle + np.random.uniform(0, angle_range)
        angles = []
        vertex = []
        for i in range(num_vertex):
            if i % 2 == 0:
                angles.append(2*math.pi - np.random.uniform(angle_min, angle_max))
            else:
                angles.append(np.random.uniform(angle_min, angle_max))

        h, w = mask.size
        vertex.append((int(np.random.randint(0, w)), int(np.random.randint(0, h))))
        for i in range(num_vertex):
            r = np.clip(
                np.random.normal(loc=average_radius, scale=average_radius//2),
                0, 2*average_radius)
            new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w)
            new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h)
            vertex.append((int(new_x), int(new_y)))

        draw = ImageDraw.Draw(mask)
        width = int(np.random.uniform(min_width, max_width))
        draw.line(vertex, fill=1, width=width)
        for v in vertex:
            draw.ellipse((v[0] - width//2,
                          v[1] - width//2,
                          v[0] + width//2,
                          v[1] + width//2),
                         fill=1)
        if np.random.random() > 0.5:
            mask.transpose(Image.FLIP_LEFT_RIGHT)
        if np.random.random() > 0.5:
            mask.transpose(Image.FLIP_TOP_BOTTOM)
    mask = np.asarray(mask, np.uint8)
    if np.random.random() > 0.5:
        mask = np.flip(mask, 0)
    if np.random.random() > 0.5:
        mask = np.flip(mask, 1)
    return mask


def RandomMask(s, hole_range=[0,1]):
    coef = min(hole_range[0] + hole_range[1], 1.0)
    while True:
        mask = np.ones((s, s), np.uint8)
        def Fill(max_size):
            w, h = np.random.randint(max_size), np.random.randint(max_size)
            ww, hh = w // 2, h // 2
            x, y = np.random.randint(-ww, s - w + ww), np.random.randint(-hh, s - h + hh)
            mask[max(y, 0): min(y + h, s), max(x, 0): min(x + w, s)] = 0
        def MultiFill(max_tries, max_size):
            for _ in range(np.random.randint(max_tries)):
                Fill(max_size)
        MultiFill(int(4 * coef), s // 2)
        MultiFill(int(2 * coef), s)
        mask = np.logical_and(mask, 1 - RandomBrush(int(8 * coef), s))  # hole denoted as 0, reserved as 1
        hole_ratio = 1 - np.mean(mask)
        if hole_range is not None and (hole_ratio <= hole_range[0] or hole_ratio >= hole_range[1]):
            continue
        return mask[np.newaxis, ...].astype(np.float32)

封装一下生成mask函数

def dir_not_exists_then_create(dir_name):
    if not os.path.exists(dir_name):
        os.makedirs(dir_name)


def generate_mask(output_mask_path, resolution=256, hole_range=[0.2, 0.4]):
    dir_name = '/'.join(output_mask_path.split('/')[:-1])
    # print(dir_name)
    dir_not_exists_then_create(dir_name)
    
    # mask = RandomBrush(s=256, max_tries=5)
    mask = RandomMask(s=resolution, hole_range=hole_range)
    mask = mask * 255
    # print(mask)
    mask = mask.repeat(3, axis=0)
    mask = mask.transpose(1, 2, 0)
    print(mask.shape)
    img = Image.fromarray(np.uint8(mask))
    img.save(output_mask_path)

调用一下

generate_mask(output_mask_path='mask/mask_{}.png'.format(1), hole_range=[0.1, 0.4])

两张图片合成

这里以celeba-256为例子

def two_image_join_by_mask(img1_path, img2_path, mask_path, output_image_path):
    dir_name = '/'.join(output_image_path.split('/')[:-1])
    # print(dir_name)
    dir_not_exists_then_create(dir_name)

    img1 = np.array(Image.open(img1_path))
    img2 = np.array(Image.open(img2_path))
    mask = np.array(Image.open(mask_path))
    mask = mask / 255.0

    img1[:, :, :] = img1[:, :, :] * mask
    img2[:, :, :] = img2[:, :, :] * (1 - mask)
    img1 = img1 + img2
    img1 = Image.fromarray(np.uint8(img1))
    img1.save(output_image_path)

调用一下

img1_path = 'pic/000255.jpg'
img2_path = 'pic/001133.jpg'
mask_path = 'mask/mask_1.png'
output_image_path = 'pic/000255_001133.jpg'
two_image_join_by_mask(img1_path, img2_path, mask_path, output_image_path)

000255.jpg

001133.jpg

000255_001133.jpg