mask rcnn 之coco.py 解读

╰半夏微凉° 2022-02-26 08:58 360阅读 0赞

代码来自：[https://github.com/matterport/Mask\_RCNN][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其实是在传递参数，一些初始化内容。

[https_github.com_matterport_Mask_RCNN]: https://github.com/matterport/Mask_RCNN