语义分割 Maskrcnn 训练COCO格式数据集

Maskrcnn 训练COCO格式数据集

使用coco-annotator标注工具

coco.py改为

"""
Mask R-CNN
Configurations and data loading code for MS COCO.

Copyright (c) 2017 Matterport, Inc.
Licensed under the MIT License (see LICENSE for details)
Written by Waleed Abdulla

------------------------------------------------------------

Usage: import the module (see Jupyter notebooks for examples), or run from
       the command line as such:

    # Train a new model starting from pre-trained COCO weights
    python3 coco.py train --dataset=/path/to/coco/ --model=coco

    # Train a new model starting from ImageNet weights. Also auto download COCO dataset
    python3 coco.py train --dataset=/path/to/coco/ --model=imagenet --download=True

    # Continue training a model that you had trained earlier
    python3 coco.py train --dataset=/path/to/coco/ --model=/path/to/weights.h5

    # Continue training the last model you trained
    python3 coco.py train --dataset=/path/to/coco/ --model=last

    # Run COCO evaluatoin on the last model you trained
    python3 coco.py evaluate --dataset=/path/to/coco/ --model=last
"""

import os
import sys
import time
import numpy as np
import imgaug  # https://github.com/aleju/imgaug (pip3 install imgaug)

# 下载和安装Python COCO工具  https://github.com/waleedka/coco
# 这个fork的原始地址 https://github.com/pdollar/coco 修复在Python 3下的bug.
# I submitted a pull request https://github.com/cocodataset/cocoapi/pull/50
# If the PR is merged then use the original repo.
# Note: Edit PythonAPI/Makefile and replace "python" with "python3".
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提供
DEFAULT_LOGS_DIR = os.path.join(ROOT_DIR, "logs")
DEFAULT_DATASET_YEAR = "2014"

############################################################
#  配置
############################################################


class CocoConfig(Config):
    """Configuration for training on MS COCO.
    Derives from the base Config class and overrides values specific
    to the COCO dataset.
    """
    # 给配置一个名字
    NAME = "coco"

    # 我们使用12G显存的GPU, 可以放两张图片.
    # 如果显存更小就调节图片数量.
    IMAGES_PER_GPU = 1

    # 取消对在8 gpus上训练的注释（默认值为1）
    # GPU_COUNT = 8

    # 分类数 (包含背景)
    NUM_CLASSES = 1 + 2  # COCO 有80个分类



    IMAGE_MIN_DIM = 512
    IMAGE_MAX_DIM = 1152
    STEPS_PER_EPOCH = 100
############################################################
#  数据集
############################################################

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: 如果提供了, 只读取给定分类id的图片.
        class_map: TODO:尚未实现。支持将不同数据集的类映射到同一个类ID。
        return_coco: 如果为真, 返回COCO对象.
        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=