Mask-rcnn训练自己的数据实践

Mask-rcnn训练自己的数据实践，这里只做了一类。

主要步骤：

1. 数据准备

（1）做标签：这里用的labelme，一张图片对应一个.json文件。数据大小1024×1024。

（2）转数据格式：在./labelme/cli/路径下找到 json_to_dataset.py，这里需要小改动一下，实现批量转格式。将步骤（1）中的.json转成训练需要的数据，一张图片对应5个文件，具体如下图。

2. mask-rcnn下载地址：https://github.com/matterport/Mask_RCNN

   pre-trained model: https://github.com/matterport/Mask_RCNN/releases

3. 写train.py文件:

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
from PIL import Image

import yaml
# Root directory of the project
ROOT_DIR = os.path.abspath("./")

# Import Mask RCNN
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn.config import Config
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize
from mrcnn.model import log

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

# iter_num=0

# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, 'pretrain_model' , "mask_rcnn_coco.h5").replace('//', '/')
# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
	utils.download_trained_weights(COCO_MODEL_PATH)

##Confiurations
class ShapesConfig(Config):
	"""Configuration for training on the toy shapes dataset.
	Derives from the base Config class and overrides values specific
	to the toy shapes dataset.
	"""
	# 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 + 1 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 = 1024
	IMAGE_MAX_DIM = 1024

	# Use smaller anchors because our image and objects are small
	RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)  # 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 = 32

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

	# use small validation steps since the epoch is small
	VALIDATION_STEPS = 100
	
config = ShapesConfig()
config.display()

#dataset
class ShapesDataset(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):

		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

	def load_shapes(self, count, height, width, imgfolder, folders_list):

		# Add classes
		self.add_class("shapes", 1, "trafficsign")
		# Add images
		for i in range(count):
			foldername = folders_list[i]
			img_path = imgfolder + foldername + '/' + 'img.png'
			mask_path = imgfolder + foldername + '/' + 'label.png'
			yaml_path = imgfolder + foldername + '/' + 'info.yaml'

			self.add_image("shapes", image_id=i, path=img_path, 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, 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)):
			if labels[i].find("trafficsign")!=-1:
				#print "box"
				labels_form.append("trafficsign")
			
		class_ids = np.array([self.class_names.index(s) for s in labels_form])

		return mask, class_ids.astype(np.int32)

# Training dataset
imgfolder_train = '/home/liesmars/maoz/Data/TStraning/train/'
folder_train_list = os.listdir(imgfolder_train)
count_train = len(os.listdir(imgfolder_train))

dataset_train = ShapesDataset()
dataset_train.load_shapes(count_train, config.IMAGE_SHAPE[0], config.IMAGE_SHAPE[1], imgfolder_train, folder_train_list)
dataset_train.prepare()

# Validation dataset
imgfolder_val = '/home/liesmars/maoz/Data/TStraning/val/'
folder_val_list = os.listdir(imgfolder_val)
count_val = len(os.listdir(imgfolder_val))

dataset_val = ShapesDataset()
dataset_val.load_shapes(count_val, config.IMAGE_SHAPE[0], config.IMAGE_SHAPE[1], imgfolder_val, folder_val_list)
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
	# are different due to the different number of classes
	# See README for instructions to download the COCO weights
	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(), by_name=True)


# Train 
model.train(dataset_train, dataset_val, 
			learning_rate=config.LEARNING_RATE,
			epochs=100, 
			layers="all")

可以用tensorboard看loss曲线，模型和event文件在logs文件夹里面。

4. 准备test.py

import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt
from PIL import Image

# Root directory of the project
ROOT_DIR = os.path.abspath("./")

# 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
from mrcnn.config import Config
# Import COCO config
# sys.path.append(os.path.join(ROOT_DIR, "samples/coco/"))  # To find local version
# import coco
import mrcnn
import pickle

# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
# Local path to trained weights file
MODEL_PATH = os.path.join(MODEL_DIR, "shapes20191012T1103/mask_rcnn_shapes_0080.h5")
# Download COCO trained weights from Releases if needed

# Directory of images to run detection on
IMAGE_DIR = os.path.join(ROOT_DIR, "img")
# Directory to run save the detection results
IMAGE_SAVE = os.path.join(ROOT_DIR, "imgresult")

if not os.path.exists(IMAGE_SAVE):
	os.makedirs(IMAGE_SAVE)

##Confiurations
class ShapesConfig(Config):
	"""Configuration for training on the toy shapes dataset.
	Derives from the base Config class and overrides values specific
	to the toy shapes dataset.
	"""
	# 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 + 1 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 = 1024
	IMAGE_MAX_DIM = 1024
	# Use smaller anchors because our image and objects are small
	RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)  # 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 = 32
	# Use a small epoch since the data is simple
	STEPS_PER_EPOCH = 238
	# use small validation steps since the epoch is small
	VALIDATION_STEPS = 100


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()
config.display()

# 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(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', 'trafficsign']


# Load a random image from the images folder
# file_names = next(os.walk(IMAGE_DIR))[2]
file_names = os.listdir(IMAGE_DIR)
# APs = []

for file in file_names:
	image = skimage.io.imread(os.path.join(IMAGE_DIR, file))

	# Run detection
	results = model.detect([image], verbose=1)
	# Visualize results
	r = results[0]

	savepath = IMAGE_SAVE + '/' + file  # the path to save the image result of testing
	name, ext = os.path.splitext(file)

	arr_mask = r['masks'] 

	## save segmentaion result of object
	np.save(IMAGE_SAVE+'/'+ name +'.npy', arr_mask)

	with open(IMAGE_SAVE + '/' + name + '.pkl', 'wb') as f:
		pickle.dump( {'rois':r['rois'], 'class_ids':r['class_ids'], 'class_names':class_names,'scores': r['scores']} , f)

	## save detection result with plt 
	visualize.display_instances(savepath, image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores'])

5. evaluation:

(1) 将输出预测的结果转为label的图片，用作后面计算IoU。这里的npy文件是预测结果的array，在test阶段生成的。

# To convert the detection results(instance segmantation result) to image for evaluation
from numpy import *
import numpy as np
import cv2
import os

npyDir = './npy/'
outputDir = './pre-label/'

if not os.path.exists(outputDir):
	os.makedirs(outputDir)

for f in os.listdir(npyDir):
	path = os.path.join(npyDir, f)
	savepath = os.path

	mask= np.load(os.path.join(npyDir, f))
	fname, ext = os.path.splitext(f)
	output = outputDir + fname +'.png'

	objnum = mask.shape[2] #object number equels the channel of mask.

	b_0 = np.zeros((1024,1024), dtype=float)

	for i in range (objnum):
		b_0 += 1*np.float32(mask[:,:,i]) #Matrix slice

	b_0 = np.reshape(b_0, (1024,1024,1))
	if not os.path.exists(output):
		cv2.imwrite(output, 1*np.float32(b_0))

(2) 计算IoU，评价训练结果。

# -*- coding: utf-8 -*-
import os
import cv2
import numpy as np

class IOUMetric:
    """
    Class to calculate mean-iou using fast_hist method
    """
    def __init__(self, num_classes):
        self.num_classes = num_classes
        self.hist = np.zeros((num_classes, num_classes))

    def _fast_hist(self, label_pred, label_true):
        # 找出标签中需要计算的类别,去掉了背景
        mask = (label_true >= 0) & (label_true < self.num_classes)

        hist = np.bincount(
            self.num_classes * label_true[mask].astype(int) +
            label_pred[mask], minlength=self.num_classes ** 2).reshape(self.num_classes, self.num_classes)
        return hist

    # input：prediction and true label.
    # 语义分割的任务是为每个像素点分配一个label
    def evaluate(self, predictions, gts):
        for lp, lt in zip(predictions, gts):
            assert len(lp.flatten()) == len(lt.flatten())
            self.hist += self._fast_hist(lp.flatten(), lt.flatten())
            
        # miou
        iou = np.diag(self.hist) / (self.hist.sum(axis=1) + self.hist.sum(axis=0) - np.diag(self.hist))
        miou = np.nanmean(iou) 
        
        # -----------------其他指标------------------------------
        # mean acc
        acc = np.diag(self.hist).sum() / self.hist.sum()
        acc_cls = np.nanmean(np.diag(self.hist) / self.hist.sum(axis=1))

        freq = self.hist.sum(axis=1) / self.hist.sum()
        fwavacc = (freq[freq > 0] * iou[freq > 0]).sum()

        return acc, acc_cls, iou, miou, fwavacc

if __name__ == '__main__':
    label_path = './Labelme2json'
    predict_path = './pre-label'
    pres = os.listdir(predict_path)
    
    labels = []
    predicts = []

    for im in pres:
        name, ext = os.path.splitext(im)
        if ext == '.png':        
            lab_path = os.path.join(label_path, name, 'label.png').replace('\\','/')
            pre_path = os.path.join(predict_path, im)

            label = cv2.imread(lab_path,0) #1024,1024

            test = label.flatten() #(1048576,)
            pre = cv2.imread(pre_path,0)

            labels.append(label)
            predicts.append(pre)

    el = IOUMetric(2) #backgroud + class
    acc, acc_cls, iou, miou, fwavacc = el.evaluate(predicts, labels)
    
    print('acc: ',acc)
    print('acc_cls: ',acc_cls)
    print('iou: ',iou)
    # print('miou: ',miou) #No mean IoU as there is only one class.
    print('fwavacc: ',fwavacc)