目标检测1-Mask RCNN模型气球检测

1.数据收集

• (1) 拍照
• (2) 网站爬虫：flickr
• (3) 公司提供
• (4) roboflow
• (5 )kaggle
• (6) 天池

2.数据标注

• (1) Labelme数据标注
• (2) Elseg数据标注
• (3) VIA数据标注

3.配置文件及数据集：balloon_dataset.py

• 文件存于balloon目录下

import sys
import os
import json
import skimage
import skimage.draw
import numpy as np
# sys.path.append('./mrcnn')
from mrcnn.config import Config
from mrcnn import utils
from mrcnn import visualize

# 3.1 指定预训练权重路径
COCO_WEIGHTS_PATH = '../mask_rcnn_coco.h5'
 
# 3.2 定义气球配置文件
class BalloonConfig(Config):
    NAME = 'balloon'
    STEPS_PER_EPOCH = 100
    NUM_CLASSES = 1 + 1  # bg + balloon
    DETECTION_MIN_CONFIDENCE = 0.9  # 置信度
class InferenceConfig(BalloonConfig):
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
# config = BalloonConfig() # 校验配置文件
# config.display()

# 3.3 气球数据集
class BalloonDataset(utils.Dataset):
    def load_balloon(self, dataset_dir, subset):
        self.add_class('balloon', 1, 'balloon')
        assert subset in ['train', 'val']
        dataset_dir = os.path.join(dataset_dir, subset)
        # 从via产生的json标注文件中提取标注信息
        with open(os.path.join(dataset_dir, 'via_region_data.json')) as fp:
            annotations = json.load(fp)   # 加载json数据
        annotations = list(annotations.values())  # 取values值
        annotations = [a for a in annotations if a['regions']]  # 取regions值
        for a in annotations:
            polygons = [r['shape_attributes'] for r in a['regions'].values()]  # 取多边形即x,y值
            image_path = os.path.join(dataset_dir, a['filename'])
            image = skimage.io.imread(image_path)
            height, width = image.shape[:2]
            self.add_image('balloon', image_id=a['filename'],
                           path=image_path,
                           width=width, height=height,
                           polygons=polygons)
    def load_mask(self, image_id):
        image_info = self.image_info[image_id]  # 根据image_id查到image的信息
        if image_info['source'] != 'balloon':
            return super().load_mask(image_id)
        # 产生一个跟原图一样大的背景图 目标为1其他为0
        mask = np.zeros([image_info['height'], image_info['width'], len(image_info['polygons'])],
                 dtype=np.uint8)
        for i, p in enumerate(image_info['polygons']):
            rr, cc = skimage.draw.polygon(p['all_points_y'], p['all_points_x'])
            mask[rr, cc, i] = 1
        return mask, np.ones([mask.shape[-1]], dtype=np.int32)
    def image_reference(self, image_id):
        info = self.image_info[image_id]
        if info['source'] == 'balloon':
            return info['path']
        else:
            super().image_reference(image_id)
            

# 3.4 JSON数据

# 3.5 测试代码
if __name__ == '__main__':
    dataset = BalloonDataset()
    dataset.load_balloon('../../balloon', 'train')
    dataset.prepare()
    print(f'Image Count: {len(dataset.image_ids)}')
    print(f'Class Count: {dataset.num_classes}')
    for i, info in enumerate(dataset.class_info):
        print(f'{i} {info["name"]}')
    image_ids = np.random.choice(dataset.image_ids, 4)
    print(image_ids)
    for image_id in image_ids:
        image = dataset.load_image(image_id)
        mask, class_ids = dataset.load_mask(image_id)
        visualize.display_top_masks(image, mask, class_ids, dataset.class_names)

4.训练及预测：train_detect.py

• 文件与mrcnn文件同级目录

import skimage
from balloon.balloon_dataset import BalloonDataset
from balloon.balloon_dataset import BalloonConfig
from balloon.balloon_dataset import InferenceConfig
from mrcnn.model import MaskRCNN
from mrcnn.visualize import display_instances
config = BalloonConfig()

# 4.1 训练
def train(model): 
    dataset_train = BalloonDataset()
    dataset_train.load_balloon('../balloon', 'train')  # 读数据
    dataset_train.prepare()  # 预处理
    dataset_val = BalloonDataset()
    dataset_val.load_balloon('../balloon', 'val')
    dataset_val.prepare()
    model.train(dataset_train, dataset_val,
                learning_rate=config.LEARNING_RATE,
                epochs=10,
                layers='heads')
                
# 4.2 测试训练代码
if __name__ == '__main__':  
    # model = MaskRCNN(mode='training', config=config, model_dir='./')
    # train(model)
    
# 4.3 命令行
    import argparse
    parser = argparse.ArgumentParser()  # 命令行解释器
    parser.add_argument('--mode', default='training', help="'train' or 'inference'")  # 增加命令行
    parser.add_argument('--dataset', required=False, default='../balloon', help='directory of balloon dataset')
    parser.add_argument('--weights', required=False, default='./mask_rcnn_coco.h5', help='path to weights')
    parser.add_argument('--image', help='Image path to apply detection', required=False)
    args = parser.parse_args()  # 解析命令行参数
    
# 4.4 命令行训练
    if args.mode == 'training':
        model = MaskRCNN(mode=args.mode, config=config, model_dir='balloon/')
        model.load_weights(args.weights, by_name=True)  # 加载初始参数
        train(model)
        
# 4.5 命令行预测
    else:
        config = InferenceConfig()  # 预测配置文件
        model = MaskRCNN(mode=args.mode, config=config, model_dir='balloon/')
        model.load_weights(args.weights, by_name=True)
        class_names = ['BG', 'balloon']
        image = skimage.io.imread(args.image)
        result = model.detect([image], verbose=1)  # 随机选择
        r = result[0]
        display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names,
                          r['scores'])

# 4.5 train_detect.py文件

5.mrcnn文件