microScan / models /create_fasterrcnn_model.py
crazyscientist1's picture
initial commit
d70f24c
raw
history blame
637 Bytes
from models import *
def return_fasterrcnn_resnet50_fpn(
num_classes, pretrained=True, coco_model=False
):
model = fasterrcnn_resnet50_fpn.create_model(
num_classes, pretrained=pretrained, coco_model=coco_model
)
return model
def return_fasterrcnn_resnet50_fpn_v2(
num_classes, pretrained=True, coco_model=False
):
model = fasterrcnn_resnet50_fpn_v2.create_model(
num_classes, pretrained=pretrained, coco_model=coco_model
)
return model
create_model = {
'fasterrcnn_resnet50_fpn': return_fasterrcnn_resnet50_fpn,
'fasterrcnn_resnet50_fpn_v2': return_fasterrcnn_resnet50_fpn_v2
}