コード例 #1
0
    def __init__(self, num_classes, pretrained=None):
        super(MaskRCNN, self).__init__()
        if pretrained is not None:
            assert pretrained in ['imagenet', 'voc2007', 'coco']
            assert pretrained not in ['voc2007', 'coco'], ("VOC2007 and COCO pretrained weights "
                                                           "will be added soon.")
        self.config = ConfigParser()
        self.config.read(os.path.join(os.path.dirname(os.path.realpath(__file__)), "config.ini"))
        self.num_classes = num_classes
        self.pooling_size_clsbbox = (7, 7)
        self.pooling_size_mask = (14, 14)
        self.use_fpn = False  # FPN backbone feature is debugging right now.
        self.train_rpn_only = False

        if self.use_fpn:
            self.backbone_fpn = ResNet_101_FPN(pretrained=pretrained)
            self.depth = 256
        else:
            self.backbone = ResNet_101(pretrained=pretrained)
            self.depth = 1024

        self.rpn = RPN(dim=self.depth, use_fpn=self.use_fpn)

        if not self.train_rpn_only:
            # RoiAlign for cls and bbox head, pooling size 7x7
            self.roi_align_clsbbox = RoiAlign(grid_size=self.pooling_size_clsbbox)
            # RoiAlign for mask head, pooling size 14x14
            self.roi_align_mask = RoiAlign(grid_size=self.pooling_size_mask)
            self.clsbbox_head = ClsBBoxHead(depth=self.depth, pool_size=self.pooling_size_clsbbox,
                                            num_classes=num_classes)
            self.mask_head = MaskHead(depth=self.depth, pool_size=self.pooling_size_mask,
                                      num_classes=num_classes)
        self.img_height = None
        self.img_width = None
        self.batch_size = None
コード例 #2
0
    def __init__(self, num_classes, pretrained=None):
        """
        
        Args:
            num_classes(int): number of classes, background should be counted in.
                e.g: there are 100 foreground objects, num_classes should be 101.                    
            pretrained(str): 'imagenet' or 'coco', set 'imagenet' indicate just
                backbone use imagenet pretrained weights, 'coco' indicate whole
                Mask R-CNN model use pretrained weights on COCO dataset.
        """

        super(MaskRCNN, self).__init__()
        if pretrained is not None:
            assert pretrained in ['imagenet', 'coco']
            assert pretrained not in [
                'coco'
            ], "COCO pretrained weights is not available yet."
        self.config = ConfigParser()
        self.config.read(
            os.path.abspath(os.path.join(__file__, "../", "config.ini")))
        self.num_classes = num_classes
        self.pooling_size_clsbbox = (7, 7)
        self.pooling_size_mask = (14, 14)
        self.validating = False  # when True output loss and predict results.
        self.fpn_on = bool(int(self.config['BACKBONE']['FPN_ON']))
        self.mask_head_on = bool(int(self.config['HEAD']['MASK_HEAD_ON']))
        self.train_rpn_only = bool(int(self.config['TRAIN']['TRAIN_RPN_ONLY']))
        resnet_layer = int(self.config['BACKBONE']['RESNET_LAYER'])
        if self.fpn_on:
            self.backbone_fpn = ResNetFPN(resnet_layer, pretrained=pretrained)
            self.depth = 256
        else:
            self.backbone = ResNet(resnet_layer, pretrained=pretrained)
            self.depth = 1024

        self.rpn = RPN(self.depth, self.fpn_on)

        if not self.train_rpn_only:
            # RoiAlign for cls and bbox head, pooling size 7x7
            self.roi_align_clsbbox = RoiAlign(
                grid_size=self.pooling_size_clsbbox)
            # RoiAlign for mask head, pooling size 14x14
            self.roi_align_mask = RoiAlign(grid_size=self.pooling_size_mask)
            self.clsbbox_head = ClsBBoxHead(
                depth=self.depth,
                pool_size=self.pooling_size_clsbbox,
                num_classes=num_classes)
            if self.mask_head_on:
                self.mask_head = MaskHead(depth=self.depth,
                                          pool_size=self.pooling_size_mask,
                                          num_classes=num_classes)
        self.img_height = None
        self.img_width = None
        self.batch_size = None
コード例 #3
0
 def __init__(self, num_classes):
     super(MaskRCNN, self).__init__()
     self.num_classes = num_classes
     self.fpn = ResNet_101_FPN()
     self.rpn = RPN(dim=512)
     self.roi_align = RoiAlign(grid_size=(14, 14))
     self.cls_box_head = ClsBBoxHead(depth=512,
                                     pool_size=14,
                                     num_classes=num_classes)
     self.mask_head = MaskHead(depth=512,
                               pool_size=14,
                               num_classes=num_classes)
コード例 #4
0
 def __init__(self, num_classes, img_size):
     super(MaskRCNN, self).__init__()
     self.config = ConfigParser()
     self.config.read(
         os.path.join(os.path.dirname(os.path.realpath(__file__)),
                      "config.ini"))
     self.num_classes = num_classes
     self.fpn = ResNet_101_FPN()
     self.rpn = RPN(dim=256)
     self.roi_align = RoiAlign(grid_size=(14, 14))
     self.cls_box_head = ClsBBoxHead(depth=256,
                                     pool_size=14,
                                     num_classes=num_classes)
     self.mask_head = MaskHead(depth=256,
                               pool_size=14,
                               num_classes=num_classes,
                               img_size=img_size)
コード例 #5
0
    def __init__(self, num_classes, pretrained=None):
        super(MaskRCNN, self).__init__()
        if pretrained is not None:
            assert pretrained in ['imagenet', 'voc2007', 'coco']
            assert pretrained not in ['voc2007', 'coco'], ("VOC2007 and COCO pretrained weights "
                                                           "will be added soon.")
        self.config = ConfigParser()
        self.config.read(os.path.abspath(os.path.join(__file__, "../",  "config.ini")))
        self.num_classes = num_classes
        self.pooling_size_clsbbox = (7, 7)
        self.pooling_size_mask = (14, 14)
        use_fpn = bool(int(self.config['TRAIN']['USE_FPN']))
        assert not use_fpn, "FPN backbone is under debugging right now."
        self.use_fpn = use_fpn
        self.train_rpn_only = bool(int(self.config['TRAIN']['TRAIN_RPN_ONLY']))

        if self.use_fpn:
            self.backbone_fpn = ResNet_101_FPN(pretrained=pretrained)
            self.depth = 256
        else:
            self.backbone = ResNet_101(pretrained=pretrained)
            self.depth = 1024

        self.rpn = RPN(self.depth, self.use_fpn)

        if not self.train_rpn_only:
            # RoiAlign for cls and bbox head, pooling size 7x7
            self.roi_align_clsbbox = RoiAlign(grid_size=self.pooling_size_clsbbox)
            # RoiAlign for mask head, pooling size 14x14
            self.roi_align_mask = RoiAlign(grid_size=self.pooling_size_mask)
            self.clsbbox_head = ClsBBoxHead(depth=self.depth, pool_size=self.pooling_size_clsbbox,
                                            num_classes=num_classes)
            self.mask_head = MaskHead(depth=self.depth, pool_size=self.pooling_size_mask,
                                      num_classes=num_classes)
        self.img_height = None
        self.img_width = None
        self.batch_size = None
コード例 #6
0
ファイル: maskrcnn.py プロジェクト: popocq/mask-rcnn-pytorch
    def __init__(self, num_classes, pretrained=None):
        """
        
        Args:
            num_classes(int): number of classes, background should be counted in.
                e.g: there are 100 foreground objects, num_classes should be 101.                    
            pretrained(str): 'imagenet' or 'coco', set 'imagenet' indicate just
                backbone use imagenet pretrained weights, 'coco' indicate whole
                Mask R-CNN model use pretrained weights on COCO dataset.
        """

        super(MaskRCNN, self).__init__()
        if pretrained is not None:
            assert pretrained in ['imagenet', 'coco']
            assert pretrained not in [
                'coco'
            ], "COCO pretrained weights is not available yet."
        self.config = ConfigParser()
        config_path = os.path.abspath(
            os.path.join(__file__, "../", "config.ini"))
        assert os.path.exists(config_path), "config.ini not exists!"
        self.config.read(config_path)
        self.num_classes = num_classes
        self.validating = False  # when validating output loss and predict results.
        self.fpn_on = bool(int(self.config['BACKBONE']['FPN_ON']))
        self.mask_head_on = bool(int(self.config['HEAD']['MASK_HEAD_ON']))
        self.train_rpn_only = bool(int(self.config['TRAIN']['TRAIN_RPN_ONLY']))

        resnet_layer = int(self.config['BACKBONE']['RESNET_LAYER'])
        self.pooling_size_clsbbox = int(
            self.config['POOLING']['POOLING_SIZE_CLSBBOX'])
        self.pooling_size_mask = int(
            self.config['POOLING']['POOLING_SIZE_MASK'])

        if self.fpn_on:
            self.backbone_fpn = ResNetFPN(resnet_layer, pretrained=pretrained)
            self.depth_for_rpn = 256
            self.depth_for_head = 2048
        else:
            backbone_type = self.config['BACKBONE']['BACKBONE_TYPE']
            assert backbone_type in ['resnet', 'vgg16']
            if backbone_type == 'resnet':
                self.backbone = ResNet(resnet_layer, pretrained=pretrained)
                self.depth_rpn = 1024
                self.depth_head = 2048
            elif backbone_type == 'vgg16':
                self.backbone = VGG16(pretrained=pretrained)
                self.depth_rpn = 512
                self.depth_head = 4096

        self.rpn = RPN(self.depth_rpn, self.fpn_on)
        self.head_common = HeadCommon(pretrained=pretrained)

        if not self.train_rpn_only:
            # RoiAlign for cls and bbox head
            self.roi_align_clsbbox = RoiAlign(
                grid_size=(self.pooling_size_clsbbox,
                           self.pooling_size_clsbbox))
            self.clsbbox_head = ClsBBoxHead(depth=self.depth_head,
                                            num_classes=num_classes)
            if self.mask_head_on:
                # RoiAlign for mask head
                self.roi_align_mask = RoiAlign(
                    grid_size=(self.pooling_size_mask, self.pooling_size_mask))
                self.mask_head = MaskHead(depth=self.depth_head,
                                          pool_size=(self.pooling_size_mask,
                                                     self.pooling_size_mask),
                                          num_classes=num_classes)
        self.img_height = None
        self.img_width = None
        self.batch_size = None