实例分割中物品遮挡数据增强实现

import os
import cv2
import numpy as np


def proces(img,mask):
    maskcopy = mask.copy()
    indexs = np.array(np.where(mask > 0)[:-1])
    x1 = np.min(indexs[1])
    y1 = np.min(indexs[0])
    x2 = np.max(indexs[1])
    y2 = np.max(indexs[0])
    w = x2 - x1
    h = y2 - y1
    if w > h:
        randw = np.random.randint(w // 15,w//3)
        startx = np.random.randint(x1 + w // 15,x2-w // 15-randw)
        mask[y1:y2,startx:startx + randw] = 0
    else:
        randh = np.random.randint(h // 15,h // 3)
        starth = np.random.randint(y1 + h // 15,y2 - h // 15 - randh)
        mask[starth:starth + randh,x1:x2] = 0
    dif = maskcopy - mask
    index = np.where(dif > 0)
    img[index] = 128


for dir in os.listdir('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/masks'):
    img = cv2.imread('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/jpg/{}.jpg'.format(dir))
    if not os.path.exists('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/newmasks/{}'.format(dir)):
        os.mkdir('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/newmasks/{}'.format(dir))
    for file in os.listdir('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/masks/{}'.format(dir)):
        mask = cv2.imread('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/masks/{}/{}'.format(dir, file))
        if 'None' in file:
            proces(img,mask)
        mmask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
        ret, mmask = cv2.threshold(mmask, 127, 255, cv2.THRESH_BINARY)
        binary = np.zeros_like(mask)
        contours, hierarchy = cv2.findContours(mmask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        contours = np.array(contours)
        if not os.path.exists('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/newtxts/{}'.format(dir)):
            os.mkdir('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/newtxts/{}'.format(dir))
        cv2.imwrite('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/newmasks/{}/{}'.format(dir,file),mask)
        np.save('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/newtxts/{}/{}.npy'.format(dir,file),contours)
    cv2.imwrite('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/newjpg/{}.jpg'.format(dir),img)