版权声明：本文为博主原创文章，未经作者允许请勿转载。 https://blog.csdn.net/heiheiya https://blog.csdn.net/heiheiya/article/details/82143967
一、导包、设置全局变量
import os
import sys
import time
import numpy as np
import imgaug # https://github.com/aleju/imgaug (pip3 install imgaug)
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from pycocotools import mask as maskUtils
import zipfile
import urllib.request
import shutil
#设置根目录
ROOT_DIR = os.path.abspath(".")
#导入Mask RCNN
sys.path.append(ROOT_DIR)
from mrcnn.config import Config
from mrcnn import model as modellib, utils
#预训练权重文件路径
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
#保存logs和model checkpoints的路径,可以通过命令行参数--logs设置
DEFAULT_LOGS_DIR = os.path.join(ROOT_DIR, "logs")
DEFAULT_DATASET_YEAR = "2014"
二、配置
class CocoConfig(Config):
""训练MS COCO数据集的配置信息。
从Config class基类继承并重写训练COCO数据集相关的值。
"""
#为该配置命名
NAME = "coco"
#每个GPU一次处理多少幅图像
IMAGES_PER_GPU = 1
#共有多少GPU参与训练(默认是1)
# GPU_COUNT = 8
#有多少类物体需要分类 (包括背景)
NUM_CLASSES = 1 + 80 # COCO数据集有80类三、数据集
3.1 加载coco数据集
class CocoDataset(utils.Dataset):
def load_coco(self, dataset_dir, subset, year=DEFAULT_DATASET_YEAR, class_ids=None,
class_map=None, return_coco=False, auto_download=False):
"""加载COCO数据集的一个子集.
dataset_dir:COCO数据集的根目录.
subset:需要下载的子集(train, val, minival, valminusminival)
year:下载哪一年的数据集(2014, 2017)，字符串而非整型
class_ids:如果给出，则只加载给定类别的图像.
class_map: TODO: 还未实现. 将不同数据集的类别映射到相同的class ID.
return_coco: If True, returns the COCO object.
auto_download: 自动下载并解压MS-COCO图像和标注文件
"""
if auto_download is True:
self.auto_download(dataset_dir, subset, year)
coco = COCO("{}/annotations/instances_{}{}.json".format(dataset_dir, subset, year))
if subset == "minival" or subset == "valminusminival":
subset = "val"
image_dir = "{}/{}{}".format(dataset_dir, subset, year)
#加载全部类别还是某个子集?
if not class_ids:
#全部类别
class_ids = sorted(coco.getCatIds())
#全部图像还是某个子集?
if class_ids:
image_ids = []
for id in class_ids:
image_ids.extend(list(coco.getImgIds(catIds=[id])))
#去掉重复
image_ids = list(set(image_ids))
else:
#全部图像
image_ids = list(coco.imgs.keys())
#添加类别
for i in class_ids:
self.add_class("coco", i, coco.loadCats(i)[0]["name"])
#添加图像
for i in image_ids:
self.add_image(
"coco", image_id=i,
path=os.path.join(image_dir, coco.imgs[i]['file_name']),
width=coco.imgs[i]["width"],
height=coco.imgs[i]["height"],
annotations=coco.loadAnns(coco.getAnnIds(
imgIds=[i], catIds=class_ids, iscrowd=None)))
if return_coco:
return coco3.2 自动下载数据集
def auto_download(self, dataDir, dataType, dataYear):
"""下载COCO数据集/标注文件.
dataDir:COCO数据集的根目录.
dataType:下载哪个子集(train, val, minival, valminusminival)
dataYear:下载哪一年的数据(2014, 2017)，string类型而非整型
Note:
For 2014, "train", "val", "minival", or "valminusminival"
For 2017,"train" and "val"
"""
#设置路径和文件名
if dataType == "minival" or dataType == "valminusminival":
imgDir = "{}/{}{}".format(dataDir, "val", dataYear)
imgZipFile = "{}/{}{}.zip".format(dataDir, "val", dataYear)
imgURL = "http://images.cocodataset.org/zips/{}{}.zip".format("val", dataYear)
else:
imgDir = "{}/{}{}".format(dataDir, dataType, dataYear)
imgZipFile = "{}/{}{}.zip".format(dataDir, dataType, dataYear)
imgURL = "http://images.cocodataset.org/zips/{}{}.zip".format(dataType, dataYear)
# print("Image paths:"); print(imgDir); print(imgZipFile); print(imgURL)
#若文件夹不存在，则创建
if not os.path.exists(dataDir):
os.makedirs(dataDir)
#如果图像不存在，则下载
if not os.path.exists(imgDir):
os.makedirs(imgDir)
print("Downloading images to " + imgZipFile + " ...")
with urllib.request.urlopen(imgURL) as resp, open(imgZipFile, 'wb') as out:
shutil.copyfileobj(resp, out)
print("... done downloading.")
print("Unzipping " + imgZipFile)
with zipfile.ZipFile(imgZipFile, "r") as zip_ref:
zip_ref.extractall(dataDir)
print("... done unzipping")
print("Will use images in " + imgDir)
#设置标注文件路径
annDir = "{}/annotations".format(dataDir)
if dataType == "minival":
annZipFile = "{}/instances_minival2014.json.zip".format(dataDir)
annFile = "{}/instances_minival2014.json".format(annDir)
annURL = "https://dl.dropboxusercontent.com/s/o43o90bna78omob/instances_minival2014.json.zip?dl=0"
unZipDir = annDir
elif dataType == "valminusminival":
annZipFile = "{}/instances_valminusminival2014.json.zip".format(dataDir)
annFile = "{}/instances_valminusminival2014.json".format(annDir)
annURL = "https://dl.dropboxusercontent.com/s/s3tw5zcg7395368/instances_valminusminival2014.json.zip?dl=0"
unZipDir = annDir
else:
annZipFile = "{}/annotations_trainval{}.zip".format(dataDir, dataYear)
annFile = "{}/instances_{}{}.json".format(annDir, dataType, dataYear)
annURL = "http://images.cocodataset.org/annotations/annotations_trainval{}.zip".format(dataYear)
unZipDir = dataDir
# print("Annotations paths:"); print(annDir); print(annFile); print(annZipFile); print(annURL)
#如果标注文件不存在，则下载
if not os.path.exists(annDir):
os.makedirs(annDir)
if not os.path.exists(annFile):
if not os.path.exists(annZipFile):
print("Downloading zipped annotations to " + annZipFile + " ...")
with urllib.request.urlopen(annURL) as resp, open(annZipFile, 'wb') as out:
shutil.copyfileobj(resp, out)
print("... done downloading.")
print("Unzipping " + annZipFile)
with zipfile.ZipFile(annZipFile, "r") as zip_ref:
zip_ref.extractall(unZipDir)
print("... done unzipping")
print("Will use annotations in " + annFile)3.3 加载mask
def load_mask(self, image_id):
"""加载给定图像的instance masks.
不同的数据集使用不同的方式存储masks。本函数将不同格式的mask转化
成同一格式的bitmap，维度是[height, width, instances].
返回:
masks:一个bool数组，尺寸是[height, width, instance count]，
一个instance有一个mask.
class_ids:一个一维数组，元素是instance masks的class IDs。
"""
#如果不是COCO数据集的图像, 则退化到其父类.
image_info = self.image_info[image_id]
if image_info["source"] != "coco":
return super(CocoDataset, self).load_mask(image_id)
instance_masks = []
class_ids = []
annotations = self.image_info[image_id]["annotations"]
#创建维度是[height, width, instance_count]的mask和
#每个通道的mask对应的class IDs的列表。
for annotation in annotations:
class_id = self.map_source_class_id(
"coco.{}".format(annotation['category_id']))
if class_id:
m = self.annToMask(annotation, image_info["height"],
image_info["width"])
#有些物体的面积小于一个像素，跳过这些物体，不予处理
if m.max() < 1:
continue
#是否是crowd?如果是，则赋给一个negative class ID.
if annotation['iscrowd']:
#crowds的使用negative class ID
class_id *= -1
#对于crowd masks, annToMask()返回的mask
#有时候会小于给定的尺寸。如果是的话，将其缩放。
if m.shape[0] != image_info["height"] or m.shape[1] != image_info["width"]:
m = np.ones([image_info["height"], image_info["width"]], dtype=bool)
instance_masks.append(m)
class_ids.append(class_id)
#将instance masks打包成数组
if class_ids:
mask = np.stack(instance_masks, axis=2).astype(np.bool)
class_ids = np.array(class_ids, dtype=np.int32)
return mask, class_ids
else:
#调用super class返回一个空的mask
return super(CocoDataset, self).load_mask(image_id)def image_reference(self, image_id):
"""返回图像在COCO Website上的链接"""
info = self.image_info[image_id]
if info["source"] == "coco":
return "http://cocodataset.org/#explore?id={}".format(info["id"])
else:
super(CocoDataset, self).image_reference(image_id)
#下面的两个函数来自pycocotools，并做了一些小的修改。
def annToRLE(self, ann, height, width):
"""
转化多边形的标注, 未压缩的RLE转成RLE.
:返回:二值mask (numpy 2D array)
"""
segm = ann['segmentation']
if isinstance(segm, list):
# 多边形 -- 一个物体可能由多个部分构成
# 将所有部分合并到一个mask rle code
rles = maskUtils.frPyObjects(segm, height, width)
rle = maskUtils.merge(rles)
elif isinstance(segm['counts'], list):
#未压缩的RLE
rle = maskUtils.frPyObjects(segm, height, width)
else:
# rle
rle = ann['segmentation']
return rle
def annToMask(self, ann, height, width):
"""
转化多边形的标注, 未压缩的RLE或者RLE转成二值的mask.
:返回: 二值的mask (numpy 2D array)
"""
rle = self.annToRLE(ann, height, width)
m = maskUtils.decode(rle)
return m四、COCO评估
def build_coco_results(dataset, image_ids, rois, class_ids, scores, masks):
"""排列结果以符合COCO的说明，http://cocodataset.org/#format
"""
#如果没有结果，则返回一个空列表
if rois is None:
return []
results = []
for image_id in image_ids:
#循环获取检测结果
for i in range(rois.shape[0]):
class_id = class_ids[i]
score = scores[i]
bbox = np.around(rois[i], 1)
mask = masks[:, :, i]
result = {
"image_id": image_id,
"category_id": dataset.get_source_class_id(class_id, "coco"),
"bbox": [bbox[1], bbox[0], bbox[3] - bbox[1], bbox[2] - bbox[0]],
"score": score,
"segmentation": maskUtils.encode(np.asfortranarray(mask))
}
results.append(result)
return results
def evaluate_coco(model, dataset, coco, eval_type="bbox", limit=0, image_ids=None):
"""运行官方的COCO验证
dataset:验证数据集
eval_type: "bbox"对应bounding box，"segm"对应segmentation evaluation
limit:如果非0，则使用全部的验证图像
"""
#从数据集中选择COCO图像
image_ids = image_ids or dataset.image_ids
#限制使用多少图像
if limit:
image_ids = image_ids[:limit]
#获得对应的COCO image IDs.
coco_image_ids = [dataset.image_info[id]["id"] for id in image_ids]
t_prediction = 0
t_start = time.time()
results = []
for i, image_id in enumerate(image_ids):
# Load image
image = dataset.load_image(image_id)
#开始检测
t = time.time()
r = model.detect([image], verbose=0)[0]
t_prediction += (time.time() - t)
#将结果转化成COCO的歌声
# 将masks转成uint8类型，因为使用bool型COCO tools会报错
image_results = build_coco_results(dataset, coco_image_ids[i:i + 1],
r["rois"], r["class_ids"],
r["scores"],
r["masks"].astype(np.uint8))
results.extend(image_results)
#加载结果.会用一些附加的属性对其做一些修改.
coco_results = coco.loadRes(results)
#评估
cocoEval = COCOeval(coco, coco_results, eval_type)
cocoEval.params.imgIds = coco_image_ids
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
print("Prediction time: {}. Average {}/image".format(
t_prediction, t_prediction / len(image_ids)))
print("Total time: ", time.time() - t_start)五、训练
if __name__ == '__main__':
import argparse
#解析命令行参数
parser = argparse.ArgumentParser(
description='Train Mask R-CNN on MS COCO.')
parser.add_argument("command",
metavar="<command>",
help="'train' or 'evaluate' on MS COCO")
parser.add_argument('--dataset', required=True,
metavar="/path/to/coco/",
help='Directory of the MS-COCO dataset')
parser.add_argument('--year', required=False,
default=DEFAULT_DATASET_YEAR,
metavar="<year>",
help='Year of the MS-COCO dataset (2014 or 2017) (default=2014)')
parser.add_argument('--model', required=True,
metavar="/path/to/weights.h5",
help="Path to weights .h5 file or 'coco'")
parser.add_argument('--logs', required=False,
default=DEFAULT_LOGS_DIR,
metavar="/path/to/logs/",
help='Logs and checkpoints directory (default=logs/)')
parser.add_argument('--limit', required=False,
default=500,
metavar="<image count>",
help='Images to use for evaluation (default=500)')
parser.add_argument('--download', required=False,
default=False,
metavar="<True|False>",
help='Automatically download and unzip MS-COCO files (default=False)',
type=bool)
args = parser.parse_args()
print("Command: ", args.command)
print("Model: ", args.model)
print("Dataset: ", args.dataset)
print("Year: ", args.year)
print("Logs: ", args.logs)
print("Auto Download: ", args.download)
#配置
if args.command == "train":
config = CocoConfig()
else:
class InferenceConfig(CocoConfig):
#因为在预测时一次只处理一幅图像，所以将batch size设为1
#Batch size = GPU_COUNT * IMAGES_PER_GPU
GPU_COUNT = 1
IMAGES_PER_GPU = 1
DETECTION_MIN_CONFIDENCE = 0
config = InferenceConfig()
config.display()
#创建model
if args.command == "train":
model = modellib.MaskRCNN(mode="training", config=config,
model_dir=args.logs)
else:
model = modellib.MaskRCNN(mode="inference", config=config,
model_dir=args.logs)
#选择加载权重文件
if args.model.lower() == "coco":
model_path = COCO_MODEL_PATH
elif args.model.lower() == "last":
#最近训练的权重
model_path = model.find_last()
elif args.model.lower() == "imagenet":
#ImageNet的权重
model_path = model.get_imagenet_weights()
else:
model_path = args.model
#加载权重
print("Loading weights ", model_path)
model.load_weights(model_path, by_name=True)
#训练或评估
if args.command == "train":
#训练数据集. 使用训练集合35K的验证集，和Mask RCNN论文一致。
dataset_train = CocoDataset()
dataset_train.load_coco(args.dataset, "train", year=args.year, auto_download=args.download)
if args.year in '2014':
dataset_train.load_coco(args.dataset, "valminusminival", year=args.year, auto_download=args.download)
dataset_train.prepare()
#验证数据集
dataset_val = CocoDataset()
val_type = "val" if args.year in '2017' else "minival"
dataset_val.load_coco(args.dataset, val_type, year=args.year, auto_download=args.download)
dataset_val.prepare()
#图像增强
#以50%的概率左右翻转图像
augmentation = imgaug.augmenters.Fliplr(0.5)
#训练 - Stage 1
print("Training network heads")
model.train(dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE,
epochs=40,
layers='heads',
augmentation=augmentation)
#训练 - Stage 2
# Finetune 四层以后的layers
print("Fine tune Resnet stage 4 and up")
model.train(dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE,
epochs=120,
layers='4+',
augmentation=augmentation)
#训练 - Stage 3
# Fine tune所有layers
print("Fine tune all layers")
model.train(dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE / 10,
epochs=160,
layers='all',
augmentation=augmentation)
elif args.command == "evaluate":
# 验证集
dataset_val = CocoDataset()
val_type = "val" if args.year in '2017' else "minival"
coco = dataset_val.load_coco(args.dataset, val_type, year=args.year, return_coco=True, auto_download=args.download)
dataset_val.prepare()
print("Running COCO evaluation on {} images.".format(args.limit))
evaluate_coco(model, dataset_val, coco, "bbox", limit=int(args.limit))
else:
print("'{}' is not recognized. "
"Use 'train' or 'evaluate'".format(args.command))