mask-rcnn的配置及训练自己的数据

最新推荐文章于 2023-07-03 13:56:10 发布

babyzbb636

最新推荐文章于 2023-07-03 13:56:10 发布

阅读量1.5k

点赞数

分类专栏：深度学习文章标签： mask_rcnn 配置训练自己的数据

本文链接：https://blog.csdn.net/babyzbb636/article/details/102631510

版权

深度学习专栏收录该内容

30 篇文章 1 订阅

订阅专栏

参考博客：

https://blog.csdn.net/qq_36810544/article/details/83582397

https://blog.csdn.net/qq_29462849/article/details/81037343

https://blog.csdn.net/weixin_42880443/article/details/93622552

https://blog.csdn.net/qq_36810544/article/details/83582397

https://blog.csdn.net/weixin_42535742/article/details/90026899

发现网上关于语义分割网络配合使用训练写的都不太详细，踩了n多坑，写了超详细。

一、配置与测试

1.下载gitclone

这两天git抽风，下载慢就算了，还老失败

git clone https://github.com/matterport/Mask_RCNN
cd Mask_RCNN
pip install -r requirement.txt
python setup.py

2.安装pycocotools

直接把faster rcnn下配置pycocotools文件夹的拷贝过来，放在mask-rcnn目录下

好像Mask_RCNN2.1要使用预训练权重mask_rcnn_balloon.h5https://github.com/matterport/Mask_RCNN/releases（使用mask_rcnn_coco.h5会报加载错误）

运行demo.py文件，一定要cd Mask_RCNN/samples 否则会报coco的错，我一直是以为没安装好pycocotools，折腾了一上午，到后来用pycharm运行之后发现不报这个错误才知道（以后不知道在那个文件下运行还是现在pycharm下试试）

cd  Mask_RCNN/samples
python demo.py

再报错误的话：

1）.又下载权重文件，这时候把权重文件路径改为绝对路径；使用mask_rcnn_balloon.h5权重文件

2）. 测试文件为balloon数据集下图片

3）cuda out of memory

os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

测试代码：

新建test.py，放在sameples文件下

# -*- coding: utf-8 -*-
import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt
import cv2
import time
from mrcnn.config import Config
from datetime import datetime 
# Root directory of the project
ROOT_DIR = os.getcwd()
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
# Import Mask RCNN
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize
# Import COCO config
sys.path.append(os.path.join(ROOT_DIR, "samples/coco/"))  # To find local version
from samples.coco import coco


# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

# Local path to trained weights file
COCO_MODEL_PATH = '/home/zbb/Mask_RCNN-master/mask_rcnn_balloon.h5' #######
# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)
    print("cuiwei***********************")

# Directory of images to run detection on
IMAGE_DIR = '/home/zbb/Mask_RCNN-master/test'

class ShapesConfig(Config):
    # Give the configuration a recognizable name
    NAME = "shapes"

    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each
    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

    # Number of classes (including background)
    NUM_CLASSES = 1 + 1  # background + 3 shapes

    # Use small images for faster training. Set the limits of the small side
    # the large side, and that determines the image shape.
    IMAGE_MIN_DIM = 320
    IMAGE_MAX_DIM = 384

    # Use smaller anchors because our image and objects are small
    RPN_ANCHOR_SCALES = (8 * 6, 16 * 6, 32 * 6, 64 * 6, 128 * 6)  # anchor side in pixels

    # Reduce training ROIs per image because the images are small and have
    # few objects. Aim to allow ROI sampling to pick 33% positive ROIs.
    TRAIN_ROIS_PER_IMAGE =100

    # Use a small epoch since the data is simple
    STEPS_PER_EPOCH = 100

    # use small validation steps since the epoch is small
    VALIDATION_STEPS = 50

#import train_tongue
#class InferenceConfig(coco.CocoConfig):
class InferenceConfig(ShapesConfig):
    # Set batch size to 1 since we'll be running inference on
    # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

config = InferenceConfig()

model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

# Create model object in inference mode.
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

# Load weights trained on MS-COCO
model.load_weights(COCO_MODEL_PATH, by_name=True)

# COCO Class names
# Index of the class in the list is its ID. For example, to get ID of
# the teddy bear class, use: class_names.index('teddy bear')
class_names = ['BG', 'balloon']  
# Load a random image from the images folder

#################批量化保存测试文件#########################
count = os.listdir(IMAGE_DIR)
for i in range(0,len(count)):
    path = os.path.join(IMAGE_DIR, count[i])
    if os.path.isfile(path):
        file_names = next(os.walk(IMAGE_DIR))[2]
        image = skimage.io.imread(os.path.join(IMAGE_DIR, count[i]))
        # Run detection
        results = model.detect([image], verbose=1)
        r = results[0]
        visualize.display_instances(image, count[i],r['rois'], r['masks'], r['class_ids'],
                            class_names, r['scores'])

同时需要修改

mrcnn/visualize下函数：

在这个函数下添加count参数

添加保存函数

二、训练自己的数据

1.labelme安装与转换

labelme安装就不多说了有很多教程

转化为json明天补上

单一类别：

import os
path = '/home/zbb/huangshang/xiugai/json/'  # path是你存放json的路径
json_file = os.listdir(path)
for file in json_file:
    os.system("python /home/zbb/anaconda3/envs/other/lib/python3.7/site-packages/labelme/cli/json_to_dataset.py %s"%(path + file))

多类别：使用上述每个类别对应颜色不一致

新建文件夹：labelme_dataset

python  json_to_picture.py  pathtojson

import os
import random
import shutil
import re


import argparse
import json
import os
import os.path as osp
import warnings
 
import PIL.Image
import yaml
 
from labelme import utils
import base64
 
def main():
    warnings.warn("This script is aimed to demonstrate how to convert the\n"
                  "JSON file to a single image dataset, and not to handle\n"
                  "multiple JSON files to generate a real-use dataset.")
    parser = argparse.ArgumentParser()
    parser.add_argument('json_file')
    parser.add_argument('-o', '--out', default=None)
    args = parser.parse_args()
 
    json_file = args.json_file
    #下面类别就是自己要分的类别数目，我的共有9类，而且下面这个字典的关键字（单引号内的）是自己在用labelme时标注的类别所以我自己又加了0-9的关键字，你的改成你自己的关键字就可以了，然后冒号后面的我猜代表的是颜色的标号，这样就会出来不同的颜色。
    label_name_to_value = {'_background_': 0,
                           'vegetation':1,
                           'farmland':2,
                           'building':3,
                           'water':4,
                           'road':5,
                         }
    if args.out is None:
        out_dir = osp.basename(json_file).replace('.', '_')
        out_dir = osp.join(osp.dirname(json_file), out_dir)
    else:
        out_dir = args.out
    if not osp.exists(out_dir):
        os.mkdir(out_dir)
 
    count = os.listdir(json_file) 
    for i in range(0, len(count)):
        path = os.path.join(json_file, count[i])
        if os.path.isfile(path):
            data = json.load(open(path))
            
            if data['imageData']:
                imageData = data['imageData']
            else:
                imagePath = os.path.join(os.path.dirname(path), data['imagePath'])
                with open(imagePath, 'rb') as f:
                    imageData = f.read()
                    imageData = base64.b64encode(imageData).decode('utf-8')
            img = utils.img_b64_to_arr(imageData)

            for shape in data['shapes']:
                label_name = shape['label']
                if label_name in label_name_to_value:
                    label_value = label_name_to_value[label_name]
                else:
                    label_value = len(label_name_to_value)
                    label_name_to_value[label_name] = label_value
            
            # label_values must be dense
            label_values, label_names = [], []
            for ln, lv in sorted(label_name_to_value.items(), key=lambda x: x[1]):
                label_values.append(lv)
                label_names.append(ln)
            assert label_values == list(range(len(label_values)))
            
            lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
            
            captions = ['{}: {}'.format(lv, ln)
                for ln, lv in label_name_to_value.items()]
            lbl_viz = utils.draw_label(lbl, img, captions)
            
            out_dir = osp.basename(count[i]).replace('.', '_')
            out_dir = osp.join(osp.dirname(count[i]), out_dir)
            if not osp.exists(out_dir):
                os.mkdir(out_dir)
            name=out_dir.split('_')
            rename=name[0]+'_'+name[1]
            PIL.Image.fromarray(img).save(osp.join(out_dir, rename+'_'+'img.png'))

            #PIL.Image.fromarray(lbl).save(osp.join(out_dir, rename+'label.png'))
            utils.lblsave(osp.join(out_dir,rename+'_'+'label.png'), lbl)
            PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir,rename+'_'+'label_viz.png'))
 
            with open(osp.join(out_dir, 'label_names.txt'), 'w') as f:
                for lbl_name in label_names:
                    f.write(lbl_name + '\n')
 
            warnings.warn('info.yaml is being replaced by label_names.txt')
            info = dict(label_names=label_names)
            with open(osp.join(out_dir, 'info.yaml'), 'w') as f:
                yaml.safe_dump(info, f, default_flow_style=False)
 
            print('Saved to: %s' % out_dir)
if __name__ == '__main__':
    main()

2.数据集的修改

所需三个文件pic，ymal，cv2_mask放置路径如下：

Mask_RCNN/data/pic

Mask_RCNN/data/cv2_mask

Mask_RCNN/data/ymal

(不需要labelme_json文件夹）

针对2中转化为json文件，先提取出pic，ymal，orig_mask(16位）

import os
import shutil
filename = "/home/zbb/huangshang/data/labelme_json"  # 存放json转化得到的文件夹名称
fileList = os.listdir(filename)
#fileList.remove(".DS_Store")

for i in range(len(fileList)):
    img_source = filename +'/' + fileList[i] + "/img.png"
    mask_source = filename+'/' + fileList[i] + "/label.png"
    yaml_source = filename+'/' + fileList[i] + "/info.yaml"

    img_target = "/home/zbb/huangshang/pic/img_{}.png".format(i)
    mask_target = "/home/zbb/huangshang/orig_mask/label_{}.png".format(i)
    yaml_target = "/home/zbb/huangshang/yaml/info_{}.yaml".format(i)

    shutil.copy(img_source, img_target)
    shutil.copy(mask_source, mask_target)
    shutil.copy(yaml_source, yaml_target)

orig_mask转化为cv2_mask（8位）

from PIL import Image
import numpy as np
import shutil
import os

src_dir = "orig_mask"
dest_dir = "cv2_mask"

for label_name in os.listdir(src_dir):
    old_mask = src_dir + "/" + label_name
    img = Image.open(old_mask)
    img = Image.fromarray(np.uint8(np.array(img)))
    new_mask = os.path.join(dest_dir, label_name)
    img.save(new_mask)

3.测试代码修改

1）修改load_shapes函数

添加自己的类别名称
修改路径文件mask，ymal（去掉主文件路径），pic folder，不需要labelme_json文件夹

或者为

 mask_path = "{}/label_{}.png".format(mask_floder, i)
 yaml_path = "{}/info_{}.yaml".format(yaml_floder, i)
 cv_img = cv2.imread("{}/img_{}.png".format(img_floder, i))

2）load_mask函数

3）数据设置修改

分训练集和测试集
对应上面吧data_root_path改为ymal folder

全部代码如下：

新建train.py文件，放在sameples文件下

# -*- coding: utf-8 -*-
import os
import sys
import random
import math
import re
import time
import numpy as np
import cv2
import matplotlib
import matplotlib.pyplot as plt
import tensorflow as tf
from mrcnn.config import Config
#import utils
from mrcnn import model as modellib,utils
from mrcnn import visualize
import yaml
from mrcnn.model import log
from PIL import Image
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 
#os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# Root directory of the project
ROOT_DIR = os.getcwd()
#ROOT_DIR = os.path.abspath("../")
# Directory to save logs and trained model
MODEL_DIR = '/home/zbb/Mask_RCNN-master/logs'
iter_num=0
 
# Local path to trained weights file
COCO_MODEL_PATH = '/home/zbb/Mask_RCNN-master/mask_rcnn_balloon.h5'
# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)
 
 
class ShapesConfig(Config):
    # Give the configuration a recognizable name
    NAME = "shapes"
    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each
    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
    # Number of classes (including background)
    NUM_CLASSES = 1 + 6  # background + 3 shapes
    # Use small images for faster training. Set the limits of the small side
    # the large side, and that determines the image shape.
    IMAGE_MIN_DIM = 320
    IMAGE_MAX_DIM = 384 
    # Use smaller anchors because our image and objects are small
    RPN_ANCHOR_SCALES = (8 * 6, 16 * 6, 32 * 6, 64 * 6, 128 * 6)  # anchor side in pixels
    # Reduce training ROIs per image because the images are small and have
    # few objects. Aim to allow ROI sampling to pick 33% positive ROIs.
    TRAIN_ROIS_PER_IMAGE = 100
    # Use a small epoch since the data is simple
    STEPS_PER_EPOCH = 100
    # use small validation steps since the epoch is small
    VALIDATION_STEPS = 50
 
 
config = ShapesConfig()
config.display()
 
class DrugDataset(utils.Dataset):
    # 得到该图中有多少个实例（物体）
    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,image_id):
        #print("draw_mask-->",image_id)
        #print("self.image_info",self.image_info)
        info = self.image_info[image_id]
        #print("info-->",info)
        #print("info[width]----->",info['width'],"-info[height]--->",info['height'])
        for index in range(num_obj):
            for i in range(info['width']):
                for j in range(info['height']):
                    #print("image_id-->",image_id,"-i--->",i,"-j--->",j)
                    #print("info[width]----->",info['width'],"-info[height]--->",info['height'])
                    at_pixel = image.getpixel((i, j))
                    if at_pixel == index + 1:
                        mask[j, i, index] = 1
        return mask
 
    # 重新写load_shapes，里面包含自己的自己的类别
    def load_shapes(self, count, img_floder, mask_floder, imglist, yaml_floder):
        ###################################### Add classes############################################
        self.add_class("shapes", 1, "building")
        self.add_class("shapes", 2, "water")
        self.add_class("shapes", 3, "road")
        self.add_class("shapes", 4, "vegetation")
        self.add_class("shapes", 5, "farmland")
        self.add_class("shapes", 6, "mountain")
        ################################################################################################
		
        for i in range(count):
            # 获取图片宽和高
 
            filestr = imglist[i].split(".")[0]
            #print(imglist[i],"-->",cv_img.shape[1],"--->",cv_img.shape[0])
            #print("id-->", i, " imglist[", i, "]-->", imglist[i],"filestr-->",filestr)
            # filestr = filestr.split("_")[1]
			
			###############对应pic，ymal，png文件夹######################################
            mask_path = mask_floder + "/" +'label'+ filestr[3:] + ".png"
            yaml_path = yaml_floder +"/" +'info'+filestr[3:] + ".yaml"
            cv_img = cv2.imread(img_floder + "/" +  filestr + ".png")
            ############################################################################
			
            self.add_image("shapes", image_id=i, path=img_floder + "/" + imglist[i],
                           width=cv_img.shape[1], height=cv_img.shape[0], mask_path=mask_path, yaml_path=yaml_path)
 
    # 重写load_mask
    def load_mask(self, 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,image_id)
        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)):
		    #################################add class##############################
            if labels[i].find("building") != -1:
                # print "car"
                labels_form.append("building")
            elif labels[i].find("road") != -1:
                # print "leg"
                labels_form.append("road")
            elif labels[i].find("water") != -1:
                # print "well"
                labels_form.append("water")
            elif labels[i].find("farmland") != -1:
                # print "leg"
                labels_form.append("farmland")
            elif labels[i].find("vegetation") != -1:
                # print "well"
                labels_form.append("vegetation")
            elif labels[i].find("mountain") != -1:
                # print "leg"
                labels_form.append("mountain")
            ########################################################################
        class_ids = np.array([self.class_names.index(s) for s in labels_form])
        return mask, class_ids.astype(np.int32)
 
def get_ax(rows=1, cols=1, size=8):
    _, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
    return ax
 
#基础设置
dataset_root_path_t="/home/zbb/Mask_RCNN-master/data/train/"
t_img_floder = dataset_root_path_t + "pic"
t_mask_floder = dataset_root_path_t + "cv2_mask"
t_yaml_floder = dataset_root_path_t +"yaml"
t_imglist = os.listdir(t_img_floder)
t_count = len(t_imglist)

dataset_root_path_v="/home/zbb/Mask_RCNN-master/data/val/"
v_img_floder = dataset_root_path_v + "pic"
v_mask_floder =dataset_root_path_v +  "cv2_mask"
v_yaml_floder = dataset_root_path_v +"yaml"
v_imglist = os.listdir(v_img_floder)
v_count = len(v_imglist)
 
#train与val数据集准备
dataset_train = DrugDataset()
dataset_train.load_shapes(t_count, t_img_floder, t_mask_floder, t_imglist,t_yaml_floder)
dataset_train.prepare()
#print("dataset_train-->",dataset_train._image_ids)
dataset_val = DrugDataset()
dataset_val.load_shapes(v_count, v_img_floder, v_mask_floder, v_imglist,v_yaml_floder)
dataset_val.prepare()

# Create model in training mode
model = modellib.MaskRCNN(mode="training", config=config,
                          model_dir=MODEL_DIR)
 
# Which weights to start with?
init_with = "coco"  # imagenet, coco, or last
 
if init_with == "imagenet":
    model.load_weights(model.get_imagenet_weights(), by_name=True)
elif init_with == "coco":
    # Load weights trained on MS COCO, but skip layers that
    model.load_weights(COCO_MODEL_PATH, by_name=True,
                       exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",
                                "mrcnn_bbox", "mrcnn_mask"])
elif init_with == "last":
    # Load the last model you trained and continue training
    model.load_weights(model.find_last()[1], by_name=True)
 
# Train the head branches
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE,
            epochs=10,
            layers='heads')
 
 
# Fine tune all layers
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE / 10,
            epochs=30,
            layers="all")

4、测试代码修改

类别数目,名称修改

babyzbb636

关注

0
点赞
踩
8

收藏

觉得还不错? 一键收藏
3
评论
mask-rcnn的配置及训练自己的数据

目录一、配置与测试1.下载gitclone2.安装pycocotools二、训练自己的数据1.labelme安装与转换2.数据集的修改3.测试代码修改1）修改load_shapes函数2）load_mask函数数据设置修改4、测试代码修改参考博客：https://blog.csdn.net/qq_36810544/article...
复制链接

扫一扫