def get_model(ie, args):
input_transform = models.InputTransform(args.reverse_input_channels, args.mean_values, args.scale_values)
common_args = (ie, args.model, input_transform)
if args.architecture_type in ('ctpn', 'yolo', 'yolov4', 'retinaface',
'retinaface-pytorch') and not input_transform.is_trivial:
raise ValueError("{} model doesn't support input transforms.".format(args.architecture_type))
if args.architecture_type == 'ssd':
return models.SSD(*common_args, labels=args.labels, keep_aspect_ratio_resize=args.keep_aspect_ratio)
elif args.architecture_type == 'ctpn':
return models.CTPN(ie, args.model, input_size=args.input_size, threshold=args.prob_threshold)
elif args.architecture_type == 'yolo':
return models.YOLO(ie, args.model, labels=args.labels,
threshold=args.prob_threshold, keep_aspect_ratio=args.keep_aspect_ratio)
elif args.architecture_type == 'yolov4':
return models.YoloV4(ie, args.model, labels=args.labels,
threshold=args.prob_threshold, keep_aspect_ratio=args.keep_aspect_ratio,
anchors=args.anchors, masks=args.masks)
elif args.architecture_type == 'faceboxes':
return models.FaceBoxes(*common_args, threshold=args.prob_threshold)
elif args.architecture_type == 'centernet':
return models.CenterNet(*common_args, labels=args.labels, threshold=args.prob_threshold)
elif args.architecture_type == 'retinaface':
return models.RetinaFace(ie, args.model, threshold=args.prob_threshold)
elif args.architecture_type == 'ultra_lightweight_face_detection':
return models.UltraLightweightFaceDetection(*common_args, threshold=args.prob_threshold)
elif args.architecture_type == 'retinaface-pytorch':
return models.RetinaFacePyTorch(ie, args.model, threshold=args.prob_threshold)
else:
raise RuntimeError('No model type or invalid model type (-at) provided: {}'.format(args.architecture_type))
def get_model(ie, args):
if args.architecture_type == 'ssd':
return models.SSD(ie, args.model, labels=args.labels, keep_aspect_ratio_resize=args.keep_aspect_ratio)
elif args.architecture_type == 'ctpn':
return models.CTPN(ie, args.model, input_size=args.input_size, threshold=args.prob_threshold)
elif args.architecture_type == 'yolo':
return models.YOLO(ie, args.model, labels=args.labels,
threshold=args.prob_threshold, keep_aspect_ratio=args.keep_aspect_ratio)
elif args.architecture_type == 'yolov4':
return models.YoloV4(ie, args.model, labels=args.labels,
threshold=args.prob_threshold, keep_aspect_ratio=args.keep_aspect_ratio)
elif args.architecture_type == 'faceboxes':
return models.FaceBoxes(ie, args.model, threshold=args.prob_threshold)
elif args.architecture_type == 'centernet':
return models.CenterNet(ie, args.model, labels=args.labels, threshold=args.prob_threshold)
elif args.architecture_type == 'retinaface':
return models.RetinaFace(ie, args.model, threshold=args.prob_threshold)
else:
raise RuntimeError('No model type or invalid model type (-at) provided: {}'.format(args.architecture_type))