def get_model_instance_segmentation_v2(num_classes,
architecture: str = 'resnet18'):
""" By modifying this function we will be able to use a large variety of
pretrained backbones but besides the backbones nothing else will be trained.
A better solution seems to be to load a pre-trained model and then to
change the mask and box predictors.
"""
# Pretrained model for num_classes=1000, but we will not use the final layers anyway.
model = pretrainedmodels.__dict__[architecture](num_classes=1000,
pretrained='imagenet')
my_backbone = MyBackbone(model.features, 512)
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
mask_roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=14,
sampling_ratio=2)
model = MaskRCNN(my_backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler,
mask_roi_pool=mask_roi_pooler)
return model
def __init__(self,
backbone=None,
architecture=None,
detector=None,
num_classes=None,
device='cpu',
*args,
**kwargs):
assert backbone is not None, ValueError('backbone can not None')
assert architecture is not None, ValueError(
'architecture can not None')
assert detector is not None, ValueError('detector can not None')
assert num_classes is not None, ValueError('num_classes can not None')
assert device is not None, ValueError('device can not None.')
self.device = device
super.__init__()
if backbone == 'efficientnet':
backbone = EfficientNet.from_pretrained(architecture)
backbone.out_channels = 1280
elif backbone == 'fishnet':
if architecture == 'fishnet99':
backbone = fishnet99()
elif architecture == 'fishnet150':
backbone = fishnet150()
else:
backbone = fishnet201()
backbone.out_channels = 1000
elif backbone == 'resnet':
backbone = resnet_fpn_backbone(architecture, pretrained=True)
self.model = MaskRCNN(backbone, num_classes=num_classes)
self.model.to(device)
def get_model(
backbone_name="resnet50",
detector_name="fasterrcnn",
trainable_layers=3,
model_ckpt=None,
):
"""Constructs a fasterrcnn or maskrcnn detector with the given backbone"""
num_classes = 2 # 1 class (wheat) + background
if model_ckpt:
# backbone = resnet_fpn_backbone('resnet101', True)
backbone = timm_resnet_fpn_backbone(backbone_name, False,
trainable_layers)
else:
backbone = timm_resnet_fpn_backbone(backbone_name, True,
trainable_layers)
if detector_name == "fasterrcnn":
model = FasterRCNN(backbone, num_classes)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
elif detector_name == "maskrcnn":
model = MaskRCNN(backbone, num_classes)
in_features_mask = (
model.roi_heads.mask_predictor.conv5_mask.in_channels)
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(
in_features_mask, hidden_layer, num_classes)
else:
raise Exception(f"{detector_name} is not supported")
if model_ckpt is not None:
model.load_state_dict(torch.load(model_ckpt)["model_state_dict"])
print("loaded ckpt")
return model
def get_instance_segmentation_model(num_classes,
model_name='maskrcnn_resnet50_fpn'):
# load a pre-trained model for classification
# and return only the features
if model_name.startswith('efficientnet'):
backbone = EfficientNet.from_pretrained(model_name,
num_classes=num_classes,
include_top=False)
# number of output channels
backbone.out_channels = int(
round_filters(1280, backbone._global_params))
model = MaskRCNN(backbone, num_classes)
else:
# load an instance segmentation model pre-trained on COCO
model = torchvision.models.detection.__dict__[model_name](
pretrained=True)
# get the number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
if model_name.startswith('mask') or model_name.startswith('efficientnet'):
# now get the number of input features for the mask classifier
in_features_mask = \
model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
# and replace the mask predictor with a new one
model.roi_heads.mask_predictor = MaskRCNNPredictor(
in_features_mask, hidden_layer, num_classes)
return model
def maskrcnn_resnet50_fpn(pretrained=False,
progress=True,
num_classes=91,
pretrained_backbone=True,
config=None,
**kwargs):
if pretrained:
# no need to download the backbone if pretrained is set
pretrained_backbone = False
backbone = resnet_fpn_backbone('resnet50', pretrained_backbone, config)
sizes = (32, 64, 128, 256, 512)
# roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0', '1', '2', '3'],
# output_size=7,
# sampling_ratio=2)
anchor_generator = AnchorGenerator(sizes=(sizes),
aspect_ratios=((0.5, 1.0, 2.0)))
model = MaskRCNN(backbone,
num_classes,
rpn_anchor_generator=anchor_generator,
**kwargs)
if pretrained:
state_dict = load_state_dict_from_url(
model_urls['maskrcnn_resnet50_fpn_coco'], progress=progress)
model.load_state_dict(state_dict)
return model
def __init__(self):
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
mask_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=14,
sampling_ratio=2)
self.net = MaskRCNN(backbone,
5,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler,
mask_roi_pool=mask_pooler)
for p in self.net.backbone.parameters():
p.requires_grad = False
params = [p for p in self.net.parameters() if p.requires_grad]
self.optim = torch.optim.SGD(params,
lr=0.001,
momentum=0.9,
weight_decay=0.0005)
self.lr_schuduler = torch.optim.lr_scheduler.StepLR(self.optim,
step_size=3,
gamma=0.1)
def get_mask_rcnn(num_classes, max_instances, backbone="resnet101"):
# load an instance segmentation model pre-trained pre-trained on COCO
if backbone == "resnet50":
print("**************Adding Resnet 50 backbone***************")
model = torchvision.models.detection.maskrcnn_resnet50_fpn(
pretrained=True, box_detections_per_img=max_instances)
else:
bb = resnet_fpn_backbone(backbone, False)
model = MaskRCNN(bb,
num_classes=91,
box_detections_per_img=max_instances)
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# now get the number of input features for the mask classifier
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
# and replace the mask predictor with a new one
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer,
num_classes)
return model
def get_model(num_classes):
# load an instance segmentation model pre-trained on COCO
# m = timm.create_model('cspresnet50', pretrained=True, num_classes=0, global_pool='')
# backbone = TimmToVision(m)
m = timm.create_model('cspresnet50', features_only=True, pretrained=True)
backbone = TimmToVisionFPN(m)
#backbone = resnet50_fpn()
model = MaskRCNN(backbone, num_classes)
'''
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
# ["0"] rather than [0]
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=["0"],
output_size=7,
sampling_ratio=2)
model = MaskRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
# get the number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# now get the number of input features for the mask classifier
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
# and replace the mask predictor with a new one
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer,
num_classes)
'''
return model
def get_backbone(num_classes):
# get backbone
backbone = torchvision.models.resnet50(pretrained=True)
# remove the fc layers
new_backbone = torch.nn.Sequential(*(list(backbone.children())[:-2]))
new_backbone.out_channels = 2048
model = MaskRCNN(new_backbone, num_classes)
return model
def get_model_instance_segmentation4(num_classes):
# COCO 에서 미리 학습된 인스턴스 분할 모델을 읽어옵니다
#model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)
#model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)
backbone = torchvision.models.squeezenet1_1(pretrained=False).features
#backbone.out_channels = 1
backbone.out_channels = 512
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
mask_roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=14,
sampling_ratio=2)
model = MaskRCNN(backbone,
num_classes=num_classes,
box_roi_pool=roi_pooler,
mask_roi_pool=mask_roi_pooler)
#print("squeezenet1_0 call2 - out_channels :1280, 18,052,473 / 72M")
#print("squeezenet1_0 call2 - out_channels :516, 4,862,777 / 19.5M")
#print("squeezenet1_1 call2 - out_channels :516, 4,849,849 4,862,777 / 19.5M")
print(
"squeezenet1_1 call2 - out_channels :256, 2,757,369 / 11M (15,000,000 / 15,000,000)"
)
print(
"squeezenet1_1 call2 - out_channels :512, 4,808,441 / 19.2M (15,000,000)"
)
print(
"squeezenet1_1 call2 - out_channels :512, 33,192,463 33,161,683 / 172M (15,000,000)"
)
#
# 분류를 위한 입력 특징 차원을 얻습니다
#in_features = backbone
# 미리 학습된 헤더를 새로운 것으로 바꿉니다
#model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
#in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
#hidden_layer = 1
# and replace the mask predictor with a new one
#model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
# hidden_layer,
# num_classes)
return model
def get_model(num_classes):
# load an instance segmentation model pre-trained on COCO
# m = timm.create_model('cspresnet50', pretrained=True, num_classes=0, global_pool='')
# backbone = TimmToVision(m)
# m = timm.create_model('cspresnet50', features_only=True, pretrained=True)
m = timm.create_model('cspresnet50',
features_only=True,
pretrained=True,
pretrained_strict=False)
backbone = TimmToVisionFPN(m)
# m = timm.create_model('cspresnet50', pretrained=True, num_classes=0, global_pool='')
# backbone = TimmToVision(m,1024)
anchor_sizes = ((32, ), (64, ), (128, ), (256, ), (512, ))
aspect_ratios = ((0.5, 1.0, 2.0), ) * len(anchor_sizes)
anchor_generator = AnchorGenerator(sizes=anchor_sizes,
aspect_ratios=aspect_ratios)
# ["0"] rather than [0]
out_channels = backbone.out_channels
num_anchors = anchor_generator.num_anchors_per_location()[0]
# CascadeRPN
rpn_head = CascadeRPNHead(out_channels,
feat_channels=out_channels,
num_anchors=num_anchors,
stage=2)
# model = FasterRCNN(backbone, num_classes=num_classes, rpn_head=rpn_head)
model = MaskRCNN(backbone, num_classes=num_classes)
# IA branch
# model = FasterRCNNIA(backbone, num_classes=num_classes, rpn_head=rpn_head)
# Box head branch
model.roi_heads.box_head = RoIFeatureExtractor(num_inputs=256,
resolution=7)
model.roi_heads.box_predictor = RoIBoxPredictor(num_classes)
return model
def get_model_instance_segmentation_mn2(num_classes):
# load a pre-trained model for classification and return
# only the features
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# MaskRCNN needs to know the number of
# output channels in a backbone. For mobilenet_v2, it's 1280
# so we need to add it here
backbone.out_channels = 1280
# let's make the RPN generate 5 x 3 anchors per spatial
# location, with 5 different sizes and 3 different aspect
# ratios. We have a Tuple[Tuple[int]] because each feature
# map could potentially have different sizes and
# aspect ratios
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
# let's define what are the feature maps that we will
# use to perform the region of interest cropping, as well as
# the size of the crop after rescaling.
# if your backbone returns a Tensor, featmap_names is expected to
# be [0]. More generally, the backbone should return an
# OrderedDict[Tensor], and in featmap_names you can choose which
# feature maps to use.
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=7,
sampling_ratio=2)
mask_roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=14,
sampling_ratio=2)
# put the pieces together inside a MaskRCNN model
model = MaskRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler,
mask_roi_pool=mask_roi_pooler)
return model
import torch
import torchvision
from torchvision.models.detection import MaskRCNN
from torchvision.models.detection.anchor_utils import AnchorGenerator
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),aspect_ratios=((0.5, 1.0, 2.0),))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=7,sampling_ratio=2)
mask_roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=14,sampling_ratio=2)
MaskRCNN_mobile_model = MaskRCNN(backbone,num_classes=2,rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler,mask_roi_pool=mask_roi_pooler)
if __name__ == '__main__':
model = MaskRCNN_mobile_model
model.eval()
x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
predictions = model(x)