用自己的数据集训练Mask R-CNN

#得到该图中有多少个实例（物体）
def get_obj_index(self, image):
        n = np.max(image)
        return n

#解析labelme中得到的yaml文件，从而得到mask每一层对应的实例标签
def from_yaml_get_class(self,image_id):
        info=self.image_info[image_id]
        with open(info['yaml_path']) as f:
            temp=yaml.load(f.read())
            labels=temp['label_names']
            del labels[0]
        return labels

#重新写draw_mask
def draw_mask(self, num_obj, mask, image):
        info = self.image_info[image_id]
        for index in range(num_obj):
            for i in range(info['width']):
                for j in range(info['height']):
                    at_pixel = image.getpixel((i, j))
                    if at_pixel == index + 1:
                        mask[j, i, index] =1
        return mask

#重新写load_shapes，里面包含自己的自己的类别
#并在self.image_info信息中添加了path、mask_path 、yaml_path
def load_shapes(self, count, height, width, img_folder, mask_folder, imglist,dataset_root_path):
        """Generate the requested number of synthetic images.
        count: number of images to generate.
        height, width: the size of the generated images.
        """
        # Add classes
        self.add_class('shapes', 1, 'your_class_1')
        self.add_class('shapes', 2, 'your_class_2')
        for i in range(count):
            filestr = imglist[i].split(".")[0]
            filestr = filestr.split("_")[1]
            mask_path = 'your mask path'
            yaml_path = 'your yaml path'
            self.add_image("shapes", image_id=i, path=img_folder + "/" + imglist[i],
                           width=width, height=height, mask_path=mask_path,yaml_path=yaml_path)

#重写load_mask
def load_mask(self, image_id):
    """
    Generate instance masks for shapes of the given image ID.
    """
    global iter_num
    info = self.image_info[image_id]
    count = 1  # number of object
    img = Image.open(info['mask_path'])
    num_obj = self.get_obj_index(img)
    mask = np.zeros([info['height'], info['width'], num_obj], dtype=np.uint8)
    mask = self.draw_mask(num_obj, mask, img)
    occlusion = np.logical_not(mask[:, :, -1]).astype(np.uint8)
    for i in range(count - 2, -1, -1):
        mask[:, :, i] = mask[:, :, i] * occlusion
        occlusion = np.logical_and(occlusion, np.logical_not(mask[:, :, i]))
    labels=[]
    labels=self.from_yaml_get_class(image_id)
    labels_form=[]
    for i in range(len(labels)):
        if labels[i] == 'your_class1':
            labels_form.append('your_class1')
        elif labels[i] == 'your_class2':
            labels_form.append('your_class2')
    class_ids = np.array([self.class_names.index(s) for s in labels_form])
    return mask, class_ids.astype(np.int32)

#基础设置
dataset_root_path = 'your dataset root path'
img_folder = dataset_root_path + 'your image folder'
mask_folder = dataset_root_path+ 'your mask folder'
imglist = listdir(img_folder)
count = len(imglist)
width = 640
height = 400

#train与val数据集准备
dataset_train = DrugDataset()
dataset_train.load_shapes(count, width, height, img_folder, mask_folder, imglist,dataset_root_path)
dataset_train.prepare()

dataset_val = DrugDataset()
dataset_val.load_shapes(count, width, height, img_folder, mask_folder, imglist,dataset_root_path)
dataset_val.prepare()