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

r囧r小猫 2022-09-16 13:18 148阅读 0赞

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)