mask rcnn 之coco.py 解讀

代碼來自：https://github.com/matterport/Mask_RCNN

使用coco.py的指令行代碼是這樣寫的：

你可以選擇谷歌翻譯。

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

# Train a new model starting from ImageNet weights
python3 samples/coco/coco.py train --dataset=/path/to/coco/ --model=imagenet

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

# Continue training the last model you trained. This will find
# the last trained weights in the model directory.
python3 samples/coco/coco.py train --dataset=/path/to/coco/ --model=last
           

coco.py這個檔案是訓練測試的入口。它一共包含了以下方法。

1.首先是與其他功能性python檔案的聯系：

config是一個統一的配置檔案，model是靜态圖模型的代碼

from mrcnn.config import Config
from mrcnn import model as modellib, utils
           

然後可以在這裡改一部分config的參數

class CocoConfig(Config):
           

2.資料集類

class CocoDataset(utils.Dataset):
           

包含了若幹方法

#加載coco資料集的子集 
def load_coco(self, dataset_dir, subset, year=DEFAULT_DATASET_YEAR, class_ids=None,
                  class_map=None, return_coco=False, auto_download=False):
#自動下載下傳
def auto_download(self, dataDir, dataType, dataYear):
#image_id會給定圖檔比如000001.jpg，這個方法會給出mask
def load_mask(self, image_id):
#對給定圖檔傳回他在coco的網址
def image_reference(self, image_id):
#下面兩個本來是pycocotools裡面的方法，坐着改動了一些，寫在這裡
def annToRLE(self, ann, height, width):
def annToMask(self, ann, height, width):
           

3.根據已經訓練好的網絡，給出evaluation結果：

#使給出的inference結果能夠符合COCO标準，這樣才能和别人比
def build_coco_results(dataset, image_ids, rois, class_ids, scores, masks):
           

在調整了我們自己輸出的結果後，就可以用官方計量tool了

#利用官方代碼計算我們網絡的準确率
def evaluate_coco(model, dataset, coco, eval_type="bbox", limit=0, image_ids=None):
           

4.訓練

訓練是coco.py的主體部分：

if __name__ == '__main__':
    import argparse
    #下面這部分會接收你傳入的各種參數，指令行語句。就是我們最初提到的使用coco.py
    # Parse command line arguments
    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)
    #配置，大家都知道網絡需要各種配置參數
    # Configurations 
    if args.command == "train":
        config = CocoConfig()
    else:
        class InferenceConfig(CocoConfig):
            # 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
            DETECTION_MIN_CONFIDENCE = 0
        config = InferenceConfig()
    config.display()

    #建立模型
    #注意，這裡用的是model.py中的MaskRCNN這個類的執行個體化模型，該模型叫做model，
    #下面就可以利用model裡面的方法和屬性了
    # Create 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)

    #選擇一個權重檔案，通常我們至少用VGG訓練的權重，很少有人用随機的權重
    # Select weights file to load
    if args.model.lower() == "coco":
        model_path = COCO_MODEL_PATH
    elif args.model.lower() == "last":
        # Find last trained weights
        model_path = model.find_last()
    elif args.model.lower() == "imagenet":
        # Start from ImageNet trained weights
        model_path = model.get_imagenet_weights()
    else:
        model_path = args.model

    #加載這個權重穩健
    # Load weights
    print("Loading weights ", model_path)
    model.load_weights(model_path, by_name=True)

    #根據你指令行輸入的語句進行訓練或者是評價
    # Train or evaluate
    if args.command == "train":
        # Training dataset. Use the training set and 35K from the
        # validation set, as as in the Mask RCNN paper.
        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()


        # Validation dataset
        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()

        # Image Augmentation
        # Right/Left flip 50% of the time
        augmentation = imgaug.augmenters.Fliplr(0.5)
        #訓練過程，你可以改裡面的參數
        # *** This training schedule is an example. Update to your needs ***

        # Training - Stage 1
        print("Training network heads")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE,
                    epochs=40,
                    layers='heads',
                    augmentation=augmentation)

        # Training - Stage 2
        # Finetune layers from ResNet stage 4 and up
        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)

        # Training - Stage 3
        # Fine tune all 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":
        # Validation dataset
        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))
           

可以看到這裡使用了model裡面的train方法來訓練。coco.py其實是在傳遞參數，一些初始化内容。