目标检测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文件