pytorch随机生成图像修复和扩展的mask

图像修复

import random
import cv2
import numpy as np
import time

t = time.time() # 用时间戳来唯一命名
def create_mask(width, height, mask_width, mask_height, x=None, y=None):
    # 黑色是0  白色缺失部分是255   实际使用时，需要将255变为1。 因为需要和原图做加减运算
    mask = np.zeros((height, width)) # 生成一个覆盖全部大小的全黑mask
    mask_x = x if x is not None else random.randint(0, width - mask_width) # 缺失部分左下角的x坐标
    mask_y = y if y is not None else random.randint(0, height - mask_height)  #缺失部分左上角的y坐标
    mask[mask_y:mask_y + mask_height, mask_x:mask_x + mask_width] = 255  # 将中间缺失白色部分标为1
    return mask


img = cv2.imread("./test_images/0.png")
imgh, imgw = img.shape[0:2]  # imgw = 256 imgh = 256
mask = create_mask(imgw, imgh, imgw // 2, imgh // 2)
cv2.imwrite(str(int(t)) + ".png", mask)
print('save mask OK!')

 

             

 

高度为从上到下，所以缺失区域为[mask_y:mask_y + mask_height]

宽度为从左到右，所以确实区域为[mask_x:mask_x + mask_width]

通过python数组中的切片操作，将缺失区域进行赋值，因为这里打印查看效果，所以赋值为255

在实际使用中，需要赋值为1.  因为需要和原图进行加减操作

 

图像扩展

import random
import cv2
import numpy as np
import time

t = time.time() # 用时间戳来唯一命名
def create_mask(width, height, mask_width, mask_height, x=None, y=None):

    mask_new = []

    # 黑色是0  白色缺失部分是255 实际使用时将255改为1
    # 右侧
    mask_right = np.zeros((height, width)) #
    mask_right_x = x if x is not None else random.randint(mask_width*2, width)
    mask_right[:, mask_right_x:width] = 255

    # 左侧
    mask_left = np.zeros((height, width))
    mask_left_x = x if x is not None else random.randint(mask_width*2, width)
    mask_left[:, 0:width-mask_left_x] = 255

    # 下侧
    mask_bottom = np.zeros((height, width))
    mask_bottom_y = y if y is not None else random.randint(mask_height*2, height)
    mask_bottom[mask_bottom_y:height, :] = 255

    # 上侧
    mask_top = np.zeros((height, width))
    mask_top_y = y if y is not None else random.randint(mask_height * 2, height)
    mask_top[0:height-mask_top_y, :] = 255

    mask_new.append(mask_right)
    mask_new.append(mask_bottom)
    mask_new.append(mask_top)
    mask_new.append(mask_left)

    return random.choice(mask_new)


# img = cv2.imread("./test_images/0.png")
# imgh, imgw = img.shape[0:2]
imgw = 256
imgh = 256
mask = create_mask(imgw, imgh, imgw // 3, imgh // 3)
cv2.imwrite(str(int(t)) + ".png", mask)

          

通过random.choice() 随机选择一个方向，进行生成

 

同样的写法，整理了一下

import random
import cv2
import numpy as np
import time

t = time.time()
def create_mask(width, height, mask_width, mask_height, x=None, y=None):

    mask_fx = ['top','right','bottom','left']
    mask = np.zeros((height, width))
    mask_random = random.choice(mask_fx)
    if mask_random == 'right' or mask_random == 'left':
        mask_x = x if x is not None else random.randint(mask_width * 2, width)
        if mask_random == 'right':
            mask[:, mask_x:width] = 255
        else:
            mask[:, 0:width - mask_x] = 255

    if mask_random == 'top' or mask_random == 'bottom':
        mask_y = y if y is not None else random.randint(mask_height*2, height)
        if mask_random == 'top':
            mask[0:height - mask_y, :] = 255
        else:
            mask[mask_y:height, :] = 255

    return mask


# img = cv2.imread("./test_images/0.png")
# imgh, imgw = img.shape[0:2]
imgw = 256
imgh = 256
mask = create_mask(imgw, imgh, imgw // 3, imgh // 3)
cv2.imwrite(str(int(t)) + ".png", mask)