mask rcnn训练自己的数据

samsara1224

已于 2023-05-24 09:52:59 修改

阅读量319

点赞数

文章标签：深度学习图像处理

于 2023-02-28 14:33:49 首次发布

本文链接：https://blog.csdn.net/d121028379/article/details/129261153

版权

主流的Mask RCNN主要有两个，一个是matterport的keras版本，另一个是facebookresearch的pytorch版本。

pytorch版本的编译报错问题较多，以下采用tensorflow-gpu版本。

检查gpu算力，低算力下高版本tf可能只调用cpu

参考：tensorflow对应CUDA版本

验证gpu

import tensorflow as tf
tf.test.is_built_with_cuda()#验证tf是否调用GPU
tf.config.list_physical_devices('GPU')#查看gpu列表

下载requirements,tensorflow,cuda,mask_rcnn_coco.h5权重

python setup.py install

划分数据集

data

——cv2_mask（掩膜）

——jason（jason）

——labelme_json(json文件夹)

——pic（原图）

其中labelme新版本在json文件夹生成info.yaml的方法

~\envs\xx\Lib\site-packages\labelme\cli\json_to_dataset.py

            f.write(lbl_name + '\n')
	# 添加部分
    logger.warning('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)

    logger.info('Saved to: {}'.format(out_dir))

训练代码参考：修改train_shapes.ipynb

其中

ROOT_DIR = os.path.abspath("../../")#根目录地址

#显存设置
GPU_COUNT = 2
IMAGES_PER_GPU = 1

COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")#权重文件地址

NUM_CLASSES = 1 + 80  # 背景加类别

#图片像素大小(用自己数据中最低的分辨率）
IMAGE_MIN_DIM = 768
IMAGE_MAX_DIM = 1024

#训练次数
model.train(dataset_train, dataset_val, 
            learning_rate=config.LEARNING_RATE, 
            epochs=300,
            layers='heads')

self.add_class("shapes", 1, "box")#逐个添加每个类别

#逐个添加每个类别 
if labels[i].find("box")!=-1:
    labels_form.append("box")

dataset_root_path=os.path.join(ROOT_DIR, "C:\\Users\\LARA\\Desktop\\Mask_RCNN-master\\train_data_test1")#数据集目录

#图片尺寸
width = 1024
height = 768
dataset_train.load_shapes(count, 768, 1024, img_floder, mask_floder, imglist,dataset_root_path)
dataset_val.load_shapes(count, 768, 1024, img_floder, mask_floder, imglist,dataset_root_path)

改写train.py中的init_with=““last””,下次训练时会接着上一次的epoch继续训练。