实例分割数据集(mask.png)转目标检测（yolo格式 .txt）+实例分割预处理

最近在做一个线粒体的实例分割数据集转目标检测数据集，记录一下

原图：

原图mask

 

尺寸4096*4096， 需要经过预处理裁成512*512，同时mask是灰度图，要先转成黑白的，顺序是先二值化rbg to mask，再统一裁成512格式，最后再转换成对应的yolo格式

二值化，师弟写的，放上来记录一下

   
import cv2
import numpy as np
import os, glob
 
 
def rgb2masks(label_name):
    lbl_id = os.path.split(label_name)[-1].split('.')[0]
    lbl = cv2.imread(label_name, 1)
    h, w = lbl.shape[:2]
    idx = 0
    num_instance = lbl[:,:,2].max()

    mask_name_list = ['Podo','ECs','GBM']

    mask = (lbl[:,:,0]!=0).astype(np.uint8)
    mask[mask==1] =255
    mask_name = save_path + lbl_id.replace('seg', 'im') + '_Mito' +'_1' + '.png'
    cv2.imwrite(mask_name, mask)
    
    # for i in range(num_instance):
    #     mask = (lbl[:,:,0]!=0).astype(np.uint8)
    #     # mask[mask==1] =255

    #     mask_name = save_path + lbl_id + '_Mito' +f'_{(i+1)}' + '.png'
    #     cv2.imwrite(mask_name, mask)

 

label_dir = './mitoem/MitoEM-H/mask_val'

save_path = './mitoem/MitoEM-H/mask_label_val/'
os.makedirs(save_path,exist_ok=True)

label_list = glob.glob(os.path.join(label_dir, '*.tif'))
for label_name in label_list:
    rgb2masks(label_name)

裁剪代码

import shutil
import os
import ast
import cv2
from PIL import Image

def crop_to_512(img_dir,save_dir,crop_list,crop_size):
    imgs = os.listdir(img_dir)  # 对应图像对应切割
    num=0
    for img in imgs:
        img_path = os.path.join(img_dir, img)
        image = cv2.imread(img_path)  # crop img
        len_x = len(crop_list)
        len_y = len(crop_list)
        for i in range(len_x):
            for j in range(len_y):
                crop_img = image[
                           crop_list[j]:crop_list[j] + crop_size[1],
                           crop_list[i]:crop_list[i] + crop_size[0],
                           ]
                new_img_name = img[:-4]
                new_img_name += '_0' + str(i) + '_0' + str(j) + '.png'
                crop_img = Image.fromarray(crop_img)
                crop_img.save(os.path.join(save_dir, new_img_name))
                num += 1
        print('finsih crop:',img)



if __name__ == '__main__':
    img_dir='./mitoem/MitoEM-H/img_train'  #png，4096*4096 to 512*512
    save_dir='./mitoem/crop_img_train'
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    img_size=[4096,4096]
    crop_size=[512,512]
    # n=0
    # list=[] #x_list
    # for i in range(0,9):
    #
    #     list.append(n)
    #     n=n+512
    #
    # print(list)
    crop_list=[0, 512, 1024, 1536, 2048, 2560, 3072, 3584]
    crop_to_512(img_dir,save_dir,crop_list,crop_size)

mask to yolo

import os
import numpy as np
from itertools import groupby
from skimage import morphology, measure
from PIL import Image
from scipy import misc
import cv2

# 因为一张图片里只有一种类别的目标，所以label图标记只有黑白两色
rgbmask = np.array([[0, 0, 0], [255, 255, 255]], dtype=np.uint8)


# 从label图得到 boundingbox 和图上连通域数量 object_num
def getboundingbox(image):
    # mask.shape = [image.shape[0], image.shape[1], classnum]
    mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
    mask[np.where(np.all(image == rgbmask[1], axis=-1))[:2]] = 1
    # 删掉小于10像素的目标
    mask_without_small = morphology.remove_small_objects(mask, min_size=10, connectivity=2)
    # 连通域标记
    label_image = measure.label(mask_without_small)
    # 统计object个数
    object_num = len(measure.regionprops(label_image))
    boundingbox = list()
    bboxs=list()
    for region in measure.regionprops(label_image):  # 循环得到每一个连通域bbox
        b=[region.bbox[1],region.bbox[3],region.bbox[0],region.bbox[2]]
        bbox=convert([4096,4096],b)
        bboxs.append(bbox)
        boundingbox.append(region.bbox)#
    return object_num, boundingbox,bboxs

def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

def mask_to_yolo(mask_dir,save_dir):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    imgs=os.listdir(mask_dir)
    for img in imgs:
        img_path=os.path.join(mask_dir,img)
        img0 = cv2.imread(img_path)
        out_file = open(save_dir+'/'+img.split('.')[0]+'.txt', 'w')
        object_num, boundingbox, bboxs = getboundingbox(img0)
        for i in bboxs:
            out_file.write('0' + " " + " ".join([str(a) for a in i]) + '\n')



if __name__ == '__main__':
    mask_to_yolo()
    # imgdir='./mitoem/MitoEM-H/mask_label_train/im0000_Mito_1.png'
    # img=cv2.imread(imgdir)
    # out_file = open('./mitoem/yolo/im0000_Mito_1.txt', 'w')
    # object_num, boundingbox,bboxs=getboundingbox(img)
    # 
    # print(object_num,boundingbox)
    # for i in bboxs:
    #     out_file.write('0' + " " + " ".join([str(a) for a in i]) + '\n')