实现物体相互遮挡下实例分割数据增强及标签制作

import os
import random
from PIL import Image
import cv2
import numpy as np

def find_files(folder, ext):
    g = os.walk(folder)
    L = []
    for path, dir_list, file_list in g:
        for file_name in file_list:
            if os.path.splitext(file_name)[1] == ext:
                L.append(file_name)
    return L

def processlip(img):
    mask_mask = np.ones_like(img) * 255
    files = find_files('/home/lixuan/snap/tiao','.jpg')
    file = files[random.randint(0,len(files) - 1)]
    backgrountd = cv2.imread('/home/lixuan/snap/tiao/{}'.format(file))
    h_img,w_img,_ = img.shape
    h_back,w_back,_ = backgrountd.shape
    if w_back > h_back:
        rate = w_img / w_back
        new_back_h = int(h_back * rate)
        backgrountd = cv2.resize(backgrountd,(w_img,new_back_h))
        starty = random.randint(0,h_img - new_back_h - 1)
        img[starty:starty + new_back_h,:] = backgrountd
        mask_mask[starty:starty + new_back_h,:] = 0
    else:
        rate = h_img / h_back
        new_back_w = int(w_back * rate)
        backgrountd = cv2.resize(backgrountd, (new_back_w,h_img))
        startx = random.randint(0, w_img - new_back_w - 1)
        img[:,startx:startx + new_back_w] = backgrountd
        mask_mask[:,startx:startx + new_back_w] = 0
    return mask_mask

def proces(img,mask,flages):
    files = find_files('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/patchs', '.png')
    file = files[random.randint(0, len(files) - 1)]
    backgrountd = Image.open('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/patchs/{}'.format(file))
    b_w,b_h = backgrountd.size
    if b_w > b_h:
        backgrountd = backgrountd.transpose(Image.ROTATE_90)
        b_w, b_h = backgrountd.size
    backgrountd = backgrountd.convert("RGBA")
    backgrountdmasks = np.zeros_like(flages)
    havebackmask = False
    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

    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    ret, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)

    if (w * h) / (img.shape[0] * img.shape[1]) > 0.005:
        cx = int(x1 + w / 2)
        cy = int(y1 + h / 2)
        if h > w:
            backgrountd = backgrountd.transpose(Image.ROTATE_90)
            rate = h / b_h
            new_back_w = h
            new_back_h = int(rate * b_w)
            backgrountd = backgrountd.resize((new_back_w, new_back_h), Image.ANTIALIAS)
            if np.sum(flages[(cy - new_back_h // 2):(cy - new_back_h // 2 + new_back_h),abs(cx - new_back_w // 2):(abs(cx - new_back_w // 2) + new_back_w)]) == 0:

                backmask = np.zeros_like(np.array(backgrountd))
                maskindex = np.where(np.array(backgrountd) > 0)
                backmask[maskindex] = 255
                backmask = cv2.cvtColor(backmask, cv2.COLOR_BGR2GRAY)
                ret, backmask = cv2.threshold(backmask, 127, 255, cv2.THRESH_BINARY)

                r, g, b, a = backgrountd.split()
                img = Image.fromarray(img)
                new_h, new_w = flages[(cy - new_back_h // 2):(cy - new_back_h // 2 + new_back_h),abs(cx - new_back_w // 2):(abs(cx - new_back_w // 2) + new_back_w)].shape
                flages[(cy - new_back_h // 2):(cy - new_back_h // 2 + new_back_h),abs(cx - new_back_w // 2):(abs(cx - new_back_w // 2) + new_back_w)] = backmask[:new_h,:new_w]
                backgrountdmasks[(cy - new_back_h // 2):(cy - new_back_h // 2 + new_back_h),abs(cx - new_back_w // 2):(abs(cx - new_back_w // 2) + new_back_w)] = backmask[:new_h,:new_w]
                backgrountdmasks_ = backgrountdmasks.astype(np.int)
                mask_and = backgrountdmasks_ & mask
                indexs_and = np.array(np.where(mask_and > 0))
                mask[indexs_and[0],indexs_and[1]] = 0
                img.paste(backgrountd, (abs(cx - new_back_w // 2),cy - new_back_h // 2), mask=a)
                havebackmask = True
        else:
            rate = w / b_h
            new_back_w = int(rate * b_w)
            new_back_h = w
            backgrountd = backgrountd.resize((new_back_w, new_back_h), Image.ANTIALIAS)
            if np.sum(flages[abs(cy - new_back_h // 2):(abs(cy - new_back_h // 2) + new_back_h),(cx - new_back_w // 2):(cx - new_back_w // 2 + new_back_w)]) == 0:

                backmask = np.zeros_like(np.array(backgrountd))
                maskindex = np.where(np.array(backgrountd) > 0)
                backmask[maskindex] = 255
                backmask = cv2.cvtColor(backmask, cv2.COLOR_BGR2GRAY)
                ret, backmask = cv2.threshold(backmask, 127, 255, cv2.THRESH_BINARY)

                r, g, b, a = backgrountd.split()
                img = Image.fromarray(img)
                new_h,new_w = flages[abs(cy - new_back_h // 2):(abs(cy - new_back_h // 2) + new_back_h),(cx - new_back_w // 2):(cx - new_back_w // 2 + new_back_w)].shape
                flages[abs(cy - new_back_h // 2):(abs(cy - new_back_h // 2) + new_back_h),(cx - new_back_w // 2):(cx - new_back_w // 2 + new_back_w)] = backmask[:new_h,:new_w]
                backgrountdmasks[abs(cy - new_back_h // 2):(abs(cy - new_back_h // 2) + new_back_h),(cx - new_back_w // 2):(cx - new_back_w // 2 + new_back_w)] = backmask[:new_h,:new_w]
                backgrountdmasks_ = backgrountdmasks.astype(np.int)
                mask_and = backgrountdmasks_ & mask
                indexs_and = np.array(np.where(mask_and > 0))
                mask[indexs_and[0],indexs_and[1]] = 0
                img.paste(backgrountd, (cx - new_back_w // 2,abs(cy - new_back_h // 2)), mask=a)
                havebackmask = True
    img = np.array(img)
    return img,mask,backgrountdmasks,havebackmask

for dir in os.listdir('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/masks'):
    try:
        count = 0
        img = cv2.imread('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/jpg/{}.jpg'.format(dir))
        masklist = []
        classlist = []
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        flages = np.zeros((img.shape[0], img.shape[1]))
        if not os.path.exists('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news1/newmasks/new_{}'.format(dir)):
            os.mkdir('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news1/newmasks/new_{}'.format(dir))
        for file in os.listdir('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/masks/{}'.format(dir)):
            try:
                mask = cv2.imread('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news/masks/{}/{}'.format(dir, file))
                if 'None' in file:
                    img, mask, backgrountdmasks, havebackmask = proces(img, mask, flages)
                    masklist.append(mask)
                    if 'not_so_clear' in file:
                        classlist.append('not_so_clear')
                    elif 'clear' in file:
                        classlist.append('clear')
                    else:
                        classlist.append('interference')
                    if havebackmask:
                        masklist.append(backgrountdmasks)
                        classlist.append('clear')
                else:
                    masklist.append(mask)
                    if 'not_so_clear' in file:
                        classlist.append('not_so_clear')
                    elif 'clear' in file:
                        classlist.append('clear')
                    else:
                        classlist.append('interference')
            except Exception as e:
                print(e,dir,file)
        mask_mask = processlip(img)
        mask_mask &= processlip(img)
        mask_mask &= processlip(img)
        mask_mask_copy = mask_mask.copy()
        mask_mask_copy = cv2.cvtColor(mask_mask_copy, cv2.COLOR_BGR2GRAY)
        ret, mask_mask_copy = cv2.threshold(mask_mask_copy, 127, 255, cv2.THRESH_BINARY)
        for index,mask in enumerate(masklist):
            try:
                mask &= mask_mask
            except:
                try:
                    mask &= mask_mask_copy
                except:
                    mask = np.array(mask,dtype='uint8')
                    mask &= mask_mask_copy
            try:
                contours, hierarchy = cv2.findContours(np.array(mask, dtype='uint8'), cv2.RETR_EXTERNAL,
                                                       cv2.CHAIN_APPROX_SIMPLE)
                contours = np.array(contours)
            except:
                mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
                ret, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
                contours, hierarchy = cv2.findContours(np.array(mask, dtype='uint8'), cv2.RETR_EXTERNAL,
                                                       cv2.CHAIN_APPROX_SIMPLE)
                contours = np.array(contours)
            if not os.path.exists(
                    '/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news1/newtxts/new_{}'.format(
                        dir)):
                os.mkdir(
                    '/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news1/newtxts/new_{}'.format(
                        dir))
            cv2.imwrite(
                '/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news1/newmasks/new_{}/{}-{}.jpg'.format(
                    dir, count,classlist[index]),np.array(mask,dtype='uint8'))
            np.save(
                '/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news1/newtxts/new_{}/{}-{}.npy'.format(
                    dir, count,classlist[index]), contours)
            count += 1
        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
        cv2.imwrite('/home/lixuan/workspace/dataset/icecreamtwo/labelme_train/news1/newjpg/new_{}.jpg'.format(dir),img)
    except Exception as e:
        print(e)