def faster_rcnn(n_class, backbone='vgg16', model_path=None):
if backbone == 'vgg16':
if model_path is None:
vgg_pretrained = True
else:
vgg_pretrained = False
extractor, head, feature_dim = get_vgg16_extractor_and_head(n_class, roip_size=7, vgg_pretrained=vgg_pretrained)
rpn_net = rpn(in_channel=feature_dim, mid_channel=512, ratio=[0.5, 1, 2], anchor_size=[128, 256, 512])
model = _Faster_RCNN_Maker(extractor, rpn_net, head)
return model
elif backbone == 'resnet50':
raise ValueError("resnet50 has not been implemented!")
else:
raise ValueError("backbone only support vgg16, resnet50!")
def faster_rcnn(n_class, backbone='vgg16', model_path=None):
if backbone == 'vgg16':
if model_path is None:
vgg_pretrained = True
else:
vgg_pretrained = False
#head is a class of vgg16 head
extractor, head, feature_dim = get_vgg16_extractor_and_head(
n_class, roip_size=7, vgg_pretrained=vgg_pretrained)
#rpn_net is also a class of rpn, remember this object oriented programing
rpn_net = rpn(in_channel=feature_dim,
mid_channel=512,
ratio=[0.5, 1, 2],
anchor_size=[128, 256, 512])
#model is class created from the instance of _faster_RCNN_maker which took rpn_net class as parent
model = _Faster_RCNN_Maker(extractor, rpn_net, head)
return model
else:
raise ValueError("backbone only support vgg16, resnet50!")