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

淡淡的烟草味﹌ 2022-09-14 14:19 60阅读 0赞

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)